chore: import upstream snapshot with attribution

python/ray/experimental/channel/intra_process_channel.py

@@ -0,0 +1,72 @@
+import uuid
+from typing import Any, Optional
+
+from ray.experimental.channel import ChannelContext
+from ray.experimental.channel.common import ChannelInterface
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="alpha")
+class IntraProcessChannel(ChannelInterface):
+    """IntraProcessChannel is a channel for communication between two tasks in the same
+    worker process. It writes data directly to the worker's _SerializationContext
+    and reads data from the _SerializationContext to avoid the serialization
+    overhead and the need for reading/writing from shared memory. Note that if the
+    readers may mutate the data, users should deep copy the data themselves to avoid
+    side effects.
+
+    Args:
+        num_readers: The number of readers that will read from this channel. Readers
+            can be the same method of the same actor.
+        _channel_id: Optional pre-generated channel identifier. If ``None``, a
+            new UUID4 is generated. Used internally for re-creating channels
+            after pickling.
+    """
+
+    def __init__(
+        self,
+        num_readers: int,
+        _channel_id: Optional[str] = None,
+    ):
+        # Generate a unique ID for the channel. The writer and reader will use
+        # this ID to store and retrieve data from the _SerializationContext.
+        self._channel_id = _channel_id
+        self._num_readers = num_readers
+        if self._channel_id is None:
+            self._channel_id = str(uuid.uuid4())
+
+    def ensure_registered_as_writer(self) -> None:
+        pass
+
+    def ensure_registered_as_reader(self) -> None:
+        pass
+
+    def __reduce__(self):
+        return IntraProcessChannel, (
+            self._num_readers,
+            self._channel_id,
+        )
+
+    def __str__(self) -> str:
+        return f"IntraProcessChannel(channel_id={self._channel_id})"
+
+    def write(self, value: Any, timeout: Optional[float] = None):
+        self.ensure_registered_as_writer()
+        # No need to check timeout as the operation is non-blocking.
+
+        # Because both the reader and writer are in the same worker process,
+        # we can directly store the data in the context instead of storing
+        # it in the channel object. This removes the serialization overhead of `value`.
+        ctx = ChannelContext.get_current().serialization_context
+        ctx.set_data(self._channel_id, value, self._num_readers)
+
+    def read(self, timeout: Optional[float] = None, deserialize: bool = True) -> Any:
+        self.ensure_registered_as_reader()
+        assert deserialize, "Data passed from the actor to itself is never serialized"
+        # No need to check timeout as the operation is non-blocking.
+        ctx = ChannelContext.get_current().serialization_context
+        return ctx.get_data(self._channel_id)
+
+    def close(self) -> None:
+        ctx = ChannelContext.get_current().serialization_context
+        ctx.reset_data(self._channel_id)