Skip to content

Commit 8d42c89

Browse files
authored
test: consolidate memory-pool test helpers and lint bindings pools (#3298)
* test: consolidate duplicated memory-pool test helpers Follow-up to #3290. The same memory-pool helpers were copied across the cuda_bindings_12, cuda_bindings_13, and cuda_core test suites. Move them into one shared module and delete the copies. No test outcome changes: the conditions behind every skipif and skip are the same. New cuda_python_test_helpers/memory.py: * POOL_SIZE (9 copies, including the cuda_core constant, which now re-exports it). * supports_memory_pools() and supports_managed_memory() (13 copies, including the two Cython test modules). * xfail_graph_mempool_oom() (3 copies of the xfail_on_graph_mempool_oom fixture). Each call site keeps its own out-of-memory check and cleanup. cuda_bindings_12/tests/test_cuda.py also uses the shared driver_version_less_than() instead of its own copy, which behaves the same. cuda_bindings_12/tests/test_cudart.py keeps its local driverVersionLessThan: it does not call cuInit, and the shared one would turn collection-time skips into errors on machines with a driver but no device. cuda_core/tests/helpers/memory.py gains skip_if_mempool_unsupported(), skip_if_managed_mempool_unsupported(), and is_pool_resource(), which replace the per-module _skip_if_no_mempool (6 copies), _skip_without_mempools, _gpu_with_mempool_or_skip, _skip_if_no_managed_mempool (2 copies), and _is_pool (3 copies). Two skip reason strings change wording; their conditions do not.
1 parent 04fe681 commit 8d42c89

28 files changed

Lines changed: 578 additions & 403 deletions

‎.pre-commit-config.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ repos:
9090
name: Check tests do not create uncapped memory pools
9191
entry: python ./ci/tools/check_mempool_hygiene.py
9292
language: python
93-
files: '^cuda_core/tests/.*\.py$'
93+
files: '^(cuda_core|cuda_bindings_1[23])/tests/.*\.py$'
9494

9595
- id: no-markdown-in-docs-source
9696
name: Prevent markdown files in docs/source directories

‎ci/tools/check_mempool_hygiene.py‎

Lines changed: 164 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,26 +9,62 @@
99
address space, after which the rest of the session fails with
1010
CUDA_ERROR_OUT_OF_MEMORY on a device with free physical memory.
1111
12-
See cuda_core/tests/AGENTS.md for the rule this enforces.
12+
Two rules are enforced:
13+
14+
* cuda_core: ``DeviceMemoryResource`` / ``PinnedMemoryResource`` and their
15+
options classes must be given ``max_size``.
16+
* cuda.bindings: a pool properties object (``MemPoolProps_v1``,
17+
``CUmemPoolProps``, ``cudaMemPoolProps``) assigned to a name must have
18+
``max_size`` / ``maxSize`` set in the same function (or at module level, for
19+
module-level code). Graph alloc-node ``params.poolProps`` is a field, not a
20+
constructor call, so it is not checked; that pool cannot be capped anyway.
21+
22+
Only ``.py`` files are checked. ``ast`` cannot parse Cython, so ``.pyx`` tests
23+
(e.g. ``cuda_bindings_*/tests/cython``) are out of scope and must be capped by
24+
hand.
25+
26+
See cuda_core/tests/AGENTS.md and cuda_bindings_13/AGENTS.md for the rules this
27+
enforces.
1328
"""
1429

1530
from __future__ import annotations
1631

1732
import argparse
1833
import ast
1934
import sys
35+
from collections.abc import Iterator
2036
from pathlib import Path
37+
from typing import NamedTuple
2138

2239
ROOT = Path(__file__).resolve().parents[2]
23-
DEFAULT_TREE = ROOT / "cuda_core" / "tests"
40+
DEFAULT_TREES = (
41+
ROOT / "cuda_core" / "tests",
42+
ROOT / "cuda_bindings_12" / "tests",
43+
ROOT / "cuda_bindings_13" / "tests",
44+
)
2445

2546
# Managed pools cannot be right-sized: cuMemPoolCreate requires maxSize == 0 for
2647
# managed pools, so ManagedMemoryResourceOptions has no max_size to set.
2748
CAPPABLE_OPTIONS = frozenset({"DeviceMemoryResourceOptions", "PinnedMemoryResourceOptions"})
2849
CAPPABLE_RESOURCES = frozenset({"DeviceMemoryResource", "PinnedMemoryResource"})
2950

51+
# Pool properties types of the cuda.bindings APIs; the size cap is a plain
52+
# attribute, spelled ``max_size`` on the v2 types and ``maxSize`` on the others.
53+
POOL_PROPS_TYPES = frozenset(
54+
{"MemPoolProps_v1", "CUmemPoolProps", "CUmemPoolProps_v1", "CUmemPoolProps_st", "cudaMemPoolProps"}
55+
)
56+
POOL_PROPS_CAP_ATTRIBUTES = frozenset({"max_size", "maxSize"})
57+
3058
OPT_OUT_MARKER = "uncapped-pool-ok"
3159

60+
RULE_RESOURCE = "resource"
61+
RULE_PROPS = "props"
62+
63+
64+
class Violation(NamedTuple):
65+
rule: str
66+
message: str
67+
3268

3369
def _callee_name(node: ast.Call) -> str:
3470
func = node.func
@@ -60,12 +96,95 @@ def _opted_out(lines: list[str], node: ast.AST) -> bool:
6096
return any(OPT_OUT_MARKER in line for line in lines[start:end])
6197

6298

63-
def violations_in(path: Path) -> list[str]:
64-
"""Return one message per uncapped pool construction in ``path``."""
65-
source = path.read_text(encoding="utf-8")
66-
lines = source.splitlines()
99+
def _walk_scope(scope: ast.AST) -> Iterator[ast.AST]:
100+
"""Yield ``scope``'s descendants in source order, not entering nested function/lambda/class bodies."""
101+
# Push children reversed so the leftmost (earliest) child pops first.
102+
pending = list(reversed(list(ast.iter_child_nodes(scope))))
103+
while pending:
104+
node = pending.pop()
105+
yield node
106+
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda, ast.ClassDef)):
107+
pending.extend(reversed(list(ast.iter_child_nodes(node))))
108+
109+
110+
def _assign_targets(node: ast.AST) -> list[ast.expr]:
111+
"""The assignment targets of ``node``, or ``[]`` if it is not an assignment."""
112+
if isinstance(node, ast.Assign):
113+
return node.targets
114+
if isinstance(node, (ast.AugAssign, ast.AnnAssign)):
115+
return [node.target]
116+
return []
117+
118+
119+
def _pool_props_violations(path: Path, scope: ast.AST, lines: list[str]) -> list[Violation]:
120+
"""Report each ``var = PoolPropsType(...)`` with no ``var.max_size`` set in the same scope.
121+
122+
Processes nodes in source order. When ``var`` is reassigned to a new
123+
construction, any cap set on the previous object is reset; only a cap that
124+
follows the new construction (in the same scope) excuses it.
125+
126+
Recurses into nested scopes so a cap set in another function does not count.
127+
128+
Handles both plain and type-annotated assignments (``ast.Assign`` and
129+
``ast.AnnAssign``).
130+
"""
131+
found: list[Violation] = []
132+
# Maps variable name -> (construction node, type name) for the latest
133+
# uncapped construction. Removed when the variable is capped; overwritten
134+
# (with a report of the old entry if it was still pending) on reassignment.
135+
pending: dict[str, tuple[ast.AST, str]] = {}
136+
137+
for node in _walk_scope(scope):
138+
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
139+
found += _pool_props_violations(path, node, lines)
140+
continue
141+
# Cap: props.max_size = ... or props.maxSize = ...
142+
for target in _assign_targets(node):
143+
if (
144+
isinstance(target, ast.Attribute)
145+
and target.attr in POOL_PROPS_CAP_ATTRIBUTES
146+
and isinstance(target.value, ast.Name)
147+
):
148+
pending.pop(target.value.id, None)
149+
# Construction: props = PoolPropsType(...) or props: T = PoolPropsType(...)
150+
con_target = con_value = None
151+
if isinstance(node, ast.Assign) and len(node.targets) == 1:
152+
con_target, con_value = node.targets[0], node.value
153+
elif isinstance(node, ast.AnnAssign) and node.value is not None:
154+
con_target, con_value = node.target, node.value
155+
if (
156+
isinstance(con_target, ast.Name)
157+
and isinstance(con_value, ast.Call)
158+
and _callee_name(con_value) in POOL_PROPS_TYPES
159+
):
160+
var, type_name = con_target.id, _callee_name(con_value)
161+
if var in pending:
162+
prev_node, prev_type = pending.pop(var)
163+
if not _opted_out(lines, prev_node):
164+
found.append(
165+
Violation(
166+
RULE_PROPS,
167+
f"{path.as_posix()}:{prev_node.lineno}: {prev_type} assigned to '{var}' "
168+
f"without {var}.max_size / {var}.maxSize",
169+
)
170+
)
171+
pending[var] = (node, type_name)
172+
173+
for var, (node, type_name) in pending.items():
174+
if not _opted_out(lines, node):
175+
found.append(
176+
Violation(
177+
RULE_PROPS,
178+
f"{path.as_posix()}:{node.lineno}: {type_name} assigned to '{var}' "
179+
f"without {var}.max_size / {var}.maxSize",
180+
)
181+
)
182+
return found
183+
184+
185+
def _resource_violations(path: Path, tree: ast.AST, lines: list[str]) -> list[Violation]:
67186
found = []
68-
for node in ast.walk(ast.parse(source, filename=str(path))):
187+
for node in ast.walk(tree):
69188
if not isinstance(node, ast.Call):
70189
continue
71190
name = _callee_name(node)
@@ -78,33 +197,62 @@ def violations_in(path: Path) -> list[str]:
78197
else:
79198
continue
80199
if uncapped and not _opted_out(lines, node):
81-
found.append(f"{path.as_posix()}:{node.lineno}: {name} without max_size")
200+
found.append(Violation(RULE_RESOURCE, f"{path.as_posix()}:{node.lineno}: {name} without max_size"))
82201
return found
83202

