fix(ci): attach the uncapped-pool opt-out marker to the call it exempts

ci/tools/check_mempool_hygiene.py

@@ -53,21 +53,58 @@ def _dict_is_capped(node: ast.Dict) -> bool:
     return False
 
 
-def _opted_out(lines: list[str], node: ast.AST) -> bool:
-    """True if the call, or the line above it, carries the opt-out marker."""
-    start = max(node.lineno - 2, 0)  # -1 for 0-based, -1 more for a preceding comment
+def _iter_calls(tree: ast.AST):
+    """Yield ``(call, statements)`` for every call in ``tree``.
+
+    ``statements`` is the chain of ``ast.stmt`` ancestors, outermost first, so
+    ``statements[-1]`` is the statement the call belongs to. The chain is what
+    bounds an opt-out marker: a marker annotates the statement it sits in (or a
+    statement containing it), not whatever the next line happens to construct.
+    """
+    stack: list[tuple[ast.AST, tuple[ast.stmt, ...]]] = [(tree, ())]
+    while stack:
+        node, stmts = stack.pop()
+        if isinstance(node, ast.stmt):
+            stmts = (*stmts, node)
+        if isinstance(node, ast.Call):
+            yield node, stmts
+        stack.extend((child, stmts) for child in ast.iter_child_nodes(node))
+
+
+def _opted_out(lines: list[str], node: ast.Call, stmts: tuple[ast.stmt, ...]) -> bool:
+    """True if an opt-out marker annotates ``node``.
+
+    A marker counts when it is on a line belonging to the call's own statement
+    (including continuation lines of a multi-line construction), on the header
+    line of a compound statement containing the call, or on a dedicated comment
+    line immediately above the call's statement.
+
+    A marker anywhere else does not count. Previously the whole line above was
+    accepted unconditionally, so a trailing marker annotating the *previous*
+    statement silently exempted the next one, and a marker appearing inside an
+    unrelated string literal exempted whatever followed it.
+    """
+    nearest = stmts[-1] if stmts else node
+    first = nearest.lineno - 1  # 0-based index of the statement's first line
     end = getattr(node, "end_lineno", node.lineno)
-    return any(OPT_OUT_MARKER in line for line in lines[start:end])
+    if any(OPT_OUT_MARKER in line for line in lines[first:end]):
+        return True
+    # Headers of the compound statements containing the call ("with", "for",
+    # "def", ...): a marker there annotates a block this call is part of.
+    if any(OPT_OUT_MARKER in lines[s.lineno - 1] for s in stmts[:-1]):
+        return True
+    if first == 0:
+        return False
+    above = lines[first - 1].strip()
+    return above.startswith("#") and OPT_OUT_MARKER in above
 
 
 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()
     found = []
-    for node in ast.walk(ast.parse(source, filename=str(path))):
-        if not isinstance(node, ast.Call):
-            continue
+    for node, stmts in _iter_calls(ast.parse(source, filename=str(path))):
         name = _callee_name(node)
         if name in CAPPABLE_OPTIONS:
             uncapped = not _is_capped(node)
@@ -77,7 +114,7 @@ def violations_in(path: Path) -> list[str]:
             uncapped = any(not _dict_is_capped(d) for d in dicts)
         else:
             continue
-        if uncapped and not _opted_out(lines, node):
+        if uncapped and not _opted_out(lines, node, stmts):
             found.append(f"{path.as_posix()}:{node.lineno}: {name} without max_size")
     return found
 

ci/tools/tests/test_check_mempool_hygiene.py

@@ -94,3 +94,63 @@ def test_the_live_test_suite_is_clean():
     # violation could ride in on a rename or a merge.
     assert DEFAULT_TREE.is_dir()
     assert main([]) == 0
+
+
+UNCAPPED_CALL = "DeviceMemoryResource(dev, DeviceMemoryResourceOptions())"
+
+# Placements where the marker does NOT annotate the offending call. Accepting
+# the whole preceding line let each of these silence a real violation.
+MARKER_DOES_NOT_CARRY = [
+    pytest.param(
+        f"{UNCAPPED_CALL}  # uncapped-pool-ok: annotates THIS line\n{UNCAPPED_CALL}\n",
+        id="trailing-marker-on-previous-statement",
+    ),
+    pytest.param(
+        'msg = "see uncapped-pool-ok in AGENTS.md"\n' + UNCAPPED_CALL + "\n",
+        id="marker-inside-an-unrelated-string",
+    ),
+    pytest.param(
+        f"# uncapped-pool-ok: detached by a blank line\n\n{UNCAPPED_CALL}\n",
+        id="blank-line-between-comment-and-call",
+    ),
+]
+
+# Placements where the marker does annotate the call and must keep working.
+MARKER_CARRIES = [
+    pytest.param(f"# uncapped-pool-ok: reason\n{UNCAPPED_CALL}\n", id="comment-line-above"),
+    pytest.param(f"def test_x():\n    # uncapped-pool-ok: reason\n    {UNCAPPED_CALL}\n", id="indented-comment-above"),
+    pytest.param(f"{UNCAPPED_CALL}  # uncapped-pool-ok: reason\n", id="inline-on-the-call"),
+    pytest.param(
+        "DeviceMemoryResource(\n    dev,  # uncapped-pool-ok: reason\n    DeviceMemoryResourceOptions(),\n)\n",
+        id="continuation-line-of-the-same-call",
+    ),
+    pytest.param(
+        "mr = DeviceMemoryResource(  # uncapped-pool-ok: reason\n    dev,\n    DeviceMemoryResourceOptions(),\n)\n",
+        id="first-line-of-a-multiline-statement",
+    ),
+    # The documented use is pytest.raises, so a marker on the block header has
+    # to keep annotating the calls inside the block.
+    pytest.param(
+        f"with pytest.raises(RuntimeError):  # uncapped-pool-ok: reason\n    {UNCAPPED_CALL}\n",
+        id="containing-with-header",
+    ),
+    pytest.param(f"def test_x():  # uncapped-pool-ok: reason\n    {UNCAPPED_CALL}\n", id="containing-def-header"),
+]
+
+
+@pytest.mark.agent_authored(model="claude-opus-5")
+@pytest.mark.parametrize("source", MARKER_DOES_NOT_CARRY)
+def test_marker_that_does_not_annotate_the_call_does_not_suppress_it(tmp_path, source):
+    """An opt-out must be attached to the call it exempts.
+
+    A trailing marker annotating one statement used to exempt the statement on
+    the next line as well -- the shape a reviewer is least likely to notice,
+    since both lines look correctly annotated.
+    """
+    assert violations_in(write(tmp_path, source))
+
+
+@pytest.mark.agent_authored(model="claude-opus-5")
+@pytest.mark.parametrize("source", MARKER_CARRIES)
+def test_marker_attached_to_the_call_still_suppresses_it(tmp_path, source):
+    assert violations_in(write(tmp_path, source)) == []

cuda_core/tests/AGENTS.md

@@ -58,6 +58,12 @@ with pytest.raises(RuntimeError, match="IPC is not available"):
     DeviceMemoryResource(mempool_device, DeviceMemoryResourceOptions(ipc_enabled=True))
 ```
 
+The marker has to be attached to the statement it exempts: on a comment line
+directly above it, anywhere within the statement itself (including the
+continuation lines of a multi-line call), or on the header of a block
+containing it. A marker somewhere else -- trailing the *previous* statement,
+or separated from the call by a blank line -- does not exempt anything.
+
 ## Release resources at test boundaries
 
 The `_init_cuda_context` fixture in `conftest.py` runs `gc.collect()` followed