figure out exactly what pickling a stream means

[cpcloud]

Yes, pickling a pointer is indeed weird, however, that is the status quo, as Stream uses the default pickling implementation and its fields are handle (stream handle), and external (bool).

Originally posted by @cpcloud in #609 (comment)

[kkraus14]

Given that we can't safely pickle a Stream, what is dependent on pickling a Stream currently and could we pull the Stream out of that in some way?

[cpcloud]

CUDADispatcher is picklable, which means _LaunchConfiguration must be picklable, which means all of its fields must be picklable and one of those fields is a cuda.core.experimental.Stream, which has a non-trivial __cinit__ and is thus not picklable (without the current underwhelming workaround).

[cpcloud]

One option is to not pickle the stream at all, and construct the default stream when unpickling.

[kkraus14]

Does CUDADispatcher need to be picklable? If it really does, then a Stream probably shouldn't be a part of it. Constructing the default stream when unpickling is really ugly behavior.

[cpcloud]

Does CUDADispatcher need to be picklable?

Not sure, but there's a test for this behavior.

If it really does, then a Stream probably shouldn't be a part of it. Constructing the default stream when unpickling is really ugly behavior.

To clarify, it's just that _LaunchConfiguration is a part of CUDADispatcher, and _LaunchConfiguration has a stream field.

[kkraus14]

Not sure, but there's a test for this behavior.

I think it's this specific test that's problematic? https://github.com/NVIDIA/numba-cuda/blob/main/numba_cuda/numba/cuda/tests/cudapy/test_serialize.py?plain=1#L36-L39

It looks like _LaunchConfiguration holds a reference to a CUDADispatcher as opposed to the opposite, and the _LaunchConfiguration is also what holds the launch dimensions and launch stream. The testing makes sense in that it's testing serializing the launch dimensions, which you could imagine wanting to serialize and reuse. I think the Stream being captured as part of the object that is being tested for serialization was an oversight and undesired behavior. Pretty sure if you pickle a LaunchConfiguration today that has a non-singleton Stream (i.e. not default stream 0x0, legacy default stream 0x1, or per-thread default stream 0x2) it ends up capturing an arbitrary pointer that will segfault on usage in another process.

Some ideas:

• In the serialization path for the CUDADispatcher object we could check if the Stream is a non-singleton Stream and then error loudly and clear if it is.
  • Given it would effectively be a guaranteed segfault if using a non-singleton Stream across processes, I can't imagine anyone is dependent on being able to pickle the Streams today.
• We could just ignore the Stream in the serialization path for the CUDADispatcher object and always use the default stream.
  • For someone to update the stream, they'd need to update the _LaunchConfiguration.stream attribute, which we'd probably want to give nicer setter / getter properties.
• We could extend _LaunchConfiguration to allow a None or some other sentinel value for its stream attribute which causes an error on launch. The stream attribute would only get this value in the pickling path and a user would be required to do something like _LaunchConfiguration.stream = ... before launching.
  • We could potentially handle singleton Streams nicely and only take this path for non-Singleton streams along with some kind of warning.

In general, the default stream (0x0) is runtime controlled via environment variable where it could change across processes and in theory cause unexpected race conditions. We could always resolve the default stream (0x0) to either the legacy default stream (0x1) or the per-thread default stream (0x2) in the serialization path to make things clearer, but I wouldn't expect users to introspect here regardless where they'd still likely potentially fall into a race.

[gmarkall]

A couple of thoughts / notes:

• IIRC CUDADispatcher needs to be picklable so it can be serialized for frameworks like Dask or Ray to send functions to remote workers
• I don't think we ought to provide a way to serialize a stream, for the reasons that have already been covered above.
• Is the _LaunchConfiguration getting cached because CUDADispatcher.configure() returns a _LaunchConfiguration and it is cached with lru_cache? If so, when pickling a CUDADispatcher I think we need to drop that cache and not serialize any launch configurations.