Update

[ghstack-poisoned]

Author: vmoens

test/test_inference_server.py

@@ -1057,7 +1057,7 @@ def test_invalid_server_backend_raises(self):
                 create_env_fn=[_counting_env_factory] * 2,
                 policy_factory=_make_counting_policy,
                 frames_per_batch=10,
-                server_backend="proces",
+                server_backend="not-a-backend",
             )
 
     def test_server_death_raises_instead_of_hanging(self):

torchrl/collectors/_async_batched.py

@@ -444,10 +444,21 @@ def server_stats(self, *, reset: bool = False) -> dict[str, float | int]:
     # Rollout: drain the result queue
     # ------------------------------------------------------------------
 
-    @staticmethod
-    def _check_worker_result(item):
-        """Re-raise exceptions propagated from coordinator threads."""
+    def _check_worker_result(self, item):
+        """Re-raise exceptions propagated from coordinator threads.
+
+        Worker threads may observe a dying server before the liveness
+        watchdog does (their transport read errors out first); attribute the
+        failure to the server in that case so the caller gets a
+        deterministic error regardless of which path wins the race.
+        """
         if isinstance(item, BaseException):
+            if not self._server.is_alive:
+                raise RuntimeError(
+                    "The inference server died while the collector was "
+                    "waiting for transitions. Check the server process "
+                    "logs (e.g. OOM kills or exceptions in the policy)."
+                ) from item
             raise RuntimeError(
                 "A collector worker thread raised an exception."
             ) from item
@@ -500,6 +511,7 @@ def _rollout_frames(self) -> TensorDictBase:
                     td = rq.get_nowait()
                 except queue.Empty:
                     break
+                self._check_worker_result(td)
                 transitions.append(td)
                 collected += td.numel()
             if self.verbose: