# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. """AsyncLLM is the main-process async frontend. Owns request intake, per-request state, scheduler IPC, and the output-dispatch loop. Inherits from ``EngineClient`` (explicit structural conformance) and ``SchedulerControlClient`` (scheduler control-plane helpers). """ import asyncio import copy import logging import os import signal import sys import threading import time import uuid from collections import deque from collections.abc import Awaitable from enum import Enum from http import HTTPStatus from typing import ( Any, Generic, TypeVar, ) import uvloop from tokenspeed.runtime.configs.model_config import ModelConfig from tokenspeed.runtime.engine.aio_rwlock import RWLock from tokenspeed.runtime.engine.collector import RequestOutputCollector from tokenspeed.runtime.engine.core_client import EngineCoreClient from tokenspeed.runtime.engine.exceptions import EngineGenerateError from tokenspeed.runtime.engine.input_processor import InputProcessor from tokenspeed.runtime.engine.io_struct import ( AbortReq, BatchEmbeddingOut, BatchStrOut, BatchTokenIDOut, CloseSessionReqInput, ConfigureLoggingReq, EmbeddingReqInput, FlushCacheReqInput, FlushCacheReqOutput, GenerateReqInput, GetLoadReqInput, HealthCheckOutput, OpenSessionReqInput, OpenSessionReqOutput, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, UpdateWeightFromDiskReqInput, UpdateWeightFromDiskReqOutput, WatchLoadUpdateReq, ) from tokenspeed.runtime.engine.output_processor import OutputProcessor, ReqState from tokenspeed.runtime.engine.parallel_sampling import ( prepare_parallel_sampling_replica, prepare_prefix_warmup, ) from tokenspeed.runtime.engine.protocol import EngineClient from tokenspeed.runtime.engine.scheduler_control_client import ( SchedulerControlClient, ) from tokenspeed.runtime.metrics.collector import RequestMetrics from tokenspeed.runtime.pd.utils import ( DisaggregationMode, KVClassType, TransferBackend, get_kv_class, ) from tokenspeed.runtime.utils import ( dataclass_to_string_truncated, get_colorful_logger, ) from tokenspeed.runtime.utils.dispatch import TypeBasedDispatcher from tokenspeed.runtime.utils.exceptions import get_exception_traceback from tokenspeed.runtime.utils.hf_transformers_utils import get_tokenizer from tokenspeed.runtime.utils.process import kill_process_tree from tokenspeed.runtime.utils.server_args import PortArgs, ServerArgs asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) logger = get_colorful_logger(__name__) def _ignore_health_check_output(_: HealthCheckOutput) -> None: return None class ServerStatus(Enum): Up = "Up" Starting = "Starting" UnHealthy = "UnHealthy" Crashed = "Crashed" class AsyncLLM(SchedulerControlClient, EngineClient): """Main-process async frontend for the tokenspeed runtime. Owns request intake, per-request state, scheduler IPC, and the output-dispatch loop. Structurally satisfies :class:`EngineClient` via the explicit inheritance declaration above. """ def __init__( self, server_args: ServerArgs, port_args: PortArgs, ): # Parse args self.server_args = server_args self.enable_metrics = server_args.enable_metrics self.log_requests = server_args.enable_log_requests self.log_requests_level = server_args.log_requests_level self.logger = logger # Init inter-process communication (scheduler IPC owned by EngineCoreClient). self.engine_core_client = EngineCoreClient(port_args) # Read model args self.model_path = server_args.model self.served_model_name = server_args.served_model_name self.model_config = ModelConfig( server_args.model, trust_remote_code=server_args.trust_remote_code, revision=server_args.revision, context_length=server_args.max_model_len, model_override_args=server_args.hf_overrides, dtype=server_args.dtype, quantization=server_args.quantization, server_args=server_args, ) self.is_generation = self.model_config.is_generation self.is_image_gen = self.model_config.is_image_gen self.context_len = self.model_config.context_len self.image_token_id = self.model_config.image_token_id # Create tokenizer. The engine never preprocesses images -- the SMG # gateway ships precomputed multimodal inputs -- so even multimodal # models only need the tokenizer, not the full HF AutoProcessor. if server_args.skip_tokenizer_init: self.tokenizer = None else: self.tokenizer = get_tokenizer( server_args.tokenizer, tokenizer_mode=server_args.tokenizer_mode, trust_remote_code=server_args.trust_remote_code, revision=server_args.revision, architectures=self.model_config.hf_config.architectures, ) if self.model_config.is_multimodal: os.environ["TOKENIZERS_PARALLELISM"] = "false" # Store states self.no_create_loop = False self.rid_to_state: dict[str, ReqState] = {} self.gracefully_exit = False self.last_receive_tstamp = 0 self.dump_requests_folder = "" # By default do not dump self.dump_requests_threshold = 1000 self.dump_request_list: list[tuple] = [] self.log_request_metadata = self.get_log_request_metadata() self.server_status = ServerStatus.Starting # The event to notify the weight sync is finished. self.model_update_lock = RWLock() self.model_update_result: Awaitable[UpdateWeightFromDiskReqOutput] | None = None self.asyncio_tasks = set() # For session info self.session_futures = {} # session_id -> asyncio event # Set after scheduler is initialized self.max_req_input_len = None self.metrics = RequestMetrics( labels={ "model_name": self.server_args.served_model_name, "app_key": self.server_args.app_key, }, enabled=( self.enable_metrics and "prometheus" in (server_args.metrics_reporters or []) ), ) self.output_processor = OutputProcessor(self) self._result_dispatcher = TypeBasedDispatcher( [ ( ( BatchStrOut, BatchEmbeddingOut, BatchTokenIDOut, ), self.output_processor.handle_batch_output, ), (OpenSessionReqOutput, self._handle_open_session_req_output), ( UpdateWeightFromDiskReqOutput, self._handle_update_weights_from_disk_req_output, ), (HealthCheckOutput, _ignore_health_check_output), ] ) self.disaggregation_mode = DisaggregationMode( self.server_args.disaggregation_mode ) self.transfer_backend = TransferBackend( self.server_args.disaggregation_transfer_backend ) # for disaggregation, start kv bootstrap server on prefill if self.disaggregation_mode == DisaggregationMode.PREFILL: # only start bootstrap server on prefill tm kv_bootstrap_server_class = get_kv_class( self.transfer_backend, KVClassType.BOOTSTRAP_SERVER ) self.bootstrap_server = kv_bootstrap_server_class( self.server_args.disaggregation_bootstrap_port ) self.init_communicators(server_args) # Tokenization lives in :class:`InputProcessor`; see # :meth:`_tokenize_one_request` for the delegation. self.input_processor = InputProcessor(self) async def generate_request( self, obj: GenerateReqInput | EmbeddingReqInput, ): created_time = time.time() self.auto_create_handle_loop() self.input_processor.validate_request(obj) obj.normalize_batch_and_arguments() if self.log_requests: max_length, skip_names, _ = self.log_request_metadata logger.info( "Receive: obj=%s", dataclass_to_string_truncated(obj, max_length, skip_names=skip_names), ) async with self.model_update_lock.reader_lock: is_single = obj.is_single if is_single: tokenized_obj = await self._tokenize_one_request(obj) self._send_one_request(obj, tokenized_obj, created_time) async for response in self._wait_one_response(obj): yield response else: async for response in self._handle_batch_request(obj, created_time): yield response async def _tokenize_one_request( self, obj: GenerateReqInput | EmbeddingReqInput, ) -> TokenizedGenerateReqInput | TokenizedEmbeddingReqInput: """Delegate to :class:`InputProcessor`. The tokenization body lives in ``InputProcessor.tokenize_one_request``. If the input-side surface ever grows (e.g. multimodal routing), it happens in that module — not here. """ return await self.input_processor.tokenize_one_request(obj) def _send_one_request( self, obj: GenerateReqInput | EmbeddingReqInput, tokenized_obj: TokenizedGenerateReqInput | TokenizedEmbeddingReqInput, created_time: float | None = None, ): state = ReqState( RequestOutputCollector(), False, asyncio.Event(), obj, created_time=created_time, tokenized_time=tokenized_obj.created_time, ) self.rid_to_state[obj.rid] = state mm_inputs = getattr(tokenized_obj, "multimodal_inputs", None) if mm_inputs is not None: mm_inputs.publish_shm_features() self.engine_core_client.send_to_scheduler.send_pyobj(tokenized_obj) def submit_encode(self, encode_request) -> None: """Send an EPD encode request to the encode-worker scheduler subprocess. Fire-and-forget over the same scheduler-input channel the LM uses (the encode subprocess runs ``run_encode_loop``, not the LM EventLoop). The encode worker runs the vision tower and ships the embeddings to prefill over Mooncake; the gateway gets no streamed response, only the gRPC ack. ``encode_request`` is an ``encode_worker.EncodeRequest`` whose multimodal tensors are CPU (gateway-reconstructed) and pickle over ZMQ directly; shm publishing of the pixels is a follow-up optimization. """ self.engine_core_client.send_to_scheduler.send_pyobj(encode_request) async def _wait_one_response( self, obj: GenerateReqInput | EmbeddingReqInput, ): """Wait for the response of one request. Cancellation contract: callers (FastAPI route handlers, the sync ``LLM`` bridge, RL-trainer drivers, etc.) signal client disconnect via ``asyncio.CancelledError`` — not via a polled ``request.is_disconnected()`` check. If the task driving this generator is cancelled mid-wait, the ``finally`` below drops the rid from ``rid_to_state`` and fires an ``AbortReq`` at the scheduler so no per-request state leaks. """ state = self.rid_to_state[obj.rid] try: while True: await state.event.wait() out = state.collector.take() if out is None: state.event.clear() continue if state.finished: if self.log_requests: max_length, skip_names, out_skip_names = ( self.log_request_metadata ) if self.model_config.is_multimodal_gen: msg = f"Finish: obj={dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}" else: if ( isinstance(obj, GenerateReqInput) and obj.input_ids is not None and obj.text is None ): if self.tokenizer is not None: obj.text = self.tokenizer.decode( obj.input_ids, skip_special_tokens=getattr( obj.sampling_params, "skip_special_tokens", False, ), ) else: obj.text = "" msg = f"Finish: obj={dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}, out={dataclass_to_string_truncated(out, max_length, skip_names=out_skip_names)}" logger.info(msg) del self.rid_to_state[obj.rid] # Check if this was an abort/error created by scheduler if isinstance(out["meta_info"].get("finish_reason"), dict): finish_reason = out["meta_info"]["finish_reason"] if finish_reason.get("type") == "abort": if ( finish_reason.get("status_code") == HTTPStatus.BAD_REQUEST ): raise EngineGenerateError(finish_reason["message"]) yield out break state.event.clear() if state.collector.has_pending(): state.event.set() if obj.stream: yield out # else: non-stream path falls through and waits for the # final chunk; external cancellation wakes us via # ``asyncio.CancelledError`` from ``state.event.wait()``. finally: # Idempotent cleanup split on ``state.finished``: # # * Normal-finish path: the yield loop above already did # ``del self.rid_to_state[obj.rid]`` at ``state.finished``. # We defensively ``pop`` with a default so a second exit # through ``finally`` (e.g. when the yield itself raised # after the del) is a no-op. # # * Abandoned path (CancelledError / unexpected exception): # the rid is still in ``rid_to_state``. Call # ``abort_request`` which **both** removes it from the # state map **and** sends ``AbortReq`` to the scheduler. # Ordering matters: ``abort_request`` early-returns if # the rid is already gone, so we must not pop first. if state.finished: self.rid_to_state.pop(obj.rid, None) else: self.abort_request(obj.rid) async def _handle_batch_request( self, obj: GenerateReqInput | EmbeddingReqInput, created_time: float | None = None, ): batch_size = obj.batch_size generators = [] rids = [] if getattr(obj, "parallel_sample_num", 1) == 1: # Send all requests for i in range(batch_size): tmp_obj = obj[i] tokenized_obj = await self._tokenize_one_request(tmp_obj) self._send_one_request(tmp_obj, tokenized_obj, created_time) generators.append(self._wait_one_response(tmp_obj)) rids.append(tmp_obj.rid) else: # Batched parallel sampling still follows a conservative path and # can be slower than duplicating requests explicitly. if batch_size > 128: logger.warning( "Sending a single large batch with parallel sampling (n > 1) has not been well optimized. " "The performance might be better if you just duplicate the requests n times or use " "many threads to send them one by one with parallel sampling (n > 1)." ) # Tokenize all requests objs = [obj[i] for i in range(batch_size)] tokenized_objs = await self.input_processor.tokenize_batch(objs) # Cache the common prefix for parallel sampling for i in range(batch_size): tmp_obj = copy.copy(objs[i]) warmup_obj = prepare_prefix_warmup(tmp_obj, tokenized_objs[i]) self._send_one_request(tmp_obj, warmup_obj, created_time) await self._wait_one_response(tmp_obj).__anext__() # Expand requests, assign new rids for them, and send them for i in range(batch_size): for _ in range(obj.parallel_sample_num): tmp_obj = copy.copy(objs[i]) replica_obj = prepare_parallel_sampling_replica( tmp_obj, tokenized_objs[i] ) self._send_one_request(tmp_obj, replica_obj, created_time) generators.append(self._wait_one_response(tmp_obj)) rids.append(tmp_obj.rid) # Wait for all requests is_stream = hasattr(obj, "stream") and obj.stream if not is_stream: outputs = await asyncio.gather(*(gen.__anext__() for gen in generators)) yield outputs else: rid_to_index = {rid: i for i, rid in enumerate(rids)} task_map = {asyncio.create_task(gen.__anext__()): gen for gen in generators} while task_map: done, _ = await asyncio.wait( task_map.keys(), return_when=asyncio.FIRST_COMPLETED ) for task in done: gen = task_map.pop(task) try: result = task.result() result["index"] = rid_to_index[result["meta_info"]["id"]] yield result new_task = asyncio.create_task(gen.__anext__()) task_map[new_task] = gen except StopAsyncIteration: pass async def flush_cache(self) -> FlushCacheReqOutput: return (await self.flush_cache_communicator(FlushCacheReqInput()))[0] def abort_request(self, rid: str): if rid not in self.rid_to_state: return del self.rid_to_state[rid] req = AbortReq(rid) self.engine_core_client.send_to_scheduler.send_pyobj(req) async def update_weights_from_disk( self, obj: UpdateWeightFromDiskReqInput, ) -> tuple[bool, str, Any]: self.auto_create_handle_loop() # default the load format to the server_args if obj.load_format is None: obj.load_format = self.server_args.load_format logger.info("Start update_weights. Load format=%s", obj.load_format) # Hold the lock if it is not async. This means that weight sync # cannot run while requests are in progress. async with self.model_update_lock.writer_lock: return await self._wait_for_model_update_from_disk(obj) async def _wait_for_model_update_from_disk( self, obj: UpdateWeightFromDiskReqInput ) -> tuple[bool, str]: self.engine_core_client.send_to_scheduler.send_pyobj(obj) self.model_update_result = asyncio.Future() if not self.server_args.mapping.attn.has_dp: result = await self.model_update_result if result.success: self.served_model_name = obj.model_path self.server_args.model = obj.model_path self.server_args.load_format = obj.load_format self.model_path = obj.model_path return result.success, result.message, result.num_paused_requests else: # self.server_args.mapping.has_attn_dp self.model_update_tmp = [] result = await self.model_update_result all_success = all([r.success for r in result]) if all_success is True: self.server_args.model = obj.model_path self.server_args.load_format = obj.load_format self.model_path = obj.model_path all_message = [r.message for r in result] all_message = " | ".join(all_message) all_paused_requests = [r.num_paused_requests for r in result] return all_success, all_message, all_paused_requests async def open_session(self, obj: OpenSessionReqInput) -> str | None: self.auto_create_handle_loop() if obj.session_id is None: obj.session_id = uuid.uuid4().hex elif obj.session_id in self.session_futures: return None self.engine_core_client.send_to_scheduler.send_pyobj(obj) self.session_futures[obj.session_id] = asyncio.Future() session_id = await self.session_futures[obj.session_id] del self.session_futures[obj.session_id] return session_id async def close_session(self, obj: CloseSessionReqInput) -> None: await self.engine_core_client.send_to_scheduler.send_pyobj(obj) async def watch_load_thread(self): # Only for dp_controller when dp_size > 1 if ( not self.server_args.mapping.attn.has_dp or self.server_args.load_balance_method == "round_robin" ): return while True: await asyncio.sleep(self.server_args.load_watch_interval) loads = await self.get_load_communicator(GetLoadReqInput()) load_udpate_req = WatchLoadUpdateReq(loads=loads) self.engine_core_client.send_to_scheduler.send_pyobj(load_udpate_req) def get_log_request_metadata(self): max_length = None skip_names = None out_skip_names = None if self.log_requests: if self.log_requests_level == 0: max_length = 1 << 30 skip_names = set( [ "text", "input_ids", "input_embeds", "image_data", "audio_data", "precomputed_multimodal_inputs", "input_multi_ids", ] ) out_skip_names = set( [ "text", "output_ids", ] ) elif self.log_requests_level == 1: max_length = 2048 elif self.log_requests_level == 2: max_length = 1 << 30 else: raise ValueError( f"Invalid --log-requests-level: {self.log_requests_level=}" ) return max_length, skip_names, out_skip_names def configure_logging(self, obj: ConfigureLoggingReq): if obj.log_requests is not None: self.log_requests = obj.log_requests if obj.log_requests_level is not None: self.log_requests_level = obj.log_requests_level if obj.dump_requests_folder is not None: self.dump_requests_folder = obj.dump_requests_folder if obj.dump_requests_threshold is not None: self.dump_requests_threshold = obj.dump_requests_threshold logging.info("Config logging: obj=%r", obj) self.log_request_metadata = self.get_log_request_metadata() # ---- Server lifecycle / health ------------------------------- # Intent-revealing wrappers around the private ``server_status`` # field. Callers (notably ``http_server.py``) drive transitions # through these methods so the ``ServerStatus`` enum and the # attribute name stay implementation-private. def is_server_starting(self) -> bool: return self.server_status == ServerStatus.Starting def mark_server_up(self) -> None: self.server_status = ServerStatus.Up def mark_server_unhealthy(self) -> None: self.server_status = ServerStatus.UnHealthy def drop_request_state(self, rid: str) -> None: """Discard the per-request state for ``rid`` if present. Used by health probes that synthesize a request, await one token through ``generate_request``, and then need to clean up the state slot regardless of whether the probe succeeded or timed out. """ self.rid_to_state.pop(rid, None) def auto_create_handle_loop(self): if self.no_create_loop: return self.no_create_loop = True loop = asyncio.get_event_loop() self.asyncio_tasks.add( loop.create_task(print_exception_wrapper(self.handle_loop)) ) # We cannot add signal handler when the tokenizer manager is not in # the main thread due to the CPython limitation. if threading.current_thread() is threading.main_thread(): signal_handler = SignalHandler(self) loop.add_signal_handler(signal.SIGTERM, signal_handler.signal_handler) else: logger.warning( "Signal handler is not added because the tokenizer manager is " "not in the main thread. This disables graceful shutdown of the " "tokenizer manager when SIGTERM is received." ) self.asyncio_tasks.add( loop.create_task(print_exception_wrapper(self.sigterm_watchdog)) ) self.asyncio_tasks.add( loop.create_task(print_exception_wrapper(self.watch_load_thread)) ) async def sigterm_watchdog(self): while not self.gracefully_exit: await asyncio.sleep(5) # Drain requests while True: remain_num_req = len(self.rid_to_state) logger.info( "Gracefully exiting... remaining number of requests %s", remain_num_req ) if remain_num_req > 0: await asyncio.sleep(5) else: break kill_process_tree(os.getpid(), include_parent=True) sys.exit(0) async def handle_loop(self): """The event loop that handles requests""" while True: recv_obj = await self.engine_core_client.recv_from_detokenizer.recv_pyobj() self._result_dispatcher(recv_obj) self.last_receive_tstamp = time.time() def _handle_open_session_req_output(self, recv_obj): self.session_futures[recv_obj.session_id].set_result( recv_obj.session_id if recv_obj.success else None ) def _handle_update_weights_from_disk_req_output(self, recv_obj): if not self.server_args.mapping.attn.has_dp: self.model_update_result.set_result(recv_obj) else: # self.server_args.mapping.has_attn_dp self.model_update_tmp.append(recv_obj) # set future if the all results are received if len(self.model_update_tmp) == self.server_args.mapping.attn.dp_size: self.model_update_result.set_result(self.model_update_tmp) async def print_exception_wrapper(func): """ Sometimes an asyncio function does not print exception. We do another wrapper to handle the exception. """ try: await func() except Exception: traceback = get_exception_traceback() logger.error("AsyncLLM hit an exception: %s", traceback) kill_process_tree(os.getpid(), include_parent=True) sys.exit(1) class SignalHandler: def __init__(self, tokenizer_manager): self.tokenizer_manager = tokenizer_manager def signal_handler(self, signum=None, frame=None): logger.warning( "SIGTERM received. signum=%r frame=%r. Draining requests and shutting down...", signum, frame, ) self.tokenizer_manager.gracefully_exit = True T = TypeVar("T") class _Communicator(Generic[T]): """Note: The communicator now only run up to 1 in-flight request at any time.""" def __init__(self, sender, fan_out: int): self._sender = sender self._fan_out = fan_out self._result_event: asyncio.Event | None = None self._result_values: list[T] | None = None self._ready_queue: deque[asyncio.Future] = deque() async def __call__(self, obj): ready_event = asyncio.Event() if self._result_event is not None or len(self._ready_queue) > 0: self._ready_queue.append(ready_event) await ready_event.wait() if self._result_event is not None or self._result_values is not None: raise RuntimeError("Communicator result state was not reset.") if obj: self._sender.send_pyobj(obj) self._result_event = asyncio.Event() self._result_values = [] await self._result_event.wait() result_values = self._result_values self._result_event = self._result_values = None if len(self._ready_queue) > 0: self._ready_queue.popleft().set() return result_values def handle_recv(self, recv_obj: T): self._result_values.append(recv_obj) if len(self._result_values) == self._fan_out: self._result_event.set()