torchrl 0.11.0__cp314-cp314-win_amd64.whl

torchrl/weight_update/_rpc.py

@@ -0,0 +1,284 @@
+from __future__ import annotations
+
+import time
+import weakref
+from typing import Any
+
+from torchrl._utils import logger as torchrl_logger
+
+from torchrl.weight_update.utils import _resolve_model
+from torchrl.weight_update.weight_sync_schemes import (
+    TransportBackend,
+    WeightStrategy,
+    WeightSyncScheme,
+)
+
+
+class RPCWeightSyncScheme(WeightSyncScheme):
+    """Weight synchronization for torch.distributed.rpc.
+
+    This scheme uses RPC calls to synchronize weights across distributed
+    workers. Each remote collector gets its own transport, following the
+    same pattern as multiprocess collectors.
+    """
+
+    def _init_on_sender_impl(
+        self,
+        *,
+        model_id: str,
+        context: Any = None,
+        num_workers: int,
+    ) -> None:
+        # Store model_id and context on scheme
+        self.model_id = model_id
+        if context is not None:
+            self.context = context
+        else:
+            raise RuntimeError(f"Expected a context for {type(self).__name__}.")
+        collector_infos = getattr(self.context, "collector_infos", None)
+        collector_rrefs = getattr(self.context, "collector_rrefs", None)
+        collector_class = getattr(self.context, "collector_class", None)
+        if (
+            collector_infos is None
+            or collector_rrefs is None
+            or collector_class is None
+        ):
+            raise RuntimeError(
+                "RPCWeightSyncScheme requires a context with the following attributes: "
+                "(context.collector_infos, context.collector_rrefs, context.collector_class)"
+            )
+
+        # Create transports for each remote collector
+        # worker_rank is i+1 because rank 0 is the main/trainer process
+        for i in range(num_workers):
+            worker_rank = i + 1
+            transport = self.create_transport(
+                collector_info=collector_infos[i],
+                collector_rref=collector_rrefs[i],
+                collector_class=collector_class,
+                worker_rank=worker_rank,
+            )
+            self._register_worker_sender(worker_idx=i, transport=transport)
+
+    def _init_on_receiver_impl(
+        self, *, model_id: str, context: Any = None, worker_idx: int | None = None
+    ) -> None:
+        """Initialize scheme on the worker (receiver) side.
+
+        Expected kwargs (as provided by collectors):
+            - model_id: str              # e.g. "policy"
+            - context: Any               # collector / inner collector
+            - worker_idx: int | None     # worker index (optional)
+        """
+        if context is None:
+            raise ValueError(
+                "RPCWeightSyncScheme.init_on_receiver requires a 'context' "
+                "providing access to the model to be synchronized."
+            )
+
+        # Store model_id and context on scheme
+        self.model_id = model_id
+        self.worker_idx = worker_idx
+        self.context = context
+        # Access weights to set up missing elements
+        self.weights  # noqa
+
+        self._receiver_transport = RPCTransport(worker_rank=worker_idx)
+
+    @property
+    def model(self) -> Any | None:
+        if self._model_ref is not None:
+            return self._model_ref()
+        if self._model_id is not None:
+            model = _resolve_model(self.context, self._model_id)
+            if model is None:
+                if self._model_id == "policy":
+                    torchrl_logger.debug("Creating policy from factory.")
+                    model = self.context.policy_factory[0]()
+                    self.context.policy = model
+                else:
+                    raise AttributeError(
+                        f"Model {self._model_id} was `None` in context {self.context}"
+                    )
+            self._model_ref = weakref.ref(model)
+            return model
+
+    @model.setter
+    def model(self, value: Any):
+        if value is None:
+            return
+        self._model_ref = weakref.ref(value)
+
+    def create_transport(
+        self,
+        *,
+        collector_info=None,
+        collector_rref=None,
+        collector_class=None,
+        worker_rank=None,
+        **kwargs,
+    ) -> TransportBackend:
+        """Create RPC-based transport for a specific remote collector.
+
+        Args:
+            collector_info: RPC worker info for the remote collector.
+            collector_rref: RPC remote reference to the collector.
+            collector_class: Class of the remote collector.
+            worker_rank: The torch.distributed rank of the remote worker.
+            **kwargs: Additional transport configuration.
+
+        Returns:
+            RPCTransport configured for this specific remote collector.
+        """
+        return RPCTransport(
+            collector_info=collector_info,
+            collector_rref=collector_rref,
+            collector_class=collector_class,
+            worker_rank=worker_rank,
+        )
+
+
+class RPCTransport:
+    """RPC transport for communicating with a single RPC remote collector.
+
+    This transport handles weight updates for ONE specific remote collector via
+    torch.distributed primitives (send/recv) with RPC used for signaling.
+    Multiple transports are created for multiple collectors, following the same
+    pattern as the DistributedDataCollector.
+    """
+
+    def __init__(
+        self,
+        collector_info=None,
+        collector_rref=None,
+        collector_class=None,
+        worker_rank=None,
+    ):
+        self._collector_info = collector_info
+        self._collector_rref = collector_rref
+        self._collector_class = collector_class
+        self._worker_rank = worker_rank  # The torch.distributed rank of this worker
+        self._pending_future = None
+        self._pending_send = None
+
+    def send_weights(self, weights: Any) -> None:
+        """Send weights to the remote collector using torch.distributed.
+
+        Uses torch.distributed.send() for the actual weight transfer and RPC
+        for signaling the remote collector to receive.
+
+        Order is critical to avoid deadlock:
+        1. Signal receiver via RPC to start recv() (non-blocking)
+        2. Send weights via torch.distributed (blocking until recv completes)
+        """
+        if self._collector_info is None or self._collector_rref is None:
+            return
+        if self._worker_rank is None:
+            raise RuntimeError("worker_rank must be set for RPC transport")
+
+        # Step 1: Signal the remote collector via RPC to start receiving (async)
+        # Use rref.rpc_async() to properly call the instance method on the remote object
+        future = self._collector_rref.rpc_async()._receive_weights_scheme()
+
+        # Step 2: Send weights via torch.distributed (blocks until receiver calls recv())
+        weights.send(self._worker_rank)
+
+        # Step 3: Wait for RPC to complete (receiver has applied weights)
+        future.wait()
+
+    def send_weights_async(self, weights: Any) -> None:
+        """Send weights to remote collector asynchronously.
+
+        Uses torch.distributed.isend() for the actual weight transfer and RPC
+        for signaling. Use wait_ack() to wait for completion.
+
+        Order is critical to avoid deadlock:
+        1. Signal receiver via RPC to start recv() (non-blocking)
+        2. Send weights via torch.distributed.isend() (non-blocking)
+        3. wait_ack() waits for both to complete
+        """
+        if self._collector_info is None or self._collector_rref is None:
+            return
+        if self._worker_rank is None:
+            raise RuntimeError("worker_rank must be set for RPC transport")
+
+        # Step 1: Signal the remote collector via RPC to start receiving (async)
+        # Use rref.rpc_async() to properly call the instance method on the remote object
+        self._pending_future = (
+            self._collector_rref.rpc_async()._receive_weights_scheme()
+        )
+
+        # Step 2: Send weights asynchronously via torch.distributed
+        # Store the Work handle for wait_ack()
+        weights.isend(self._worker_rank)
+
+    def wait_ack(self) -> None:
+        """Wait for both the RPC call and the distributed send to complete."""
+        # Wait for the RPC call to complete
+        if hasattr(self, "_pending_future") and self._pending_future is not None:
+            self._pending_future.wait()
+            del self._pending_future
+
+    def receive_weights(
+        self,
+        timeout: float | None = None,
+        *,
+        weights: Any = None,
+        model: Any = None,
+        strategy: WeightStrategy | None = None,
+    ) -> Any | None:
+        """Receive weights from sender using torch.distributed.
+
+        Args:
+            timeout: Maximum time to wait for weights (seconds). If None,
+                blocks until weights are received.
+            weights: Pre-allocated weight buffer to receive into.
+            model: The model to apply weights to.
+            strategy: Strategy for applying weights to the model.
+
+        Returns:
+            The received weights, or None if timeout expires.
+        """
+        if weights is None:
+            return None
+
+        if timeout is None:
+            # Blocking receive
+            weights.recv(0)
+        else:
+            # Non-blocking receive with timeout support
+            futures = weights.irecv(src=0, return_premature=True)
+            if futures:
+                start_time = time.monotonic()
+                while True:
+                    # Check if all futures are complete
+                    all_complete = all(f.is_completed() for f in futures)
+                    if all_complete:
+                        break
+                    # Check timeout
+                    elapsed = time.monotonic() - start_time
+                    if elapsed >= timeout:
+                        # Timeout expired before receiving all weights
+                        return None
+                    # Small sleep to avoid busy-waiting
+                    time.sleep(0.001)
+
+        # Apply the received weights to the model
+        if model is not None and strategy is not None:
+            strategy.apply_weights(model, weights)
+
+        return weights
+
+    def setup_connection_and_weights_on_sender(self) -> None:
+        """No-op for RPCTransport - weights are sent via send_weights()."""
+
+    def setup_connection_and_weights_on_receiver(
+        self,
+        *,
+        worker_idx: int,
+        weights: Any = None,
+        model: Any = None,
+        strategy: WeightStrategy | None = None,
+    ) -> Any:
+        """No-op for RPCTransport - weights are received via receive()."""
+        return None