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

Author: cpcloud

cuda_core/cuda/core/experimental/_device.pyx

@@ -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,
@@ -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,

cuda_core/cuda/core/experimental/_memory.pyx

@@ -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,
@@ -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":

cuda_core/cuda/core/experimental/_stream.pxd

@@ -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()

cuda_core/cuda/core/experimental/_stream.pyx

@@ -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
 
@@ -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
@@ -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")
+
+    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

cuda_core/tests/test_stream.py

@@ -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
 
 
@@ -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)
-
-
 def test_stream_subclassing(init_cuda):
     class MyStream(Stream):
         pass