intermediate-fix

Author: vmoens

benchmarks/storage/benchmark_sample_latency_over_rpc.py

@@ -144,7 +144,7 @@ def __init__(self, capacity: int):
     rank = args.rank
     storage_type = args.storage
 
-    torchrl_logger.info(f"Rank: {rank}; Storage: {storage_type}")
+    torchrl_logger.debug(f"RANK: {rank}; Storage: {storage_type}")
 
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "29500"

examples/distributed/replay_buffers/distributed_replay_buffer.py

@@ -172,7 +172,7 @@ def __init__(self, capacity: int):
 if __name__ == "__main__":
     args = parser.parse_args()
     rank = args.rank
-    torchrl_logger.info(f"Rank: {rank}")
+    torchrl_logger.debug(f"RANK: {rank}")
 
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "29500"

test/test_distributed.py

@@ -10,35 +10,21 @@
 
 import abc
 import argparse
+import importlib
 import os
+import socket
 import sys
 import time
 from functools import partial
 
 import pytest
-from tensordict import TensorDict
-from tensordict.nn import TensorDictModuleBase
-from torchrl._utils import logger as torchrl_logger
-from torchrl.data import (
-    LazyTensorStorage,
-    RandomSampler,
-    RayReplayBuffer,
-    RoundRobinWriter,
-    SamplerWithoutReplacement,
-)
-
-try:
-    import ray
-
-    _has_ray = True
-    RAY_ERR = None
-except ModuleNotFoundError as err:
-    _has_ray = False
-    RAY_ERR = err
 
 import torch
+from tensordict import TensorDict
+from tensordict.nn import TensorDictModuleBase
 
 from torch import multiprocessing as mp, nn
+from torchrl._utils import logger as torchrl_logger
 
 from torchrl.collectors import (
     MultiaSyncDataCollector,
@@ -52,8 +38,17 @@
     RPCDataCollector,
 )
 from torchrl.collectors.distributed.ray import DEFAULT_RAY_INIT_CONFIG
+from torchrl.data import (
+    LazyTensorStorage,
+    RandomSampler,
+    RayReplayBuffer,
+    RoundRobinWriter,
+    SamplerWithoutReplacement,
+)
 from torchrl.envs.utils import RandomPolicy
 
+_has_ray = importlib.util.find_spec("ray") is not None
+
 if os.getenv("PYTORCH_TEST_FBCODE"):
     from pytorch.rl.test.mocking_classes import ContinuousActionVecMockEnv, CountingEnv
 else:
@@ -115,7 +110,6 @@ def _test_distributed_collector_basic(cls, queue, frames_per_batch):
                 **cls.distributed_kwargs(),
             )
             total = 0
-            torchrl_logger.info("getting data...")
             for data in collector:
                 total += data.numel()
                 assert data.numel() == frames_per_batch
@@ -289,7 +283,9 @@ def _test_distributed_collector_updatepolicy(cls, queue, collector_class, sync):
                 n_collectors = 1
             else:
                 n_collectors = 2
-            collector = cls.distributed_class()(
+            dcls = cls.distributed_class()
+            torchrl_logger.info(f"Using distributed collector {dcls}")
+            collector = dcls(
                 [env] * n_collectors,
                 policy,
                 collector_class=collector_class,
@@ -307,6 +303,7 @@ def _test_distributed_collector_updatepolicy(cls, queue, collector_class, sync):
                 if i == 0:
                     first_batch = data
                     policy.weight.data += 1
+                    torchrl_logger.info("TEST -- Calling update_policy_weights_()")
                     collector.update_policy_weights_()
                 elif total == total_frames - frames_per_batch:
                     last_batch = data
@@ -338,7 +335,8 @@ def test_distributed_collector_updatepolicy(self, collector_class, sync):
         proc.start()
         try:
             out = queue.get(timeout=TIMEOUT)
-            assert out == "passed"
+            if out != "passed":
+                raise AssertionError(out)
         finally:
             proc.join(10)
             if proc.is_alive():
@@ -353,7 +351,13 @@ def distributed_class(cls) -> type:
 
     @classmethod
     def distributed_kwargs(cls) -> dict:
-        return {"launcher": "mp", "tcp_port": "4324"}
+        # Pick an ephemeral free TCP port on localhost for each test process to
+        # avoid address-in-use errors when tests are run repeatedly or in quick
+        # succession.
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("localhost", 0))
+            port = s.getsockname()[1]
+        return {"launcher": "mp", "tcp_port": str(port)}
 
     @classmethod
     def _start_worker(cls):
@@ -367,7 +371,10 @@ def distributed_class(cls) -> type:
 
     @classmethod
     def distributed_kwargs(cls) -> dict:
-        return {"launcher": "mp", "tcp_port": "4324"}
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("localhost", 0))
+            port = s.getsockname()[1]
+        return {"launcher": "mp", "tcp_port": str(port)}
 
     @classmethod
     def _start_worker(cls):
