[Feature] Add max-inflight guard for remote policy clients

test/test_inference_server.py

@@ -437,6 +437,7 @@ def test_forward_as_tensordict_module(self):
                 transport,
                 in_keys=["observation"],
                 out_keys=["action"],
+                max_inflight=1,
             )
             td = TensorDict({"observation": torch.randn(4)})
             result = remote_policy(td)
@@ -508,6 +509,136 @@ def test_update_policy_weights_cascade_bumps_version(self):
         server.update_policy_weights_()
         assert server.policy_version == 1
 
+    def test_max_inflight_validation(self):
+        with pytest.raises(ValueError, match="max_inflight must be at least 1"):
+            PolicyClientModule(lambda td: td, max_inflight=0)
+        with pytest.raises(ValueError, match="max_inflight must be at least 1"):
+            PolicyClientModule(lambda td: td, max_inflight=-3)
+        with pytest.raises(ValueError, match="max_inflight_per_env"):
+            InferenceServerConfig(max_inflight_per_env=0)
+
+    def test_max_inflight_survives_pickling(self):
+        """The guard is rebuilt on unpickle; clients stay picklable."""
+        remote_policy = PolicyClientModule(_echo_client, max_inflight=2)
+        restored = pickle.loads(pickle.dumps(remote_policy))
+        assert restored.max_inflight == 2
+        assert restored._inflight_sem is not None
+        # The rebuilt guard still enforces the limit
+        release_one = restored._acquire_inflight()
+        release_two = restored._acquire_inflight()
+        acquired = restored._inflight_sem.acquire(blocking=False)
+        assert not acquired
+        release_one()
+        release_two()
+        result = restored(TensorDict({"observation": torch.randn(4)}))
+        assert "observation" in result.keys()
+
+    def test_max_inflight_blocks_until_completion(self):
+        """The guard blocks a second submit and frees on completion (not result())."""
+
+        class _ManualClient:
+            def __init__(self):
+                self.futures = []
+
+            def submit(self, td):
+                fut = concurrent.futures.Future()
+                self.futures.append(fut)
+                return fut
+
+        client = _ManualClient()
+        remote = PolicyClientModule(client, max_inflight=1)
+        remote.submit(TensorDict({}))
+        second_done = threading.Event()
+
+        def _second_submit():
+            remote.submit(TensorDict({}))
+            second_done.set()
+
+        t = threading.Thread(target=_second_submit, daemon=True)
+        t.start()
+        time.sleep(0.2)
+        assert not second_done.is_set()  # guard enforced
+        # Completing the first request frees the slot even though result()
+        # is never called on it (done-callback release).
+        client.futures[0].set_result(TensorDict({}))
+        assert second_done.wait(timeout=5.0)
+        t.join(timeout=5.0)
+
+    def test_max_inflight_releases_on_error(self):
+        """A failed request frees its slot; the guard cannot deadlock."""
+
+        class _ManualClient:
+            def __init__(self):
+                self.futures = []
+
+            def submit(self, td):
+                fut = concurrent.futures.Future()
+                self.futures.append(fut)
+                return fut
+
+        client = _ManualClient()
+        remote = PolicyClientModule(client, max_inflight=1)
+        fut = remote.submit(TensorDict({}))
+        client.futures[0].set_exception(ValueError("boom"))
+        with pytest.raises(ValueError, match="boom"):
+            fut.result(timeout=5.0)
+        # Slot was released; this would deadlock otherwise.
+        fut2 = remote.submit(TensorDict({}))
+        client.futures[1].set_result(TensorDict({}))
+        assert fut2.result(timeout=5.0) is not None
+
+    def test_max_inflight_timeout_keeps_slot(self):
+        """A result() timeout must not free the slot of a running request."""
+
+        class _PullFuture:
+            """Pull-based future without add_done_callback support."""
+
+            def __init__(self):
+                self._event = threading.Event()
+                self._result = None
+
+            def done(self):
+                return self._event.is_set()
+
+            def result(self, timeout=None):
+                if not self._event.wait(timeout=timeout):
+                    raise TimeoutError("still running")
+                return self._result
+
+            def set_result(self, result):
+                self._result = result
+                self._event.set()
+
+        class _PullClient:
+            def __init__(self):
+                self.futures = []
+
+            def submit(self, td):
+                fut = _PullFuture()
+                self.futures.append(fut)
+                return fut
+
+        client = _PullClient()
+        remote = PolicyClientModule(client, max_inflight=1)
+        fut = remote.submit(TensorDict({}))
+        with pytest.raises(TimeoutError):
+            fut.result(timeout=0.05)
+        second_done = threading.Event()
+
+        def _second_submit():
+            remote.submit(TensorDict({}))
+            second_done.set()
+
+        t = threading.Thread(target=_second_submit, daemon=True)
+        t.start()
+        time.sleep(0.2)
+        # The request is still inflight: its slot must still be held.
+        assert not second_done.is_set()
+        client.futures[0].set_result(TensorDict({}))
+        assert fut.result(timeout=5.0) is not None
+        assert second_done.wait(timeout=5.0)
+        t.join(timeout=5.0)
+
 
 # =============================================================================
 # Tests: ThreadingTransport (Commit 2)

torchrl/collectors/_async_batched.py

@@ -399,6 +399,7 @@ def __init__(
                 policy_version_key=policy_version_key,
             )
         self._policy_version_key = policy_version_key
+        self._max_inflight_per_env = _server_defaults.max_inflight_per_env
 
         # ---- collector settings -----------------------------------------------
         self.requested_frames_per_batch = frames_per_batch
@@ -446,7 +447,10 @@ def _ensure_started(self) -> None:
         # inherited by the child process.
         if self._clients is None:
             self._clients = [
-                PolicyClientModule(self._transport.client())
+                PolicyClientModule(
+                    self._transport.client(),
+                    max_inflight=self._max_inflight_per_env,
+                )
                 for _ in range(self._num_envs)
             ]
 

torchrl/modules/inference_server/_client.py

@@ -4,8 +4,11 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
+import queue
+import threading
 from collections.abc import Callable, Sequence
 from concurrent.futures import Future
+from typing import Any
 
 from tensordict.base import TensorDictBase
 from tensordict.nn import TensorDictModuleBase
@@ -27,6 +30,72 @@ def result(self, timeout: float | None = None) -> TensorDictBase:
         return self._result
 
 
+class _InflightGuardedFuture:
+    """Future proxy that frees an inflight slot when the request *completes*.
+
+    Callback-capable futures (:class:`concurrent.futures.Future`) release
+    through ``add_done_callback``, so a dropped or cancelled future cannot
+    leak its slot. Pull-based futures (e.g. queue transports) release when a
+    completed result is first observed through :meth:`result` or
+    :meth:`done`. In both cases a ``result(timeout=...)`` that times out does
+    **not** free the slot -- the request is still running on the server, and
+    releasing early would let the number of genuinely inflight requests
+    exceed ``max_inflight``. Garbage collection of the proxy releases the
+    slot as a last resort so an abandoned pull-based future cannot
+    permanently exhaust the guard.
+    """
+
+    def __init__(self, future, release: Callable[[], None]) -> None:
+        self.future = future
+        self._release_cb = release
+        self._released = False
+        self._release_lock = threading.Lock()
+        add_done_callback = getattr(future, "add_done_callback", None)
+        if add_done_callback is not None:
+            add_done_callback(lambda _fut: self._release_once())
+
+    def _release_once(self) -> None:
+        with self._release_lock:
+            if self._released:
+                return
+            self._released = True
+        self._release_cb()
+
+    def done(self) -> bool:
+        is_done = self.future.done()
+        if is_done:
+            self._release_once()
+        return is_done
+
+    def result(self, timeout: float | None = None) -> TensorDictBase:
+        try:
+            result = self.future.result(timeout=timeout)
+        except (queue.Empty, TimeoutError):
+            # The request is still inflight on the server; keep the slot.
+            raise
+        except BaseException:
+            self._release_once()
+            raise
+        self._release_once()
+        return result
+
+    def __getattr__(self, name: str) -> Any:
+        # __getattr__ only fires for missing attributes; route through
+        # __dict__ explicitly so a partially-initialised proxy (e.g. during
+        # unpickling) raises AttributeError instead of recursing.
+        try:
+            future = object.__getattribute__(self, "future")
+        except AttributeError:
+            raise AttributeError(name) from None
+        return getattr(future, name)
+
+    def __del__(self) -> None:
+        try:
+            self._release_once()
+        except Exception:
+            pass
+
+
 class PolicyClientModule(TensorDictModuleBase):
     """TensorDict policy wrapper for remote inference-server clients.
 
@@ -59,6 +128,12 @@ class PolicyClientModule(TensorDictModuleBase):
             module. The full input TensorDict is still sent to the server.
         out_keys (sequence of NestedKey, optional): output keys advertised by
             the module.
+        max_inflight (int, optional): maximum number of unresolved
+            asynchronous requests submitted through this module; further
+            :meth:`submit` calls block until a slot frees up. A slot is
+            freed when its request *completes* (including errors), not when
+            ``result()`` is first called; a timed-out ``result()`` keeps the
+            slot. Must be at least ``1``. ``None`` means unbounded.
 
     .. note::
         Version tracking is an instance of the generic *service-stamped
@@ -114,13 +189,45 @@ def __init__(
         *,
         in_keys: Sequence[NestedKey] | None = None,
         out_keys: Sequence[NestedKey] | None = None,
+        max_inflight: int | None = None,
     ) -> None:
         super().__init__()
         if isinstance(client, InferenceTransport):
             client = client.client()
+        if max_inflight is not None and max_inflight < 1:
+            raise ValueError(
+                f"max_inflight must be at least 1 (got {max_inflight}); "
+                "use None to disable the guard."
+            )
         self.client = client
         self.in_keys = list(in_keys or [])
         self.out_keys = list(out_keys or [])
+        self.max_inflight = max_inflight
+        self._inflight_sem = (
+            threading.BoundedSemaphore(max_inflight)
+            if max_inflight is not None
+            else None
+        )
+
+    def __getstate__(self):
+        # Semaphores are not picklable; the guard is a per-process resource,
+        # so a fresh one (with a full complement of slots) is rebuilt on
+        # unpickling. This keeps clients picklable per the Client contract.
+        state = super().__getstate__()
+        state = dict(state)
+        state["_inflight_sem"] = None
+        return state
+
+    def __setstate__(self, state) -> None:
+        super().__setstate__(state)
+        if self.max_inflight is not None:
+            self._inflight_sem = threading.BoundedSemaphore(self.max_inflight)
+
+    def _acquire_inflight(self) -> Callable[[], None]:
+        if self._inflight_sem is None:
+            return lambda: None
+        self._inflight_sem.acquire()
+        return self._inflight_sem.release
 
     def submit(self, tensordict: TensorDictBase) -> Future | _ImmediateFuture:
         """Submit a TensorDict request and return a future-like object.
