Update

[ghstack-poisoned]

Author: vmoens

docs/source/reference/modules_inference_server.rst

@@ -28,6 +28,7 @@ Transport Backends
     :template: rl_template_noinherit.rst
 
     ThreadingTransport
+    SlotTransport
     MPTransport
     RayTransport
     MonarchTransport

test/test_inference_server.py

@@ -981,3 +981,68 @@ def tracking_collate(items):
 
         # At least one batch should have >= min_batch_size items
         assert any(s >= min_bs for s in seen_sizes)
+
+
+# =============================================================================
+# Tests: bugfix regressions
+# =============================================================================
+
+
+class TestShutdownPendingFutures:
+    def test_shutdown_resolves_pending_futures(self):
+        """Pending futures receive an exception on shutdown (no hang)."""
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        server = InferenceServer(policy, transport, max_batch_size=1024)
+        server.start()
+        futures = [
+            transport.submit(TensorDict({"observation": torch.randn(4)}))
+            for _ in range(5)
+        ]
+        time.sleep(0.05)
+        server.shutdown(timeout=5.0)
+        for f in futures:
+            try:
+                f.result(timeout=2.0)
+            except Exception:
+                pass  # exception is acceptable; hanging is not
+
+
+class TestThreadingTransportNoLostSignals:
+    def test_rapid_submit_no_lost_signals(self):
+        """Rapid submits from many threads don't lose signals."""
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        n = 100
+        with InferenceServer(policy, transport, max_batch_size=4, timeout=0.001):
+            client = transport.client()
+            with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
+                futs = [
+                    pool.submit(
+                        lambda: client(TensorDict({"observation": torch.randn(4)}))
+                    )
+                    for _ in range(n)
+                ]
+                results = [f.result(timeout=10.0) for f in futs]
+        assert len(results) == n
+        for r in results:
+            assert "action" in r.keys()
+
+
+class TestWorkerCrashPropagation:
+    def test_worker_crash_propagates(self):
+        """If the model always fails, the collector propagates the error."""
+
+        def bad_model(td):
+            raise RuntimeError("model crash")
+
+        collector = AsyncBatchedCollector(
+            create_env_fn=[_counting_env_factory] * 2,
+            policy=bad_model,
+            frames_per_batch=10,
+            total_frames=100,
+        )
+        with pytest.raises(RuntimeError, match="worker thread"):
+            for _ in collector:
+                pass
+        collector.shutdown()

torchrl/collectors/_async_batched.py

