diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index ca3a791763..dfc9b5cec8 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -287,13 +287,22 @@ def _prepare_multimodal_inputs( return inputs["input_ids"][0].tolist(), multimodal_data +# ── Batched stream dispatch ────────────────────────────────────────────── +# Per-seq `call_soon_threadsafe` floods the API event loop at high batch size +# (one call per token). Instead the callback only decodes + buffers; the mgr +# flushes a whole step with one scheduled call (see `flush_stream_batch`). +import threading as _threading # noqa: E402 + +_stream_batch_tls = _threading.local() + + def _send_stream_chunk_direct( request_output: RequestOutput, request_id: str, stream_queue: asyncio.Queue, loop: AbstractEventLoop, ) -> None: - """Send stream chunk directly to the queue.""" + """Decode the chunk and buffer it; dispatch happens in flush_stream_batch.""" global tokenizer new_text = tokenizer.decode(request_output.output_tokens, skip_special_tokens=True) @@ -309,7 +318,34 @@ def _send_stream_chunk_direct( } if getattr(request_output, "kv_transfer_params_output", None): chunk_data["kv_transfer_params"] = request_output.kv_transfer_params_output - loop.call_soon_threadsafe(stream_queue.put_nowait, chunk_data) + + buf = getattr(_stream_batch_tls, "buf", None) + if buf is None: + buf = _stream_batch_tls.buf = [] + buf.append((loop, stream_queue, chunk_data)) + + +def _drain_batch_into_queues(items: list) -> None: + """Runs ON the event loop: push each chunk into its per-request queue. + One scheduled call handles a whole step's worth of chunks.""" + for _loop, q, chunk in items: + q.put_nowait(chunk) + + +def flush_stream_batch() -> None: + """Flush a step's buffered chunks: one call_soon_threadsafe per loop + (normally one — all requests on a rank share the API loop).""" + buf = getattr(_stream_batch_tls, "buf", None) + if not buf: + return + _stream_batch_tls.buf = [] + # Group by loop (normally a single loop). dict preserves insertion order + # so per-request chunk ordering within the step is maintained. + by_loop: Dict[AbstractEventLoop, list] = {} + for loop, q, chunk in buf: + by_loop.setdefault(loop, []).append((loop, q, chunk)) + for loop, items in by_loop.items(): + loop.call_soon_threadsafe(_drain_batch_into_queues, items) def _send_stream_chunk_tagged( diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index 51b053c123..df7ae294e8 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -45,6 +45,8 @@ def __init__(self, config: Config): self.stream_outputs_queue = queue.Queue() self.utility_response_queue = queue.Queue() self._seq_id_to_callback = {} + # Batched stream-flush hook, resolved lazily (avoids import cycle). + self._flush_stream_batch_fn = None self.engine_core_processes = [] self.input_sockets = [] self.output_sockets = [] @@ -230,18 +232,16 @@ def process_outputs_socket(): f"{self.label}: Received STREAM message with {len(stream_outputs)} outputs" ) self.stream_outputs_queue.put_nowait(stream_outputs) - # Also call callbacks if registered + # Run per-seq callbacks (decode + buffer), then flush + # the whole step in one scheduled call — avoids a + # per-token call_soon_threadsafe storm on the API loop. + any_callback = False for seq_id, request_output in stream_outputs: callback = self._seq_id_to_callback.get(seq_id) - logger.debug( - f"{self.label}: seq_id={seq_id}, callback={'found' if callback is not None else 'NOT FOUND'}, tokens={request_output.output_tokens}" - ) if callback is not None: + any_callback = True try: callback(request_output) - logger.debug( - f"{self.label}: Successfully called callback for seq_id={seq_id}" - ) except Exception as e: logger.warning( f"Error calling stream_callback for sequence {seq_id}: {e}", @@ -249,8 +249,12 @@ def process_outputs_socket(): ) if request_output.finished: self._seq_id_to_callback.pop(seq_id, None) - logger.debug( - f"{self.label}: Cleaned up callback for finished sequence {seq_id}" + if any_callback: + try: + self._flush_stream_batch() + except Exception: + logger.exception( + f"{self.label}: flush_stream_batch failed" ) elif request_type == EngineCoreRequestType.UTILITY_RESPONSE: self.utility_response_queue.put_nowait(data) @@ -413,6 +417,20 @@ def add_request(self, seqs: List[Sequence]): copy=False, ) + def _flush_stream_batch(self): + """Flush this step's buffered stream chunks (see flush_stream_batch). + Resolved lazily to avoid an import cycle; no-op without the entrypoint.""" + fn = self._flush_stream_batch_fn + if fn is None: + try: + from atom.entrypoints.openai.api_server import flush_stream_batch + + fn = self._flush_stream_batch_fn = flush_stream_batch + except Exception: + self._flush_stream_batch_fn = lambda: None # resolve to no-op + return + fn() + def get_stream_outputs(self): try: return self.stream_outputs_queue.get_nowait()