API server: wake the message loop on results instead of a 1 ms nap

mstar/api_server/data_worker.py

@@ -149,7 +149,8 @@ def _video_frame_metadata(
 class InlineTokens:
     """One RESULT_TOKENS frame's values for the requests still known here,
     queued behind the tensor results so each request's chunks keep arrival
-    order. One item per frame, not per request. No loop indices: a chunk's
+    order. One item per frame, not per request, so the data worker wakes the
+    message thread once per step (``ResultQueue``). No loop indices: a chunk's
     order is its arrival sequence, and nothing reads them."""
     request_ids: list[str]
     values: list[int]
@@ -182,6 +183,33 @@ def _kv_page_sizes(model: Model, model_config: dict) -> dict[str, int]:
     }
 
 
+class ResultQueue(queue.Queue):
+    """What the data worker hands the API message thread: output chunks, or
+    profile updates.
+
+    A put doesn't wake that thread: ``flush`` signals ``wakeup`` once for
+    everything put since the last flush. The data worker thread flushes after
+    each item it takes and after its reads. A RESULT_TOKENS frame is one item
+    whose chunks are put as it is taken, so a token step costs the message
+    thread two passes however many rows it had, one for the frame's arrival and
+    one for its chunks. A result tensor is one request's edge, and taking it
+    only starts a read, whose chunks go out with the flush after the reads."""
+
+    def __init__(self, wakeup: EventWakeup):
+        super().__init__()
+        self.wakeup = wakeup
+        self._unsignaled = False
+
+    def put(self, item, block=True, timeout=None):
+        super().put(item, block, timeout)
+        self._unsignaled = True
+
+    def flush(self) -> None:
+        if self._unsignaled:
+            self._unsignaled = False
+            self.wakeup.signal()
+
+
 class PreprocessWorker:
     def __init__(
         self,
@@ -202,8 +230,12 @@ def __init__(
         self.abort_request_queue = queue.Queue()
         self.reads_done_queue = queue.Queue()
         self.discard_tensor_queue = queue.Queue()
-        self.output_queue = queue.Queue()
-        self.profile_queue = queue.Queue()
+        # Ends the API message thread's poll: the flushes of these two queues,
+        # and the API server when a request handler changes what that
+        # thread's next pass reads, signal it.
+        self.results_wakeup = EventWakeup()
+        self.output_queue = ResultQueue(self.results_wakeup)
+        self.profile_queue = ResultQueue(self.results_wakeup)
         # Set by report_preprocess_finish(); read by the worker thread.
         self.timing_requested = threading.Event()
         # Set by enable_worker_metrics(); the worker thread tells the conductor.
@@ -517,8 +549,8 @@ def __init__(
         self,
         in_queue: queue.Queue, # for preprocessing
         result_tensor_queue: queue.Queue, # for output streaming
-        out_queue: queue.Queue,
-        profile_queue: queue.Queue,
+        out_queue: ResultQueue,
+        profile_queue: ResultQueue,
         cleanup_request_queue: queue.Queue,
         abort_request_queue: queue.Queue,
         reads_done_queue: queue.Queue,
@@ -1641,6 +1673,8 @@ def run(self):
         if not isinstance(self.communicator, OutboxCommunicator):
             self.communicator = OutboxCommunicator(self.communicator)
         idle_wait = self._idle_wait()
+        flush_results = self.out_queue.flush
+        flush_profiles = self.profile_queue.flush
         while not self.stop_event.is_set():
             did_work = False
             sends_pending = False
@@ -1660,23 +1694,25 @@ def run(self):
                         # Nothing to read or ack: postprocess in place, in
                         # arrival order with the tensor results.
                         self._emit_inline_tokens(result)
-                        continue
-                    # Draining for teardown: don't start new reads — ack the
-                    # tensors back so the producing worker can free its buffers.
-                    if result.request_id in self._draining_rids:
+                    elif result.request_id in self._draining_rids:
+                        # Draining for teardown: don't start new reads — ack
+                        # the tensors back so the producing worker can free
+                        # its buffers.
                         self._discard_result_tensor(result)
-                        continue
-                    try:
-                        self._read_result_tensor(result)
-                    except Exception as exc:  # noqa: BLE001 — must reach the client
-                        # The read never started, so none of this edge's
-                        # tensors will ever produce a chunk; answer for all of
-                        # them at once.
-                        self._fail_request(
-                            result.request_id, exc,
-                            f"{result.modality} output transfer",
-                            count=len(result.graph_edge.tensor_info),
-                        )
+                    else:
+                        try:
+                            self._read_result_tensor(result)
+                        except Exception as exc:  # noqa: BLE001 — must reach the client
+                            # The read never started, so none of this edge's
+                            # tensors will ever produce a chunk; answer for all
+                            # of them at once.
+                            self._fail_request(
+                                result.request_id, exc,
+                                f"{result.modality} output transfer",
+                                count=len(result.graph_edge.tensor_info),
+                            )
+                    # one wake per item, and so per step
+                    flush_results()
                 did_work = self._process_releases() or did_work
                 while not self.discard_tensor_queue.empty():
                     did_work = True
@@ -1687,6 +1723,7 @@ def run(self):
                     self._cleanup_request_state(req_id)
                 did_work = self._process_read_tensors() or did_work
                 did_work = self._release_encoded() or did_work
+                flush_results()
                 # Reads may have just resolved; ACK any drains now free of them.
                 for rid in list(self._draining_rids):
                     self._complete_drain_if_ready(rid)
@@ -1707,6 +1744,10 @@ def run(self):
                         logger.exception("Input for request %s failed", pre_input.request_id)
             except Exception:
                 logger.exception("PreprocessWorkerThread error")
+            # a new input's profile update or error chunks, and anything a
+            # raise skipped
+            flush_results()
+            flush_profiles()
 
             # getattr: tests drive this loop with stand-in tensor managers
             transfers = getattr(self.tensor_manager, "transfer_stats", None)

mstar/api_server/entrypoint.py

@@ -7,6 +7,7 @@
 import itertools
 import json
 import logging
+import math
 import multiprocessing as mp
 import os
 import re
@@ -817,7 +818,12 @@ async def wait_for_request_id(self, request_id: str | None) -> None:
     def _retire(self, key: str) -> None:
         """Forget ``key``'s request_id once neither dict holds the key.
         Caller holds ``request_lock``."""
-        if key in self.pending_requests or key in self.recently_completed:
+        if key in self.recently_completed:
+            if key not in self.pending_requests:
+                # its client let go: the message thread's prune releases it
+                self._wake_message_loop()
+            return
+        if key in self.pending_requests:
             return
         request_id = self._client_ids.pop(key, None)
         if request_id is not None and self._live.get(request_id) == key:
@@ -924,6 +930,7 @@ def reject_input(self, key: str, error: str, status: int) -> bool:
             # it; the result thread reads the data worker's bookkeeping for
             # it until then.
             self.recently_completed[key] = time.time()
+            self._wake_message_loop()  # its TTL starts now
         req.notify()
         self.preprocess_worker.forward_abort(key)
         return True
@@ -932,9 +939,13 @@ def reject_input(self, key: str, error: str, status: int) -> bool:
     # Result collection (background thread)
     # ----------------------------------------------------------
 
-    def _prune_recently_completed(self) -> None:
+    def _prune_recently_completed(self) -> float | None:
+        """Release the finished requests whose outputs are all delivered, or
+        whose client let go, and fail the ones past the delivery TTL. Returns
+        when the TTL of the oldest one kept runs out (``time.time``)."""
         now = time.time()
         stale = []
+        oldest = None
         for rid, ts in self.recently_completed.items():
             # The client handler pops pending_requests on its own timeout (and
             # a streaming client after its final flush), so the entry can be
@@ -952,6 +963,8 @@ def _prune_recently_completed(self) -> None:
                 stale.append((rid, False, drained))
             elif (now - ts) >= self._recently_completed_ttl:
                 stale.append((rid, True, drained))
+            elif oldest is None or ts < oldest:
+                oldest = ts
         for rid, lost_outputs, drained in stale:
             # only set the event when there are no more pending chunks
             req = self.pending_requests.get(rid)
@@ -989,13 +1002,15 @@ def _prune_recently_completed(self) -> None:
             self.preprocess_worker.finished_reading(rid, drained=drained)
             self.recently_completed.pop(rid, None)
             self._retire(rid)
+        return None if oldest is None else oldest + self._recently_completed_ttl
 
     def _process_messages(self) -> None:
         """Drain the ZMQ pull socket and route results to pending requests.
         Also watches the conductor process. Once it exits, every pending
         request is failed and the HTTP server is told to stop."""
         self.pin_result_thread()
         next_liveness_check = 0.0
+        self.communicator.register_event_for_poll(self.preprocess_worker.results_wakeup)
         while self.running:
             now = time.monotonic()
             if now >= next_liveness_check:
@@ -1004,11 +1019,8 @@ def _process_messages(self) -> None:
                 if exited is not None:
                     self._fail_pending_for_dead_conductor(exited)
                     return
+            next_expiry = None
             try:
-                with self.request_lock:
-                    if len(self.recently_completed) > 0:
-                        self._prune_recently_completed()
-
                 for message in self.communicator.get_all_new_messages():
                     if not isinstance(message, APIServerMessage):
                         logger.warning("Unexpected message type: %s", type(message))
@@ -1225,11 +1237,30 @@ def _process_messages(self) -> None:
                         sink.loop.call_soon_threadsafe(self._relay_chunks, sink, reqs)
                     except RuntimeError:
                         pass  # the loop is closed: its clients are gone
+                # Last, so a request whose final chunk or completion came in
+                # this pass is released in this pass.
+                with self.request_lock:
+                    if self.recently_completed:
+                        next_expiry = self._prune_recently_completed()
             except Exception:
                 if self.running:
                     logger.exception("Error in message processing loop")
                     time.sleep(0.01)
-            time.sleep(0.001)
+                # the prune may not have run: the next pass goes after the back-off
+                next_expiry = time.time()
+            # Everything else that changes what a pass reads signals the poll:
+            # a message, the data worker's flush (``ResultQueue``) and a
+            # request handler (``_wake_message_loop``). Only the liveness
+            # check and the delivery TTL are on a timer.
+            wait_s = next_liveness_check - time.monotonic()
+            if next_expiry is not None:
+                wait_s = min(wait_s, next_expiry - time.time())
+            self.communicator.poll_for_messages(timeout_ms=max(0, math.ceil(wait_s * 1000)))
+
+    def _wake_message_loop(self) -> None:
+        """End the message thread's wait: a request handler changed what its
+        next pass reads."""
+        self.preprocess_worker.results_wakeup.signal()
 
     # ----------------------------------------------------------
     # Streaming helper
@@ -1714,6 +1745,7 @@ def release_submission(self, key: str | None) -> None:
     def cleanup(self) -> None:
         self.preprocess_worker.shutdown()
         self.running = False
+        self._wake_message_loop()
         if hasattr(self, "_msg_thread") and self._msg_thread.is_alive():
             self._msg_thread.join(timeout=2)
 

mstar/communication/communicator.py

@@ -52,6 +52,17 @@ def try_send_bytes(self, entity_id: str, data: bytes) -> bool:
     def get_all_new_messages(self) -> list:
         pass
 
+    @abstractmethod
+    def register_event_for_poll(self, event: EventWakeup) -> None:
+        """Have ``poll_for_messages`` (and ``wait_for_work``) also end when
+        ``event`` is signalled."""
+
+    @abstractmethod
+    def poll_for_messages(self, timeout_ms: int = 20) -> bool:
+        """Block until a message is readable, the registered event fires, or
+        ``timeout_ms`` elapses, whichever comes first. True when a message is
+        waiting for ``get_all_new_messages``; a fired event is drained."""
+
     # -- endpoint scheme (shared by every ZMQ-based communicator) ------------
     # Subclasses set ``self.protocol`` and ``self.ipc_socket_path_prefix``.
 

test/distributed/test_tensor_transport.py

@@ -39,6 +39,12 @@ def send(self, entity_id: str, msg):
     def get_all_new_messages(self):
         return []
 
+    def register_event_for_poll(self, event):
+        pass
+
+    def poll_for_messages(self, timeout_ms=20):
+        return False
+
 
 _TCP_AVAILABLE = MOONCAKE_IMPORT_ERROR is None
 _PROTOCOLS = [

test/modular/api_server_helpers.py

@@ -10,10 +10,13 @@
 Import it by module name (``from api_server_helpers import ...``): on CI the
 stdlib ``test`` package shadows the repo's ``test/`` directory.
 """
+import select
 import tempfile
+import time
 from pathlib import Path
 
 from mstar.api_server.entrypoint import APIServer
+from mstar.communication.event import EventWakeup
 
 
 def make_test_api_server(**overrides) -> APIServer:
@@ -45,3 +48,32 @@ def make_test_api_server(**overrides) -> APIServer:
             raise TypeError(f"APIServer has no field {name!r}")
         setattr(server, name, value)
     return server
+
+
+class StandInSocket:
+    """The communicator surface ``APIServer._process_messages`` uses, over a
+    test's own ``messages`` (or an override of ``get_all_new_messages``). Its
+    poll ends on the registered wakeup as the real communicators' does, and
+    otherwise after at most ``tick_ms``: messages a test queues signal
+    nothing."""
+
+    def __init__(self, messages=lambda: [], tick_ms: float = 1.0):
+        self._messages = messages
+        self.tick_ms = tick_ms
+        self.event: EventWakeup | None = None
+        self.timeouts: list[int] = []  # each poll's timeout_ms, as the loop asked
+
+    def get_all_new_messages(self) -> list:
+        return self._messages()
+
+    def register_event_for_poll(self, event: EventWakeup) -> None:
+        self.event = event
+
+    def poll_for_messages(self, timeout_ms: int = 20) -> bool:
+        self.timeouts.append(timeout_ms)
+        timeout_s = min(timeout_ms, self.tick_ms) / 1000
+        if self.event is None:
+            time.sleep(timeout_s)
+        elif select.select([self.event.fd], [], [], timeout_s)[0]:
+            self.event.drain()
+        return False

test/modular/test_api_completion_guard.py

@@ -12,8 +12,9 @@
 import threading
 from types import SimpleNamespace
 
-from mstar.api_server.data_worker import PreprocessWorker, PreprocessWorkerThread
+from mstar.api_server.data_worker import PreprocessWorker, PreprocessWorkerThread, ResultQueue
 from mstar.api_server.request_types import PreprocessInput
+from mstar.communication.event import EventWakeup
 
 
 def _worker_with(received_chunk_counts):
@@ -46,7 +47,8 @@ def process_prompt(self, *args, **kwargs):
     wt = PreprocessWorkerThread.__new__(PreprocessWorkerThread)
     wt.in_queue = queue.Queue()
     wt._is_pending = lambda key: True
-    wt.out_queue = queue.Queue()
+    wt.out_queue = ResultQueue(EventWakeup())
+    wt.profile_queue = ResultQueue(EventWakeup())
     wt.result_tensor_queue = queue.Queue()
     wt.cleanup_request_queue = queue.Queue()
     wt.abort_request_queue = queue.Queue()