[core][gpu-objects] GPU Objects POC

[kevin85421]

Why are these changes needed?

High-level

This PR implements most parts of the diagram above.

• Currently, this PR supports only Tensors, not TensorDicts.
• This PR supports data that mixes both CPU and GPU.
• This PR implements a very simple GC mechanism: all data in the in-actor object store is garbage-collected once it has been consumed.

Details

• Step 1: Users annotate the sender’s actor method with @ray.method(tensor_transport=...). The valid, case-insensitive values for tensor_transport are nccl, gloo, and object_store (default).

  @ray.remote
  class GPUTestActor:
      @ray.method(tensor_transport="GLOO")
      def echo(self, data):
          return data
      ...

• Step 2: Users create a communication group, such as an NCCL group, for actors that need to communicate with each other. In addition, each actor needs to register the custom serializer.

  • This step can be optimized for better UX in the future.

• Step 3: Pass the tensor_transport information through the stack—Python → Cython → C++ → Cython —when submitting a task to the sender actor.

  • Python (actor.py) → Cython → C++ → Cython
  • If tensor_transport is not OBJECT_STORE, serialize_and_store_gpu_objects will be called to extract tensors from the task output and store them in the GPUObjectManager.

• Step 4: When the driver process resolves the dependencies of the receiver actor’s task argument, if that argument is an ObjectRef pointing to an object created by an actor method annotated with @ray.method(tensor_transport="...") (NCCL or GLOO), it submits a __ray_send__ task to the sender actor to initiate the send operation (e.g., NCCL send) and a __ray_recv__ task to the receiver actor to initiate the receive operation (e.g., NCCL recv).

• Step 5: Pass the object ID through the stack—C++ (driver) → C++ (receiver actor) → Cython → Python (def deserialize).

  • In step 4, the data’s tensors are transferred via out-of-band NCCL or GLOO communication, then stored in the receiver actor’s in-actor object store using the object ID as the key.
  • In def deserialize, use the object ID to retrieve tensors from the in-actor object store, add them to the serialization context, and then deserialize to obtain the argument.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[stephanie-wang]

Looking good so far!

[kevin85421]

Hi @edoakes, I’ve updated the type from a string to a Protobuf enum. Would you have a chance to take another look? I’ll be adding some tests today.

[kevin85421]

I added a test to cover both non-trivial public functions in GPUObjectManager.

2884f4f

[kevin85421]

I decide not to rebase to fix the sign off issue because:

[crypdick]

comment for linking