Skip to content

fix: ensure that there are no module global lookups when calling default_stream#1107

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
leofang merged 1 commit into from
Oct 8, 2025
Merged

fix: ensure that there are no module global lookups when calling default_stream#1107

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 2 additions & 1 deletion cuda_core/cuda/core/experimental/_device.pyx
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -16,7 +16,7 @@ from cuda.core.experimental._context import Context, ContextOptions
from cuda.core.experimental._event import Event, EventOptions
from cuda.core.experimental._graph import GraphBuilder
from cuda.core.experimental._memory import Buffer, DeviceMemoryResource, MemoryResource, _SynchronousMemoryResource
from cuda.core.experimental._stream import IsStreamT, Stream, StreamOptions, default_stream
from cuda.core.experimental._stream import IsStreamT, Stream, StreamOptions
from cuda.core.experimental._utils.clear_error_support import assert_type
from cuda.core.experimental._utils.cuda_utils import (
ComputeCapability,
Expand All@@ -25,6 +25,7 @@ from cuda.core.experimental._utils.cuda_utils import (
handle_return,
runtime,
)
from cuda.core.experimental._stream cimport default_stream


# TODO: I prefer to type these as "cdef object" and avoid accessing them from within Python,
Expand Down
3 changes: 2 additions & 1 deletion cuda_core/cuda/core/experimental/_memory.pyx
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -12,6 +12,7 @@ from libc.string cimport memset, memcpy
from cuda.bindings cimport cydriver

from cuda.core.experimental._stream cimport Stream as cyStream
from cuda.core.experimental._stream cimport default_stream
from cuda.core.experimental._utils.cuda_utils cimport (
_check_driver_error as raise_if_driver_error,
check_or_create_options,
Expand All@@ -30,7 +31,7 @@ import platform
import weakref

from cuda.core.experimental._dlpack import DLDeviceType, make_py_capsule
from cuda.core.experimental._stream import Stream, default_stream
from cuda.core.experimental._stream import Stream
from cuda.core.experimental._utils.cuda_utils import ( driver, Transaction, get_binding_version )

if platform.system() == "Linux":
Expand Down
3 changes: 3 additions & 0 deletions cuda_core/cuda/core/experimental/_stream.pxd
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -22,3 +22,6 @@ cdef class Stream:
cpdef close(self)
cdef int _get_context(self) except?-1 nogil
cdef int _get_device_and_context(self) except?-1


cdef Stream default_stream()
24 changes: 18 additions & 6 deletions cuda_core/cuda/core/experimental/_stream.pyx
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -5,6 +5,7 @@
from __future__ import annotations

from libc.stdint cimport uintptr_t, INT32_MIN
from libc.stdlib cimport strtol, getenv

from cuda.bindings cimport cydriver

Expand DownExpand Up@@ -388,11 +389,16 @@ cdef class Stream:
return GraphBuilder._init(stream=self, is_stream_owner=False)


LEGACY_DEFAULT_STREAM = Stream._legacy_default()
PER_THREAD_DEFAULT_STREAM = Stream._per_thread_default()
# c-only python objects, not public
cdef Stream C_LEGACY_DEFAULT_STREAM = Stream._legacy_default()
cdef Stream C_PER_THREAD_DEFAULT_STREAM = Stream._per_thread_default()

# standard python objects, public
LEGACY_DEFAULT_STREAM = C_LEGACY_DEFAULT_STREAM
PER_THREAD_DEFAULT_STREAM = C_PER_THREAD_DEFAULT_STREAM

def default_stream():

cdef Stream default_stream():
"""Return the default CUDA :obj:`~_stream.Stream`.

The type of default stream returned depends on if the environment
Expand All@@ -403,8 +409,14 @@ def default_stream():

"""
# TODO: flip the default
use_ptds = int(os.environ.get("CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM", 0))
cdef const char* use_ptds_raw = getenv("CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM")

@kkraus14kkraus14Oct 8, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Instead of looking this up in the function, we could probably resolve this as a global. Since we resolve the underlying function symbols only once we don't handle the situation where someone changes the environment variable value during process execution.


cdef int use_ptds = 0
if use_ptds_raw != NULL:
use_ptds = strtol(use_ptds_raw, NULL, 10)

# value is non-zero, including for weird stuff like 123foo
if use_ptds:
return PER_THREAD_DEFAULT_STREAM
return C_PER_THREAD_DEFAULT_STREAM
else:
return LEGACY_DEFAULT_STREAM
return C_LEGACY_DEFAULT_STREAM
7 changes: 1 addition & 6 deletions cuda_core/tests/test_stream.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -4,7 +4,7 @@
import pytest
from cuda.core.experimental import Device, Stream, StreamOptions
from cuda.core.experimental._event import Event
from cuda.core.experimental._stream import LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM, default_stream
from cuda.core.experimental._stream import LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM
from cuda.core.experimental._utils.cuda_utils import driver


Expand DownExpand Up@@ -107,11 +107,6 @@ def test_per_thread_default_stream():
assert isinstance(PER_THREAD_DEFAULT_STREAM, Stream)


def test_default_stream():
stream = default_stream()
assert isinstance(stream, Stream)
Comment thread
leofang marked this conversation as resolved.


def test_stream_subclassing(init_cuda):
class MyStream(Stream):
pass
Expand Down
Loading