From 83df73a9fdce9fc668239c1ee8025b66a93c6bd2 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 1 Nov 2026 07:24:43 +0000
Subject: [PATCH 1/3] data worker: prepare requests on a pool

---
 docs/adding_models.rst                        |  16 +
 mstar/api_server/data_worker.py               | 137 +++++++-
 mstar/model/base.py                           |   7 +
 mstar/model/qwen3_omni/qwen3_omni_model.py    |   2 +-
 mstar/utils/cpus.py                           |   7 +
 test/modular/test_data_worker_prepare_pool.py | 321 ++++++++++++++++++
 6 files changed, 474 insertions(+), 16 deletions(-)
 create mode 100644 test/modular/test_data_worker_prepare_pool.py

diff --git a/docs/adding_models.rst b/docs/adding_models.rst
index eb6d13f8f0..ebbac0b99c 100644
--- a/docs/adding_models.rst
+++ b/docs/adding_models.rst
@@ -2047,6 +2047,22 @@ torch's default. It does not size video decoding: each torchcodec decode that as
 ``media_limits.video_decode_threads`` (``Model.load_video`` does) runs FFmpeg with up to 8
 threads, no more than the CPUs the process's CPU plan leaves unpinned (``CpuPlan.left``).
 
+The data worker prepares one request at a time on its own thread, so a request's media
+decode holds up every request behind it and the output chunks of the ones running. A model
+whose ``load_image`` / ``load_audio`` / ``load_video`` and ``process_prompt`` keep no shared
+state sets ``PREPROCESS_PARALLEL = True``, and requests then prepare on a pool sized to the
+CPUs left (``CpuPlan.preprocess_workers``: at most 8, each with its own intra-op threads).
+Tensors are still stored and requests still reach the conductor from the data worker's thread,
+in arrival order. A deployment sets the pool at the top level of its config; 1 turns it off:
+
+.. code-block:: yaml
+
+   model: "qwen3_omni"
+   preprocess_workers: 2
+
+The pool only helps work that releases the GIL, so measure each stage first. A model with
+client streams always prepares inline, since an append has to follow its request.
+
 A pool that small is slow for a large copy, and a 121-frame 832x480 video is 145 MB. A
 video decoder should therefore emit its clip in the layout the encoder takes, transposed
 on the GPU inside its forward, so ``postprocess`` hands the tensor it reads straight to the
diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py
index 7f649a2850..f8cfa918be 100644
--- a/mstar/api_server/data_worker.py
+++ b/mstar/api_server/data_worker.py
@@ -9,6 +9,7 @@
 import time
 from collections import deque
 from collections.abc import Callable, Collection
+from concurrent.futures import Future
 from dataclasses import dataclass, field
 from functools import cached_property
 from typing import NamedTuple
@@ -594,6 +595,21 @@ class ClientCall:
     charged: dict[str, int] = field(default_factory=dict)
 
 
+class PreparedInput(NamedTuple):
+    """What a request's preparation hands back: its tensors, their metadata
+    and its kwargs, none of them stored or sent yet."""
+    tensors: NameToTensorList
+    input_metadata: dict[str, list[dict]]
+    model_kwargs: dict
+
+
+class PreparingInput(NamedTuple):
+    """A request on the prepare pool, queued in arrival order."""
+    input: PreprocessInput
+    profile: DataWorkerProfile | None
+    future: Future[PreparedInput]
+
+
 class PreprocessWorkerThread:
     # next time the thread publishes its transfer counts; class-level so a
     # thread built without __init__ has one
@@ -602,6 +618,9 @@ class PreprocessWorkerThread:
     worker_metrics_requested: threading.Event | None = None
     # None (a thread built without __init__) postprocesses every output inline
     _encode_pool: ThreadPool | None = None
+    # None (no model that declares it, one CPU's worth of workers, or a thread
+    # built without __init__) prepares every request inline
+    _prepare_pool: ThreadPool | None = None
     # Appends handled per pass before output delivery gets its turn again: one
     # for each of up to 32 live calls appending at once. An append costs a file
     # decode, far less than a request's preprocess, which is one per pass.
