Update

[ghstack-poisoned]

Author: vmoens

test/test_inference_server.py

@@ -6,6 +6,7 @@
 
 import concurrent.futures
 import threading
+import time
 
 import pytest
 import torch
@@ -21,6 +22,7 @@
     InferenceTransport,
     MPTransport,
     RayTransport,
+    SlotTransport,
     ThreadingTransport,
 )
 from torchrl.modules.inference_server._monarch import MonarchTransport
@@ -606,8 +608,6 @@ def test_weight_sync_init_called(self):
 
         with InferenceServer(policy, transport, weight_sync=ws):
             # Give the worker thread a moment to start
-            import time
-
             time.sleep(0.1)
             assert ws.initialized_on_receiver
             assert ws.synchronized_on_receiver
@@ -634,8 +634,6 @@ def test_weight_update_applied(self):
             ws.push(new_weights)
 
             # Give the server loop a chance to apply the update
-            import time
-
             time.sleep(0.2)
 
             # Now inference should reflect zero weights
@@ -662,8 +660,6 @@ def test_inference_continues_after_weight_update(self):
             new_weights = TensorDict.from_module(policy)
             ws.push(new_weights)
 
-            import time
-
             time.sleep(0.1)
 
             # Continue making requests
@@ -868,3 +864,120 @@ def postproc(td):
             pass
         collector.shutdown()
         assert called["count"] >= 1
+
+
+# =============================================================================
+# Tests: SlotTransport
+# =============================================================================
+
+
+class TestSlotTransport:
+    def test_single_request(self):
+        transport = SlotTransport(num_slots=4)
+        policy = _make_policy()
+        with InferenceServer(policy, transport, max_batch_size=4):
+            client = transport.client()
+            td = TensorDict({"observation": torch.randn(4)})
+            result = client(td)
+            assert "action" in result.keys()
+            assert result["action"].shape == (2,)
+
+    def test_concurrent_actors(self):
+        """Multiple threads submit concurrently via slot clients."""
+        n_actors = 4
+        n_requests = 30
+        transport = SlotTransport(num_slots=n_actors)
+        policy = _make_policy()
+
+        results_per_actor: list[list[TensorDictBase]] = [[] for _ in range(n_actors)]
+        clients = [transport.client() for _ in range(n_actors)]
+
+        def actor_fn(actor_id):
+            for _ in range(n_requests):
+                td = TensorDict({"observation": torch.randn(4)})
+                result = clients[actor_id](td)
+                results_per_actor[actor_id].append(result)
+
+        with InferenceServer(policy, transport, max_batch_size=n_actors):
+            with concurrent.futures.ThreadPoolExecutor(max_workers=n_actors) as pool:
+                futs = [pool.submit(actor_fn, i) for i in range(n_actors)]
+                concurrent.futures.wait(futs)
+                for f in futs:
+                    f.result()
+
+        for actor_results in results_per_actor:
+            assert len(actor_results) == n_requests
+            for r in actor_results:
+                assert "action" in r.keys()
+                assert r["action"].shape == (2,)
+
+    def test_too_many_clients_raises(self):
+        """Creating more clients than slots raises RuntimeError."""
+        transport = SlotTransport(num_slots=2)
+        transport.client()
+        transport.client()
+        with pytest.raises(RuntimeError, match="slots"):
+            transport.client()
+
+    def test_submit_raises(self):
+        """Direct submit() on SlotTransport is not supported."""
+        transport = SlotTransport(num_slots=1)
+        td = TensorDict({"observation": torch.randn(4)})
+        with pytest.raises(NotImplementedError):
+            transport.submit(td)
+
+    def test_exception_propagates(self):
+        """Model exceptions propagate through SlotTransport."""
+
+        def bad_model(td):
+            raise ValueError("slot model error")
+
+        transport = SlotTransport(num_slots=1)
+        with InferenceServer(bad_model, transport, max_batch_size=4):
+            client = transport.client()
+            td = TensorDict({"observation": torch.randn(4)})
+            with pytest.raises(ValueError, match="slot model error"):
+                client(td)
+
+
+# =============================================================================
+# Tests: min_batch_size
+# =============================================================================
+
+
+class TestMinBatchSize:
+    def test_min_batch_size_accumulates(self):
+        """With min_batch_size > 1, the server waits for enough items."""
+        min_bs = 4
+        seen_sizes = []
+
+        def tracking_collate(items):
+            seen_sizes.append(len(items))
+            return lazy_stack(items)
+
+        transport = ThreadingTransport()
+        policy = _make_policy()
+        n = 8
+
+        with InferenceServer(
+            policy,
+            transport,
+            max_batch_size=16,
+            min_batch_size=min_bs,
+            collate_fn=tracking_collate,
+            timeout=1.0,
+        ):
+            client = transport.client()
+            # Submit items from threads to give the server time to accumulate
+            with concurrent.futures.ThreadPoolExecutor(max_workers=n) as pool:
+                futs = [
+                    pool.submit(
+                        lambda: client(TensorDict({"observation": torch.randn(4)}))
+                    )
+                    for _ in range(n)
+                ]
+                for f in futs:
+                    f.result(timeout=10.0)
+
+        # At least one batch should have >= min_batch_size items
+        assert any(s >= min_bs for s in seen_sizes)

