From 3a582c9f37a6d12d6e19fa380b70168b30dec3f8 Mon Sep 17 00:00:00 2001
From: Rui Luo <ruluo@nvidia.com>
Date: Fri, 23 Oct 2026 14:42:26 +0800
Subject: [PATCH 1/2] ci: run coverage pytest on main thread with larger stack

---
 .github/workflows/coverage.yml    |  2 +-
 ci/tools/run_pytest_with_stack.py | 97 ++++++++++++++++++++++++++-----
 2 files changed, 83 insertions(+), 16 deletions(-)

diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml
index aee89d41f2d..448e13ad2ff 100644
--- a/.github/workflows/coverage.yml
+++ b/.github/workflows/coverage.yml
@@ -440,7 +440,7 @@ jobs:
           "$GITHUB_WORKSPACE/.venv/Scripts/pytest" -v --cov=./cuda --cov-append --cov-context=test --cov-config="$GITHUB_WORKSPACE/.coveragerc" "$GITHUB_WORKSPACE/cuda_pathfinder/tests"
 
       # Cython linetrace under coverage on Windows needs more stack than the
-      # default 1 MB thread size.  The helper runs pytest on an 8 MB thread.
+      # default 1 MB. The helper runs pytest on the main thread with 8 MB.
       - name: Run cuda.bindings tests (with 8MB stack)
         continue-on-error: true
         run: |
diff --git a/ci/tools/run_pytest_with_stack.py b/ci/tools/run_pytest_with_stack.py
index e1e1f782ba1..12abb49d766 100644
--- a/ci/tools/run_pytest_with_stack.py
+++ b/ci/tools/run_pytest_with_stack.py
@@ -4,24 +4,88 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Run pytest on a thread with a larger stack size.
+"""Run pytest on Python's main thread with a larger stack.
 
-Cython linetrace instrumentation under coverage on Windows can exceed the
-default 1 MB thread stack.  This helper spawns a single worker thread with
-a configurable stack (default 8 MB) so the rest of the CI workflow stays
-readable.
+On Windows, start a temporary copy of the base Python executable with an
+8 MB PE stack reservation. This preserves the virtual environment without
+modifying the installed executable. Running pytest on the main thread lets
+CPython process pending calls used by CUDA graph cleanup.
 
 Usage:
     python run_pytest_with_stack.py [--stack-mb N] [--cwd DIR] [pytest args ...]
 """
 
 import argparse
-import concurrent.futures
 import os
+import struct
+import subprocess
 import sys
-import threading
+import tempfile
+from pathlib import Path
 