@@ -709,6 +728,23 @@ def __init__(
         # Finished encodes, appended by the pool's done callbacks and released
         # by this thread, which alone touches the per-request output state.
         self._encoded: deque = deque()
+        self._prepare_workers = self._num_prepare_workers()
+        if self._prepare_workers > 1:
+            self._prepare_pool = ThreadPool(
+                max_workers=self._prepare_workers, thread_name_prefix="mstar-prepare",
+            )
+        # Requests on the prepare pool, oldest first. They reach the conductor
+        # in this order, whichever finishes first.
+        self._preparing: deque[PreparingInput] = deque()
+
+    def _num_prepare_workers(self) -> int:
+        """Requests prepared at once. A model has to declare that its loaders
+        and ``process_prompt`` may overlap, and client streams need their
+        appends handled behind the request, so anything else runs inline."""
+        if not getattr(self.model, "PREPROCESS_PARALLEL", False) or self._client_edges:
+            return 1
+        workers = (self._model_config or {}).get("preprocess_workers")
+        return cpus.plan().preprocess_workers if workers is None else int(workers)
 
     def _cleanup_request_state(self, request_id: str, *, force: bool = False) -> None:
         """Release transport and postprocessing state owned by this thread.
@@ -774,10 +810,21 @@ def _load_media(
     def _process_input(
         self, input: PreprocessInput, profile: DataWorkerProfile | None = None,
     ):
-        file_states: dict[str, list[tuple[int, int, int]]] = {}
+        self._start_input(input)
+        self._finish_input(input, self._prepare_input(input), profile)
+
+    def _start_input(self, input: PreprocessInput) -> None:
+        """Count the request as started here. On this thread, as it leaves the
+        queue, so a release later in the queue's order finds it."""
         self._live_requests[input.request_id] = self._live_requests.get(input.request_id, 0) + 1
         self.in_flight_requests.add(input.request_id)
 
+    def _prepare_input(self, input: PreprocessInput) -> PreparedInput:
+        """Decode the request's media and run the model's prompt processing.
+        Runs on the prepare pool when there is one, so it reads only what
+        ``__init__`` fixed and writes nothing the thread owns."""
+        file_states: dict[str, list[tuple[int, int, int]]] = {}
+
         # First, load raw modality tensors from file_paths (images, audio, video)
         # so they can be passed to process_prompt() below.
         if input.file_paths and self._keyed_labels:
@@ -834,6 +881,13 @@ def _process_input(
             tensors["text_inputs"] = [torch.tensor(
                 list(byte_data), dtype=torch.uint8, device=self.device
             )]
+        return PreparedInput(tensors, input_metadata, model_kwargs)
+
+    def _finish_input(
+        self, input: PreprocessInput, prepared: PreparedInput, profile: DataWorkerProfile | None,
+    ) -> None:
+        """Store a prepared request's tensors and hand it to the conductor."""
+        tensors, input_metadata, model_kwargs = prepared
 
         # A client stream's first chunk, not a prompt signal: a one-shot upload
         # is a stream of one chunk. Stored before the request reaches the
@@ -1765,8 +1819,11 @@ def _publish_transfers(self, transfers: EdgeTransfers) -> None:
     def _next_inputs(self) -> list[PreprocessInput | ClientAppend]:
         """This pass's input: up to MAX_APPENDS_PER_PASS items in arrival
         order, ending at the first new request (a preprocess each, seconds for
-        media), so output delivery gets its turn between requests."""
+        media), so output delivery gets its turn between requests. None while
+        the prepare pool is full, so requests wait in the queue undecoded."""
         batch: list[PreprocessInput | ClientAppend] = []
+        if self._prepare_pool is not None and len(self._preparing) >= self._prepare_workers:
+            return batch
         while len(batch) < self.MAX_APPENDS_PER_PASS:
             try:
                 item = self.in_queue.get_nowait()
@@ -1797,20 +1854,66 @@ def _process_waiting_input(self, pre_input: PreprocessInput | ClientAppend) -> N
         try:
             if self.timing_requested.is_set():
                 profile = self._start_profile(pre_input)
-            self._process_input(pre_input, profile)
+            if self._prepare_pool is not None:
+                self._submit_prepare(pre_input, profile)
+            else:
+                self._process_input(pre_input, profile)
+        except Exception as exc:  # noqa: BLE001 — any failure must reach the client
+            self._input_failed(pre_input, profile, exc)
+
+    def _input_failed(
+        self, pre_input: PreprocessInput, profile: DataWorkerProfile | None, exc: Exception,
+    ) -> None:
+        if profile is not None:
+            profile.failed = True
+            self.profile_queue.put(profile)
+        # A request whose media load or prompt processing fails never
+        # reaches the conductor, so nothing downstream would ever complete
+        # it; surface the failure as an error chunk instead of leaving the
+        # client to wait out the server timeout.
+        self._fail_request(pre_input.request_id, exc, "preprocessing", abort=False)
+        # Never reached the conductor, so there are no remote readers to
+        # race: hard-drop the (possibly persisted) input signals directly.
+        self._cleanup_request_state(pre_input.request_id, force=True)
+        self._removed(pre_input.request_id)
+
+    def _submit_prepare(self, pre_input: PreprocessInput, profile: DataWorkerProfile | None) -> None:
+        """Prepare on the pool; :meth:`_release_prepared` finishes the request."""
+        self._start_input(pre_input)
+        try:
+            future = self._prepare_pool.submit(self._prepare_input, pre_input)
+        except RuntimeError as exc:  # shut down under a stopping server
+            self._input_failed(pre_input, profile, exc)
+            return
+        self._preparing.append(PreparingInput(pre_input, profile, future))
+        if self.wakeup is not None:
+            self.wakeup.register_future(future)
+
+    def _release_prepared(self) -> bool:
+        """Finish the oldest prepared request, one per pass so output delivery
+        gets its turn, as it does between inline requests. The request waits
+        on the pool for its turn behind older ones, so a conductor sees
+        requests in arrival order.
+
+        An abort that landed during preparation found nothing to stop. The API
+        server lets go of a request before it queues the abort, so the request
+        is dropped here unsent, as one still in the queue is; the conductor
+        takes an abort and a READS_DONE for a rid it never saw."""
+        if self._prepare_pool is None or not self._preparing or not self._preparing[0].future.done():
+            return False
+        pre_input, profile, future = self._preparing.popleft()
+        try:
+            if self._is_pending(pre_input.request_id):
+                self._finish_input(pre_input, future.result(), profile)
+            else:
+                self._cleanup_request_state(pre_input.request_id, force=True)
+                self._removed(pre_input.request_id)
+                if profile is not None:
+                    profile.preprocess_finish_time = time.perf_counter()
+                    self.profile_queue.put(profile)
         except Exception as exc:  # noqa: BLE001 — any failure must reach the client
-            if profile is not None:
-                profile.failed = True
-                self.profile_queue.put(profile)
-            # A request whose media load or prompt processing fails never
-            # reaches the conductor, so nothing downstream would ever complete
-            # it; surface the failure as an error chunk instead of leaving the
-            # client to hit the server timeout.
-            self._fail_request(pre_input.request_id, exc, "preprocessing", abort=False)
-            # Never reached the conductor, so there are no remote readers to
-            # race: hard-drop the (possibly persisted) input signals directly.
-            self._cleanup_request_state(pre_input.request_id, force=True)
-            self._removed(pre_input.request_id)
+            self._input_failed(pre_input, profile, exc)
+        return True
 
     def _idle_wait(self):
         """How a pass that found nothing waits: on the communicator's poll,
@@ -1912,6 +2015,8 @@ def run(self):
                         self._process_waiting_input(pre_input)
                     except Exception:  # noqa: BLE001 — the pass's other input still goes
                         logger.exception("Input for request %s failed", pre_input.request_id)
+                if not sends_pending:
+                    did_work = self._release_prepared() or did_work
             except Exception:
                 logger.exception("PreprocessWorkerThread error")
             # a new input's profile update or error chunks, and anything a
@@ -1935,3 +2040,5 @@ def run(self):
             self._hard_cleanup(request_id)
         if self._encode_pool is not None:
             self._encode_pool.shutdown(wait=False, cancel_futures=True)
+        if self._prepare_pool is not None:
+            self._prepare_pool.shutdown(wait=False, cancel_futures=True)
diff --git a/mstar/model/base.py b/mstar/model/base.py
index d18d59b98e..a46cd7d53e 100644
--- a/mstar/model/base.py
+++ b/mstar/model/base.py
@@ -462,6 +462,13 @@ class Model(ABC):
     # the model's config.yaml overrides either. See
     # ``entrypoint._set_preprocess_threads``.
     PREPROCESS_TORCH_THREADS: int | None = None
+    # ``load_image`` / ``load_audio`` / ``load_video`` and ``process_prompt``
+    # may run for several requests at once. The data worker then prepares
+    # requests on a pool (``CpuPlan.preprocess_workers``; ``preprocess_workers``
+    # in the model's config.yaml overrides it) and still hands them to the
+    # conductor in arrival order. Declare it only once those calls are checked
+    # for shared state, and for the GIL: work that holds it gains nothing here.
+    PREPROCESS_PARALLEL: bool = False
 
     def _get_worker_graphs_for_graph_walk(
         self,
diff --git a/mstar/model/qwen3_omni/qwen3_omni_model.py b/mstar/model/qwen3_omni/qwen3_omni_model.py
index d1ea5739e7..77a06deb75 100644
--- a/mstar/model/qwen3_omni/qwen3_omni_model.py
+++ b/mstar/model/qwen3_omni/qwen3_omni_model.py
@@ -184,7 +184,7 @@ class Qwen3OmniModel(Model):
     """Qwen3-Omni: Thinker + Talker + Code2Wav 3-partition streaming model."""
 
     LENGTH_KWARGS = ("max_output_tokens", "talker_max_output_tokens")
-
+    PREPROCESS_PARALLEL = True
 
     def __init__(
         self,
diff --git a/mstar/utils/cpus.py b/mstar/utils/cpus.py
index 860d28cc30..e80cd90fa6 100644
--- a/mstar/utils/cpus.py
+++ b/mstar/utils/cpus.py
@@ -153,6 +153,13 @@ def preprocess_threads(self) -> int:
         MB."""
         return max(1, min(4, self.left // 4))
 
+    @property
+    def preprocess_workers(self) -> int:
+        """Requests the data worker prepares at once, for a model that
+        declares ``PREPROCESS_PARALLEL``. Each runs ``preprocess_threads`` of
+        its own, so the pool times that stays within the CPUs left."""
+        return max(1, min(8, self.left // self.preprocess_threads))
+
     @property
     def png_threads(self) -> int:
         """Threads one PNG deflates on (``image_io.uint8_to_png``)."""
diff --git a/test/modular/test_data_worker_prepare_pool.py b/test/modular/test_data_worker_prepare_pool.py
new file mode 100644
index 0000000000..645588634d
--- /dev/null
+++ b/test/modular/test_data_worker_prepare_pool.py
@@ -0,0 +1,321 @@
+"""A model that declares PREPROCESS_PARALLEL prepares requests on the data worker's pool:
+they reach the conductor in arrival order, output delivery runs meanwhile, and a failed or
+abandoned request ends as it does inline. A model that doesn't declare it runs as before."""
+from __future__ import annotations
+
+import queue
+import threading
+import time
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from mstar.api_server.data_worker import InlineTokens, PreprocessWorkerThread, ResultQueue
+from mstar.api_server.request_types import PreprocessInput
+from mstar.communication.event import EventWakeup
+from mstar.model.base import ProcessPromptOutput
+from mstar.streaming.topology import PartitionTopology
+from mstar.utils import cpus
+from mstar.utils.ipc_format import ConductorMessageType
+
+
+class _Model:
+    """process_prompt blocks on the gate named by the prompt, so a test decides
+    when each preparation finishes."""
+    PREPROCESS_PARALLEL = True
+
+    def __init__(self):
+        self.gates: dict[str, threading.Event] = {}
+        self.threads: dict[str, str] = {}
+        self.overlap = 0
+        self._running = 0
+        self._lock = threading.Lock()
+
+    def get_partition_topology(self):
+        return PartitionTopology(partitions=["default"])
+
+    def kv_labels(self):
+        return {}
+
+    def gate(self, text: str) -> threading.Event:
+        return self.gates.setdefault(text, threading.Event())
+
+    def process_prompt(self, text, input_modalities, output_modalities, **kwargs):
+        self.threads[text] = threading.current_thread().name
+        with self._lock:
+            self._running += 1
+            self.overlap = max(self.overlap, self._running)
+        try:
+            if not self.gate(text).wait(timeout=10):
+                raise TimeoutError(text)
+        finally:
+            with self._lock:
+                self._running -= 1
+        if text.startswith("bad"):
+            raise ValueError("bad prompt")
+        return ProcessPromptOutput({"text_inputs": [torch.tensor([1])]}, {"seen": text})
+
+
+class _Serial(_Model):
+    PREPROCESS_PARALLEL = False
+
+
+class _Tensors:
+    def __init__(self):
+        self.stored: list[tuple[str, str]] = []
+        self.dropped: list[str] = []
+
+    def store_and_return_tensor_info(self, request_id, tensors):
+        self.stored.append((request_id, threading.current_thread().name))
+        return {}
+
+    def register_for_send(self, request_id, tensor_infos):
+        pass
+
+    def cleanup_request(self, request_id):
+        pass
+
+    def force_cleanup_request(self, request_id):
+        self.dropped.append(request_id)
+
+    def has_inflight_reads(self, request_id):
+        return False
+
+    def get_ready_tensors(self):
+        return {}
+
+    def take_failed_reads(self):
+        return []
+
+
+class _Comm:
+    # the outbox the run loop wraps this in encodes with it: identity, so sends arrive as objects
+    codec = SimpleNamespace(encode=lambda msg: msg)
+
+    def __init__(self):
+        self.sent: list = []
+
+    def get_all_new_messages(self):
+        return []
+
+    def send(self, entity, msg):
+        self.sent.append(msg)
+
+    def try_send_bytes(self, entity, msg):
+        self.sent.append(msg)
+        return True
+
+
+def _thread(model, config=None, pending=None) -> PreprocessWorkerThread:
+    return PreprocessWorkerThread(
+        in_queue=queue.Queue(), result_tensor_queue=queue.Queue(), out_queue=ResultQueue(EventWakeup()),
+        profile_queue=ResultQueue(EventWakeup()), cleanup_request_queue=queue.Queue(),
+        abort_request_queue=queue.Queue(), reads_done_queue=queue.Queue(),
+        discard_tensor_queue=queue.Queue(), stop_event=threading.Event(), communicator=_Comm(),
+        tensor_manager=_Tensors(), model=model, model_config=config, is_pending=pending,
+    )
+
+
+@pytest.fixture
+def worker(monkeypatch):
+    monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
+    wt = _thread(_Model())
+    yield wt
+    for gate in wt.model.gates.values():
+        gate.set()
+    wt._prepare_pool.shutdown(wait=True)
+    wt._encode_pool.shutdown(wait=True)
+
+
+def _input(request_id: str, text: str | None = None) -> PreprocessInput:
+    return PreprocessInput(
+        request_id=request_id, text=text or request_id, file_paths=None,
+        input_modalities=["text"], output_modalities=["text"], model_kwargs={},
+    )
+
+
+def _sent(wt) -> list[str]:
+    comm = getattr(wt.communicator, "_inner", wt.communicator)  # run() wraps it in an outbox
+    return [
+        m.body.request_id for m in comm.sent
+        if m.message_type == ConductorMessageType.NEW_REQUEST
+    ]
+
+
+def _release_until(wt, n: int, timeout_s: float = 5.0) -> list[str]:
+    deadline = threading.Event()
+    timer = threading.Timer(timeout_s, deadline.set)
+    timer.start()
+    try:
+        while len(_sent(wt)) < n and not deadline.is_set():
+            wt._release_prepared()
+    finally:
+        timer.cancel()
+    return _sent(wt)
+
+
+@pytest.mark.parametrize(("usable", "workers"), [(1, 1), (3, 3), (8, 4), (16, 4), (40, 8), (256, 8)])
+def test_the_pool_times_its_intra_op_threads_fits_the_cpus(usable, workers):
+    plan = cpus.CpuPlan(frozenset(range(256)), usable)
+    assert plan.preprocess_workers == workers
+    assert workers * plan.preprocess_threads <= usable
+
+
+def test_only_a_model_that_declares_it_gets_a_pool(monkeypatch):
+    monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
+    serial, parallel = _thread(_Serial()), _thread(_Model())
+    assert serial._prepare_pool is None
+    assert parallel._prepare_workers == 4 and parallel._prepare_pool is not None
+    parallel._prepare_pool.shutdown()
+
+
+def test_one_cpu_prepares_inline(monkeypatch):
+    monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset({0}), usable=1))
+    assert _thread(_Model())._prepare_pool is None
+
+
+def test_a_deployment_sets_the_pool_or_turns_it_off(monkeypatch):
+    monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
+    assert _thread(_Model(), {"preprocess_workers": 1})._prepare_pool is None
+    wt = _thread(_Model(), {"preprocess_workers": 2})
+    assert wt._prepare_workers == 2
+    wt._prepare_pool.shutdown()
+
+
+def test_a_model_with_client_streams_prepares_inline(monkeypatch):
+    class Streams(_Model):
+        def get_partition_topology(self):
+            topology = super().get_partition_topology()
+            topology.client_connections = lambda: [
+                type("Conn", (), {"edge_name": "mic", "to_partition": "default"})(),
+            ]
+            return topology
+
+    monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
+    assert _thread(Streams())._prepare_pool is None
+
+
+def test_a_model_without_the_declaration_runs_on_the_thread(monkeypatch):
+    monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
+    wt = _thread(_Serial())
+    wt.model.gate("a").set()
+    wt._process_waiting_input(_input("a"))
+    assert wt.model.threads["a"] == threading.current_thread().name
+    assert _sent(wt) == ["a"]
+
+
+def test_requests_overlap_and_reach_the_conductor_in_arrival_order(worker):
+    worker._process_waiting_input(_input("a"))
+    worker._process_waiting_input(_input("b"))
+    worker._process_waiting_input(_input("c"))
+    # c and b finish first
+    worker.model.gate("c").set()
+    worker.model.gate("b").set()
+    worker._preparing[2].future.result(timeout=5)
+    worker._preparing[1].future.result(timeout=5)
+    assert worker._release_prepared() is False and _sent(worker) == []
+    worker.model.gate("a").set()
+    assert _release_until(worker, 3) == ["a", "b", "c"]
+    assert all(name.startswith("mstar-prepare") for name in worker.model.threads.values())
+
+
+def test_the_request_is_stored_and_sent_on_the_data_workers_thread(worker):
+    worker.model.gate("a").set()
+    worker._process_waiting_input(_input("a"))
+    _release_until(worker, 1)
+    assert worker.tensor_manager.stored == [("a", threading.current_thread().name)]
+    assert worker.request_model_kwargs["a"]["seen"] == "a"
+
+
+def test_a_full_pool_leaves_requests_in_the_queue_undecoded(worker):
+    for i in range(6):
+        worker.in_queue.put(_input(str(i)))
+    while (batch := worker._next_inputs()):
+        worker._process_waiting_input(batch[0])
+    assert len(worker._preparing) == worker._prepare_workers == 4
+    assert worker.in_queue.qsize() == 2
+    for i in range(6):
+        worker.model.gate(str(i)).set()
+    _release_until(worker, 1)
+    assert len(worker._next_inputs()) == 1
+
+
+def test_a_failed_preparation_answers_in_its_turn_and_leaves_nothing(worker):
+    worker.model.gate("bad").set()
+    worker.model.gate("b").set()
+    worker._process_waiting_input(_input("a", "bad"))
+    worker._process_waiting_input(_input("b"))
+    assert _release_until(worker, 1) == ["b"]
+    chunk = worker.out_queue.get_nowait()
+    assert (chunk.request_id, chunk.modality, chunk.metadata["status"]) == ("a", "error", 400)
+    assert "a" not in worker._live_requests and "a" not in worker.in_flight_requests
+    assert worker.tensor_manager.dropped == ["a"]
+
+
+def test_a_request_abandoned_while_it_prepares_is_never_sent(monkeypatch):
+    monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
+    let_go: set[str] = set()
+    wt = _thread(_Model(), pending=lambda key: key not in let_go)
+    wt._process_waiting_input(_input("a"))
+    # the API server lets go first, then queues the abort
+    let_go.add("a")
+    wt.abort_request_queue.put("a")
+    wt._process_releases()
+    wt.model.gate("a").set()
+    wt._preparing[0].future.result(timeout=5)
+    assert wt._release_prepared() is True
+    assert _sent(wt) == []
+    assert "a" not in wt._live_requests and "a" not in wt._draining_rids and "a" not in wt._reads_done_sent
+    assert wt.tensor_manager.stored == [] and wt.tensor_manager.dropped == ["a"]
+    wt._prepare_pool.shutdown(wait=True)
+
+
+def test_a_release_queued_before_the_request_applies_to_the_earlier_one(worker):
+    worker._live_requests["a"] = 1  # an earlier request under the rid
+    worker.reads_done_queue.put(("a", True))
+    worker._process_releases()
+    worker._process_waiting_input(_input("a"))
+    assert worker._live_requests["a"] == 2
+    assert "a" in worker._reads_done_sent
+
+
+def test_the_run_loop_delivers_outputs_while_a_request_prepares(monkeypatch):
+    monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
+    wt = _thread(_Model())
+    delivered = threading.Event()
+    wt._emit_inline_tokens = lambda result: delivered.set()
+    wt.in_queue.put(_input("slow"))
+    thread = threading.Thread(target=wt.run, daemon=True)
+    thread.start()
+    try:
+        deadline = time.monotonic() + 5
+        while "slow" not in wt.model.threads and time.monotonic() < deadline:
+            time.sleep(0.001)
+        assert "slow" in wt.model.threads
+        wt.result_tensor_queue.put(InlineTokens(["other"], [1], "text"))
+        assert delivered.wait(timeout=5), "an output waited behind a request's preparation"
+        assert _sent(wt) == []
+        wt.model.gate("slow").set()
+        deadline = time.monotonic() + 5
+        while not _sent(wt) and time.monotonic() < deadline:
+            time.sleep(0.001)
+        assert _sent(wt) == ["slow"]
+    finally:
+        wt.model.gate("slow").set()
+        wt.stop_event.set()
+        thread.join(timeout=5)
+    assert not thread.is_alive()
+
+
+def test_a_stopping_thread_refuses_a_request_it_can_no_longer_prepare(worker):
+    worker._prepare_pool.shutdown(wait=True)
+    worker._process_waiting_input(_input("a"))
+    chunk = worker.out_queue.get_nowait()
+    assert (chunk.request_id, chunk.modality) == ("a", "error")
+    assert "a" not in worker._live_requests
+
+
+def test_a_thread_built_without_init_prepares_inline():
+    wt = PreprocessWorkerThread.__new__(PreprocessWorkerThread)
+    assert wt._prepare_pool is None and wt._release_prepared() is False

From 645a0a6001590f8cfc1bd7b266c7bbee02f911f7 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 1 Nov 2026 18:26:27 -0800
Subject: [PATCH 2/3] data worker: size the prepare pool by the threads this
 process runs

A config's preprocess_torch_threads above the plan's default put the pool over the CPUs left. The pool now divides by torch's thread count, which is set before the API server is built, and serving.rst gets a preprocess_workers row.
---
 docs/adding_models.rst                        |  4 +++-
 docs/serving.rst                              |  5 +++++
 mstar/api_server/data_worker.py               |  6 +++++-
 mstar/model/base.py                           |  5 +++--
 mstar/utils/cpus.py                           | 13 +++++++++----
 test/modular/test_data_worker_prepare_pool.py | 19 +++++++++++++++++++
 6 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/docs/adding_models.rst b/docs/adding_models.rst
index ebbac0b99c..12c9fa4814 100644
--- a/docs/adding_models.rst
+++ b/docs/adding_models.rst
@@ -2051,7 +2051,9 @@ The data worker prepares one request at a time on its own thread, so a request's
 decode holds up every request behind it and the output chunks of the ones running. A model
 whose ``load_image`` / ``load_audio`` / ``load_video`` and ``process_prompt`` keep no shared
 state sets ``PREPROCESS_PARALLEL = True``, and requests then prepare on a pool sized to the
-CPUs left (``CpuPlan.preprocess_workers``: at most 8, each with its own intra-op threads).
+CPUs left (``CpuPlan.prepare_workers``: at most 8, and the pool times the process's intra-op
+thread count stays within the CPUs left, so a higher ``preprocess_torch_threads`` means fewer
+workers).
 Tensors are still stored and requests still reach the conductor from the data worker's thread,
 in arrival order. A deployment sets the pool at the top level of its config; 1 turns it off:
 
diff --git a/docs/serving.rst b/docs/serving.rst
index be13f258a8..992d40e1c6 100644
--- a/docs/serving.rst
+++ b/docs/serving.rst
@@ -350,6 +350,11 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key
      - *(optional)* Async scheduling for the parallel nodes: ``true``
        (default), ``false`` or a list of node names (see ``MSTAR_TP_ASYNC_SCHED``, which
        overrides it).
+   * - ``preprocess_workers``
+     - *(optional)* Requests the data worker prepares at once, for a model that
+       declares ``PREPROCESS_PARALLEL``. Default: as many as fit the CPUs left
+       beside the API server's intra-op threads (``preprocess_torch_threads``),
+       at most 8; ``1`` prepares inline. See :doc:`adding_models`.
    * - ``knobs``
      - *(optional)* Engine knobs for this deployment, by name: ``device_loopback``,
        ``sampler_ingraph_scatter``, ``sampler_slots``, ``ingraph_decode_rope``,
diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py
index f8cfa918be..a6d33cf84c 100644
--- a/mstar/api_server/data_worker.py
+++ b/mstar/api_server/data_worker.py
@@ -744,7 +744,11 @@ def _num_prepare_workers(self) -> int:
         if not getattr(self.model, "PREPROCESS_PARALLEL", False) or self._client_edges:
             return 1
         workers = (self._model_config or {}).get("preprocess_workers")
-        return cpus.plan().preprocess_workers if workers is None else int(workers)
+        if workers is not None:
+            return int(workers)
+        # each request runs on this process's intra-op pool, which the config
+        # or the model may size above the plan's default
+        return cpus.plan().prepare_workers(torch.get_num_threads())
 
     def _cleanup_request_state(self, request_id: str, *, force: bool = False) -> None:
         """Release transport and postprocessing state owned by this thread.
diff --git a/mstar/model/base.py b/mstar/model/base.py
index a46cd7d53e..dd11a48c8d 100644
--- a/mstar/model/base.py
+++ b/mstar/model/base.py
@@ -464,8 +464,9 @@ class Model(ABC):
     PREPROCESS_TORCH_THREADS: int | None = None
     # ``load_image`` / ``load_audio`` / ``load_video`` and ``process_prompt``
     # may run for several requests at once. The data worker then prepares
-    # requests on a pool (``CpuPlan.preprocess_workers``; ``preprocess_workers``
-    # in the model's config.yaml overrides it) and still hands them to the
+    # requests on a pool (``CpuPlan.prepare_workers`` at this process's intra-op
+    # thread count; ``preprocess_workers`` in the model's config.yaml overrides
+    # it) and still hands them to the
     # conductor in arrival order. Declare it only once those calls are checked
     # for shared state, and for the GIL: work that holds it gains nothing here.
     PREPROCESS_PARALLEL: bool = False
diff --git a/mstar/utils/cpus.py b/mstar/utils/cpus.py
index e80cd90fa6..f110dcfbd6 100644
--- a/mstar/utils/cpus.py
+++ b/mstar/utils/cpus.py
@@ -153,12 +153,17 @@ def preprocess_threads(self) -> int:
         MB."""
         return max(1, min(4, self.left // 4))
 
+    def prepare_workers(self, threads: int) -> int:
+        """Requests the data worker prepares at once, for a model that
+        declares ``PREPROCESS_PARALLEL``, when each runs on ``threads``
+        intra-op threads: at most 8, and the pool times the threads stays
+        within the CPUs left."""
+        return max(1, min(8, self.left // max(1, threads)))
+
     @property
     def preprocess_workers(self) -> int:
-        """Requests the data worker prepares at once, for a model that
-        declares ``PREPROCESS_PARALLEL``. Each runs ``preprocess_threads`` of
-        its own, so the pool times that stays within the CPUs left."""
-        return max(1, min(8, self.left // self.preprocess_threads))
+        """``prepare_workers`` at the default thread count."""
+        return self.prepare_workers(self.preprocess_threads)
 
     @property
     def png_threads(self) -> int:
diff --git a/test/modular/test_data_worker_prepare_pool.py b/test/modular/test_data_worker_prepare_pool.py
index 645588634d..e464f7c512 100644
--- a/test/modular/test_data_worker_prepare_pool.py
+++ b/test/modular/test_data_worker_prepare_pool.py
@@ -162,14 +162,33 @@ def test_the_pool_times_its_intra_op_threads_fits_the_cpus(usable, workers):
     assert workers * plan.preprocess_threads <= usable
 
 
+@pytest.mark.parametrize(("usable", "threads", "workers"), [(8, 4, 2), (8, 8, 1), (16, 1, 8), (40, 4, 8), (3, 0, 3)])
+def test_a_higher_thread_count_shrinks_the_pool(usable, threads, workers):
+    plan = cpus.CpuPlan(frozenset(range(256)), usable)
+    assert plan.prepare_workers(threads) == workers
+    assert workers * max(1, threads) <= usable
+
+
 def test_only_a_model_that_declares_it_gets_a_pool(monkeypatch):
     monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
+    monkeypatch.setattr(torch, "get_num_threads", lambda: 2)
     serial, parallel = _thread(_Serial()), _thread(_Model())
     assert serial._prepare_pool is None
     assert parallel._prepare_workers == 4 and parallel._prepare_pool is not None
     parallel._prepare_pool.shutdown()
 
 
+def test_the_pool_follows_the_threads_this_process_runs(monkeypatch):
+    # ``preprocess_torch_threads`` may set the intra-op pool above the plan's default
+    monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
+    monkeypatch.setattr(torch, "get_num_threads", lambda: 4)
+    wt = _thread(_Model())
+    assert wt._prepare_workers == 2
+    wt._prepare_pool.shutdown()
+    monkeypatch.setattr(torch, "get_num_threads", lambda: 8)
+    assert _thread(_Model())._prepare_pool is None
+
+
 def test_one_cpu_prepares_inline(monkeypatch):
     monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset({0}), usable=1))
     assert _thread(_Model())._prepare_pool is None

From 2b8047bf9f91703430c6e0aa564d0843e3840229 Mon Sep 17 00:00:00 2001
From: merceod <atj10@stanford.edu>
Date: Sun, 1 Nov 2026 18:40:07 -0800
Subject: [PATCH 3/3] tests: the prepare pool tests pin the intra-op thread
 count

The pool now divides by torch's thread count, so the module fixes it at 2: 8 usable CPUs give 4 workers on any runner.
---
 test/modular/test_data_worker_prepare_pool.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/test/modular/test_data_worker_prepare_pool.py b/test/modular/test_data_worker_prepare_pool.py
index e464f7c512..a781db4c0a 100644
--- a/test/modular/test_data_worker_prepare_pool.py
+++ b/test/modular/test_data_worker_prepare_pool.py
@@ -117,6 +117,13 @@ def _thread(model, config=None, pending=None) -> PreprocessWorkerThread:
     )
 
 
+@pytest.fixture(autouse=True)
+def _two_intra_op_threads(monkeypatch):
+    # the pool divides the CPUs left by this process's intra-op threads, so
+    # 8 usable CPUs give 4 workers wherever the tests run
+    monkeypatch.setattr(torch, "get_num_threads", lambda: 2)
+
+
 @pytest.fixture
 def worker(monkeypatch):
     monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
@@ -171,7 +178,6 @@ def test_a_higher_thread_count_shrinks_the_pool(usable, threads, workers):
 
 def test_only_a_model_that_declares_it_gets_a_pool(monkeypatch):
     monkeypatch.setattr(cpus, "_plan", cpus.CpuPlan(frozenset(range(8)), usable=8))
-    monkeypatch.setattr(torch, "get_num_threads", lambda: 2)
     serial, parallel = _thread(_Serial()), _thread(_Model())
     assert serial._prepare_pool is None
     assert parallel._prepare_workers == 4 and parallel._prepare_pool is not None