Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions .github/workflows/coverage.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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: |
Expand Down Expand Up @@ -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
Expand Down
97 changes: 82 additions & 15 deletions ci/tools/run_pytest_with_stack.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

🔎 Supported by static analysis

🏁 Script executed:

sed -n '1,135p' ci/tools/run_pytest_with_stack.py
rg -n 'run_pytest_with_stack|stack-mb|RLIMIT_STACK' .github ci tests

Repository: NVIDIA/cuda-python

Length of output: 5595


🌐 Web query:

official Python documentation resource.setrlimit ValueError soft limit exceeds hard limit

💡 Result:

<source_evidence>
<source kind="search">
<title>resource — Resource usage information — Python 3.14.7 documentation</title>
<location>https://docs.python.org/3/library/resource.html</location>
<excerpt>citeturn0search0 [wordlim: 200] Crawled: last month; The soft limit can never exceed the hard limit. ... resource.setrlimit(resource, limits)¶ ... Raises `ValueError` if an invalid resource is specified, if the new soft limit exceeds the hard limit, or if a process tries to raise its hard limit.
</excerpt>
</source>
<source kind="search">
<title>resource — Resource usage information — Python 3.15.1 documentation</title>
<location>https://docs.python.org/3.15/library/resource.html</location>
<excerpt>citeturn0search1 [wordlim: 200] Crawled: yesterday; resource.setrlimit(resource, limits)¶ ... Raises `ValueError` if an invalid resource is specified, if the new soft limit exceeds the hard limit, or if a process tries to raise its hard limit.
</excerpt>
</source>
<source kind="search">
<title>resource — Resource usage information — Python 3.14.7 documentation</title>
<location>https://docs.python.org/uk/3/library/resource.html</location>
<excerpt>citeturn0search2 [wordlim: 200] Crawled: last month; resource.setrlimit(resource, limits)¶ ... Raises `ValueError` if an invalid resource is specified, if the new soft limit exceeds the hard limit, or if a process tries to raise its hard limit.
</excerpt>
</source>
<source kind="search">
<title>resource — Resource usage information — Python 3.15.0 documentation</title>
<location>https://docs.python.org/3/library/resource.html?highlight=option</location>
<excerpt>citeturn0search3 [wordlim: 200] Crawled: 4 days ago; Each resource is controlled by a pair of limits: a soft limit and a hard limit. ... The soft limit can never exceed the hard limit. ... Raises `ValueError` if an invalid resource is specified, or `error` if the underlying system call fails unexpectedly.resource.setrlimit(resource, limits)¶
</excerpt>
</source>
<source kind="search">
<title>cpython/Doc/library/resource.rst at main · python/cpython · GitHub</title>
<location>https://github.com/python/cpython/blob/main/Doc/library/resource.rst?plain=1</location>
<excerpt>citeturn0search4 [wordlim: 200] Crawled: 4 days ago; Each resource is controlled by a pair of limits: a soft limit and a hard ... Raises :exc:`ValueError` if an invalid resource is specified, or ... .. function:: setrlimit(resource, limits)
</excerpt>
</source>
<omitted_sources count="26">Additional source records were omitted from this bounded result.</omitted_sources>
</source_evidence>

Citations:

- 1: https://docs.python.org/3/library/resource.html?utm_source=openai
- 2: https://docs.python.org/3/library/resource.html?utm_source=openai

suggestion: Reject requests above the hard limit.

resource.setrlimit raises ValueError when the requested soft limit exceeds a finite hard limit. Silently clamping would violate the wrapper's contract because it would not provide the requested stack.

--- "a/ci/tools/run_pytest_with_stack.py"
+++ "b/ci/tools/run_pytest_with_stack.py"
@@ -83,9 +83,14 @@
     else:
         import resource
 
         soft, hard = resource.getrlimit(resource.RLIMIT_STACK)
         if soft != resource.RLIM_INFINITY and soft < stack_bytes:
+            if hard != resource.RLIM_INFINITY and stack_bytes > hard:
+                raise RuntimeError(
+                    f"Cannot ensure requested stack size: {stack_bytes} exceeds "
+                    f"RLIMIT_STACK hard limit {hard}"
+                )
             resource.setrlimit(resource.RLIMIT_STACK, (stack_bytes, hard))
 
 
 def main():
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
resource.setrlimit(resource.RLIMIT_STACK, (stack_bytes, hard))
if hard != resource.RLIM_INFINITY and stack_bytes > hard:
raise RuntimeError(
f"Cannot ensure requested stack size: {stack_bytes} exceeds "
f"RLIMIT_STACK hard limit {hard}"
)
resource.setrlimit(resource.RLIMIT_STACK, (stack_bytes, hard))



def main():
Expand All @@ -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__":
Expand Down
Loading