Skip to content

Commit 7695e91

Browse files
rwgkrwgk
committed
test: patch pytest-run-parallel worker context setup
1 parent dad6a42 commit 7695e91

4 files changed

Lines changed: 456 additions & 12 deletions

File tree

‎ci/tools/tests/test_pytest_run_parallel_patch.py‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import sys
5+
import threading
6+
import types
7+
from contextlib import contextmanager
8+
from pathlib import Path
9+
10+
import pytest
11+
12+
_TEST_HELPERS_ROOT = Path(__file__).resolve().parents[3] / "cuda_python_test_helpers"
13+
sys.path.insert(0, str(_TEST_HELPERS_ROOT))
14+
15+
from cuda_python_test_helpers.pytest_run_parallel import (
16+
install_run_parallel_worker_context_patch,
17+
mark_item_for_worker_context,
18+
)
19+
20+
21+
def _install_fake_pytest_run_parallel(monkeypatch):
22+
package = types.ModuleType("pytest_run_parallel")
23+
package.__path__ = []
24+
plugin = types.ModuleType("pytest_run_parallel.plugin")
25+
26+
def wrap_function_parallel(fn, n_workers, n_iterations):
27+
raise AssertionError("unpatched fake wrapper should not be called")
28+
29+
plugin.wrap_function_parallel = wrap_function_parallel
30+
monkeypatch.setitem(sys.modules, "pytest_run_parallel", package)
31+
monkeypatch.setitem(sys.modules, "pytest_run_parallel.plugin", plugin)
32+
return plugin
33+
34+
35+
@pytest.mark.agent_authored(model="gpt-5")
36+
def test_install_run_parallel_worker_context_patch_is_idempotent(monkeypatch):
37+
plugin = _install_fake_pytest_run_parallel(monkeypatch)
38+
39+
assert install_run_parallel_worker_context_patch() is True
40+
patched = plugin.wrap_function_parallel
41+
assert patched._cuda_python_patched_run_parallel_worker_context
42+
43+
assert install_run_parallel_worker_context_patch() is True
44+
assert plugin.wrap_function_parallel is patched
45+
46+
47+
@pytest.mark.agent_authored(model="gpt-5")
48+
def test_patched_wrapper_runs_context_with_isolated_kwargs(monkeypatch):
49+
plugin = _install_fake_pytest_run_parallel(monkeypatch)
50+
install_run_parallel_worker_context_patch()
51+
52+
lock = threading.Lock()
53+
context_events = []
54+
calls = []
55+
56+
@contextmanager
57+
def worker_context(*, thread_index, iteration_index, kwargs):
58+
token = object()
59+
kwargs["token"] = (thread_index, iteration_index, id(token))
60+
kwargs["context_kwargs_id"] = id(kwargs)
61+
with lock:
62+
context_events.append(("enter", thread_index, iteration_index, id(kwargs)))
63+
try:
64+
yield
65+
finally:
66+
with lock:
67+
context_events.append(("exit", thread_index, iteration_index, id(kwargs)))
68+
69+
def test_body(*, thread_index, iteration_index, token, context_kwargs_id, static_value):
70+
with lock:
71+
calls.append(
72+
{
73+
"thread_index": thread_index,
74+
"iteration_index": iteration_index,
75+
"token": token,
76+
"context_kwargs_id": context_kwargs_id,
77+
"static_value": static_value,
78+
}
79+
)
80+
81+
item = types.SimpleNamespace(obj=test_body)
82+
assert mark_item_for_worker_context(item, worker_context) is True
83+
84+
wrapped = plugin.wrap_function_parallel(item.obj, n_workers=3, n_iterations=2)
85+
wrapped(thread_index=-1, iteration_index=-1, static_value="fixture-value")
86+
87+
expected_pairs = {(thread_index, iteration_index) for thread_index in range(3) for iteration_index in range(2)}
88+
actual_pairs = {(call["thread_index"], call["iteration_index"]) for call in calls}
89+
assert actual_pairs == expected_pairs
90+
assert {call["token"][:2] for call in calls} == expected_pairs
91+
assert {call["static_value"] for call in calls} == {"fixture-value"}
92+
93+
kwargs_ids = {call["context_kwargs_id"] for call in calls}
94+
assert len(kwargs_ids) == 6
95+
assert len(context_events) == 12
96+
assert {event[3] for event in context_events} == kwargs_ids
97+
98+
99+
@pytest.mark.agent_authored(model="gpt-5")
100+
def test_mark_item_for_worker_context_wraps_callables_without_attrs():
101+
class CallableWithoutAttrs:
102+
__slots__ = ("calls",)
103+
104+
def __init__(self):
105+
self.calls = []
106+
107+
def __call__(self, **kwargs):
108+
self.calls.append(kwargs)
109+
110+
@contextmanager
111+
def worker_context(*, thread_index, iteration_index, kwargs):
112+
kwargs["patched"] = True
113+
yield
114+
115+
original = CallableWithoutAttrs()
116+
item = types.SimpleNamespace(obj=original)
117+
118+
assert mark_item_for_worker_context(item, worker_context) is True
119+
assert item.obj is not original
120+
item.obj()
121+
assert original.calls == [{}]

‎cuda_bindings/tests/conftest.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import pathlib
55
import sys
6+
from contextlib import contextmanager
67
from importlib.metadata import PackageNotFoundError, distribution
78

89
import pytest
@@ -25,6 +26,59 @@
2526
sys.path.insert(0, test_helpers_root)
2627

2728