84203

204+
def _violations(path: Path) -> list[Violation]:
205+
source = path.read_text(encoding="utf-8")
206+
lines = source.splitlines()
207+
tree = ast.parse(source, filename=str(path))
208+
return [*_resource_violations(path, tree, lines), *_pool_props_violations(path, tree, lines)]
209+
210+
211+
def violations_in(path: Path) -> list[str]:
212+
"""Return one message per uncapped pool construction in ``path``."""
213+
return [violation.message for violation in _violations(path)]
214+
215+
216+
def _report(header: str, messages: list[str], hint: str) -> None:
217+
"""Print ``header``, the bulleted ``messages``, and ``hint`` to stderr; a no-op if empty."""
218+
if not messages:
219+
return
220+
print(header, file=sys.stderr)
221+
for message in messages:
222+
print(f" - {message}", file=sys.stderr)
223+
print(hint, file=sys.stderr)
224+
225+
85226
def main(argv: list[str] | None = None) -> int:
86227
parser = argparse.ArgumentParser(description=__doc__)
87228
parser.add_argument(
88229
"paths",
89230
nargs="*",
90231
type=Path,
91-
help=f"Files to check. Defaults to every .py under {DEFAULT_TREE.relative_to(ROOT).as_posix()}.",
232+
help="Files to check. Defaults to every .py under "
233+
+ ", ".join(tree.relative_to(ROOT).as_posix() for tree in DEFAULT_TREES)
234+
+ ".",
92235
)
93236
args = parser.parse_args(argv)
94237