@@ -92,9 +92,9 @@ def _env_loop(
             if shutdown_event.is_set():
                 break
             action_td = client(next_obs)
-    except Exception:
+    except Exception as exc:
         if not shutdown_event.is_set():
-            raise
+            result_queue.put(exc)
 
 
 class AsyncBatchedCollector(BaseCollector):
@@ -367,6 +367,14 @@ def policy(self) -> Callable:
     # Rollout: drain the result queue
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _check_worker_result(item):
+        """Re-raise exceptions propagated from coordinator threads."""
+        if isinstance(item, BaseException):
+            raise RuntimeError(
+                "A collector worker thread raised an exception."
+            ) from item
+
     def _rollout_frames(self) -> TensorDictBase:
         """Drain ``frames_per_batch`` transitions from the workers."""
         rq = self._result_queue
@@ -376,6 +384,7 @@ def _rollout_frames(self) -> TensorDictBase:
         while collected < self.frames_per_batch:
             # Block for at least one transition
             td = rq.get()
+            self._check_worker_result(td)
             transitions.append(td)
             collected += td.numel()
             # Batch-drain any additional items already in the queue
@@ -399,6 +408,7 @@ def _rollout_yield_trajs(self) -> TensorDictBase:
 
         while not self._trajectory_queue:
             td = rq.get()
+            self._check_worker_result(td)
             env_id = 0
             eid = td.get(_ENV_IDX_KEY, default=None)
             if eid is not None:

torchrl/modules/inference_server/_queue_transport.py

@@ -112,8 +112,10 @@ def _get_result(self, req_id: int, timeout: float | None = None) -> Any:
                     raise queue.Empty(f"Timeout waiting for result of request {req_id}")
             try:
                 rid, result = self._response_queue.get(timeout=remaining)
-            except Exception:
-                raise queue.Empty(f"Timeout waiting for result of request {req_id}")
+            except Exception as e:
+                raise queue.Empty(
+                    f"Timeout waiting for result of request {req_id}"
+                ) from e
             if rid == req_id:
                 return result
             self._buffered[rid] = result
@@ -135,15 +137,15 @@ class QueueBasedTransport(InferenceTransport):
     * :meth:`_make_response_queue` -- factory for creating a new response queue.
 
     .. note::
-        ``wait_for_work`` uses a blocking ``get`` followed by ``put`` to peek
-        at the request queue.  This is safe because a single server thread
-        calls both ``wait_for_work`` and ``drain`` sequentially -- there is no
-        concurrent consumer that could miss the re-enqueued item.
+        ``wait_for_work`` uses a blocking ``get`` to detect new work.  The
+        retrieved item is stored in ``_peeked`` and consumed by the next
+        ``drain`` call, preserving FIFO ordering.
     """
 
     def __init__(self):
         self._lock = threading.Lock()
         self._next_actor_id = 0
+        self._peeked = None
 
     # -- to be implemented by subclasses --------------------------------------
 
@@ -181,7 +183,13 @@ def drain(
         """Dequeue up to *max_items* pending requests (non-blocking)."""
         items: list[TensorDictBase] = []
         callbacks: list[tuple[int, int]] = []
-        for _ in range(max_items):
+        peeked = self._peeked
+        if peeked is not None:
+            self._peeked = None
+            actor_id, req_id, td = peeked
+            items.append(td)
+            callbacks.append((actor_id, req_id))
+        for _ in range(max_items - len(items)):
             try:
                 actor_id, req_id, td = self._request_queue.get(block=False)
             except Exception:
@@ -192,10 +200,10 @@ def drain(
 
     def wait_for_work(self, timeout: float) -> None:
         """Block until at least one request is available or *timeout* elapses."""
+        if self._peeked is not None:
+            return
         try:
-            item = self._request_queue.get(timeout=timeout)
-            # Put it back so drain() can consume it.
-            self._request_queue.put(item)
+            self._peeked = self._request_queue.get(timeout=timeout)
         except Exception:
             pass
 

torchrl/modules/inference_server/_server.py

@@ -160,47 +160,59 @@ def _poll_weight_update(self) -> None:
     def _run(self) -> None:
         self._init_weight_sync()
 
-        while not self._shutdown_event.is_set():
-            # Poll for weight updates between batches (non-blocking)
-            self._poll_weight_update()
-
-            self.transport.wait_for_work(timeout=self.timeout)
-
+        try:
+            while not self._shutdown_event.is_set():
+                self._poll_weight_update()
+
+                self.transport.wait_for_work(timeout=self.timeout)
+
+                items, callbacks = self.transport.drain(self.max_batch_size)
+                if not items:
+                    continue
+
+                # Accumulate up to min_batch_size (or until timeout expires)
+                if len(items) < self.min_batch_size:
+                    deadline = time.monotonic() + self.timeout
+                    while len(items) < self.min_batch_size:
+                        remaining = deadline - time.monotonic()
+                        if remaining <= 0:
+                            break
+                        self.transport.wait_for_work(timeout=remaining)
+                        more_items, more_cbs = self.transport.drain(
+                            self.max_batch_size - len(items)
+                        )
+                        items.extend(more_items)
+                        callbacks.extend(more_cbs)
+
+                batch = self.collate_fn(items)
+                if self.device is not None:
+                    batch = batch.to(self.device)
+
+                try:
+                    with self._model_lock:
+                        results = self.model(batch).unbind(0)
+                    if len(results) != len(callbacks):
+                        raise RuntimeError(
+                            f"Model returned {len(results)} results for a "
+                            f"batch of {len(callbacks)} inputs."
+                        )
+                    for cb, res in zip(callbacks, results):
+                        self.transport.resolve(cb, res)
+                except Exception as exc:
+                    for cb in callbacks:
+                        self.transport.resolve_exception(cb, exc)
+        finally:
+            self._drain_pending_on_shutdown()
+
+    def _drain_pending_on_shutdown(self) -> None:
+        """Resolve all pending requests with an error during shutdown."""
+        shutdown_exc = RuntimeError("InferenceServer is shutting down.")
+        while True:
             items, callbacks = self.transport.drain(self.max_batch_size)
             if not items:
-                continue
-
-            # Accumulate up to min_batch_size (or until timeout expires)
-            if len(items) < self.min_batch_size:
-                deadline = time.monotonic() + self.timeout
-                while len(items) < self.min_batch_size:
-                    remaining = deadline - time.monotonic()
-                    if remaining <= 0:
-                        break
-                    self.transport.wait_for_work(timeout=remaining)
-                    more_items, more_cbs = self.transport.drain(
-                        self.max_batch_size - len(items)
-                    )
-                    items.extend(more_items)
-                    callbacks.extend(more_cbs)
-
-            batch = self.collate_fn(items)
-            if self.device is not None:
-                batch = batch.to(self.device)
-
-            try:
-                with self._model_lock:
-                    results = self.model(batch).unbind(0)
-                if len(results) != len(callbacks):
-                    raise RuntimeError(
-                        f"Model returned {len(results)} results for a "
-                        f"batch of {len(callbacks)} inputs."
-                    )
-                for cb, res in zip(callbacks, results):
-                    self.transport.resolve(cb, res)
-            except Exception as exc:
-                for cb in callbacks:
-                    self.transport.resolve_exception(cb, exc)
+                break
+            for cb in callbacks:
+                self.transport.resolve_exception(cb, shutdown_exc)
 
     # -- context manager ------------------------------------------------------
 

torchrl/modules/inference_server/_threading.py

@@ -15,8 +15,9 @@
 class ThreadingTransport(InferenceTransport):
     """In-process transport for actors that are threads.
 
-    Uses a shared list protected by a :class:`threading.Lock` as the request
-    queue and :class:`~concurrent.futures.Future` objects for response routing.
+    Uses a shared list protected by a :class:`threading.Condition` as the
+    request queue and :class:`~concurrent.futures.Future` objects for response
+    routing.
 
     This is the simplest backend and is appropriate when all actors live in the
     same process (e.g. running in a :class:`~concurrent.futures.ThreadPoolExecutor`).
@@ -25,21 +26,20 @@ class ThreadingTransport(InferenceTransport):
     def __init__(self):
         self._queue: list[TensorDictBase] = []
         self._futures: list[Future] = []
-        self._lock = threading.Lock()
-        self._event = threading.Event()
+        self._cond = threading.Condition(threading.Lock())
 
     def submit(self, td: TensorDictBase) -> Future[TensorDictBase]:
         """Enqueue a request and return a Future for the result."""
         fut: Future[TensorDictBase] = Future()
-        with self._lock:
+        with self._cond:
             self._queue.append(td)
             self._futures.append(fut)
-        self._event.set()
+            self._cond.notify()
         return fut
 
     def drain(self, max_items: int) -> tuple[list[TensorDictBase], list[Future]]:
         """Dequeue up to *max_items* pending requests."""
-        with self._lock:
+        with self._cond:
             n = min(len(self._queue), max_items)
             items = self._queue[:n]
             futs = self._futures[:n]
@@ -49,8 +49,9 @@ def drain(self, max_items: int) -> tuple[list[TensorDictBase], list[Future]]:
 
     def wait_for_work(self, timeout: float) -> None:
         """Block until at least one request is enqueued or *timeout* elapses."""
-        self._event.wait(timeout=timeout)
-        self._event.clear()
+        with self._cond:
+            if not self._queue:
+                self._cond.wait(timeout=timeout)
 
     def resolve(self, callback: Future, result: TensorDictBase) -> None:
         """Set the result on the actor's Future."""