IPC Mempool Serialization and multiprocessing Module Support

[Andy-Jost]

Add serialization to memory IPC. Memory resources can be passed directly to new processes. Expands testing. Removes IPCChannel.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Andy-Jost]

/ok to test 6be4686

[Andy-Jost]

/ok to test e0d0bf4

[Andy-Jost]

It looks like this change is good to go. I'm not aware of any remaining issues.

[leofang]

Awesome, thanks Andy! I'll re-review in an hour or two and then merge. I don't expect any change is needed, just wanna read the new tests once more.

[Andy-Jost]

@ksimpson-work I ended up creating unique identifiers to track memory pools across processes. Is this something we could add to the Driver API as a mempool attribute?

[leofang]

Sorry Andy, I re-read this again and still have a few questions 😓

btw we should fix a few doc issues

• add DeviceMemoryResourceOptions to cuda_core/docs/source/api.rst
• add IPCAllocationHandle to cuda_core/docs/source/api_private.rst
• address other comments below

The doc preview CI can be used to check what's missing.

[leofang]

Sorry I can't resist but circle back to ask this: We wanted Device to be serializable when only Buffer objects are passed over to child processes through multiprocessing APIs. We called .set_current() because the child function arguments are restored/deserialized prior to dropping into the function (xref: #1020 (comment)).

Questions:

1. What happens in a multi-GPU environment? If I send over buf0 allocated on GPU 0 and buf1 allocated on GPU 1 to the child, which device is current by the time the child function starts executing? What if we change the order of args from (buf0, buf1) to (buf1, buf0)? We don't have CI for multi-GPU tests for now, but the design should be multi-GPU-safe.
2. If we really need a device ID, can't we just look up from Buffer.mr.device_id, instead of serializing Device?

I still feel it is unwise to set the global state behind users' back.

[Andy-Jost]

The problem here is that I don't see a way to initialize CUDA (?) without calling set_current on a device. If I remove set_current from _reconstruct_device then I get CUDA_ERROR_INVALID_CONTEXT errors everywhere.

Device._has_inited is only set in set_current. To avoid messing with the global state, I'd rather say something like Device(n).init() to be sure that device n is ready to go. Then, later, when mapping a buffer, I could use mr.device_id to get that device. I just don't see any way to do that with the current _device code.

[Andy-Jost]

OK, I think I have something here that we can accept. I was able to remove the call to set_current in Device reconstruction since it turns out that memory resources can be mapped without setting a device context.

The only thing that really needs to be done implicitly when a process spawns is to call cuInit. The needs to be the very first thing or else all CUDA API calls (including mapping a memory resource or buffer) fails. That can be accomplished by just constructing a Device.

I tried removing device serialization completely but that actually makes the code quite a bit more cumbersome in several places. It's just easier to send an applicative of the form (Device, (device_id,)) rather than send a device ID, import Device, and then reconstruct manually on the other end. I also think it is conceptually sensible for Device objects to be serializable.

[leofang]

Sounds great, I'll take a look asap! A few quick notes before I forget:

1. CUDA mempools are by design not bound to any CUDA context. We (CUDA) used to have a context as the powerhouse for everything, memory, queue, module, ..., but now we are making things "context-independent," including mempools and kernel modules.
2. I read the multiprocessing docs more carefully, and I think we spawned the processes wrong (we should document this after this PR, in a Q&A or Tips & Tricks like in cuda-bindings). As you said calling cuInit() is the minimal:
   • For processes spawned via Process, we can inherit from the Process class, in the constructor __init__() we call Device().set_current() which would under the hood call cuInit(), and then call super().__init__(). My take is that this will ensure CUDA is initialized before the arguments are deserialized
   • For processes spawned via a Pool of workers, the same can be done by passing initializer and initargs.

Prove me wrong! 🙂

[Andy-Jost]

I think subclassing Process as you describe would work. Are you suggesting we ship a subclass of Process to simplify this for our users? The current code places set_current calls within the child main functions. That's actually pretty nice, since it mirrors what we ask users to do in any process. As you point out, the memory pool and buffer mappings don't require a context, which simplifies this a lot.

When using worker pools, initializer and initargs do work (we have examples of this in test_workpool.py). Still, I think there's a ton of value in supporting the ultra-simple use case of mapping a workerpool over buffers without doing anything special at start up. The latest code just puts the call to set_current within each worker function, which seems great, actually. That way, one could mix-and-match buffers from different devices and it ought to work.

One limitation I'm facing now is that I don't know whether the test runners have multipe GPUs, so I'm not sure how to move forward with developing the multi-device scenario.

[leofang]

Are you suggesting we ship a subclass of Process to simplify this for our users?

Certainly not, at least not now (or never). We should only do so in our own test suite, and teach users about this.

We did not have to teach this in most of Python GPU libraries (wearing my CuPy hat), because most libraries only interact with CUDA through runtime instead of driver, and runtime does the implicit initialization in every single cudart API. We don't and by design we always want users to call dev.set_current() first.

Now, with our own test suite, your other question is valid: Device().set_current() sets by default GPU 0 to current if there is no prior current device, this is a bit unfortunate if the child main wants to use say GPU 1. I think for now we need to call set_current() in both places, the process initializer (as a proxy to cuInit() because we definitely don't want to teach this) as well as child main (to actually select the right GPU, following the usual cuda.core requirement). We should do this for now and discuss if there is a better approach on Friday meeting. I believe cccl-rt will hit the same issue: NVIDIA/cccl#6073.

[Andy-Jost]

/ok to test

[copy-pr-bot]

/ok to test

@Andy-Jost, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

[Andy-Jost]

/ok to test e5b8542

[Andy-Jost]

/ok to test 64f154c

[leofang]

❗

[leofang]

/ok to test d584ea9

[leofang]

I suspect this comment is outdated.

[github-actions]

Doc Preview CI
Preview removed because the pull request was closed or merged.