-import pytest
+
+def _python_with_stack(executable, stack_bytes):
+    """Return a PE executable image with at least the requested stack reserve."""
+    image = bytearray(executable.read_bytes())
+    if len(image) < 64 or image[:2] != b"MZ":
+        raise ValueError(f"Not a PE executable: {executable}")
+    pe_offset = struct.unpack_from("<I", image, 0x3C)[0]
+    if image[pe_offset : pe_offset + 4] != b"PE\0\0":
+        raise ValueError(f"Invalid PE signature: {executable}")
+    optional_offset = pe_offset + 24
+    if len(image) < optional_offset + 88:
+        raise ValueError(f"Truncated PE optional header: {executable}")
+    magic = struct.unpack_from("<H", image, optional_offset)[0]
+    if magic not in (0x10B, 0x20B):
+        raise ValueError(f"Unsupported PE optional header: {magic:#x}")
+    size_format = "<Q" if magic == 0x20B else "<I"
+    reserve_offset = optional_offset + 72
+    reserve = struct.unpack_from(size_format, image, reserve_offset)[0]
+    struct.pack_into(size_format, image, reserve_offset, max(reserve, stack_bytes))
+    return image
+
+
+def _run_windows(stack_mb, cwd, pytest_args):
+    # A venv's python.exe may be a redirector, so patch the base interpreter.
+    base_executable = Path(sys._base_executable)
+    image = _python_with_stack(base_executable, stack_mb * 1024 * 1024)
+    with tempfile.TemporaryDirectory(prefix="pytest-main-stack-") as directory:
+        executable = Path(directory) / "python.exe"
+        executable.write_bytes(image)
+        env = os.environ.copy()
+        env["__PYVENV_LAUNCHER__"] = sys.executable
+        env["PATH"] = str(base_executable.parent) + os.pathsep + env.get("PATH", "")
+        cmd = [str(executable), os.path.abspath(__file__), f"--stack-mb={stack_mb}"]
+        cmd.append("--stack-child")
+        if cwd:
+            cmd.extend(("--cwd", cwd))
+        cmd.extend(pytest_args)
+        return subprocess.call(cmd, env=env)  # noqa: S603 - trusted Python copy
+
+
+def _ensure_main_stack(stack_mb):
+    stack_bytes = stack_mb * 1024 * 1024
+    if sys.platform == "win32":
+        import ctypes
+
+        low, high = ctypes.c_size_t(), ctypes.c_size_t()
+        get_limits = ctypes.WinDLL("kernel32", use_last_error=True).GetCurrentThreadStackLimits
+        get_limits.argtypes = [
+            ctypes.POINTER(ctypes.c_size_t),
+            ctypes.POINTER(ctypes.c_size_t),
+        ]
+        get_limits.restype = None
+        get_limits(ctypes.byref(low), ctypes.byref(high))
+        actual = high.value - low.value
+        if actual < stack_bytes:
+            raise RuntimeError(f"Main thread stack is {actual} bytes, requested {stack_bytes}")
+        print(f"[stack-wrapper] main thread stack: {actual} bytes", flush=True)
+    else:
+        import resource
+
+        soft, hard = resource.getrlimit(resource.RLIMIT_STACK)
+        if soft != resource.RLIM_INFINITY and soft < stack_bytes:
+            resource.setrlimit(resource.RLIMIT_STACK, (stack_bytes, hard))
 
 
 def main():
@@ -30,24 +94,27 @@ def main():
         "--stack-mb",
         type=int,
         default=8,
-        help="Thread stack size in megabytes (default: 8)",
+        help="Main thread stack size in megabytes (default: 8)",
     )
     parser.add_argument(
         "--cwd",
         default=None,
         help="Working directory for the test run",
     )
+    parser.add_argument("--stack-child", action="store_true", help=argparse.SUPPRESS)
     args, pytest_args = parser.parse_known_args()
+    if args.stack_mb <= 0:
+        parser.error("--stack-mb must be positive")
+
+    if sys.platform == "win32" and not args.stack_child:
+        sys.exit(_run_windows(args.stack_mb, args.cwd, pytest_args))
 
+    _ensure_main_stack(args.stack_mb)
     if args.cwd:
         os.chdir(args.cwd)
+    import pytest
 
-    threading.stack_size(args.stack_mb * 1024 * 1024)
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
-        code = pool.submit(pytest.main, pytest_args).result()
-
-    sys.exit(code)
+    sys.exit(pytest.main(pytest_args))
 
 
 if __name__ == "__main__":

From 9a73fa83140eb0b72f8f9620ad2df34fb48914fb Mon Sep 17 00:00:00 2001
From: Rui Luo <ruluo@nvidia.com>
Date: Fri, 23 Oct 2026 15:17:43 +0800
Subject: [PATCH 2/2] ci: keep manual PR coverage runs from deploying

---
 .github/workflows/coverage.yml | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml
index 448e13ad2ff..f0bd243f24e 100644
--- a/.github/workflows/coverage.yml
+++ b/.github/workflows/coverage.yml
@@ -581,7 +581,11 @@ jobs:
     # GitHub discards prerequisite artifacts on a new run attempt. If only this
     # job fails, use "Re-run all jobs" so coverage recreates its payload.
     needs: combine-and-deploy
-    if: ${{ needs.combine-and-deploy.result == 'success' && github.repository_owner == 'nvidia' }}
+    # Manual runs on PR mirror branches should not publish over main's report.
+    if: >-
+      needs.combine-and-deploy.result == 'success' &&
+      github.repository_owner == 'nvidia' &&
+      github.ref == 'refs/heads/main'
     runs-on: ubuntu-latest
     permissions:
       contents: write