torch.cuda.init() unstacks existing CUDA contexts

[QuiteAFoxtrot]

🐛 Describe the bug

I am using the Nvidia's 'cuda-python' bindings, and I have discovered that torch expects to be the only CUDA resident in a given program. In my code I was doing something as follows:

import torch
from cuda import cuda #nvidia package
device_num = 5
err, = cuda.cuInit(0)
err, device = cuda.cuGetDevice(device_num)
err, cuda_context = cuda.cuCtxCreate(device)
...
#some stuff with the current active context which works fine
err, = some_func_with_my(cuda_context)
err, cuda_context = cuda.cuCtxPopCurrent() #remove my context from the stack so its clean for PyTorch's turn
...
#Then I need to do some pytorch stuff
torch_dev = torch.device("cuda:"+str(device_num))
torch.cuda.init() #Or any torch.function which initializes CUDA

...
#Then later I go to use my cuda context
err, = cuda.cuCtxPushCurrent(cuda_context) #push my context back on the stack, since other programs need not be aware of my program
err, = some_other_func_which_uses_my(cuda_context) #Crashes with cuda error Context Destroyed

As far as I can tell, there is only one programming paradigm which Nvidia/CUDA makes safe for having multiple residents within a single process both able to use CUDA, as described here:
https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#context

Under the hood, torch.cuda.init() must be calling something like cuDeviceReset() or doing something otherwise funky, but the point is after calling torch.cuda.init(), the context I created was destroyed. The workaround to this is to call torch.cuda.init() BEFORE I create my context - but it would be nice if torch could adhere to the Nvidia provided paradigm and "play nice in the sandbox" with others, instead of destroying things.

Update: As per example below, context is not destroyed, just removed from the CUDA stack.

At minimum, this behavior should be explicitly documented in torch.cuda.init() (though that wouldn't be playing nicely in the sandbox). The current documentation notes "Does nothing if the CUDA state is already initialized." which is quite false.

Versions

python 3.8
pytorch 1.10.1
cuda_toolkit 11.2
cuda-python 11.6.1

cc @ngimel

[ngimel]

cc @ptrblck

[ptrblck]

@QuiteAFoxtrot I think this is expected behavior based on the Initialization section of the same linked docs.
In particular these section might be relevant:

The runtime creates a CUDA context for each device in the system (see Context for more details on CUDA contexts). This context is the primary context for this device and is initialized at the first runtime function which requires an active context on this device.

An application can mix runtime API code with driver API code.
If a context is created and made current via the driver API, subsequent runtime calls will pick up this context instead of creating a new one.
If the runtime is initialized (implicitly as mentioned in CUDA Runtime), cuCtxGetCurrent() can be used to retrieve the context created during initialization. This context can be used by subsequent driver API calls.

The driver API is interoperable with the runtime and it is possible to access the primary context (see Initialization) managed by the runtime from the driver API via cuDevicePrimaryCtxRetain().

I tried to reproduce the issue you are seeing using cuda-python as well as a recent PyTorch nightly binary to check this behavior, but wasn't able to trigger an error.
Since your methods are undefined I used e.g. cuda.cuCtxSetCurrent(cuda_context) or cuda.cuCtxGetApiVersion(cuda_context) to consume the context without any failures.

Could you provide a minimal, executable code snippet to reproduce the issue?

[QuiteAFoxtrot]

Ah - I was mistaken it seems. Torch doesn't actually destroy the user-created context, it just removes it completely from the stack (which I think I had mis-interpreted as 'destroyed'), but it seems to still be valid, according to the snippet below:

from cuda import cuda, cudart
import torch

def ASSERT_DRV(err):
    """
    This is a helper function to turn CUDA messages into errors when
    appropriate, since by default the CUDA package doesn't raise
    Python errors, it returns error messages
    """
    if isinstance(err, cuda.CUresult):
        if err != cuda.CUresult.CUDA_SUCCESS:
            raise RuntimeError("Cuda Error: {}".format(err))
    elif isinstance(err, cudart.cudaError_t):
        if err != cudart.cudaError_t.cudaSuccess:
            raise RuntimeError("Cudart Error: {}".format(err))
    else:
        raise RuntimeError("Unknown error type: {}".format(err))


def print_existing_contexts():
    
    valid_contexts = []
    
    while True:
    
        err, cuda_context = cuda.cuCtxPopCurrent()
        
        try:
            ASSERT_DRV(err)
        except RuntimeError:
            break
        else:
            valid_contexts.append(cuda_context)
            
    print("Existing, valid contexts: ", valid_contexts)
            
    for curr_ctx in reversed(valid_contexts):
    
        err, = cuda.cuCtxPushCurrent(curr_ctx)
        ASSERT_DRV(err)




if __name__ == "__main__":
    
    device_num = 2
    
    print("Initializing CUDA via Driver API")
    
    err, = cuda.cuInit(0) #always 0
    
    err, device = cuda.cuDeviceGet(device_num)

    err, cuda_context = cuda.cuCtxCreate(0, device)
    
    print_existing_contexts()
    
    print("Initializing Torch/CUDA")
    
    torch.cuda.init()
    
    print_existing_contexts()
    
    print("Attempting to use User-Created Context", cuda_context)
    
    #First, must push back on the stack
    err, = cuda.cuCtxPushCurrent(cuda_context)
    print(err)
    ASSERT_DRV(err)
    
    #Attempt to allocate 1kb on the GPU device
    err, ptr = cuda.cuMemAlloc(1024)
    print(err)
    ASSERT_DRV(err)
    
    print_existing_contexts()

I think this would still be good to have an update in the torch.cuda.init() docs, since it does still modify the CUDA state and not quite "does nothing" as the docstring suggests.