Update

[ghstack-poisoned]

Author: vmoens

test/test_weightsync.py

@@ -702,7 +702,29 @@ def test_shared_mem_scheme_serialize_after_init(self):
             strategy="tensordict",
             auto_register=False,
         )
-        scheme.init_on_sender(model_id="policy", pipes=[parent_pipe])
+
+        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

torchrl/weight_update/weight_sync_schemes.py

@@ -202,7 +202,9 @@ def register_weights(self, model_id: str, weights: TensorDictBase) -> None:
                     f"Model '{model_id}' has already been registered with workers."
                 )
 
-    def _send_buffer_to_workers(self, model_id: str, buffer: TensorDictBase) -> None:
+    def _send_buffer_to_workers(
+        self, model_id: str, buffer: TensorDictBase, timeout: float = 10.0
+    ) -> None:
         """Send shared memory buffer reference to all workers via pipes.
 
         This is called once per model_id when lazy registration occurs.
@@ -219,6 +221,8 @@ def _send_buffer_to_workers(self, model_id: str, buffer: TensorDictBase) -> None
 
         # Wait for acknowledgments from all workers
         for pipe in self._pipes:
+            if not pipe.poll(timeout):
+                raise TimeoutError(f"Timeout waiting for acknowledgment from worker")
             _, msg = pipe.recv()
             if msg != "registered":
                 raise RuntimeError(f"Expected 'registered' acknowledgment, got '{msg}'")