Update

[ghstack-poisoned]

Author: vmoens

test/llm/test_updaters.py

@@ -41,7 +41,7 @@ def get_open_port():
 if _has_ray:
     import ray
 
-    from ray_helpers import (
+    from torchrl.testing import (
         WorkerTransformerDoubleBuffer,
         WorkerTransformerNCCL,
         WorkerVLLMDoubleBuffer,

test/test_rb.py

@@ -4188,6 +4188,25 @@ def test_ray_rb_iter(self):
         finally:
             rb.close()
 
+    def test_ray_rb_serialization(self):
+        import ray
+
+        class Worker:
+            def __init__(self, rb):
+                self.rb = rb
+
+            def run(self):
+                self.rb.extend(TensorDict({"x": torch.ones(100)}, batch_size=100))
+
+        rb = RayReplayBuffer(
+            storage=partial(LazyTensorStorage, 100), ray_init_config={"num_cpus": 1}
+        )
+        try:
+            remote_worker = ray.remote(Worker).remote(rb)
+            ray.get(remote_worker.run.remote())
+        finally:
+            rb.close()
+
 
 class TestSharedStorageInit:
     def worker(self, rb, worker_id, queue):

torchrl/data/replay_buffers/ray_buffer.py

@@ -147,6 +147,7 @@ def __init__(
         else:
             self.has_gpu = False
         self._rb = remote_cls(*args, **kwargs)
+        self._delayed_init = False
 
     def close(self):
         """Terminates the Ray actor associated with this replay buffer."""

torchrl/testing/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Testing utilities for TorchRL.
+
+This module provides helper classes and utilities for testing TorchRL functionality,
+particularly for distributed and Ray-based tests that require importable classes.
+"""
+
+from torchrl.testing.ray_helpers import (
+    WorkerTransformerDoubleBuffer,
+    WorkerTransformerNCCL,
+    WorkerVLLMDoubleBuffer,
+    WorkerVLLMNCCL,
+)
+
+__all__ = [
+    "WorkerVLLMNCCL",
+    "WorkerTransformerNCCL",
+    "WorkerVLLMDoubleBuffer",
+    "WorkerTransformerDoubleBuffer",
+]

test/llm/ray_helpers.py renamed to torchrl/testing/ray_helpers.py

@@ -6,8 +6,8 @@
 """Helper classes for Ray-based weight synchronization tests.
 
 This module contains Ray actor classes that need to be importable by Ray workers.
-These classes are used in test_updaters.py but must be defined at module level
-so Ray can serialize and import them on remote workers.
+These classes are used in tests but must be defined at module level in a proper
+Python package (not in test files) so Ray can serialize and import them on remote workers.
 """
 
 import torch