Skip to content

Commit 2f40730

Browse files
aryanputtaaryanputta
committed
cuda.core: resolve default-stream context per call instead of caching it
LEGACY_DEFAULT_STREAM and PER_THREAD_DEFAULT_STREAM wrap default-stream tokens, which denote whatever context is current. Both are module-level singletons, and Stream_ensure_ctx / Stream_ensure_ctx_device stored the first context and device they observed on the object and never cleared them, so a process-wide object became permanently bound to one context. Replace the two helpers with resolvers that return the context and device through out-parameters and cache on the object only when the stream is not a default-stream token. Stream.context, .device, .resources, .record(), and __repr__ now follow the current context, a query no longer pins a context reference for the lifetime of the process, and the shared singletons are no longer written to from multiple threads. Object identity is preserved, so __eq__ and __hash__ keying off the handle are unaffected. Fixes #2485
1 parent e5c7b06 commit 2f40730

4 files changed

Lines changed: 146 additions & 30 deletions

File tree

‎cuda_core/cuda/core/_stream.pyi‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,16 @@ class Stream:
165165

166166
@property
167167
def context(self) -> Context:
168-
"""Return the :obj:`~_context.Context` associated with this stream."""
168+
"""Return the :obj:`~_context.Context` associated with this stream.
169+
170+
Note
171+
----
172+
A default stream (:obj:`LEGACY_DEFAULT_STREAM` or
173+
:obj:`PER_THREAD_DEFAULT_STREAM`) carries no context of its own; it
174+
refers to whichever context is current, so this returns the current
175+
context at call time.
176+
177+
"""
169178

170179
@property
171180
def resources(self) -> DeviceResources:

‎cuda_core/cuda/core/_stream.pyx‎

Lines changed: 67 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -216,8 +216,9 @@ cdef class Stream:
216216
return as_intptr(self._h_stream) == as_intptr((<Stream>other)._h_stream)
217217

218218
def __repr__(self) -> str:
219-
Stream_ensure_ctx(self)
220-
return f"<Stream handle={as_intptr(self._h_stream):#x} context={as_intptr(self._h_context):#x}>"
219+
cdef ContextHandle h_context
220+
Stream_get_ctx(self, &h_context)
221+
return f"<Stream handle={as_intptr(self._h_stream):#x} context={as_intptr(h_context):#x}>"
221222

