test: consolidate memory-pool test helpers and lint bindings pools

.pre-commit-config.yaml

@@ -90,7 +90,7 @@ repos:
         name: Check tests do not create uncapped memory pools
         entry: python ./ci/tools/check_mempool_hygiene.py
         language: python
-        files: '^cuda_core/tests/.*\.py$'
+        files: '^(cuda_core|cuda_bindings_1[23])/tests/.*\.py$'
 
       - id: no-markdown-in-docs-source
         name: Prevent markdown files in docs/source directories

ci/tools/check_mempool_hygiene.py

@@ -9,26 +9,62 @@
 address space, after which the rest of the session fails with
 CUDA_ERROR_OUT_OF_MEMORY on a device with free physical memory.
 
-See cuda_core/tests/AGENTS.md for the rule this enforces.
+Two rules are enforced:
+
+* cuda_core: ``DeviceMemoryResource`` / ``PinnedMemoryResource`` and their
+  options classes must be given ``max_size``.
+* cuda.bindings: a pool properties object (``MemPoolProps_v1``,
+  ``CUmemPoolProps``, ``cudaMemPoolProps``) assigned to a name must have
+  ``max_size`` / ``maxSize`` set in the same function (or at module level, for
+  module-level code). Graph alloc-node ``params.poolProps`` is a field, not a
+  constructor call, so it is not checked; that pool cannot be capped anyway.
+
+Only ``.py`` files are checked. ``ast`` cannot parse Cython, so ``.pyx`` tests
+(e.g. ``cuda_bindings_*/tests/cython``) are out of scope and must be capped by
+hand.
+
+See cuda_core/tests/AGENTS.md and cuda_bindings_13/AGENTS.md for the rules this
+enforces.
 """
 
 from __future__ import annotations
 
 import argparse
 import ast
 import sys
+from collections.abc import Iterator
 from pathlib import Path
+from typing import NamedTuple
 
 ROOT = Path(__file__).resolve().parents[2]
-DEFAULT_TREE = ROOT / "cuda_core" / "tests"
+DEFAULT_TREES = (
+    ROOT / "cuda_core" / "tests",
+    ROOT / "cuda_bindings_12" / "tests",
+    ROOT / "cuda_bindings_13" / "tests",
+)
 
 # Managed pools cannot be right-sized: cuMemPoolCreate requires maxSize == 0 for
 # managed pools, so ManagedMemoryResourceOptions has no max_size to set.
 CAPPABLE_OPTIONS = frozenset({"DeviceMemoryResourceOptions", "PinnedMemoryResourceOptions"})
 CAPPABLE_RESOURCES = frozenset({"DeviceMemoryResource", "PinnedMemoryResource"})
 
+# Pool properties types of the cuda.bindings APIs; the size cap is a plain
+# attribute, spelled ``max_size`` on the v2 types and ``maxSize`` on the others.
+POOL_PROPS_TYPES = frozenset(
+    {"MemPoolProps_v1", "CUmemPoolProps", "CUmemPoolProps_v1", "CUmemPoolProps_st", "cudaMemPoolProps"}
+)
+POOL_PROPS_CAP_ATTRIBUTES = frozenset({"max_size", "maxSize"})
+
 OPT_OUT_MARKER = "uncapped-pool-ok"
 
+RULE_RESOURCE = "resource"
+RULE_PROPS = "props"
+
+
+class Violation(NamedTuple):
+    rule: str
+    message: str
+
 
 def _callee_name(node: ast.Call) -> str:
     func = node.func
@@ -60,12 +96,95 @@ def _opted_out(lines: list[str], node: ast.AST) -> bool:
     return any(OPT_OUT_MARKER in line for line in lines[start:end])
 
 
-def violations_in(path: Path) -> list[str]:
-    """Return one message per uncapped pool construction in ``path``."""
-    source = path.read_text(encoding="utf-8")
-    lines = source.splitlines()
+def _walk_scope(scope: ast.AST) -> Iterator[ast.AST]:
+    """Yield ``scope``'s descendants in source order, not entering nested function/lambda/class bodies."""
+    # Push children reversed so the leftmost (earliest) child pops first.
+    pending = list(reversed(list(ast.iter_child_nodes(scope))))
+    while pending:
+        node = pending.pop()
+        yield node
+        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda, ast.ClassDef)):
+            pending.extend(reversed(list(ast.iter_child_nodes(node))))
+
+
+def _assign_targets(node: ast.AST) -> list[ast.expr]:
+    """The assignment targets of ``node``, or ``[]`` if it is not an assignment."""
+    if isinstance(node, ast.Assign):
+        return node.targets
+    if isinstance(node, (ast.AugAssign, ast.AnnAssign)):
+        return [node.target]
+    return []
+
+
+def _pool_props_violations(path: Path, scope: ast.AST, lines: list[str]) -> list[Violation]:
+    """Report each ``var = PoolPropsType(...)`` with no ``var.max_size`` set in the same scope.
+
+    Processes nodes in source order.  When ``var`` is reassigned to a new
+    construction, any cap set on the previous object is reset; only a cap that
+    follows the new construction (in the same scope) excuses it.
+
+    Recurses into nested scopes so a cap set in another function does not count.
+
+    Handles both plain and type-annotated assignments (``ast.Assign`` and
+    ``ast.AnnAssign``).
+    """
+    found: list[Violation] = []
+    # Maps variable name -> (construction node, type name) for the latest
+    # uncapped construction.  Removed when the variable is capped; overwritten
+    # (with a report of the old entry if it was still pending) on reassignment.
+    pending: dict[str, tuple[ast.AST, str]] = {}
+
+    for node in _walk_scope(scope):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+            found += _pool_props_violations(path, node, lines)
+            continue
+        # Cap: props.max_size = ... or props.maxSize = ...
+        for target in _assign_targets(node):
+            if (
+                isinstance(target, ast.Attribute)
+                and target.attr in POOL_PROPS_CAP_ATTRIBUTES
+                and isinstance(target.value, ast.Name)
+            ):
+                pending.pop(target.value.id, None)
+        # Construction: props = PoolPropsType(...) or props: T = PoolPropsType(...)
+        con_target = con_value = None
+        if isinstance(node, ast.Assign) and len(node.targets) == 1:
+            con_target, con_value = node.targets[0], node.value
+        elif isinstance(node, ast.AnnAssign) and node.value is not None:
+            con_target, con_value = node.target, node.value
+        if (
+            isinstance(con_target, ast.Name)
+            and isinstance(con_value, ast.Call)
+            and _callee_name(con_value) in POOL_PROPS_TYPES
+        ):
+            var, type_name = con_target.id, _callee_name(con_value)
+            if var in pending:
+                prev_node, prev_type = pending.pop(var)
+                if not _opted_out(lines, prev_node):
+                    found.append(
+                        Violation(
+                            RULE_PROPS,
+                            f"{path.as_posix()}:{prev_node.lineno}: {prev_type} assigned to '{var}' "
+                            f"without {var}.max_size / {var}.maxSize",
+                        )
+                    )
+            pending[var] = (node, type_name)
+
+    for var, (node, type_name) in pending.items():
+        if not _opted_out(lines, node):
+            found.append(
+                Violation(
+                    RULE_PROPS,
+                    f"{path.as_posix()}:{node.lineno}: {type_name} assigned to '{var}' "
+                    f"without {var}.max_size / {var}.maxSize",
+                )
+            )
+    return found
+
+
+def _resource_violations(path: Path, tree: ast.AST, lines: list[str]) -> list[Violation]:
     found = []
-    for node in ast.walk(ast.parse(source, filename=str(path))):
+    for node in ast.walk(tree):
         if not isinstance(node, ast.Call):
             continue
         name = _callee_name(node)
@@ -78,33 +197,62 @@ def violations_in(path: Path) -> list[str]:
         else:
             continue
         if uncapped and not _opted_out(lines, node):
-            found.append(f"{path.as_posix()}:{node.lineno}: {name} without max_size")
+            found.append(Violation(RULE_RESOURCE, f"{path.as_posix()}:{node.lineno}: {name} without max_size"))
     return found
 
 
+def _violations(path: Path) -> list[Violation]:
+    source = path.read_text(encoding="utf-8")
+    lines = source.splitlines()
+    tree = ast.parse(source, filename=str(path))
+    return [*_resource_violations(path, tree, lines), *_pool_props_violations(path, tree, lines)]
+
+
+def violations_in(path: Path) -> list[str]:
+    """Return one message per uncapped pool construction in ``path``."""
+    return [violation.message for violation in _violations(path)]
+
+
+def _report(header: str, messages: list[str], hint: str) -> None:
+    """Print ``header``, the bulleted ``messages``, and ``hint`` to stderr; a no-op if empty."""
+    if not messages:
+        return
+    print(header, file=sys.stderr)
+    for message in messages:
+        print(f"  - {message}", file=sys.stderr)
+    print(hint, file=sys.stderr)
+
+
 def main(argv: list[str] | None = None) -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
         "paths",
         nargs="*",
         type=Path,