29+
from cuda_python_test_helpers.pytest_run_parallel import (
30+
install_run_parallel_worker_context_patch,
31+
mark_item_for_worker_context,
32+
)
33+
34+
35+
def pytest_configure(config):
36+
install_run_parallel_worker_context_patch()
37+
38+
39+
@contextmanager
40+
def _thread_context():
41+
(err,) = cuda.cuInit(0)
42+
assert err == cuda.CUresult.CUDA_SUCCESS
43+
err, device = cuda.cuDeviceGet(0)
44+
assert err == cuda.CUresult.CUDA_SUCCESS
45+
err, ctx = cuda.cuCtxCreate(None, 0, device)
46+
assert err == cuda.CUresult.CUDA_SUCCESS
47+
try:
48+
yield device, ctx
49+
finally:
50+
(err,) = cuda.cuCtxDestroy(ctx)
51+
assert err == cuda.CUresult.CUDA_SUCCESS
52+
53+
54+
@contextmanager
55+
def _cuda_bindings_worker_context(*, thread_index, iteration_index, kwargs):
56+
with _thread_context() as (device, ctx):
57+
if "device" in kwargs:
58+
kwargs["device"] = device
59+
if "ctx" in kwargs:
60+
kwargs["ctx"] = ctx
61+
yield
62+
63+
64+
def _is_cudla_item(item):
65+
nodeid = item.nodeid.replace("\\", "/")
66+
return nodeid.startswith("tests/cudla/") or "cuda_bindings/tests/cudla/" in nodeid
67+
68+
69+
def _item_needs_thread_ctx(item):
70+
if _is_cudla_item(item):
71+
return False
72+
fixturenames = set(getattr(item, "fixturenames", ()))
73+
return bool(fixturenames & {"device", "ctx", "driver", "cufile_env_json"})
74+
75+
76+
def pytest_collection_modifyitems(config, items):
77+
for item in items:
78+
if _item_needs_thread_ctx(item):
79+
mark_item_for_worker_context(item, _cuda_bindings_worker_context)
80+
81+
2882
@pytest.fixture(scope="module")
2983
def cuda_driver():
3084
(err,) = cuda.cuInit(0)

‎cuda_core/tests/conftest.py‎

Lines changed: 72 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,76 @@ def xfail_if_mempool_oom(err_or_exc, api_name=None, device=0):
9191
sys.path.insert(0, test_helpers_root)
9292

9393

94+
from cuda_python_test_helpers.pytest_run_parallel import (
95+
install_run_parallel_worker_context_patch,
96+
mark_item_for_worker_context,
97+
)
98+
99+
100+
def pytest_configure(config):
101+
install_run_parallel_worker_context_patch()
102+
103+
104+
@contextmanager
105+
def _init_cuda_context():
106+
# TODO: rename this to e.g. init_context
107+
device = Device(0)
108+
device.set_current()
109+
110+
# Set option to avoid spin-waiting on synchronization.
111+
if int(os.environ.get("CUDA_CORE_TEST_BLOCKING_SYNC", 0)) != 0:
112+
handle_return(
113+
driver.cuDevicePrimaryCtxSetFlags(device.device_id, driver.CUctx_flags.CU_CTX_SCHED_BLOCKING_SYNC)
114+
)
115+
116+
try:
117+
yield device
118+
finally:
119+
_ = _device_unset_current()
120+
121+
122+
@contextmanager
123+
def _cuda_core_worker_context(*, thread_index, iteration_index, kwargs):
124+
with _init_cuda_context() as device:
125+
if "init_cuda" in kwargs:
126+
kwargs["init_cuda"] = device
127+
if "mempool_device" in kwargs:
128+
kwargs["mempool_device"] = device
129+
if "ipc_device" in kwargs:
130+
kwargs["ipc_device"] = device
131+
if "mempool_device_x2" in kwargs:
132+
kwargs["mempool_device_x2"] = _mempool_device_impl(2)
133+
if "mempool_device_x3" in kwargs:
134+
kwargs["mempool_device_x3"] = _mempool_device_impl(3)
135+
if "ipc_mempool_device_x2" in kwargs:
136+
kwargs["ipc_mempool_device_x2"] = _require_ipc_mempool_devices(_mempool_device_impl(2))
137+
yield
138+
139+
140+
_CUDA_CONTEXT_FIXTURES = frozenset(
141+
{
142+
"init_cuda",
143+
"ipc_device",
144+
"ipc_memory_resource",
145+
"mempool_device",
146+
"mempool_device_x2",
147+
"mempool_device_x3",
148+
"ipc_mempool_device_x2",
149+
"memory_resource_factory",
150+
}
151+
)
152+
153+
154+
def _item_needs_thread_ctx(item):
155+
return bool(_CUDA_CONTEXT_FIXTURES & set(getattr(item, "fixturenames", ())))
156+
157+
158+
def pytest_collection_modifyitems(config, items):
159+
for item in items:
160+
if _item_needs_thread_ctx(item):
161+
mark_item_for_worker_context(item, _cuda_core_worker_context)
162+
163+
94164
def skip_if_pinned_memory_unsupported(device):
95165
try:
96166
if not device.properties.host_memory_pools_supported:
@@ -194,18 +264,8 @@ def session_setup():
194264

195265
@pytest.fixture
196266
def init_cuda():
197-
# TODO: rename this to e.g. init_context
198-
device = Device(0)
199-
device.set_current()
200-
201-
# Set option to avoid spin-waiting on synchronization.
202-
if int(os.environ.get("CUDA_CORE_TEST_BLOCKING_SYNC", 0)) != 0:
203-
handle_return(
204-
driver.cuDevicePrimaryCtxSetFlags(device.device_id, driver.CUctx_flags.CU_CTX_SCHED_BLOCKING_SYNC)
205-
)
206-
207-
yield device
208-
_ = _device_unset_current()
267+
with _init_cuda_context() as device:
268+
yield device
209269

210270

211271
def _device_unset_current() -> bool:

0 commit comments

Comments
 (0)