99address space, after which the rest of the session fails with
1010CUDA_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
1530from __future__ import annotations
1631
1732import argparse
1833import ast
1934import sys
35+ from collections .abc import Iterator
2036from pathlib import Path
37+ from typing import NamedTuple
2138
2239ROOT = 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.
2748CAPPABLE_OPTIONS = frozenset ({"DeviceMemoryResourceOptions" , "PinnedMemoryResourceOptions" })
2849CAPPABLE_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+
3058OPT_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
3369def _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+
85226def 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
0 commit comments