Update (base update)

[ghstack-poisoned]

Author: vmoens

examples/collectors/async_batched_collector.py

@@ -4,12 +4,20 @@
 run many environments in parallel while automatically batching policy inference
 through an :class:`~torchrl.modules.InferenceServer`.
 
+Architecture:
+  - An :class:`~torchrl.envs.AsyncEnvPool` runs environments in parallel using
+    the chosen backend (``"multiprocessing"`` by default for true parallelism,
+    or ``"threading"``/``"asyncio"``).
+  - An :class:`~torchrl.modules.InferenceServer` batches incoming observations
+    and runs a single forward pass.
+  - A lightweight coordinator thread bridges the two: when an env finishes
+    stepping its observation is submitted to the server, and when an action is
+    ready the env is sent back for stepping -- all without synchronisation
+    barriers.
+
 The user only supplies:
   - A list of environment factories
   - A policy (or policy factory)
-
-The collector creates the ``AsyncEnvPool``, ``InferenceServer``, and
-``ThreadingTransport`` internally -- no manual wiring required.
 """
 import torch.nn as nn
 from tensordict.nn import TensorDictModule

test/test_inference_server.py

@@ -733,6 +733,7 @@ def test_basic_collection(self):
             frames_per_batch=frames_per_batch,
             total_frames=total_frames,
             max_batch_size=num_envs,
+            env_backend="threading",
         )
         total_collected = 0
         for batch in collector:
@@ -750,6 +751,7 @@ def test_policy_factory(self):
             frames_per_batch=10,
             total_frames=20,
             max_batch_size=num_envs,
+            env_backend="threading",
         )
         total_collected = 0
         for batch in collector:
@@ -789,6 +791,7 @@ def test_yield_completed_trajectories(self):
             total_frames=30,
             yield_completed_trajectories=True,
             max_batch_size=num_envs,
+            env_backend="threading",
         )
         count = 0
         for batch in collector:
@@ -806,6 +809,7 @@ def test_shutdown_idempotent(self):
             policy=policy,
             frames_per_batch=10,
             total_frames=10,
+            env_backend="threading",
         )
         # Consume one batch to start
         for _batch in collector:
@@ -821,6 +825,7 @@ def test_endless_collector(self):
             policy=policy,
             frames_per_batch=10,
             total_frames=-1,
+            env_backend="threading",
         )
         collected = 0
         for batch in collector:
@@ -830,18 +835,16 @@ def test_endless_collector(self):
         collector.shutdown()
         assert collected >= 50
 
-    def test_env_property(self):
-        """The env property returns an AsyncEnvPool."""
-        from torchrl.envs import AsyncEnvPool
-
+    def test_num_envs(self):
+        """The collector knows the number of environments."""
         policy = _make_counting_policy()
         collector = AsyncBatchedCollector(
             create_env_fn=[_counting_env_factory] * 2,
             policy=policy,
             frames_per_batch=10,
             total_frames=10,
         )
-        assert isinstance(collector.env, AsyncEnvPool)
+        assert collector._num_envs == 2
         collector.shutdown()
 
     def test_postproc(self):
@@ -859,6 +862,7 @@ def postproc(td):
             frames_per_batch=10,
             total_frames=20,
             postproc=postproc,
+            env_backend="threading",
         )
         for _ in collector:
             pass

torchrl/modules/inference_server/_monarch.py

@@ -29,13 +29,22 @@ def __init__(self, client: _MonarchInferenceClient, req_id: int):
         self._req_id = req_id
         self._result: Any = _SENTINEL
 
+    def done(self) -> bool:
+        """Return ``True`` if the result is available without blocking."""
+        if self._result is not _SENTINEL:
+            return True
+        try:
+            self._result = self._client._get_result(self._req_id, timeout=0)
+        except queue.Empty:
+            return False
+        return True
+
     def result(self, timeout: float | None = None) -> TensorDictBase:
         """Block until the result is available."""
         if self._result is _SENTINEL:
-            item = self._client._get_result(self._req_id, timeout=timeout)
-            if isinstance(item, BaseException):
-                raise item
-            self._result = item
+            self._result = self._client._get_result(self._req_id, timeout=timeout)
+        if isinstance(self._result, BaseException):
+            raise self._result
         return self._result
 
 

torchrl/modules/inference_server/_mp.py

@@ -33,6 +33,16 @@ def __init__(self, client: _MPInferenceClient, req_id: int):
         self._req_id = req_id
         self._result: Any = _SENTINEL
 
+    def done(self) -> bool:
+        """Return ``True`` if the result is available without blocking."""
+        if self._result is not _SENTINEL:
+            return True
+        try:
+            self._result = self._client._get_result(self._req_id, timeout=0)
+        except queue.Empty:
+            return False
+        return True
+
     def result(self, timeout: float | None = None) -> TensorDictBase:
         """Block until the result is available.
 
@@ -44,10 +54,9 @@ def result(self, timeout: float | None = None) -> TensorDictBase:
             Exception: if the server set an exception instead of a result.
         """
         if self._result is _SENTINEL:
-            item = self._client._get_result(self._req_id, timeout=timeout)
-            if isinstance(item, BaseException):
-                raise item
-            self._result = item
+            self._result = self._client._get_result(self._req_id, timeout=timeout)
+        if isinstance(self._result, BaseException):
+            raise self._result
         return self._result
 
 

torchrl/modules/inference_server/_ray.py

@@ -32,13 +32,22 @@ def __init__(self, client: _RayInferenceClient, req_id: int):
         self._req_id = req_id
         self._result: Any = _SENTINEL
 
+    def done(self) -> bool:
+        """Return ``True`` if the result is available without blocking."""
+        if self._result is not _SENTINEL:
+            return True
+        try:
+            self._result = self._client._get_result(self._req_id, timeout=0)
+        except queue.Empty:
+            return False
+        return True
+
     def result(self, timeout: float | None = None) -> TensorDictBase:
         """Block until the result is available."""
         if self._result is _SENTINEL:
-            item = self._client._get_result(self._req_id, timeout=timeout)
-            if isinstance(item, BaseException):
-                raise item
-            self._result = item
+            self._result = self._client._get_result(self._req_id, timeout=timeout)
+        if isinstance(self._result, BaseException):
+            raise self._result
         return self._result