95-
paths = args.paths or sorted(DEFAULT_TREE.rglob("*.py"))
96-
violations = sorted(v for path in paths if path.suffix == ".py" for v in violations_in(path))
238+
paths = args.paths or sorted(path for tree in DEFAULT_TREES for path in tree.rglob("*.py"))
239+
violations = sorted(v for path in paths if path.suffix == ".py" for v in _violations(path))
97240
if not violations:
98241
return 0
99242

100-
print("error: memory pools created by tests must set max_size:", file=sys.stderr)
101-
for violation in violations:
102-
print(f" - {violation}", file=sys.stderr)
103-
print(
243+
_report(
244+
"error: memory pools created by tests must set max_size:",
245+
[v.message for v in violations if v.rule == RULE_RESOURCE],
104246
f"Use the suite-wide POOL_SIZE from cuda_core/tests/helpers/constants.py, or annotate a\n"
105247
f"deliberate exception with a '# {OPT_OUT_MARKER}: <reason>' comment.\n"
106248
f"See cuda_core/tests/AGENTS.md.",
107-
file=sys.stderr,
249+
)
250+
_report(
251+
"error: pool properties created by tests must set max_size / maxSize:",
252+
[v.message for v in violations if v.rule == RULE_PROPS],
253+
f"Set max_size (v2 bindings) or maxSize (other bindings) on the props object you constructed,\n"
254+
f"using POOL_SIZE from cuda_python_test_helpers.memory, or annotate a deliberate exception\n"
255+
f"with a '# {OPT_OUT_MARKER}: <reason>' comment. See cuda_bindings_13/AGENTS.md.",
108256
)
109257
return 1
110258

‎ci/tools/tests/test_check_mempool_hygiene.py‎

Lines changed: 125 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import pytest
1010

1111
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
12-
from check_mempool_hygiene import DEFAULT_TREE, main, violations_in
12+
from check_mempool_hygiene import DEFAULT_TREES, main, violations_in
1313

1414

1515
def write(tmp_path, source):
@@ -92,5 +92,128 @@ def test_main_ignores_non_python_files(tmp_path):
9292
def test_the_live_test_suite_is_clean():
9393
# Without a default the hook would only ever see changed files, so a
9494
# violation could ride in on a rename or a merge.
95-
assert DEFAULT_TREE.is_dir()
95+
assert all(tree.is_dir() for tree in DEFAULT_TREES)
9696
assert main([]) == 0
97+
98+
99+
POOL_PROPS_CAPPED = [
100+
pytest.param("props = driver.MemPoolProps_v1()\nprops.max_size = POOL_SIZE\n", id="v2-max_size"),
101+
pytest.param("props = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\n", id="driver-maxSize"),
102+
pytest.param("props = cudart.cudaMemPoolProps()\nprops.maxSize = POOL_SIZE\n", id="runtime-maxSize"),
103+
pytest.param("def f():\n props = cuda.CUmemPoolProps_v1()\n props.maxSize = POOL_SIZE\n", id="in-function"),
104+
# A field of an alloc-node params object is not a constructor call, so it is not checked.
105+
pytest.param("params.poolProps.allocType = TYPE\nparams.pool_props = other\n", id="graph-params-fields"),
106+
# Type-annotated assignment (ast.AnnAssign) must also be detected and treated as capped.
107+
pytest.param(
108+
"props: cuda.CUmemPoolProps = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\n",
109+
id="annotated-capped",
110+
),
111+
# Second construction caps correctly; the first was already capped.
112+
pytest.param(
113+
"props = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\nprops = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\n",
114+
id="two-constructions-both-capped",
115+
),
116+
]
117+
118+
POOL_PROPS_UNCAPPED = [
119+
pytest.param("props = driver.MemPoolProps_v1()\nprops.location = loc\n", id="v2"),
120+
pytest.param("props = cuda.CUmemPoolProps()\n", id="driver"),
121+
pytest.param("props = cudart.cudaMemPoolProps()\n", id="runtime"),
122+
pytest.param(
123+
"def f():\n props = cuda.CUmemPoolProps()\n other.maxSize = POOL_SIZE\n",
124+
id="other-object",
125+
),
126+
# The cap must be set in the function that constructed the props.
127+
pytest.param(
128+
"def make():\n return cuda.CUmemPoolProps()\n\n\ndef f():\n props = cuda.CUmemPoolProps()\n g(props)\n\n\ndef g(props):\n props.maxSize = POOL_SIZE\n",
129+
id="capped-in-another-function",
130+
),
131+
# Type-annotated assignment without a cap must also be reported (ast.AnnAssign, Finding 1).
132+
pytest.param("props: cuda.CUmemPoolProps = cuda.CUmemPoolProps()\n", id="annotated-uncapped"),
133+
# A cap on the first object does not cover a second construction of the same name (Finding 2).
134+
pytest.param(
135+
"props = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\nprops = cuda.CUmemPoolProps()\n",
136+
id="reassignment-second-uncapped",
137+
),
138+
]
139+
140+
141+
@pytest.mark.agent_authored(model="claude-sonnet-5.5")
142+
@pytest.mark.parametrize("source", POOL_PROPS_CAPPED)
143+
def test_capped_pool_props_are_not_reported(tmp_path, source):
144+
assert violations_in(write(tmp_path, source)) == []
145+
146+
147+
@pytest.mark.agent_authored(model="claude-sonnet-5.5")
148+
@pytest.mark.parametrize("source", POOL_PROPS_UNCAPPED)
149+
def test_uncapped_pool_props_are_reported(tmp_path, source):
150+
assert violations_in(write(tmp_path, source))
151+
152+
153+
@pytest.mark.agent_authored(model="claude-sonnet-5.5")
154+
def test_pool_props_marker_opts_out(tmp_path):
155+
source = "# uncapped-pool-ok: pool is never created\nprops = cuda.CUmemPoolProps()\n"
156+
157+
assert violations_in(write(tmp_path, source)) == []
158+
159+
160+
@pytest.mark.agent_authored(model="claude-sonnet-5.5")
161+
def test_pool_props_message_names_file_line_and_variable(tmp_path):
162+
path = write(tmp_path, "x = 1\nprops = driver.MemPoolProps_v1()\n")
163+
164+
(violation,) = violations_in(path)
165+
166+
assert violation.startswith(path.as_posix())
167+
assert ":2:" in violation
168+
assert "MemPoolProps_v1 assigned to 'props' without props.max_size / props.maxSize" in violation
169+
170+
171+
@pytest.mark.agent_authored(model="claude-sonnet-5.5")
172+
def test_main_reports_pool_props_with_their_own_message(tmp_path, capsys):
173+
path = write(tmp_path, "props = cuda.CUmemPoolProps()\n")
174+
175+
assert main([str(path)]) == 1
176+
177+
err = capsys.readouterr().err
178+
assert "pool properties created by tests must set max_size / maxSize" in err
179+
assert "must set max_size:" not in err
180+
181+
182+
@pytest.mark.agent_authored(model="claude-sonnet-5.5")
183+
def test_main_reports_both_rules_separately(tmp_path, capsys):
184+
path = write(
185+
tmp_path,
186+
"DeviceMemoryResource(dev, DeviceMemoryResourceOptions())\nprops = cuda.CUmemPoolProps()\n",
187+
)
188+
189+
assert main([str(path)]) == 1
190+
191+
err = capsys.readouterr().err
192+
assert "memory pools created by tests must set max_size:" in err
193+
assert "pool properties created by tests must set max_size / maxSize" in err
194+
195+
196+
@pytest.mark.agent_authored(model="claude-sonnet-4.6")
197+
def test_annotated_pool_props_assignment_reports_correct_line(tmp_path):
198+
# Finding 1: props: T = CUmemPoolProps() is ast.AnnAssign; it must be detected.
199+
path = write(tmp_path, "x = 1\nprops: cuda.CUmemPoolProps = cuda.CUmemPoolProps()\n")
200+
201+
(violation,) = violations_in(path)
202+
203+
assert ":2:" in violation
204+
205+
206+
@pytest.mark.agent_authored(model="claude-sonnet-4.6")
207+
def test_reassignment_reports_only_the_uncapped_construction(tmp_path):
208+
# Finding 2: a cap on the first object must not excuse a second construction
209+
# of the same variable. Only the second (uncapped) object should be reported.
210+
source = (
211+
"props = cuda.CUmemPoolProps()\n" # line 1
212+
"props.maxSize = POOL_SIZE\n" # line 2 — caps line 1
213+
"props = cuda.CUmemPoolProps()\n" # line 3 — new object, no cap
214+
)
215+
path = write(tmp_path, source)
216+
217+
(violation,) = violations_in(path)
218+
219+
assert ":3:" in violation

0 commit comments

Comments
 (0)