perf(api-server): stop keeping finished sequences, keep prompt rendering and tokenizing off the event loop

atom/config.py

@@ -2103,6 +2103,12 @@ class Config:
     compilation_config: CompilationConfig = field(default_factory=CompilationConfig)
     quant_config: QuantizationConfig = field(init=False)
     asyncio_mode: bool = False
+    # Whether CoreManager also queues finished sequences on ``outputs_queue``
+    # for ``LLMEngine.step`` / ``generate`` and the asyncio path to read. Online
+    # front ends deliver every output through per-sequence callbacks and build
+    # the engine with this off; left on there, every finished Sequence stays
+    # on a queue nothing reads, for the life of the process.
+    keep_finished_outputs: bool = True
     mark_trace: bool = False
     load_dummy: str | None = None
     enable_expert_parallel: bool = False

atom/entrypoints/atomesh/server.py

@@ -101,7 +101,8 @@ def initialize_engine(args: argparse.Namespace) -> tuple[Any, Any]:
 
     logger.info("Initializing engine with model %s...", args.model)
     engine_args = EngineArgs.from_cli_args(args)
-    engine = engine_args.create_engine(tokenizer=tokenizer)
+    # Outputs reach requests through the per-sequence callbacks.
+    engine = engine_args.create_engine(tokenizer=tokenizer, keep_finished_outputs=False)
     return engine, tokenizer
 
 

atom/model_engine/arg_utils.py

@@ -945,6 +945,17 @@ def _get_engine_kwargs(self) -> dict:
 
         return kwargs
 
-    def create_engine(self, tokenizer=None) -> LLMEngine:
-        """Create and return an LLMEngine instance with the configured parameters."""
-        return LLMEngine(self.model, tokenizer=tokenizer, **self._get_engine_kwargs())
+    def create_engine(
+        self, tokenizer=None, *, keep_finished_outputs: bool = True
+    ) -> LLMEngine:
+        """Create and return an LLMEngine instance with the configured parameters.
+
+        ``keep_finished_outputs=False`` is for online front ends, which take
+        every output through per-sequence callbacks (``Config.keep_finished_outputs``).
+        """
+        return LLMEngine(
+            self.model,
+            tokenizer=tokenizer,
+            keep_finished_outputs=keep_finished_outputs,
+            **self._get_engine_kwargs(),
+        )

atom/model_engine/engine_core_mgr.py

@@ -231,6 +231,10 @@ def _init_shared_state(
         )
         self.ctx = zmq.Context(io_threads=2)
         self.outputs_queue = queue.Queue[list[Sequence]]()
+        # Off for online front ends, which read every output through the
+        # per-sequence callbacks: nothing would drain the queue
+        # (`Config.keep_finished_outputs`).
+        self.keep_finished_outputs = getattr(config, "keep_finished_outputs", True)
         # Utility replies are matched to their caller by request id, not by
         # position on a shared queue, where a late or unasked-for reply became
         # the next caller's answer.
@@ -834,7 +838,8 @@ def process_outputs_socket():
                                     f"{self.label}: flush_stream_batch failed: {e}",
                                     exc_info=True,
                                 )
-                        self.outputs_queue.put_nowait(seqs)
+                        if self.keep_finished_outputs:
+                            self.outputs_queue.put_nowait(seqs)
                         # These completions freed slots on their ranks — re-feed
                         # from the pending pool so a fast rank never drains to
                         # empty while peers still have work. No-op when
@@ -1306,6 +1311,11 @@ def _bp_fail_stalled_run(self) -> None:
         ``LLMEngine.generate`` waits while the engine is alive and any request
         is unfinished. A dropped seq never finishes, so that wait never ends.
         The seqs are still in ``_bp_pending``; this only surfaces the error.
