Discussion: Restore (or document how to restore) the old Device context manager behavior (pre-v10)?

Follow-up of #5705, #6867.

For downstream libraries depending on CuPy, having a context manager that respects any cudaSetDevice/cudaGetDevice calls inside or outside CuPy is useful. However, this is not possible since v10.

We are investigating this following snippet as a workaround:

@contextlib.contextmanager
def device_ctx(new_device_id):
    """
    old_device_id = cp.cuda.runtime.getDevice()
    try:
        if old_device_id != new_device_id:
            cp.cuda.runtime.setDevice(new_device_id)
        yield
    finally:
        # We should always restore the old device at exit.
        cp.cuda.runtime.setDevice(old_device_id)

For comparison this is what CuPy is currently doing:

cupy/cupy/cuda/device.pyx

Lines 180 to 191 in fc04c94

    
           def __enter__(self): 
        
               # N.B. for maintainers: do not use this context manager in CuPy 
        
               # codebase. See #5943 and #5963. 
        
               if self.id != runtime.getDevice(): 
        
                   runtime.setDevice(self.id) 
        
               _ThreadLocalStack.get().push_device(self.id) 
        
               return self 
        
           def __exit__(self, *args): 
        
               cdef int prev_device = _ThreadLocalStack.get().pop_device() 
        
               if prev_device != runtime.getDevice(): 
        
                   runtime.setDevice(prev_device)

Compare it with CuPy's ctx mgr, we solely rely on the facts that

CUDA device context has a per-thread state, so we don't need a thread-local stack.
The device_ctx is implemented such that it is stateless (which is critical when leaving the context). Each device_ctx instance keeps its own old_device_id thanks to the contextmanager decorator.

This snippet has passed CuPy's device tests. I am wondering what have we missed in CuPy earlier? I assume it is not possible for CuPy to adopt this approach given that CuPy's Device is stateful. If so, perhaps we can at least document it?

We discuss this and think this is reasonable to do.

@leofang

@leofang, thanks for raising this. I had discussion with core team members and concluded it's better to revert the v9 behavior in v12.

The motivation of the change in CuPy v10 was to align the behavior with that of cupy.cuda.Stream API (#5705), which was also changed in CuPy v10 to allow per-device & per-thread stream management (#3980). At that time I thought cupy.cuda.Device API change is fine, but it seems we overlooked that Device states are shared between libraries, in contrary to the concept of "default stream" which is clsoed to each library (CuPy, PyTorch, etc.). Supporting the interoperability between libraries is obviously important requirement for us, and v9 behavior is actually better in terms of that.

FYI, PyTorch's device context manager also uses the same strategy as in CuPy v9: https://github.com/pytorch/pytorch/blob/0ac2986d3334f8f9b35ca2fa7a30c20022c26fa6/torch/cuda/__init__.py#L265-L289

@kmaehashi

Thanks a lot @kmaehashi. Adding a note for the expected semantics:

The device context manager makes the specified device current from the point of entry until the point of exit.
When the context manager exits, the current device is reset to what it was when the context manager was entered.
Any explicit setting of the device within the context manager (using cupy.cuda.Device().use(), torch.cuda.set_device(), cudaSetDevice(), etc) will overrule the device set by the context manager from that point onwards till the context manager exits. In other words, the context manager provides a local device scope and the current device can be explicitly reset for the remainder of that scope. An analogy is changing the value of an argument that is passed by value to a function.
The device context managers can be arbitrarily nested.

I'll find some time to work on it 🙂

@kmaehashi

@kmaehashi could you assign this task to me so I don't drop the ball?

@kmaehashi

@kmaehashi @emcastillo I took some time to refresh my memory and did some investigation. My understanding is that the fix is very simple, but it must be a breaking change, mainly users should get their code updated, see below.

Based on my original post, the expected semantics requires that the previous ID to not be stored as a state of the object. However, this is not possible when creating a context manager with the __enter__ and __exit__ methods of a class, as we need a way to pass this info between the methods. This rules out the single-use context manager, as defined here.

The next viable, also breaking approach is to store the previous ID and to document it as a "reusable, but not reentrant" context manager (which PyTorch also follows). What it means is that storing the created Device objects and reusing them in the subsequent code

dev0 = cp.cuda.Device(0)
dev1 = cp.cuda.Device(1)

# do arbitrary work with the same objects dev0 & dev1
with dev0:
    with dev1: 
        ...

will likely be unsafe (breaking the expected semantics). Any user code must rewrite to use each instance exactly once, as if they are single-use context managers:

# always recreate the Device object
with cp.cuda.Device(0) as dev0:
    with cp.cuda.Device(1) as dev1: 
        ...

This is how PyTorch got around of this awkwardness IIUC. If this is OK I'll send a PR.

Sorry for the delay here Leo, it seems I got lost track of the discussion. Isn't it insufficient to revert #5853 to meet #6965 (comment)? I think the implementation in v9 is reentrant.

@kmaehashi

@kmaehashi I don't remember the detail, you might be right but I don't have time to evaluate it carefully this week 😅 Perhaps you can just send a revert PR and close #7315?

Navigation Menu

Repository navigation

Discussion: Restore (or document how to restore) the old `Device` context manager behavior (pre-v10)? #6965

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	def __enter__(self):
	# N.B. for maintainers: do not use this context manager in CuPy
	# codebase. See #5943 and #5963.
	if self.id != runtime.getDevice():
	runtime.setDevice(self.id)
	_ThreadLocalStack.get().push_device(self.id)
	return self

	def __exit__(self, *args):
	cdef int prev_device = _ThreadLocalStack.get().pop_device()
	if prev_device != runtime.getDevice():
	runtime.setDevice(prev_device)

Discussion: Restore (or document how to restore) the old Device context manager behavior (pre-v10)? #6965

Description

Activity

emcastillo commented on Aug 4, 2022

kmaehashi commented on Aug 29, 2022

leofang commented on Aug 30, 2022

leofang commented on Jan 15, 2023

leofang commented on Jan 16, 2023

kmaehashi commented on Feb 13, 2023

leofang commented on Feb 15, 2023

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Discussion: Restore (or document how to restore) the old `Device` context manager behavior (pre-v10)? #6965