-        help=f"Files to check. Defaults to every .py under {DEFAULT_TREE.relative_to(ROOT).as_posix()}.",
+        help="Files to check. Defaults to every .py under "
+        + ", ".join(tree.relative_to(ROOT).as_posix() for tree in DEFAULT_TREES)
+        + ".",
     )
     args = parser.parse_args(argv)
 
-    paths = args.paths or sorted(DEFAULT_TREE.rglob("*.py"))
-    violations = sorted(v for path in paths if path.suffix == ".py" for v in violations_in(path))
+    paths = args.paths or sorted(path for tree in DEFAULT_TREES for path in tree.rglob("*.py"))
+    violations = sorted(v for path in paths if path.suffix == ".py" for v in _violations(path))
     if not violations:
         return 0
 
-    print("error: memory pools created by tests must set max_size:", file=sys.stderr)
-    for violation in violations:
-        print(f"  - {violation}", file=sys.stderr)
-    print(
+    _report(
+        "error: memory pools created by tests must set max_size:",
+        [v.message for v in violations if v.rule == RULE_RESOURCE],
         f"Use the suite-wide POOL_SIZE from cuda_core/tests/helpers/constants.py, or annotate a\n"
         f"deliberate exception with a '# {OPT_OUT_MARKER}: <reason>' comment.\n"
         f"See cuda_core/tests/AGENTS.md.",
-        file=sys.stderr,
+    )
+    _report(
+        "error: pool properties created by tests must set max_size / maxSize:",
+        [v.message for v in violations if v.rule == RULE_PROPS],
+        f"Set max_size (v2 bindings) or maxSize (other bindings) on the props object you constructed,\n"
+        f"using POOL_SIZE from cuda_python_test_helpers.memory, or annotate a deliberate exception\n"
+        f"with a '# {OPT_OUT_MARKER}: <reason>' comment. See cuda_bindings_13/AGENTS.md.",
     )
     return 1
 

ci/tools/tests/test_check_mempool_hygiene.py

@@ -9,7 +9,7 @@
 import pytest
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
-from check_mempool_hygiene import DEFAULT_TREE, main, violations_in
+from check_mempool_hygiene import DEFAULT_TREES, main, violations_in
 
 
 def write(tmp_path, source):
@@ -92,5 +92,128 @@ def test_main_ignores_non_python_files(tmp_path):
 def test_the_live_test_suite_is_clean():
     # Without a default the hook would only ever see changed files, so a
     # violation could ride in on a rename or a merge.
-    assert DEFAULT_TREE.is_dir()
+    assert all(tree.is_dir() for tree in DEFAULT_TREES)
     assert main([]) == 0
+
+
+POOL_PROPS_CAPPED = [
+    pytest.param("props = driver.MemPoolProps_v1()\nprops.max_size = POOL_SIZE\n", id="v2-max_size"),
+    pytest.param("props = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\n", id="driver-maxSize"),
+    pytest.param("props = cudart.cudaMemPoolProps()\nprops.maxSize = POOL_SIZE\n", id="runtime-maxSize"),
+    pytest.param("def f():\n    props = cuda.CUmemPoolProps_v1()\n    props.maxSize = POOL_SIZE\n", id="in-function"),
+    # A field of an alloc-node params object is not a constructor call, so it is not checked.
+    pytest.param("params.poolProps.allocType = TYPE\nparams.pool_props = other\n", id="graph-params-fields"),
+    # Type-annotated assignment (ast.AnnAssign) must also be detected and treated as capped.
+    pytest.param(
+        "props: cuda.CUmemPoolProps = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\n",
+        id="annotated-capped",
+    ),
+    # Second construction caps correctly; the first was already capped.
+    pytest.param(
+        "props = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\nprops = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\n",
+        id="two-constructions-both-capped",
+    ),
+]
+
+POOL_PROPS_UNCAPPED = [
+    pytest.param("props = driver.MemPoolProps_v1()\nprops.location = loc\n", id="v2"),
+    pytest.param("props = cuda.CUmemPoolProps()\n", id="driver"),
+    pytest.param("props = cudart.cudaMemPoolProps()\n", id="runtime"),
+    pytest.param(
+        "def f():\n    props = cuda.CUmemPoolProps()\n    other.maxSize = POOL_SIZE\n",
+        id="other-object",
+    ),
+    # The cap must be set in the function that constructed the props.
+    pytest.param(
+        "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",
+        id="capped-in-another-function",
+    ),
+    # Type-annotated assignment without a cap must also be reported (ast.AnnAssign, Finding 1).
+    pytest.param("props: cuda.CUmemPoolProps = cuda.CUmemPoolProps()\n", id="annotated-uncapped"),
+    # A cap on the first object does not cover a second construction of the same name (Finding 2).
+    pytest.param(
+        "props = cuda.CUmemPoolProps()\nprops.maxSize = POOL_SIZE\nprops = cuda.CUmemPoolProps()\n",
+        id="reassignment-second-uncapped",
+    ),
+]
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-5.5")
+@pytest.mark.parametrize("source", POOL_PROPS_CAPPED)
+def test_capped_pool_props_are_not_reported(tmp_path, source):
+    assert violations_in(write(tmp_path, source)) == []
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-5.5")
+@pytest.mark.parametrize("source", POOL_PROPS_UNCAPPED)
+def test_uncapped_pool_props_are_reported(tmp_path, source):
+    assert violations_in(write(tmp_path, source))
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-5.5")
+def test_pool_props_marker_opts_out(tmp_path):
+    source = "# uncapped-pool-ok: pool is never created\nprops = cuda.CUmemPoolProps()\n"
+
+    assert violations_in(write(tmp_path, source)) == []
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-5.5")
+def test_pool_props_message_names_file_line_and_variable(tmp_path):
+    path = write(tmp_path, "x = 1\nprops = driver.MemPoolProps_v1()\n")
+
+    (violation,) = violations_in(path)
+
+    assert violation.startswith(path.as_posix())
+    assert ":2:" in violation
+    assert "MemPoolProps_v1 assigned to 'props' without props.max_size / props.maxSize" in violation
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-5.5")
+def test_main_reports_pool_props_with_their_own_message(tmp_path, capsys):
+    path = write(tmp_path, "props = cuda.CUmemPoolProps()\n")
+
+    assert main([str(path)]) == 1
+
+    err = capsys.readouterr().err
+    assert "pool properties created by tests must set max_size / maxSize" in err
+    assert "must set max_size:" not in err
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-5.5")
+def test_main_reports_both_rules_separately(tmp_path, capsys):
+    path = write(
+        tmp_path,
+        "DeviceMemoryResource(dev, DeviceMemoryResourceOptions())\nprops = cuda.CUmemPoolProps()\n",
+    )
+
+    assert main([str(path)]) == 1
+
+    err = capsys.readouterr().err
+    assert "memory pools created by tests must set max_size:" in err
+    assert "pool properties created by tests must set max_size / maxSize" in err
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-4.6")
+def test_annotated_pool_props_assignment_reports_correct_line(tmp_path):
+    # Finding 1: props: T = CUmemPoolProps() is ast.AnnAssign; it must be detected.
+    path = write(tmp_path, "x = 1\nprops: cuda.CUmemPoolProps = cuda.CUmemPoolProps()\n")
+
+    (violation,) = violations_in(path)
+
+    assert ":2:" in violation
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-4.6")
+def test_reassignment_reports_only_the_uncapped_construction(tmp_path):
+    # Finding 2: a cap on the first object must not excuse a second construction
+    # of the same variable.  Only the second (uncapped) object should be reported.
+    source = (
+        "props = cuda.CUmemPoolProps()\n"  # line 1
+        "props.maxSize = POOL_SIZE\n"  # line 2 — caps line 1
+        "props = cuda.CUmemPoolProps()\n"  # line 3 — new object, no cap
+    )
+    path = write(tmp_path, source)
+
+    (violation,) = violations_in(path)
+
+    assert ":3:" in violation