Why are these changes needed?

Issue statement

https://gist.github.com/kevin85421/a7f14ea38d64420b105fbd79fd31fb8a

Without this PR, both SynchronousReader._read_list and AwaitableBackgroundReader._run call the read function of each input channel sequentially. If the first input channel involves a long-running task and the second one fails immediately, the reader still has to wait until all channels have been read before failing.

Based on #46337,

This is a problem for tensor-parallel inference, because all of the workers execute in lockstep, and if one actor throws an exception, the others may hang. Depending on the order of the actors, ray.get() may never return.

Implementation details

• experimental_wait_and_get_mutable_objects / CoreWorker::WaitAndGetExperimentalMutableObjects

  • WaitAndGetExperimentalMutableObjects iterates through a list of mutable object references and retrieves the data when the objects are ready. The function returns when either num_objects mutable objects are acquired or the operation times out.

• SynchronousReader._read_list / AwaitableBackgroundReader._run

  • _get_all_waitables_to_num_consumers: Iterate through self._input_channels and call get_ray_waitables to retrieve all underlying mutable object references.
  • worker.experimental_wait_and_get_mutable_objects: Attempt to retrieve a single mutable object from a list of mutable object references. If the return value is a RayTaskError, immediately return and raise an exception. Otherwise, write the return value into ChannelContext.
  • After all mutable objects have been retrieved, iterate through self._input_channels and call channel.read().

• Channel.read (shared_memory_channel.py)

  • Because the mutable objects are retrieved in _read_list, Channel.read retrieves the data from ChannelContext instead of object store.

• get_ray_waitables

  • Retrieve the underlying mutable object references that the "next" read operation plans to access. Therefore, get_ray_waitables may return different results for different read function calls on the same channel. For example, BufferedSharedMemoryChannel's get_ray_waitables:

    def get_ray_waitables(self) -> List[Tuple[ObjectRef, bool]]:
        self.ensure_registered_as_reader()
        return self._buffers[self._next_read_index].get_ray_waitables()

• Special case: TorchTensorNcclChannel: See the comments in the file for more details.

  • Call CPU write before NCCL write. Since the channel's read is only called after all required mutable objects have been retrieved, NCCL read will only be invoked after the reader has already retrieved the mutable object. Therefore, CPU write must occur before NCCL write to avoid a deadlock.
  • _read_list and _run will skip deserialization for the TorchTensorNcclChannel's mutable object because TorchTensorNcclChannel relies on a custom serializer, which replaces placeholders in the CPU data with tensors read from the NCCL channel during deserialization.
    • If we deserialize the mutable object in _read_list or _run before the reader has retrieved the GPU tensors via the NCCL channel and placed the out-of-band tensors into the serialization context, issues may arise.
    • Instead, the reader will deserialize the CPU data after the out-of-band tensors are ready in the channel's read operation.

Related issue number

Closes #46337

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 :(