@@ -381,7 +388,10 @@ def distributed_class(cls) -> type:
 
     @classmethod
     def distributed_kwargs(cls) -> dict:
-        return {"launcher": "mp", "tcp_port": "4324"}
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("localhost", 0))
+            port = s.getsockname()[1]
+        return {"launcher": "mp", "tcp_port": str(port)}
 
     @classmethod
     def _start_worker(cls):
@@ -459,14 +469,17 @@ def test_distributed_collector_updatepolicy(self, collector_class, update_interv
             queue.close()
 
 
-@pytest.mark.skipif(not _has_ray, reason=f"Ray not found (error: {RAY_ERR})")
+@pytest.mark.skipif(
+    not _has_ray, reason="Ray not found. Ray may be badly configured or not installed."
+)
 class TestRayCollector(DistributedCollectorBase):
     """A testing distributed data collector class that runs tests without using a Queue,
     to avoid potential deadlocks when combining Ray and multiprocessing.
     """
 
     @pytest.fixture(autouse=True, scope="class")
     def start_ray(self):
+        import ray
         from torchrl.collectors.distributed.ray import DEFAULT_RAY_INIT_CONFIG
 
         ray.init(**DEFAULT_RAY_INIT_CONFIG)
@@ -480,6 +493,8 @@ def distributed_class(cls) -> type:
 
     @classmethod
     def distributed_kwargs(cls) -> dict:
+        import ray
+
         ray.shutdown()  # make sure ray is not running
         ray_init_config = DEFAULT_RAY_INIT_CONFIG
         ray_init_config["runtime_env"] = {

test/test_weightsync.py

@@ -638,7 +638,7 @@ def test_multiprocess_scheme_serialize_before_init(self):
         assert restored._sender is None
         assert restored._receiver is None
         assert not restored._initialized_on_sender
-        assert not restored._initialized_on_worker
+        assert not restored._initialized_on_receiver
 
     def test_multiprocess_scheme_serialize_after_sender_init(self):
         """Test that initialized sender can be pickled (excluding runtime state)."""
@@ -660,7 +660,7 @@ def test_multiprocess_scheme_serialize_after_sender_init(self):
         assert restored._sender is None  # Runtime state excluded
         assert restored._receiver is None
         assert not restored._initialized_on_sender  # Reset
-        assert not restored._initialized_on_worker
+        assert not restored._initialized_on_receiver
 
         # Clean up
         parent_pipe.close()

torchrl/_utils.py

@@ -52,7 +52,7 @@ def strtobool(val: Any) -> bool:
 
 LOGGING_LEVEL = os.environ.get("RL_LOGGING_LEVEL", "INFO")
 logger = logging.getLogger("torchrl")
-logger.setLevel(getattr(logging, LOGGING_LEVEL))
+logger.setLevel(LOGGING_LEVEL)
 logger.propagate = False
 # Clear existing handlers
 while logger.hasHandlers():
@@ -85,7 +85,9 @@ def format(self, record):
 console_handler = logging.StreamHandler(stream=stream_handler)
 console_handler.setFormatter(_CustomFormatter())
 logger.addHandler(console_handler)
-console_handler.setLevel(logging.INFO)
+
+console_handler.setLevel(LOGGING_LEVEL)
+logger.debug(f"Logging level: {logger.getEffectiveLevel()}")
 
 VERBOSE = strtobool(os.environ.get("VERBOSE", str(logger.isEnabledFor(logging.DEBUG))))
 _os_is_windows = sys.platform == "win32"

torchrl/collectors/__init__.py

@@ -5,12 +5,13 @@
 
 from torchrl.envs.utils import RandomPolicy
 
+from ._base import DataCollectorBase
+
 from ._multi_async import MultiaSyncDataCollector
 from ._multi_sync import MultiSyncDataCollector
 from ._single import SyncDataCollector
 
 from ._single_async import aSyncDataCollector
-from .base import DataCollectorBase
 from .weight_update import (
     MultiProcessedWeightUpdater,
     RayWeightUpdater,

torchrl/collectors/base.py renamed to torchrl/collectors/_base.py

@@ -16,10 +16,11 @@
 from tensordict.nn import TensorDictModule, TensorDictModuleBase
 from torch import nn as nn
 from torch.utils.data import IterableDataset
+from torchrl._utils import logger as torchrl_logger
 from torchrl.collectors.utils import _map_weight
 
 from torchrl.collectors.weight_update import WeightUpdaterBase
-from torchrl.weight_update import WeightReceiver, WeightSender, WeightSyncScheme
+from torchrl.weight_update import WeightSyncScheme
 
 
 class DataCollectorBase(IterableDataset, metaclass=abc.ABCMeta):
@@ -35,8 +36,6 @@ class DataCollectorBase(IterableDataset, metaclass=abc.ABCMeta):
     cudagraphed_policy: bool
     _weight_updater: WeightUpdaterBase | None = None
     _weight_sync_schemes: dict[str, WeightSyncScheme] | None = None
-    _weight_senders: dict[str, WeightSender] | None = None
-    _weight_receivers: dict[str, WeightReceiver] | None = None
     verbose: bool = False
 
     @property
@@ -320,40 +319,81 @@ def _weight_update_impl(
         if policy_or_weights is not None:
             weights_dict = {"policy": policy_or_weights}
 
-        # Priority: new weight sync schemes > old weight updater system
-        if self._weight_senders:
-            if model_id is not None:
+        if self._weight_sync_schemes:
+            if model_id is None:
+                model_id = "policy"
+            if weights_dict is None:
                 # Compose weight_dict
                 weights_dict = {model_id: policy_or_weights}
-            if weights_dict is None:
-                if "policy" in self._weight_senders:
-                    weights_dict = {"policy": policy_or_weights}
-                elif len(self._weight_senders) == 1:
-                    single_model_id = next(iter(self._weight_senders.keys()))
-                    weights_dict = {single_model_id: policy_or_weights}
-                else:
-                    raise ValueError(
-                        "Cannot determine the model to update. Please provide a weights_dict."
-                    )
             for target_model_id, weights in weights_dict.items():
-                if target_model_id not in self._weight_senders:
+                if target_model_id not in self._weight_sync_schemes:
                     raise KeyError(
-                        f"Model '{target_model_id}' not found in registered weight senders. "
-                        f"Available models: {list(self._weight_senders.keys())}"
+                        f"Model '{target_model_id}' not found in registered weight sync schemes. "
+                        f"Available models: {list(self._weight_sync_schemes.keys())}"
                     )
                 processed_weights = self._extract_weights_if_needed(
                     weights, target_model_id
                 )
                 # Use new send() API with worker_ids support
-                self._weight_senders[target_model_id].send(
-                    weights=processed_weights, worker_ids=worker_ids
+                torchrl_logger.debug("weight update -- getting scheme")
+                scheme = self._weight_sync_schemes.get(target_model_id)
+                if not isinstance(scheme, WeightSyncScheme):
+                    raise TypeError(f"Expected WeightSyncScheme, got {target_model_id}")
+                torchrl_logger.debug(
+                    f"calling send() on scheme {type(scheme).__name__}"
                 )
+                scheme.send(weights=processed_weights, worker_ids=worker_ids)
         elif self._weight_updater is not None:
             # unreachable
             raise RuntimeError
         else:
             return self.receive_weights(policy_or_weights)
 
+    def _receive_weights_scheme(self):
+        """Receive weights via registered receiver schemes and cascade to nested collectors.
+
+        This method enables cascading weight updates across multiple collector layers:
+        - RPCDataCollector -> MultiSyncDataCollector -> SyncDataCollector
+        - DistributedDataCollector -> MultiSyncDataCollector -> SyncDataCollector
+
+        Process:
+        1. Receive weights for all registered receiver schemes (_receiver_schemes)
+        2. If this collector has nested collectors (_weight_sync_schemes), propagate
+           the updates by calling update_policy_weights_()
+
+        """
+        # Receive weights for all registered schemes
+        updates = {}
+        if not hasattr(self, "_receiver_schemes"):
+            raise RuntimeError("No receiver schemes registered.")
+
+        for model_id, scheme in self._receiver_schemes.items():
+            # scheme.receive() pulls weights from the transport and applies them locally
+            # For RPC/Ray: weights are already passed as argument, receive() is a no-op
+            # For Distributed: receive() pulls from TCPStore
+            # For MultiProcess: receive() checks the pipe
+            received_weights = scheme.receive()
+            if received_weights is not None:
+                updates[model_id] = received_weights
+
+        # If we have nested collectors (e.g., MultiSyncDataCollector with inner workers)
+        # AND we actually received updates, propagate them down via their senders
+        if (
+            updates
+            and hasattr(self, "_weight_sync_schemes")
+            and self._weight_sync_schemes
+        ):
+            # Build weights_dict for all models that need propagation to nested collectors
+            weights_dict = {}
+            for model_id in updates:
+                if model_id in self._weight_sync_schemes:
+                    # This model has a sender scheme - propagate to nested workers
+                    weights_dict[model_id] = updates[model_id]
+
+            if weights_dict:
+                # Propagate to nested collectors via their sender schemes
+                self.update_policy_weights_(weights_dict=weights_dict)
+
     def receive_weights(self, policy_or_weights: TensorDictBase | None = None):
         # No weight updater configured
         # For single-process collectors, apply weights locally if explicitly provided
@@ -389,6 +429,42 @@ def receive_weights(self, policy_or_weights: TensorDictBase | None = None):
             strategy.apply_weights(self.policy, weights)
         # Otherwise, no action needed - policy is local and changes are immediately visible
 
+    def _set_scheme_receiver(self, weight_sync_schemes: dict[str, WeightSyncScheme]):
+        """Set up receiver schemes for this collector.
+
+        This method initializes receiver schemes and stores them in _receiver_schemes
+        for later use by _receive_weights_scheme() and receive_weights().
+
+        Args:
+            weight_sync_schemes: Dictionary of {model_id: WeightSyncScheme} to set up as receivers
+        """
+        # Initialize _receiver_schemes if not already present
+        if not hasattr(self, "_receiver_schemes"):
+            self._receiver_schemes = {}
+
+        # Initialize each scheme on the receiver side
+        for model_id, scheme in weight_sync_schemes.items():
+            if not scheme.initialized_on_receiver:
+                if scheme.initialized_on_sender:
+                    raise RuntimeError(
+                        "Weight sync scheme cannot be initialized on both sender and receiver."
+                    )
+                scheme.init_on_receiver(
+                    model_id=model_id,
+                    context=self,
+                    worker_idx=getattr(self, "_worker_idx", None),
+                )
+
+            # Store the scheme for later use in receive_weights()
+            self._receiver_schemes[model_id] = scheme
+
+        # Perform initial synchronization
+        for scheme in weight_sync_schemes.values():
+            if not scheme.synchronized_on_receiver:
+                scheme.synchronize_weights(
+                    worker_idx=getattr(self, "_worker_idx", None)
+                )
+
     def __iter__(self) -> Iterator[TensorDictBase]:
         try:
             yield from self.iterator()