+
+        Only offline seqs are ever in the pool (``add_request`` keeps online
+        requests out of it), and offline ``generate`` reads ``outputs_queue``.
+        If online seqs could wait there, this error would have to reach them
+        through their callbacks instead.
         """
         pending = len(getattr(self, "_bp_pending", ()))
         logger.error(

tests/entrypoints/test_api_compatibility.py

@@ -1052,6 +1052,8 @@ async def generate(*args, **kwargs):
             "text": "partial" if partial else "",
             "finish_reason": reason,
             "num_cached_tokens": 0,
+            "num_tokens_input": 2,
+            "num_tokens_output": 1 if partial else 0,
         }
 
     monkeypatch.setattr(api_server, "generate_async", generate)

tests/entrypoints/test_api_server_helpers.py

@@ -75,7 +75,7 @@ def add_cli_args(cls, parser):
         def from_cli_args(cls, args):
             return cls()
 
-        def create_engine(self, tokenizer=None):
+        def create_engine(self, tokenizer=None, **kwargs):
             return None
 
     _try_import_else_stub(
@@ -876,6 +876,85 @@ async def drive():
     assert events == ["initialize", "serving", "shutdown"]
 
 
+class _FakeSeq:
+    def __init__(self, seq_id, token_ids):
+        self.id = seq_id
+        self.prompt_token_ids = list(token_ids)
+        self.num_prompt_tokens = len(token_ids)
+
+
+def test_prompt_is_checked_and_truncated_in_the_preprocess_call(monkeypatch):
+    # /v1/messages and /v1/responses: no encode on the event loop. The length
+    # check reads the engine's own token count from the one preprocess call,
+    # and only an over-long prompt is built again, from its first ids.
+    requests = {}
+    monkeypatch.setattr(
+        api_server,
+        "engine",
+        SimpleNamespace(io_processor=SimpleNamespace(requests=requests)),
+    )
+    monkeypatch.setattr(
+        api_server,
+        "tokenizer",
+        SimpleNamespace(decode=lambda ids, **_kw: f"<{len(ids)} ids>"),
+    )
+    built = []
+
+    def preprocess(prompt):
+        tokens = list(range(10)) if isinstance(prompt, str) else prompt
+        seq = _FakeSeq(len(built), tokens)
+        built.append(prompt)
+        requests[seq.id] = seq
+        return seq
+
+    cap = api_server._PromptCap(16)
+    seq = api_server._preprocess_prompt(preprocess, "short", cap)
+    assert built == ["short"] and seq.num_prompt_tokens == 10
+    assert cap.cut_text is None
+    seq = api_server._preprocess_prompt(preprocess, "short", None)
+    assert seq.num_prompt_tokens == 10
+
+    built.clear()
+    requests.clear()
+    cap = api_server._PromptCap(4)
+    seq = api_server._preprocess_prompt(preprocess, "long prompt", cap)
+    assert built == ["long prompt", [0, 1, 2, 3]]
+    assert seq.prompt_token_ids == [0, 1, 2, 3]
+    # The text the engine runs, for the reasoning channel's seed.
+    assert cap.cut_text == "<4 ids>"
+    # The discarded first build is not left registered.
+    assert list(requests) == [seq.id]
+
+
+def test_responses_channel_follows_the_text_the_engine_runs(monkeypatch):
+    # Structured: cut before the engine call (off the loop), so the channel
+    # the sampling params are built from reads the kept text. Unstructured:
+    # re-seeded from `cut_text` after the preprocess call cut it.
+    monkeypatch.setattr(
+        api_server,
+        "tokenizer",
+        SimpleNamespace(
+            encode=lambda text, **_kw: list(range(len(text))),
+            decode=lambda ids, **_kw: "x" * len(ids),
+        ),
+    )
+    cap = api_server._PromptCap(4)
+    assert api_server._cut_prompt("abc", cap) == ("abc", "abc")
+    assert api_server._cut_prompt("abcdefgh", cap) == ([0, 1, 2, 3], "xxxx")
+
+    seeds = []
+    monkeypatch.setattr(
+        api_server, "prompt_starts_in_reasoning", lambda text: seeds.append(text)
+    )
+    monkeypatch.setattr(
+        api_server, "reasoning_channel", lambda start, **_kw: ("reseeded", start)
+    )
+    assert api_server._reseed_after_cut("orig", cap, {}) == "orig"
+    cap.cut_text = "kept"
+    assert api_server._reseed_after_cut("orig", cap, {})[0] == "reseeded"
+    assert seeds == ["kept"]
+
+
 class TestDeploymentDefaultOutputCap:
     @pytest.mark.parametrize("explicit", [False, True])
     def test_small_deployment_distinguishes_default_from_explicit(

tests/entrypoints/test_pd_messages_responses.py

@@ -49,19 +49,35 @@ def server(monkeypatch):
     """Stub the engine and record what each endpoint hands to it."""
     calls = SimpleNamespace(generate=[], setup=[], template="<rendered prompt>")
 
+    def encode(text, **_kw):
+        return [0] * (LONG_PROMPT_TOKENS if text == "<long prompt>" else 2)
+
+    def engine_prompt(prompt, kwargs):
+        # What preprocess builds: the prompt, cut to `prompt_cap.tokens`
+        # (shown as `<n>` for an n-token cut) when it is longer.
+        cap = kwargs.get("prompt_cap")
+        limit = None if cap is None else cap.tokens
+        n = len(encode(prompt))
+        return (f"<{limit}>", limit) if limit is not None and n > limit else (prompt, n)
+
     async def fake_generate_async(prompt, _sampling, _request_id, **kwargs):
-        calls.generate.append((prompt, kwargs))
-        output = {"text": "hello", "finish_reason": "eos", "num_cached_tokens": 0}
+        sent, n = engine_prompt(prompt, kwargs)
+        calls.generate.append((sent, kwargs))
+        output = {
+            "text": "hello",
+            "finish_reason": "eos",
+            "num_cached_tokens": 0,
+            "num_tokens_input": n,
+            "num_tokens_output": 1,
+        }
         if kwargs.get("kv_transfer_params") is not None:
             output["kv_transfer_output_meta_info"] = PREFILL_OUTPUT
         yield output
 
     async def fake_setup_streaming_request(prompt, _sampling, _request_id, **kwargs):
-        calls.setup.append((prompt, kwargs))
-        return 1, _FinishedCollector(), 2
-
-    def encode(text, **_kw):
-        return [0] * (LONG_PROMPT_TOKENS if text == "<long prompt>" else 2)
+        sent, n = engine_prompt(prompt, kwargs)
+        calls.setup.append((sent, kwargs))
+        return 1, _FinishedCollector(), n
 
     monkeypatch.setattr(api_server, "model_name", "m")
     monkeypatch.setattr(api_server, "_request_logger", None)

tests/test_engine_core_mgr_shared_state.py

@@ -76,6 +76,7 @@ def _bare_manager(cls, **kwargs):
             dp_load_balance="round_robin",
             tp_world_size=1,
             prefill_context_parallel_size=1,
+            keep_finished_outputs=kwargs.get("keep_finished_outputs", True),
         ),
         label=cls.__name__,
         local_engine_count=kwargs.get("local_engine_count", 2),
@@ -130,6 +131,33 @@ def test_shutdown_frame_marks_unhealthy_before_cleanup(self):
         finally:
             sender.close(linger=0)
 
+    def test_finished_sequences_are_queued_only_if_the_config_keeps_them(self):
+        # Online front ends build the engine with keep_finished_outputs=False
+        # and take outputs through the callbacks; otherwise every finished
+        # Sequence would stay queued for the life of the process.
+        for keep in (True, False):
+            mgr = _bare_manager(CoreManager, keep_finished_outputs=keep)
+            self.addCleanup(mgr.ctx.term)
+            mgr._deliver_terminal_output = lambda seq: False
+            mgr._release_seq_load = lambda seq_id: None
+            mgr._shutdown_engine_core_rank = lambda rank: None
+            url = f"inproc://outputs-{uuid.uuid4().hex}"
+            receiver = mgr.ctx.socket(zmq.PULL)
+            receiver.bind(url)
+            sender = mgr.ctx.socket(zmq.PUSH)
+            sender.connect(url)
+            thread = mgr._create_output_thread(0, receiver, url + "-shutdown")
+            thread.start()
+            try:
+                seqs = [SimpleNamespace(id=7)]
+                sender.send(pickle.dumps((EngineCoreRequestType.ADD, seqs)))
+                sender.send(pickle.dumps((EngineCoreRequestType.SHUTDOWN, None)))
+                thread.join(timeout=3)
+                self.assertFalse(thread.is_alive())
+                self.assertEqual(mgr.outputs_queue.qsize(), 1 if keep else 0)
+            finally:
+                sender.close(linger=0)
+
     def _terminate(self):
         for mgr in self.managers:
             mgr.ctx.term()
@@ -189,3 +217,54 @@ def test_subclass_does_not_override_the_shared_initialiser(self):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+def test_create_engine_hands_keep_finished_outputs_to_the_config(monkeypatch):
+    # Decided at construction: create_engine hands the choice to Config, and
+    # the default keeps the offline / asyncio behaviour.
+    import dataclasses
+
+    from atom.config import Config
+    from atom.model_engine import arg_utils
+
+    default = {f.name: f.default for f in dataclasses.fields(Config)}
+    assert default["keep_finished_outputs"] is True
+    seen = {}
+
+    def fake_engine(model, tokenizer=None, **kwargs):
+        seen.update(kwargs)
+
+    monkeypatch.setattr(arg_utils, "LLMEngine", fake_engine)
+    monkeypatch.setattr(arg_utils.EngineArgs, "_get_engine_kwargs", lambda self: {})
+    args = arg_utils.EngineArgs.__new__(arg_utils.EngineArgs)
+    args.model = "m"
+    args.create_engine(keep_finished_outputs=False)
+    assert seen == {"keep_finished_outputs": False}
+    args.create_engine()
+    assert seen == {"keep_finished_outputs": True}
+
+
+def test_atomesh_front_end_builds_the_engine_without_kept_outputs(monkeypatch):
+    # An online front end: every output reaches its request through the
+    # callbacks, so it must not ask the engine to keep finished sequences.
+    from atom.entrypoints.atomesh import server as atomesh_server
+    from atom.model_engine import arg_utils, llm_engine
+
+    seen = {}
+
+    class FakeArgs:
+        @classmethod
+        def from_cli_args(cls, args):
+            return cls()
+
+        def create_engine(self, **kwargs):
+            seen.update(kwargs)
+            return "engine"
+
+    monkeypatch.setattr(arg_utils, "EngineArgs", FakeArgs)
+    monkeypatch.setattr(llm_engine, "load_tokenizer", lambda *a: "tok")
+    engine, _ = atomesh_server.initialize_engine(
+        SimpleNamespace(model="m", trust_remote_code=False)
+    )
+    assert engine == "engine"
+    assert seen == {"tokenizer": "tok", "keep_finished_outputs": False}