222223
@property
223224
def handle(self) -> cuda.bindings.driver.CUstream:
@@ -277,9 +278,11 @@ cdef class Stream:
277278
# Create an Event object (or reusing the given one) by recording
278279
# on the stream. Event flags such as disabling timing, nonblocking,
279280
# and CU_EVENT_RECORD_EXTERNAL, can be set in EventOptions.
281+
cdef ContextHandle h_context
282+
cdef int device_id
280283
if event is None:
281-
Stream_ensure_ctx_device(self)
282-
event = cyEvent._init(cyEvent, self._device_id, self._h_context, options, False)
284+
Stream_get_ctx_device(self, &h_context, &device_id)
285+
event = cyEvent._init(cyEvent, device_id, h_context, options, False)
283286
elif event.is_ipc_enabled:
284287
raise TypeError(
285288
"IPC-enabled events should not be re-recorded, instead create a "
@@ -350,15 +353,27 @@ cdef class Stream:
350353

351354
"""
352355
from cuda.core._device import Device # avoid circular import
353-
Stream_ensure_ctx_device(self)
354-
return Device(self._device_id)
356+
cdef ContextHandle h_context
357+
cdef int device_id
358+
Stream_get_ctx_device(self, &h_context, &device_id)
359+
return Device(device_id)
355360

356361
@property
357362
def context(self) -> Context:
358-
"""Return the :obj:`~_context.Context` associated with this stream."""
359-
Stream_ensure_ctx(self)
360-
Stream_ensure_ctx_device(self)
361-
return Context._from_handle(Context, self._h_context, self._device_id)
363+
"""Return the :obj:`~_context.Context` associated with this stream.
364+
365+
Note
366+
----
367+
A default stream (:obj:`LEGACY_DEFAULT_STREAM` or
368+
:obj:`PER_THREAD_DEFAULT_STREAM`) carries no context of its own; it
369+
refers to whichever context is current, so this returns the current
370+
context at call time.
371+
372+
"""
373+
cdef ContextHandle h_context
374+
cdef int device_id
375+
Stream_get_ctx_device(self, &h_context, &device_id)
376+
return Context._from_handle(Context, h_context, device_id)
362377

363378
@property
364379
def resources(self) -> DeviceResources:
@@ -368,9 +383,10 @@ cdef class Stream:
368383
that context was provisioned with. For streams on the primary
369384
context, returns the full device resources.
370385
"""
371-
Stream_ensure_ctx(self)
372-
Stream_ensure_ctx_device(self)
373-
return DeviceResources._init_from_ctx(self._h_context, self._device_id)
386+
cdef ContextHandle h_context
387+
cdef int device_id
388+
Stream_get_ctx_device(self, &h_context, &device_id)
389+
return DeviceResources._init_from_ctx(h_context, device_id)
374390

375391
@staticmethod
376392
def from_handle(handle) -> Stream:
@@ -447,39 +463,61 @@ cpdef Stream default_stream():
447463
return LEGACY_DEFAULT_STREAM
448464

449465

450-
cdef inline int Stream_ensure_ctx(Stream self) except?-1 nogil:
451-
"""Ensure the stream's context handle is populated."""
466+
cdef inline bint Stream_is_default_token(Stream self) noexcept nogil:
467+
"""Return True for CU_STREAM_LEGACY and CU_STREAM_PER_THREAD.
468+
469+
These tokens carry no context of their own; they refer to whatever context
470+
is current, so nothing resolved from one may be cached on the object.
471+
"""
472+
cdef uintptr_t h = <uintptr_t>as_cu(self._h_stream)
473+
return h == <uintptr_t>cydriver.CU_STREAM_LEGACY or h == <uintptr_t>cydriver.CU_STREAM_PER_THREAD
474+
475+
476+
cdef inline int Stream_get_ctx(Stream self, ContextHandle* h_context) except?-1 nogil:
477+
"""Resolve the stream's context handle into ``h_context``."""
452478
cdef cydriver.CUcontext ctx
453-
if not self._h_context:
454-
self._h_context = get_stream_context(self._h_stream)
479+
cdef bint is_default = Stream_is_default_token(self)
480+
455481
if self._h_context:
482+
h_context[0] = self._h_context
456483
return 0
457-
HANDLE_RETURN(cydriver.cuStreamGetCtx(as_cu(self._h_stream), &ctx))
458-
if ctx != NULL:
459-
with gil:
460-
self._h_context = create_context_handle_ref(ctx)
484+
485+
h_context[0] = get_stream_context(self._h_stream)
486+
if not h_context[0]:
487+
HANDLE_RETURN(cydriver.cuStreamGetCtx(as_cu(self._h_stream), &ctx))
488+
if ctx != NULL:
489+
with gil:
490+
h_context[0] = create_context_handle_ref(ctx)
491+
492+
if h_context[0] and not is_default:
493+
self._h_context = h_context[0]
461494
return 0
462495

463496

464-
cdef inline int Stream_ensure_ctx_device(Stream self) except?-1:
465-
"""Ensure the stream's context and device_id are populated."""
497+
cdef inline int Stream_get_ctx_device(Stream self, ContextHandle* h_context, int* device_id) except?-1:
498+
"""Resolve the stream's context handle and device ID."""
466499
cdef cydriver.CUcontext ctx
467500
cdef cydriver.CUdevice target_dev
468501
cdef ContextHandle current_context
469502
cdef bint switch_context
503+
cdef bint is_default = Stream_is_default_token(self)
470504

471-
if self._device_id < 0:
472-
with nogil:
505+
with nogil:
506+
Stream_get_ctx(self, h_context)
507+
if self._device_id >= 0 and not is_default:
508+
device_id[0] = self._device_id
509+
else:
473510
# Get device ID from context, switching context temporarily if needed
474-
Stream_ensure_ctx(self)
475511
current_context = get_current_context()
476-
switch_context = (as_cu(current_context) != as_cu(self._h_context))
512+
switch_context = (as_cu(current_context) != as_cu(h_context[0]))
477513
if switch_context:
478-
HANDLE_RETURN(cydriver.cuCtxPushCurrent(as_cu(self._h_context)))
514+
HANDLE_RETURN(cydriver.cuCtxPushCurrent(as_cu(h_context[0])))
479515
HANDLE_RETURN(cydriver.cuCtxGetDevice(&target_dev))
480516
if switch_context:
481517
HANDLE_RETURN(cydriver.cuCtxPopCurrent(&ctx))
482-
self._device_id = <int>target_dev
518+
device_id[0] = <int>target_dev
519+
if not is_default:
520+
self._device_id = device_id[0]
483521
return 0
484522

485523

‎cuda_core/docs/source/release/1.2.0-notes.rst‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,15 @@ Fixes and enhancements
2525
versions 12.2 or newer.
2626
(`#2352 <https://github.com/NVIDIA/cuda-python/issues/2352>`__)
2727

28+
- The default-stream singletons ``LEGACY_DEFAULT_STREAM`` and
29+
``PER_THREAD_DEFAULT_STREAM`` no longer cache the first context and device
30+
they observe. A default-stream token refers to whatever context is current,
31+
so ``Stream.context``, ``Stream.device``, ``Stream.resources``, and
32+
``Stream.record()`` now resolve against the current context on every call.
33+
Previously the first query pinned the singleton to one context for the
34+
lifetime of the process, which also kept that context alive.
35+
(`#2485 <https://github.com/NVIDIA/cuda-python/issues/2485>`__)
36+
2837
Deprecation Notices
2938
-------------------
3039

‎cuda_core/tests/test_stream.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -405,3 +405,63 @@ def test_default_stream_per_thread_when_env_set(monkeypatch):
405405
assert default_stream() is PER_THREAD_DEFAULT_STREAM
406406
monkeypatch.delenv("CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM", raising=False)
407407
assert default_stream() is LEGACY_DEFAULT_STREAM
408+
409+
410+
def _skip_unless_multi_gpu():
411+
from cuda.core import system
412+
413+
if system.get_num_devices() < 2:
414+
pytest.skip("requires 2+ GPUs")
415+
416+
417+
@pytest.mark.agent_authored(model="claude-opus-5")
418+
@pytest.mark.parametrize("stream", [LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM])
419+
def test_default_stream_follows_current_context(stream):
420+
"""A default-stream token denotes whatever context is current, so its
421+
queries follow a context switch instead of reporting the first context
422+
they ever saw (issue #2485)."""
423+
_skip_unless_multi_gpu()
424+
425+
for device_id in (0, 1, 0):
426+
dev = Device(device_id)
427+
dev.set_current()
428+
assert stream.device.device_id == device_id
429+
assert stream.context == dev.context
430+
assert stream.record().context == dev.context
431+
432+
433+
@pytest.mark.agent_authored(model="claude-opus-5")
434+
def test_default_stream_first_touch_does_not_pin_context():
435+
"""Any query resolves the context, including repr(), so logging a default
436+
stream must not bind the singleton to whichever context happened to be
437+
current at the time (issue #2485)."""
438+
_skip_unless_multi_gpu()
439+
440+
Device(1).set_current()
441+
repr(LEGACY_DEFAULT_STREAM)
442+
443+
dev0 = Device(0)
444+
dev0.set_current()
445+
assert LEGACY_DEFAULT_STREAM.device.device_id == 0
446+
assert LEGACY_DEFAULT_STREAM.context == dev0.context
447+
448+
449+
@pytest.mark.agent_authored(model="claude-opus-5")
450+
def test_created_stream_keeps_its_own_context():
451+
"""A created stream has a context fixed at creation and keeps reporting it
452+
across a switch; only default-stream tokens follow the current context
453+
(issue #2485)."""
454+
_skip_unless_multi_gpu()
455+
456+
dev0 = Device(0)
457+
dev0.set_current()
458+
stream = dev0.create_stream()
459+
assert stream.context == dev0.context
460+
461+
try:
462+
Device(1).set_current()
463+
assert stream.device.device_id == 0
464+
assert stream.context == dev0.context
465+
finally:
466+
dev0.set_current()
467+
stream.close()

0 commit comments

Comments
 (0)