[BugFix] Make async recv timeout bound the whole call (#4184)

Author: vmoens

test/envs/test_special.py

@@ -945,9 +945,12 @@ def test_shared_memory_lazy_result_lifetime(self, make_envs):
     def test_shared_memory_deadline_batching(self, make_envs):
         env = AsyncEnvPool(make_envs, backend="multiprocessing", exchange="shm")
         try:
-            env.async_reset_send(env_index=list(range(env.num_envs)))
-            first = env.async_reset_recv(min_get=1, max_get=4, timeout=0.0)
-            remaining = env.async_reset_recv(min_get=3, max_get=3, timeout=0.1)
+            # Only env 0 is commanded, so the first recv returns a partial
+            # batch (1 of max_get=4) when its deadline expires.
+            env.async_reset_send(env_index=[0])
+            first = env.async_reset_recv(min_get=1, max_get=4, timeout=1.0)
+            env.async_reset_send(env_index=[1, 2, 3])
+            remaining = env.async_reset_recv(min_get=3, max_get=3, timeout=10.0)
             assert first.shape[0] == 1
             assert remaining.shape[0] == 3
             stats = env.stats(reset=True)
@@ -967,13 +970,47 @@ def test_deadline_batching(self, make_envs, backend):
         env = AsyncEnvPool(make_envs, backend=backend)
         try:
             env.async_reset_send(env_index=list(range(env.num_envs)))
-            first = env.async_reset_recv(min_get=1, max_get=1, timeout=0.0)
-            remaining = env.async_reset_recv(min_get=3, max_get=3, timeout=0.1)
+            first = env.async_reset_recv(min_get=1, max_get=1, timeout=10.0)
+            remaining = env.async_reset_recv(min_get=3, max_get=3, timeout=10.0)
             assert first.shape[0] == 1
             assert remaining.shape[0] == 3
         finally:
             env._maybe_shutdown()
 
+    @pytest.mark.parametrize(
+        "backend,exchange",
+        [
+            ("multiprocessing", "queue"),
+            ("multiprocessing", "shm"),
+            ("threading", "queue"),
+        ],
+    )
+    @set_capture_non_tensor_stack(False)
+    def test_recv_timeout_bounds_whole_call(self, backend, exchange):
+        """The recv deadline covers the wait for min_get, without losing results.
+
+        One env steps immediately, the other takes ~1s. A recv asking for both
+        with a 0.2s deadline must raise TimeoutError instead of blocking until
+        the slow env finishes, and the fast env's result must remain available
+        to the next call.
+        """
+        makers = [
+            partial(_DelayedCountingEnv, delay=0.0, max_steps=1000),
+            partial(_DelayedCountingEnv, delay=1.0, max_steps=1000),
+        ]
+        kwargs = {"exchange": exchange} if backend == "multiprocessing" else {}
+        env = AsyncEnvPool(makers, backend=backend, **kwargs)
+        try:
+            reset = env.reset()
+            reset.set("action", torch.ones(reset.shape + (1,)))
+            env.async_step_send(reset)
+            with pytest.raises(TimeoutError, match="timed out"):
+                env.async_step_recv(min_get=2, timeout=0.2)
+            result = env.async_step_recv(min_get=2, timeout=60.0)
+            assert result.shape[0] == 2
+        finally:
+            env._maybe_shutdown()
+
     @pytest.mark.parametrize("backend", ["multiprocessing", "threading"])
     def test_deadline_batch_validation(self, make_envs, backend):
         env = AsyncEnvPool(make_envs, backend=backend)

torchrl/envs/_async_exchange.py

@@ -32,23 +32,48 @@ def _receive_batch(
     if timeout is not None and timeout < 0:
         raise ValueError(f"timeout must be non-negative, got {timeout}.")
 
-    items = [result_queue.get()]
+    # The deadline is anchored at call entry and bounds the entire call,
+    # including the wait for the first result.
     deadline_timer = (
         None if timeout is None else timeit("async_env_batch_deadline").start()
     )
 
+    items: list[Any] = []
     while len(items) < min_get:
-        items.append(result_queue.get())
+        if deadline_timer is None:
+            items.append(result_queue.get())
+            continue
+        remaining = timeout - deadline_timer.elapsed()
+        try:
+            if remaining <= 0:
+                # Deadline passed: drain what is already available without
+                # waiting before giving up.
+                items.append(result_queue.get_nowait())
+            else:
+                items.append(result_queue.get(timeout=remaining))
+        except queue.Empty:
+            # Requeue the partial harvest so no result is lost, then signal
+            # the missed deadline. The pool's pending-result accounting is
+            # only updated by the caller on success, so state stays
+            # consistent.
+            for item in items:
+                result_queue.put(item)
+            raise TimeoutError(
+                f"async recv timed out: {len(items)}/{min_get} results after "
+                f"{timeout}s; partial results were requeued and remain "
+                f"available to the next call."
+            ) from None
 
     limit = max_get if max_get is not None else float("inf")
     while len(items) < limit:
         try:
             if deadline_timer is None:
                 items.append(result_queue.get_nowait())
+                continue
+            remaining = timeout - deadline_timer.elapsed()
+            if remaining <= 0:
+                items.append(result_queue.get_nowait())
             else:
-                remaining = timeout - deadline_timer.elapsed()
-                if remaining <= 0:
-                    break
                 items.append(result_queue.get(timeout=remaining))
         except queue.Empty:
             break

torchrl/envs/async_envs.py

@@ -480,6 +480,32 @@ def async_step_recv(
         max_get: int | None = None,
         timeout: float | None = None,
     ) -> TensorDictBase:
+        """Collects step results from the pool.
+
+        Args:
+            min_get (int, optional): minimum number of results to collect
+                before returning. Defaults to ``self.min_get``.
+            env_index (int, optional): if provided, waits for the result of
+                this specific env instead of forming a batch; ``min_get``,
+                ``max_get`` and ``timeout`` are ignored.
+
+        Keyword Args:
+            max_get (int, optional): maximum number of results to return.
+                ``None`` (default) drains everything available once
+                ``min_get`` is satisfied.
+            timeout (float, optional): deadline in seconds for the entire
+                call, measured from call entry. If fewer than ``min_get``
+                results are available when the deadline expires, a
+                :class:`TimeoutError` is raised; results collected so far are
+                put back and remain available to the next call. Once
+                ``min_get`` is satisfied, the call returns at the deadline
+                with whatever is ready (up to ``max_get``). ``timeout=0``
+                polls without blocking. ``None`` (default) waits
+                indefinitely.
+
+        Returns:
+            TensorDictBase: the stacked results, carrying ``"env_index"``.
+        """
         raise NotImplementedError
 
     @abc.abstractmethod
@@ -497,6 +523,18 @@ def async_step_and_maybe_reset_recv(
         max_get: int | None = None,
         timeout: float | None = None,
     ) -> tuple[TensorDictBase, TensorDictBase]:
+        """Collects step-and-maybe-reset results from the pool.
+
+        ``min_get``, ``env_index``, ``max_get`` and ``timeout`` behave as in
+        :meth:`async_step_recv`; in particular ``timeout`` bounds the entire
+        call and raises :class:`TimeoutError` (without losing results) if
+        fewer than ``min_get`` results arrive in time.
+
+        Returns:
+            tuple of (TensorDictBase, TensorDictBase): the stacked step
+            results and the stacked roots of the next steps (reset where
+            needed), both carrying ``"env_index"``.
+        """
         raise NotImplementedError
 
     @abc.abstractmethod
@@ -516,6 +554,17 @@ def async_reset_recv(
         max_get: int | None = None,
         timeout: float | None = None,
     ) -> TensorDictBase:
+        """Collects reset results from the pool.
+
+        ``min_get``, ``env_index``, ``max_get`` and ``timeout`` behave as in
+        :meth:`async_step_recv`; in particular ``timeout`` bounds the entire
+        call and raises :class:`TimeoutError` (without losing results) if
+        fewer than ``min_get`` results arrive in time.
+
+        Returns:
+            TensorDictBase: the stacked reset results, carrying
+            ``"env_index"``.
+        """
         raise NotImplementedError
 
     def stats(self, *, reset: bool = False) -> dict[str, float | int]:
@@ -1118,12 +1167,33 @@ def _receive_futures(futures, min_get, max_get, timeout):
         limit = len(futures) if max_get is None else max_get
         pending = set(futures)
         completed = []
-        deadline_timer = None
+        # The deadline is anchored at call entry and bounds the entire call,
+        # including the wait for the first result.
+        deadline_timer = (
+            None
+            if timeout is None
+            else timeit("async_env_future_batch_deadline").start()
+        )
         while pending and len(completed) < min_get:
-            done, pending = wait(pending, return_when=FIRST_COMPLETED)
+            if deadline_timer is None:
+                done, pending = wait(pending, return_when=FIRST_COMPLETED)
+            else:
+                remaining = timeout - deadline_timer.elapsed()
+                # A non-positive remaining still polls: wait(timeout=0)
+                # returns whatever is already done without blocking.
+                done, pending = wait(
+                    pending, timeout=max(remaining, 0.0), return_when=FIRST_COMPLETED
+                )
+                if not done:
+                    # Completed futures stay in the pool's pending lists (the
+                    # caller only removes returned ones), so no result is
+                    # lost and state stays consistent.
+                    raise TimeoutError(
+                        f"async recv timed out: {len(completed)}/{min_get} "
+                        f"results after {timeout}s; completed futures remain "
+                        f"available to the next call."
+                    )
             completed.extend(list(done)[: limit - len(completed)])
-            if deadline_timer is None and timeout is not None:
-                deadline_timer = timeit("async_env_future_batch_deadline").start()
         while pending and len(completed) < limit:
             done = {future for future in pending if future.done()}
             if not done: