Add dask serialization of CUDA objects #3482
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This should make room for Dask serializers to also be specified and added.
To make TCP a bit more performant with RMM, provide Dask serializers to allow going to and from host memory.
|
Thanks @jakirkham ! For testing, what do you think about calling serialize on a few objects, like a cupy array, and verifying that some of the frames are |
|
Yep, agree this needs some tests. Those sound like reasonable suggestions. Will work on those and update when they are up. Thanks Matt! 😄 |
To make sure that different CUDA objects can use different serialization protocols, test with each one individual and ensure it completes. In particular test both "cuda" and "dask". Where supported also test "pickle", but skip it when it is not (like with Numba).
To make sure Dask can handle transmission of the frames serialized, test they match the type expected by the protocol used. With "cuda" ensure we get something that supports `__cuda_array_interface__`. With "dask" make sure we get a `memoryview`.
jakirkham
force-pushed
the
add_dask_cuda_serializers
branch
from
e8ec4c6 to
f73b614
Compare
|
Have parametrized tests around serializers so this tests Planning on merging in 24hrs if no comments. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This implements
"dask"serialization for CUDA objects. All serializers work by first performing"cuda"serialization on the different objects and then performing a device-to-host transfer. The deserializers perform a host-to-device transfer before running the respective"cuda"deserializer for the object considered. Should provide an alternative way to do TCP transfers with CUDA objects that does not rely on pickling.Note: Just like with UCX transfers, we preferentially add objects to the RMM pool if RMM is available. Otherwise we fallback to using Numba.