torchrl/modules/inference_server/_monarch.py

@@ -4,103 +4,50 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-import queue
 import threading
-import time
-from typing import Any
 
-from tensordict.base import TensorDictBase
+from torchrl.modules.inference_server._queue_transport import (
+    _QueueInferenceClient,
+    QueueBasedTransport,
+)
 
-from torchrl.modules.inference_server._transport import InferenceTransport
 
-_SENTINEL = object()
+class _MonarchRequestQueue:
+    """Wrapper around ``MonarchQueue`` that signals a :class:`threading.Event` on put.
 
+    Also adapts the Monarch queue API (``get(block=False)``) to the standard
+    ``get_nowait()`` expected by :class:`QueueBasedTransport`.
+    """
 
-class _MonarchFuture:
-    """Future-like object for Monarch transport results.
+    def __init__(self, monarch_queue, has_work: threading.Event):
+        self._queue = monarch_queue
+        self._has_work = has_work
 
-    Args:
-        client: the :class:`_MonarchInferenceClient` that created this future.
-        req_id: the unique request identifier within that client.
-    """
+    def put(self, item):
+        self._queue.put(item)
+        self._has_work.set()
 
-    def __init__(self, client: _MonarchInferenceClient, req_id: int):
-        self._client = client
-        self._req_id = req_id
-        self._result: Any = _SENTINEL
+    def get(self, timeout=None):
+        return self._queue.get(timeout=timeout)
 
-    def done(self) -> bool:
-        """Return ``True`` if the result is available without blocking."""
-        if self._result is not _SENTINEL:
-            return True
-        try:
-            self._result = self._client._get_result(self._req_id, timeout=0)
-        except queue.Empty:
-            return False
-        return True
-
-    def result(self, timeout: float | None = None) -> TensorDictBase:
-        """Block until the result is available."""
-        if self._result is _SENTINEL:
-            self._result = self._client._get_result(self._req_id, timeout=timeout)
-        if isinstance(self._result, BaseException):
-            raise self._result
-        return self._result
-
-
-class _MonarchInferenceClient:
-    """Actor-side client for :class:`MonarchTransport`.
-
-    Each client owns a dedicated response queue and routes results by
-    request-id.
-
-    Args:
-        request_queue: the shared Monarch queue for requests.
-        response_queue: this client's dedicated response queue.
-        actor_id: the unique identifier assigned by the transport.
-    """
+    def get_nowait(self):
+        return self._queue.get(block=False)
+
+
+class _MonarchResponseQueue:
+    """Thin wrapper adapting the MonarchQueue get API."""
+
+    def __init__(self, monarch_queue):
+        self._queue = monarch_queue
 