@@ -137,14 +244,24 @@ def submit(self, tensordict: TensorDictBase) -> Future | _ImmediateFuture:
             and errors are deferred to ``result()`` on a reduced future that
             only implements ``done()`` and ``result()``.
         """
+        release = self._acquire_inflight()
         submit = getattr(self.client, "submit", None)
         if submit is None:
+            # The plain-callable path runs eagerly, so the request has
+            # already completed here: free the slot immediately.
             try:
                 result = self.client(tensordict)
                 return _ImmediateFuture(result)
             except Exception as exc:
                 return _ImmediateFuture(exc)
-        return submit(tensordict)
+            finally:
+                release()
+        try:
+            future = submit(tensordict)
+        except BaseException:
+            release()
+            raise
+        return _InflightGuardedFuture(future, release)
 
     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
         return self.submit(tensordict).result()

torchrl/modules/inference_server/_config.py

@@ -99,6 +99,12 @@ class InferenceServerConfig:
             throughput and latency stats. Defaults to ``True``.
         stats_window_size (int, optional): number of recent timing samples kept
             for percentile stats. Defaults to ``1024``.
+        max_inflight_per_env (int, optional): maximum unresolved remote-policy
+            requests each environment coordinator may have inflight (consumed
+            by :class:`~torchrl.collectors.AsyncBatchedCollector` when
+            building its clients). Defaults to ``None`` (unbounded), so the
+            guard never throttles by surprise; set an explicit bound when
+            backpressure is wanted.
 
     Examples:
         >>> import torch
@@ -130,10 +136,16 @@ class InferenceServerConfig:
     timeout: float = 0.01
     collect_stats: bool = True
     stats_window_size: int = 1024
+    max_inflight_per_env: int | None = None
 
     def __post_init__(self) -> None:
         if self.backend not in ("thread", "process"):
             raise ValueError(
                 f"backend={self.backend!r} is not supported. "
                 "Expected 'thread' or 'process'."
             )
+        if self.max_inflight_per_env is not None and self.max_inflight_per_env < 1:
+            raise ValueError(
+                f"max_inflight_per_env must be at least 1 (got "
+                f"{self.max_inflight_per_env}); use None to disable the guard."
+            )