Stream context manager does not preserve previous cupy.cuda.Stream.use() call
Description
We have a code where we set the stream at program startup to a non-default CUDA stream (via the use option). We discovered that if somewhere later on in the code a with statement with a different stream is called, the stream resets to the CUDA default stream once we leave the scope of the context manager! This is unexpected and unintuitive behavior.
Some comments in https://github.com/cupy/cupy/blob/main/cupy/cuda/stream.pyx state that mixing use and with is an "anti-pattern" but I think it is reasonable for a user to use a stream to set a new default stream and then use with to temporarily change streams during the program duration. I would agree that using a use within a with block makes much less sense.
To Reproduce
import cupy as cp
s1 = cp.cuda.Stream()
s2 = cp.cuda.Stream()
print(f"s1 = {s1}")
print(f"s2 = {s2}")
print(cp.cuda.get_current_stream())
# Set current stream to s1
s1.use()
print(cp.cuda.get_current_stream())
# Temporarily set stream to s2
with s2:
print(cp.cuda.get_current_stream())
# Stream resets to default stream but should reset to s1
print(cp.cuda.get_current_stream())Installation
Wheel (pip install cupy-***)
Environment
OS : Linux-5.4.0-169-generic-x86_64-with-glibc2.35
Python Version : 3.10.12
CuPy Version : 13.1.0
CuPy Platform : NVIDIA CUDA
NumPy Version : 1.26.4
SciPy Version : 1.13.0
Cython Build Version : 0.29.36
Cython Runtime Version : None
CUDA Root : /opt/nvidia/hpc_sdk/Linux_x86_64/24.1/compilers
nvcc PATH : /opt/nvidia/hpc_sdk/Linux_x86_64/24.1/compilers/bin/nvcc
CUDA Build Version : 12040
CUDA Driver Version : 12020
CUDA Runtime Version : 12040 (linked to CuPy) / 12030 (locally installed)
cuBLAS Version : (available)
cuFFT Version : 11012
cuRAND Version : 10304
cuSOLVER Version : (11, 5, 4)
cuSPARSE Version : (available)
NVRTC Version : (12, 3)
Thrust Version : 200200
CUB Build Version : 200200
Jitify Build Version : <unknown>
cuDNN Build Version : (not loaded; try `import cupy.cuda.cudnn` first)
cuDNN Version : (not loaded; try `import cupy.cuda.cudnn` first)
NCCL Build Version : 21602
NCCL Runtime Version : 21805
cuTENSOR Version : 20000
cuSPARSELt Build Version : None
Device 0 Name : Tesla V100-DGXS-16GB
Device 0 Compute Capability : 70
Device 0 PCI Bus ID : 0000:07:00.0
Device 1 Name : Tesla V100-DGXS-16GB
Device 1 Compute Capability : 70
Device 1 PCI Bus ID : 0000:08:00.0
Device 2 Name : Tesla V100-DGXS-16GB
Device 2 Compute Capability : 70
Device 2 PCI Bus ID : 0000:0E:00.0
Device 3 Name : Tesla V100-DGXS-16GB
Device 3 Compute Capability : 70
Device 3 PCI Bus ID : 0000:0F:00.0Additional Information
The problem seems to lie in the handling of current_stream_stack in the implementation of Stream. The initial/base value of the stack is always initialized to the CUDA default stream. Only calls using with push or pop streams to this stack so once we leave the scope of the with statement, the stream gets reset to the base value in the stack which is the CUDA default stream, not the one currently set with use.
Source: cupy/cupy