-    def __init__(self, request_queue, response_queue, actor_id: int):
-        self._request_queue = request_queue
-        self._response_queue = response_queue
-        self._actor_id = actor_id
-        self._next_req_id = 0
-        self._buffered: dict[int, Any] = {}
-
-    def __call__(self, td: TensorDictBase) -> TensorDictBase:
-        """Submit a request and block until the result is ready."""
-        return self.submit(td).result()
-
-    def submit(self, td: TensorDictBase) -> _MonarchFuture:
-        """Submit a request and return a :class:`_MonarchFuture`."""
-        req_id = self._next_req_id
-        self._next_req_id += 1
-        self._request_queue.put((self._actor_id, req_id, td))
-        return _MonarchFuture(self, req_id)
-
-    # -- internal -------------------------------------------------------------
-
-    def _get_result(self, req_id: int, timeout: float | None = None) -> Any:
-        """Return the result for *req_id*, buffering any earlier arrivals."""
-        if req_id in self._buffered:
-            return self._buffered.pop(req_id)
-        deadline = None if timeout is None else time.monotonic() + timeout
-        while True:
-            remaining = None
-            if deadline is not None:
-                remaining = deadline - time.monotonic()
-                if remaining <= 0:
-                    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}")
-            if rid == req_id:
-                return result
-            self._buffered[rid] = result
-
-
-class MonarchTransport(InferenceTransport):
+    def put(self, item):
+        self._queue.put(item)
+
+    def get(self, timeout=None):
+        return self._queue.get(timeout=timeout)
+
+
+class MonarchTransport(QueueBasedTransport):
     """Transport using Monarch for distributed inference on GPU clusters.
 
     Uses Monarch's actor model and RDMA-capable channels for efficient
@@ -118,6 +65,7 @@ class MonarchTransport(InferenceTransport):
     """
 
     def __init__(self, *, max_queue_size: int = 1000):
+        super().__init__()
         try:
             import monarch  # noqa: F401
             from monarch.tools.queue import MonarchQueue
@@ -126,66 +74,21 @@ def __init__(self, *, max_queue_size: int = 1000):
                 "Monarch is required for MonarchTransport. "
                 "Install it following the Monarch documentation."
             )
-        self._request_queue = MonarchQueue(maxsize=max_queue_size)
-        self._response_queues: dict[int, Any] = {}
-        self._lock = threading.Lock()
-        self._next_actor_id = 0
+        self._has_work = threading.Event()
+        self._request_queue = _MonarchRequestQueue(
+            MonarchQueue(maxsize=max_queue_size), self._has_work
+        )
+        self._response_queues: dict[int, _MonarchResponseQueue] = {}
         self._MonarchQueue = MonarchQueue
 
-    # -- actor API ------------------------------------------------------------
+    def _make_response_queue(self) -> _MonarchResponseQueue:
+        return _MonarchResponseQueue(self._MonarchQueue(maxsize=1000))
 
-    def client(self) -> _MonarchInferenceClient:
+    def client(self) -> _QueueInferenceClient:
         """Create an actor-side client with a dedicated response queue.
 
         Returns:
-            A :class:`_MonarchInferenceClient` that can be passed to a Monarch
+            A :class:`_QueueInferenceClient` that can be passed to a Monarch
             actor.
         """
-        with self._lock:
-            actor_id = self._next_actor_id
-            self._next_actor_id += 1
-        response_queue = self._MonarchQueue(maxsize=1000)
-        self._response_queues[actor_id] = response_queue
-        return _MonarchInferenceClient(self._request_queue, response_queue, actor_id)
-
-    def submit(self, td: TensorDictBase):
-        """Not supported -- use :meth:`client` to obtain an actor handle."""
-        raise RuntimeError(
-            "MonarchTransport.submit() is not supported. "
-            "Call transport.client() to create a _MonarchInferenceClient."
-        )
-
-    # -- server API -----------------------------------------------------------
-
-    def drain(
-        self, max_items: int
-    ) -> tuple[list[TensorDictBase], list[tuple[int, int]]]:
-        """Dequeue up to *max_items* pending requests (non-blocking)."""
-        items: list[TensorDictBase] = []
-        callbacks: list[tuple[int, int]] = []
-        for _ in range(max_items):
-            try:
-                actor_id, req_id, td = self._request_queue.get(block=False)
-                items.append(td)
-                callbacks.append((actor_id, req_id))
-            except Exception:
-                break
-        return items, callbacks
-
-    def wait_for_work(self, timeout: float) -> None:
-        """Block until at least one request is available or *timeout* elapses."""
-        try:
-            item = self._request_queue.get(timeout=timeout)
-            self._request_queue.put(item)
-        except Exception:
-            pass
-
-    def resolve(self, callback: tuple[int, int], result: TensorDictBase) -> None:
-        """Route the result to the correct actor's response queue."""
-        actor_id, req_id = callback
-        self._response_queues[actor_id].put((req_id, result))
-
-    def resolve_exception(self, callback: tuple[int, int], exc: BaseException) -> None:
-        """Route an exception to the correct actor's response queue."""
-        actor_id, req_id = callback
-        self._response_queues[actor_id].put((req_id, exc))
+        return super().client()