final!

Author: vmoens

docs/source/reference/collectors_weightsync.rst

@@ -49,7 +49,7 @@ Weight update schemes can be used outside of collectors for custom synchronizati
 The new simplified API provides four core methods for weight synchronization:
 
 - ``init_on_sender(model_id, **kwargs)`` - Initialize on the main process (trainer) side
-- ``init_on_worker(model_id, **kwargs)`` - Initialize on worker process side
+- ``init_on_receiver(model_id, **kwargs)`` - Initialize on worker process side
 - ``get_sender()`` - Get the configured sender instance
 - ``get_receiver()`` - Get the configured receiver instance
 
@@ -85,7 +85,7 @@ Here's a basic example:
     # or sender.send_async(weights); sender.wait_async()  # Asynchronous send
 
     # On the worker process side:
-    # scheme.init_on_worker(model_id="policy", pipe=child_pipe, model=policy)
+    # scheme.init_on_receiver(model_id="policy", pipe=child_pipe, model=policy)
     # receiver = scheme.get_receiver()
     # # Non-blocking check for new weights
     # if receiver.receive(timeout=0.001):

test/test_weightsync.py

@@ -6,7 +6,9 @@
 
 import argparse
 import importlib.util
+
 import pickle
+import threading
 import time
 
 import pytest
@@ -26,12 +28,10 @@
     RayWeightSyncScheme,
     RPCWeightSyncScheme,
     SharedMemTransport,
-)
-from torchrl.weight_update.utils import _resolve_model
-from torchrl.weight_update.weight_sync_schemes import (
     SharedMemWeightSyncScheme,
     WeightStrategy,
 )
+from torchrl.weight_update.utils import _resolve_model
 
 _has_ray = importlib.util.find_spec("ray") is not None
 
@@ -43,7 +43,7 @@ def worker_update_policy(pipe, timeout=5.0):
         policy.bias.fill_(0.0)
 
     scheme = MultiProcessWeightSyncScheme(strategy="state_dict")
-    scheme.init_on_worker(model_id="policy", pipe=pipe, model=policy)
+    scheme.init_on_receiver(model_id="policy", pipe=pipe, model=policy)
     receiver = scheme.get_receiver()
 
     if receiver._transport.pipe.poll(timeout):
@@ -62,7 +62,7 @@ def worker_update_policy_tensordict(pipe, timeout=5.0):
         policy.bias.fill_(0.0)
 
     scheme = MultiProcessWeightSyncScheme(strategy="tensordict")
-    scheme.init_on_worker(model_id="policy", pipe=pipe, model=policy)
+    scheme.init_on_receiver(model_id="policy", pipe=pipe, model=policy)
     receiver = scheme.get_receiver()
 
     if receiver._transport.pipe.poll(timeout):
@@ -100,7 +100,7 @@ def test_mp_transport_basic(self):
         proc.start()
 
         test_weights = {"weight": torch.ones(2, 4), "bias": torch.ones(2)}
-        transport.send_weights("policy", test_weights)
+        transport.send_weights(test_weights)
 
         proc.join(timeout=10.0)
         assert not proc.is_alive()
@@ -113,7 +113,7 @@ def test_mp_transport_async(self):
         proc.start()
 
         test_weights = {"weight": torch.ones(2, 4), "bias": torch.ones(2)}
-        transport.send_weights_async("policy", test_weights)
+        transport.send_weights_async(test_weights)
         transport.wait_ack()
 
         proc.join(timeout=10.0)
@@ -124,13 +124,16 @@ def test_shared_mem_transport(self):
             {"weight": torch.zeros(2, 4), "bias": torch.zeros(2)}, batch_size=[]
         ).share_memory_()
 
-        transport = SharedMemTransport({"policy": shared_buffer})
+        transport = SharedMemTransport()
+        transport.register_weights(
+            params_map={0: shared_buffer}, init_queues={0: mp.Queue()}
+        )
 
         new_weights = TensorDict(
             {"weight": torch.ones(2, 4), "bias": torch.ones(2)}, batch_size=[]
         )
 
-        transport.send_weights("policy", new_weights)
+        transport.send_weights(new_weights)
 
         assert torch.allclose(shared_buffer["weight"], torch.ones(2, 4))
         assert torch.allclose(shared_buffer["bias"], torch.ones(2))
@@ -255,7 +258,10 @@ def test_shared_mem_scheme(self):
             {"weight": torch.ones(2, 4), "bias": torch.ones(2)}, batch_size=[]
         )
 
-        transport.send_weights("policy", new_weights)
+        transport.register_weights(
+            params_map={0: shared_buffer}, init_queues={0: mp.Queue()}
+        )
+        transport.send_weights(new_weights)
 
         assert torch.allclose(shared_buffer["weight"], torch.ones(2, 4))
         assert torch.allclose(shared_buffer["bias"], torch.ones(2))
@@ -265,7 +271,7 @@ def test_no_weight_sync_scheme(self):
         transport = scheme.create_transport(None)
 
         weights = {"weight": torch.ones(2, 4), "bias": torch.ones(2)}
-        transport.send_weights("policy", weights)
+        transport.send_weights(weights)
 
     @classmethod
     def _worker_with_receive(cls, pipe, scheme):
@@ -274,7 +280,7 @@ def _worker_with_receive(cls, pipe, scheme):
             policy.weight.fill_(0.0)
             policy.bias.fill_(0.0)
 
-        scheme.init_on_worker(model_id="policy", pipe=pipe, model=policy)
+        scheme.init_on_receiver(model_id="policy", pipe=pipe, model=policy)
         receiver = scheme.get_receiver()
 
         # Non-blocking receive should return False when no data
@@ -354,7 +360,7 @@ def test_syncdatacollector_multiprocess_scheme(self, simple_policy):
         collector.shutdown()
 
     def test_multisyncdatacollector_multiprocess_scheme(self, simple_policy):
-        scheme = MultiProcessWeightSyncScheme(strategy="state_dict")
+        scheme = MultiProcessWeightSyncScheme()
 
         collector = MultiSyncDataCollector(
             create_env_fn=[
@@ -660,73 +666,76 @@ def test_multiprocess_scheme_serialize_after_sender_init(self):
         parent_pipe.close()
         child_pipe.close()
 
-    def test_shared_mem_scheme_serialize_before_init(self):
-        """Test that uninitialized SharedMemWeightSyncScheme can be pickled."""
-        scheme = SharedMemWeightSyncScheme(strategy="tensordict")
-
-        # Serialize and deserialize
-        pickled = pickle.dumps(scheme)
-        restored = pickle.loads(pickled)
-
-        # Check that configuration is preserved
-        assert restored.strategy == "tensordict"
-        assert restored._sender is None
-        assert restored._receiver is None
+    # Serialize and deserialize
+    @staticmethod
+    def _get_scheme_from_queue(q, scheme):
+        try:
+            restored = scheme
+            # Check that configuration is preserved but runtime state is cleared
+            assert restored.strategy == "tensordict"
+            assert restored._sender is None
+            assert not restored._initialized_on_sender
+
+            q.put("success")
+        except Exception as err:
+            q.put(f"failure: {err}")
+        finally:
+            q.close()
 
+    @pytest.mark.timeout(10)
     def test_shared_mem_scheme_serialize_after_init(self):
         """Test that initialized SharedMemWeightSyncScheme can be pickled."""
         parent_pipe, child_pipe = mp.Pipe()
+        q = mp.Queue()
+        try:
+            # Create shared buffer
+            shared_buffer = TensorDict(
+                {"weight": torch.zeros(2, 4), "bias": torch.zeros(2)}, batch_size=[]
+            ).share_memory_()
+
+            scheme = SharedMemWeightSyncScheme()
+
+            def init_on_sender(scheme, pipe):
+                scheme.init_on_sender(params_map={0: shared_buffer})
+                scheme.synchronize_weights()
+                msg = pipe.recv()
+                assert msg == "registered"
+
+            def init_on_receiver(scheme: SharedMemWeightSyncScheme, child_pipe):
+                scheme.init_on_receiver(
+                    worker_idx=0, model=nn.Linear(4, 2, device="meta")
+                )
+                scheme.synchronize_weights()
+                child_pipe.send("registered")
+
+            future_sender = threading.Thread(
+                target=init_on_sender,
+                kwargs={"scheme": scheme, "pipe": parent_pipe},
+            )
+            future_receiver = threading.Thread(
+                target=init_on_receiver,
+                kwargs={"scheme": scheme, "child_pipe": child_pipe},
+            )
+            future_receiver.start()
+            future_sender.start()
+            future_receiver.join(timeout=10.0)
+            future_sender.join(timeout=10.0)
 
-        # Create shared buffer
-        shared_buffer = TensorDict(
-            {"weight": torch.zeros(2, 4), "bias": torch.zeros(2)}, batch_size=[]
-        ).share_memory_()
-
-        scheme = SharedMemWeightSyncScheme(
-            strategy="tensordict",
-        )
-
-        def init_on_sender(scheme, child_pipe):
-            (model_id, data), msg = child_pipe.recv()
-            if msg == "register_shared_weights":
-                child_pipe.send((None, "registered"))
-            else:
-                raise ValueError(f"Expected 'register_shared_weights' but got {msg}")
-
-        # Initialize the scheme with the pipes, in 2 separate threads because init requires acknowledgement from the worker
-        import threading
-
-        future_sender = threading.Thread(
-            target=scheme.init_on_sender,
-            kwargs={"model_id": "policy", "pipes": [parent_pipe]},
-        )
-        future_receiver = threading.Thread(
-            target=init_on_sender,
-            kwargs={"scheme": scheme, "child_pipe": child_pipe},
-        )
-        future_receiver.start()
-        future_sender.start()
-        future_receiver.join()
-        future_sender.join()
-
-        # Scheme now has _sender with non-serializable state
-        assert scheme._sender is not None
-
-        # Serialize and deserialize
-        pickled = pickle.dumps(scheme)
-        restored = pickle.loads(pickled)
-
-        # Check that configuration is preserved but runtime state is cleared
-        assert restored.strategy == "tensordict"
-        assert restored._sender is None
-        assert not restored._initialized_on_sender
-
-        # Note: policy_weights dict is preserved (but may need re-sharing)
-        assert "policy" in restored.policy_weights
+            # Scheme now has _sender with non-serializable state
+            assert scheme._sender is not None
 
-        # Clean up
-        parent_pipe.close()
-        child_pipe.close()
+            proc = mp.Process(target=self._get_scheme_from_queue, args=(q, scheme))
+            proc.start()
+            try:
+                msg = q.get(timeout=10.0)
+                assert msg == "success", msg
+            finally:
+                proc.join()
+        finally:
+            q.close()
+            # Clean up
+            parent_pipe.close()
+            child_pipe.close()
 
     def test_no_weight_sync_scheme_serialize(self):
         """Test that NoWeightSyncScheme can be pickled."""
@@ -809,7 +818,7 @@ def test_scheme_reinitialization_after_unpickle(self):
         """Test that a scheme can be re-initialized after unpickling.
 
         This is the expected workflow: pickle a scheme, unpickle it in a worker,
-        then call init_on_worker() to establish new runtime resources.
+        then call init_on_receiver() to establish new runtime resources.
         """
         # Initialize and pickle a scheme
         parent_pipe, child_pipe = mp.Pipe()

torchrl/collectors/_runner.py

@@ -39,7 +39,7 @@ def _make_policy_factory(
 
     if weight_sync_scheme is not None:
         # Initialize the receiver on the worker side
-        weight_sync_scheme.init_on_worker(
+        weight_sync_scheme.init_on_receiver(
             model=policy, model_id="policy", worker_idx=worker_idx, pipe=pipe
         )
         # Get the receiver and synchronize initial weights
@@ -147,7 +147,7 @@ def _main_async_collector(
                     inner_collector._weight_receivers[model_id] = receiver
                 else:
                     # Initialize receivers for other models
-                    scheme.init_on_worker(model_id=model_id, context=inner_collector)
+                    scheme.init_on_receiver(model_id=model_id, context=inner_collector)
                     receiver = scheme.get_receiver()
                     receiver.synchronize_weights(worker_idx=worker_idx)
                     inner_collector._weight_receivers[model_id] = receiver

torchrl/weight_update/_mp.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import weakref
-from typing import Any
+from typing import Any, overload
 
 from torchrl.weight_update.weight_sync_schemes import (
     TransportBackend,
@@ -22,7 +22,7 @@ class MultiProcessWeightSyncScheme(WeightSyncScheme):
     Synchronization flow:
     - init_on_sender() creates a MPWeightSender and registers all worker pipes
     - synchronize_weights() triggers the initial weight distribution via pipes
-    - init_on_worker() creates a MPWeightReceiver that receives from its pipe
+    - init_on_receiver() creates a MPWeightReceiver that receives from its pipe
     - Subsequent updates use send() which extracts, sends, and waits for ACKs
 
     Args:
@@ -55,6 +55,27 @@ def synchronize_weights(self):
             )
         self._sender.synchronize_weights()
 
+    @overload
+    def init_on_sender(
+        self,
+        model_id: str,
+        context: Any,
+        **kwargs,
+    ) -> None:
+        ...
+
+    @overload
+    def init_on_sender(
+        self,
+        model_id: str,
+        context: None = None,
+        *,
+        pipes: list = ...,
+        num_workers: int | None = None,
+        **kwargs,
+    ) -> None:
+        ...
+
     def init_on_sender(
         self,
         model_id: str,
@@ -93,7 +114,28 @@ def init_on_sender(
         self._sender = sender
         self._initialized_on_sender = True
 
-    def init_on_worker(
+    @overload
+    def init_on_receiver(
+        self,
+        model_id: str,
+        context: Any,
+        **kwargs,
+    ) -> None:
+        ...
+
+    @overload
+    def init_on_receiver(
+        self,
+        model_id: str,
+        context: None = None,
+        *,
+        pipe: Any = ...,
+        model: Any | None = None,
+        **kwargs,
+    ) -> None:
+        ...
+
+    def init_on_receiver(
         self,
         model_id: str,
         context: Any = None,
@@ -138,7 +180,7 @@ def create_transport(self, pipe: Any) -> TransportBackend:
         """Create an MPTransport using the provided pipe.
 
         Note:
-            This is used internally by init_on_sender/init_on_worker.
+            This is used internally by init_on_sender/init_on_receiver.
         """
         return MPTransport(pipe)