Update

[ghstack-poisoned]

Author: vmoens

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

@@ -152,33 +152,45 @@ 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
+
+                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
-
-            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."""