diff --git a/atom/config.py b/atom/config.py index cc5ca0f5e2..dd3add9415 100644 --- a/atom/config.py +++ b/atom/config.py @@ -2108,6 +2108,9 @@ class Config: compilation_config: CompilationConfig = field(default_factory=CompilationConfig) quant_config: QuantizationConfig = field(init=False) asyncio_mode: bool = False + # Queue finished sequences on outputs_queue; only its readers turn it on + # (LLMEngine.add_request for step / generate, and asyncio_mode). + keep_finished_outputs: bool = False mark_trace: bool = False load_dummy: str | None = None enable_expert_parallel: bool = False diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 864e4c97ed..d988ba65a1 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -51,6 +51,7 @@ from atom.model_engine.arg_utils import EngineArgs from atom.model_engine.llm_engine import load_tokenizer from atom.model_engine.request import RequestOutput +from atom.model_engine.segment_tokenizer import PrefixEncoding from atom.model_engine.sequence import new_token_ids from atom.multimodal.processing import prepare_multimodal_inputs from atom.multimodal.registry import get_native_multimodal_processor @@ -853,7 +854,7 @@ async def _submit_engine_requests(seqs): async def generate_async( - prompt_or_tokens: str | list[int], + prompt_or_tokens: str | list[int] | PrefixEncoding, sampling_params: SamplingParams, request_id: str, kv_transfer_params: dict[str, Any] | None = None, @@ -1306,6 +1307,50 @@ def _prompt_headroom( return headroom +def _prepare_prompt( + messages, + template_kwargs: dict[str, Any], + *, + tools, + preserve_tool_call_arguments: bool, + max_tokens: int, + kv_transfer_params: dict[str, Any] | None, +) -> tuple[PrefixEncoding | list[int], str]: + """Render, tokenize the engine's way and cut to the context, in one + executor step for ``/v1/messages`` and ``/v1/responses``. + + Returns ``(engine_input, text)``: what preprocess takes in place of the + prompt (no second tokenize), and the text the engine runs, which seeds the + reasoning channel. A prompt longer than the context minus the output + headroom is cut from the end; a context with no room for any prompt is + refused (ValueError, a 400).""" + prompt = apply_chat_template( + tokenizer, + custom_message_encoder, + messages, + tools=tools, + preserve_tool_call_arguments=preserve_tool_call_arguments, + **template_kwargs, + ) + encoded = engine.io_processor.encode_prompt(prompt) + max_ctx = int(_get_engine_max_model_len() or 30720) + cap = max_ctx - _prompt_headroom(max_tokens, max_ctx, kv_transfer_params) + if cap < 1: + raise ValueError( + f"max_tokens={max_tokens} leaves no room for a prompt in this " + f"model's {max_ctx}-token context" + ) + token_ids = encoded.token_ids if isinstance(encoded, PrefixEncoding) else encoded + if len(token_ids) <= cap: + return encoded, prompt + logger.warning("Prompt too long (%d > %d), truncating", len(token_ids), cap) + token_ids = list(token_ids[:cap]) + text = tokenizer.decode(token_ids, skip_special_tokens=False) + if not text: + raise ValueError("The prompt is empty after truncation to the context") + return token_ids, text + + def _with_kv_transfer_params(response: dict, final_output: dict) -> dict: """Return the prefill's KV location to the PD router, as chat does.""" if "kv_transfer_output_meta_info" in final_output: @@ -1314,7 +1359,7 @@ def _with_kv_transfer_params(response: dict, final_output: dict) -> dict: async def setup_streaming_request( - prompt_or_tokens: str | list[int], + prompt_or_tokens: str | list[int] | PrefixEncoding, sampling_params: SamplingParams, request_id: str, kv_transfer_params: dict[str, Any] | None = None, @@ -1924,7 +1969,8 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) prompt_or_tokens = pretokenized kv_transfer_params = _engine_kv_transfer_params(kv_transfer_params) else: - prompt_or_tokens = apply_chat_template( + prompt_or_tokens = await asyncio.to_thread( + apply_chat_template, tokenizer, custom_message_encoder, messages, @@ -2351,15 +2397,6 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req # no such spelling, and are not answered here or on the chat path. if forbids_tool_calls(request.tool_choice): merged_kwargs["tool_choice"] = "none" - prompt = apply_chat_template( - tokenizer, - custom_message_encoder, - messages, - tools=anthropic_to_openai_tools(request.tools), - preserve_tool_call_arguments=preserve_tool_call_arguments, - **merged_kwargs, - ) - generation_config = engine.config.generation_config model_temperature = getattr(generation_config, "temperature", None) model_top_p = getattr(generation_config, "top_p", None) @@ -2388,47 +2425,23 @@ async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Req request_id = uuid.uuid4().hex[:24] set_request_id(request_id) kv_transfer_params = _engine_kv_transfer_params(request.kv_transfer_params) - input_tokens = len(tokenizer.encode(prompt)) - - max_ctx = None - for _path in ( - lambda: engine.config.max_model_len, - lambda: engine.model_config.max_model_len, - lambda: engine.scheduler.max_model_len, - lambda: engine.max_model_len, - ): - try: - _v = _path() - if _v: - max_ctx = int(_v) - break - except Exception: - logger.debug("max_ctx candidate rejected", exc_info=True) - continue - if not max_ctx: - max_ctx = 30720 - logger.warning(f"[anthropic] resolved max_ctx={max_ctx}") - headroom = _prompt_headroom( - sampling_params.max_tokens, max_ctx, kv_transfer_params + engine_input, prompt = await asyncio.to_thread( + _prepare_prompt, + messages, + merged_kwargs, + tools=anthropic_to_openai_tools(request.tools), + preserve_tool_call_arguments=preserve_tool_call_arguments, + max_tokens=sampling_params.max_tokens, + kv_transfer_params=kv_transfer_params, ) - max_input = max_ctx - headroom - if input_tokens > max_input: - logger.warning( - f"Prompt too long ({input_tokens} > {max_input}), truncating" - ) - token_ids = tokenizer.encode(prompt)[:max_input] - prompt = tokenizer.decode(token_ids, skip_special_tokens=False) - input_tokens = max_input if request.stream: # Streaming response - seq_id, stream_collector, _num_prompt_tokens = ( - await setup_streaming_request( - prompt, - sampling_params, - request_id, - kv_transfer_params=kv_transfer_params, - ) + seq_id, stream_collector, input_tokens = await setup_streaming_request( + engine_input, + sampling_params, + request_id, + kv_transfer_params=kv_transfer_params, ) async def generate_anthropic_stream(): @@ -2600,7 +2613,7 @@ async def generate_anthropic_stream(): # Non-streaming response final_output = await _run_nonstream_with_disconnect( generate_async( - prompt, + engine_input, sampling_params, request_id, kv_transfer_params=kv_transfer_params, @@ -2610,6 +2623,7 @@ async def generate_anthropic_stream(): ) if final_output is None: raise RuntimeError("No output generated") + input_tokens = final_output["num_tokens_input"] raw_text = final_output["text"] # Separating is unconditional -- the same call the chat path makes, @@ -2633,7 +2647,7 @@ async def generate_anthropic_stream(): if drop_reasoning: events = [e for e in events if e[0] != "reasoning"] _, tool_calls = flatten_tool_events(events) - output_tokens = len(tokenizer.encode(raw_text)) + output_tokens = final_output["num_tokens_output"] cache_read_input_tokens = final_output.get("num_cached_tokens", 0) response = build_anthropic_response( request_id=request_id, @@ -2783,15 +2797,6 @@ async def responses_create(raw_request: Request): rules=rules, ) - prompt = apply_chat_template( - tokenizer, - custom_message_encoder, - messages, - tools=openai_tools, - preserve_tool_call_arguments=preserve_tool_call_arguments, - **merged_kwargs, - ) - generation_config = engine.config.generation_config model_temperature = getattr(generation_config, "temperature", None) model_top_p = getattr(generation_config, "top_p", None) @@ -2828,30 +2833,15 @@ async def responses_create(raw_request: Request): schema = schema_from_format((text_config or {}).get("format"), responses=True) request_id = f"resp_{uuid.uuid4().hex[:24]}" set_request_id(request_id) - input_tokens = len(tokenizer.encode(prompt)) - - max_ctx = 30720 - for obj in ( - getattr(engine, "config", None), - getattr(engine, "model_config", None), - getattr(engine, "scheduler", None), - engine, - ): - value = getattr(obj, "max_model_len", None) if obj is not None else None - if value: - max_ctx = int(value) - break - headroom = _prompt_headroom( - sampling_params.max_tokens, max_ctx, kv_transfer_params + engine_input, prompt = await asyncio.to_thread( + _prepare_prompt, + messages, + merged_kwargs, + tools=openai_tools, + preserve_tool_call_arguments=preserve_tool_call_arguments, + max_tokens=sampling_params.max_tokens, + kv_transfer_params=kv_transfer_params, ) - max_input = max_ctx - headroom - if input_tokens > max_input: - logger.warning( - f"Prompt too long ({input_tokens} > {max_input}), truncating" - ) - token_ids = tokenizer.encode(prompt)[:max_input] - prompt = tokenizer.decode(token_ids, skip_special_tokens=False) - input_tokens = max_input created_at = int(time.time()) _reasoning = reasoning_channel( @@ -2869,13 +2859,11 @@ async def responses_create(raw_request: Request): _reasoning = structured_response_channel(_reasoning, schema) if request.stream: - seq_id, stream_collector, _num_prompt_tokens = ( - await setup_streaming_request( - prompt, - sampling_params, - request_id, - kv_transfer_params=kv_transfer_params, - ) + seq_id, stream_collector, input_tokens = await setup_streaming_request( + engine_input, + sampling_params, + request_id, + kv_transfer_params=kv_transfer_params, ) async def generate_responses_stream(): @@ -2984,7 +2972,7 @@ async def generate_responses_stream(): final_output = await _run_nonstream_with_disconnect( generate_async( - prompt, + engine_input, sampling_params, request_id, kv_transfer_params=kv_transfer_params, @@ -2994,6 +2982,7 @@ async def generate_responses_stream(): ) if final_output is None: raise RuntimeError("No output generated") + input_tokens = final_output["num_tokens_input"] raw_text = final_output["text"] events = read_whole_blocks( @@ -3003,7 +2992,7 @@ async def generate_responses_stream(): openai_tools, suppress_calls=forbids_tool_calls(request.tool_choice), ) - output_tokens = len(tokenizer.encode(raw_text)) + output_tokens = final_output["num_tokens_output"] cache_read_input_tokens = final_output.get("num_cached_tokens", 0) response = build_responses_response( request_id=request_id, diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index 63d8def114..44329db0ed 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -231,6 +231,7 @@ def _init_shared_state( ) self.ctx = zmq.Context(io_threads=2) self.outputs_queue = queue.Queue[list[Sequence]]() + self.keep_finished_outputs = config.keep_finished_outputs or config.asyncio_mode # 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 +835,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,7 +1308,23 @@ 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. + A pending seq with a stream callback is not reading that queue, so it + is taken out and ended through its callback, as abort does. """ + with self._lb_lock: + waiting = [s for s in self._bp_pending if s.id in self._seq_id_to_callback] + for seq in waiting: + self._bp_pending.remove(seq) + delivered = False + for seq in waiting: + if not getattr(seq, "leave_reason", None): + seq.leave_reason = "aborted" + delivered |= self._deliver_terminal_output(seq) + if delivered and self._flush_stream_batch_fn is not None: + try: + self._flush_stream_batch_fn() + except Exception as e: # noqa: BLE001 - the error below must still go out + logger.warning("%s: flush_stream_batch failed: %s", self.label, e) pending = len(getattr(self, "_bp_pending", ())) logger.error( "%s: backpressure send failed with nothing in flight; " diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index d3d2426714..50df10485f 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -14,7 +14,7 @@ from atom.config import Config from atom.model_engine.engine_core_mgr import CoreManager, DisaggCoreManager from atom.model_engine.prefix_hash import hash_block_size_for, ship_prompt_hash_chain -from atom.model_engine.segment_tokenizer import SegmentTokenizer +from atom.model_engine.segment_tokenizer import PrefixEncoding, SegmentTokenizer from atom.model_engine.sequence import Sequence from atom.multimodal.registry import get_mrope_input_positions from atom.sampling_params import RequiredChannel, SamplingParams @@ -330,6 +330,7 @@ def add_request( parent_request_id=request_id, ) reqs.extend(fanout) + self.core_mgr.keep_finished_outputs = True self.core_mgr.add_request(reqs) def step(self) -> list[Sequence]: @@ -813,7 +814,7 @@ def _per_req_cache_model_types() -> frozenset[str]: def preprocess( self, - prompt_or_tokens: str | list[int], + prompt_or_tokens: str | list[int] | PrefixEncoding, sampling_params: SamplingParams, stream_callback=None, kv_transfer_params=None, @@ -874,9 +875,29 @@ def _engine_channel(self, channel: RequiredChannel) -> RequiredChannel: ) return self._required_channel + def encode_prompt( + self, prompt: str, *, multimodal: bool = False + ) -> PrefixEncoding | list[int]: + """Tokenize ``prompt`` the way ``preprocess`` does (segment cache and, + when enabled, the prompt's block-hash chain). ``preprocess`` takes the + result in place of the text without tokenizing again.""" + segment_tokenizer = getattr(self, "segment_tokenizer", None) + if segment_tokenizer is None: + return self.tokenizer.encode(prompt) + if ( + getattr(self, "_reuse_block_hashes", False) + and self.prefix_hash_block_size + and envs.ATOM_FRONTEND_PREFIX_HASH + and not multimodal + ): + return segment_tokenizer.encode_with_prefix( + prompt, self.prefix_hash_block_size + ) + return segment_tokenizer.encode(prompt) + def preprocess_fanout( self, - prompt_or_tokens: str | list[int], + prompt_or_tokens: str | list[int] | PrefixEncoding, sampling_params: SamplingParams, stream_callback=None, stream_callbacks: list | None = None, @@ -918,25 +939,16 @@ def preprocess_fanout( ) n = max(1, int(getattr(sampling_params, "n", 1))) - prefix_encoding = None - segment_tokenizer = getattr(self, "segment_tokenizer", None) - if not isinstance(prompt_or_tokens, str): - tokens = prompt_or_tokens - elif segment_tokenizer is not None: - if ( - getattr(self, "_reuse_block_hashes", False) - and self.prefix_hash_block_size - and envs.ATOM_FRONTEND_PREFIX_HASH - and multimodal_data is None - ): - prefix_encoding = segment_tokenizer.encode_with_prefix( - prompt_or_tokens, self.prefix_hash_block_size - ) - tokens = prefix_encoding.token_ids - else: - tokens = segment_tokenizer.encode(prompt_or_tokens) + if isinstance(prompt_or_tokens, str): + prompt_or_tokens = self.encode_prompt( + prompt_or_tokens, multimodal=multimodal_data is not None + ) + if isinstance(prompt_or_tokens, PrefixEncoding): + prefix_encoding = prompt_or_tokens + tokens = prefix_encoding.token_ids else: - tokens = self.tokenizer.encode(prompt_or_tokens) + prefix_encoding = None + tokens = prompt_or_tokens mrope_positions = None mrope_position_delta = 0 if multimodal_data is not None: diff --git a/tests/entrypoints/test_api_compatibility.py b/tests/entrypoints/test_api_compatibility.py index 13b38ef79a..6233ef01c8 100644 --- a/tests/entrypoints/test_api_compatibility.py +++ b/tests/entrypoints/test_api_compatibility.py @@ -72,7 +72,8 @@ def encode(messages, **kwargs): "engine": SimpleNamespace( config=SimpleNamespace( generation_config=SimpleNamespace(), max_model_len=8192 - ) + ), + io_processor=SimpleNamespace(encode_prompt=lambda text: [1, 2]), ), }.items(): monkeypatch.setattr(api_server, name, value) @@ -1052,6 +1053,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) diff --git a/tests/entrypoints/test_api_server_helpers.py b/tests/entrypoints/test_api_server_helpers.py index 32a2895418..b6b8a73a87 100644 --- a/tests/entrypoints/test_api_server_helpers.py +++ b/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,84 @@ async def drive(): assert events == ["initialize", "serving", "shutdown"] +class TestPreparePrompt: + """`_prepare_prompt`: render, tokenize the engine's way, cut to the + context, on the executor, for /v1/messages and /v1/responses.""" + + @pytest.fixture + def prepare(self, monkeypatch): + from atom.model_engine.segment_tokenizer import PrefixEncoding + + encoded = {} + + def encode_prompt(text): + # One id per character; with a block-hash chain when asked. + ids = list(range(len(text))) + return PrefixEncoding(ids, "chain") if encoded.get("chain") else ids + + monkeypatch.setattr( + api_server, + "engine", + SimpleNamespace( + config=SimpleNamespace(max_model_len=100), + io_processor=SimpleNamespace(encode_prompt=encode_prompt), + ), + ) + monkeypatch.setattr( + api_server, + "tokenizer", + SimpleNamespace(decode=lambda ids, **_kw: "k" * len(ids)), + ) + monkeypatch.setattr(api_server, "custom_message_encoder", None) + rendered = {} + monkeypatch.setattr( + api_server, + "apply_chat_template", + lambda _tok, _enc, messages, **kwargs: rendered.setdefault( + "text", messages + ), + ) + + def run(prompt, max_tokens, chain=False, kv=None): + encoded["chain"] = chain + rendered.clear() + return api_server._prepare_prompt( + prompt, + {}, + tools=None, + preserve_tool_call_arguments=False, + max_tokens=max_tokens, + kv_transfer_params=kv, + ) + + return run + + def test_a_prompt_that_fits_goes_to_the_engine_as_encoded(self, prepare): + engine_input, text = prepare("x" * 50, max_tokens=10, chain=True) + # The engine's own encoding, chain included, so preprocess reuses it. + assert engine_input.token_ids == list(range(50)) + assert engine_input.block_hashes == "chain" + assert text == "x" * 50 + + def test_an_over_long_prompt_is_cut_to_the_context(self, prepare): + # 100-token context, 10 reserved for output: 90 kept. + engine_input, text = prepare("x" * 150, max_tokens=10, chain=True) + assert engine_input == list(range(90)) + assert text == "k" * 90 + + @pytest.mark.parametrize("max_tokens", [100, 1024]) + def test_no_room_for_a_prompt_is_refused(self, prepare, max_tokens): + with pytest.raises(ValueError, match="no room for a prompt"): + prepare("x" * 5, max_tokens=max_tokens) + + def test_an_empty_cut_is_refused(self, prepare, monkeypatch): + monkeypatch.setattr( + api_server, "tokenizer", SimpleNamespace(decode=lambda ids, **_kw: "") + ) + with pytest.raises(ValueError, match="empty after truncation"): + prepare("x" * 150, max_tokens=10) + + class TestDeploymentDefaultOutputCap: @pytest.mark.parametrize("explicit", [False, True]) def test_small_deployment_distinguishes_default_from_explicit( diff --git a/tests/entrypoints/test_kimi_k3_request_compat.py b/tests/entrypoints/test_kimi_k3_request_compat.py index e7030eaad8..4de97961d7 100644 --- a/tests/entrypoints/test_kimi_k3_request_compat.py +++ b/tests/entrypoints/test_kimi_k3_request_compat.py @@ -407,7 +407,8 @@ async def fake_generate_async(_prompt, sampling, _request_id, **_kwargs): SimpleNamespace( config=SimpleNamespace( generation_config=SimpleNamespace(), max_model_len=8192 - ) + ), + io_processor=SimpleNamespace(encode_prompt=lambda _text: [1, 2]), ), ) monkeypatch.setattr(api_server, "generate_async", fake_generate_async) diff --git a/tests/entrypoints/test_pd_messages_responses.py b/tests/entrypoints/test_pd_messages_responses.py index b77c8ac683..7fd8ddc04e 100644 --- a/tests/entrypoints/test_pd_messages_responses.py +++ b/tests/entrypoints/test_pd_messages_responses.py @@ -49,19 +49,27 @@ def server(monkeypatch): """Stub the engine and record what each endpoint hands to it.""" calls = SimpleNamespace(generate=[], setup=[], template="") + def encode(text, **_kw): + return [0] * (LONG_PROMPT_TOKENS if text == "" else 2) + + # Only what the engine does: count what it is handed. Any cut is the + # production code's (`_prepare_prompt`). 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} + output = { + "text": "hello", + "finish_reason": "eos", + "num_cached_tokens": 0, + "num_tokens_input": len(prompt), + "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 == "" else 2) + return 1, _FinishedCollector(), len(prompt) monkeypatch.setattr(api_server, "model_name", "m") monkeypatch.setattr(api_server, "_request_logger", None) @@ -83,7 +91,8 @@ def encode(text, **_kw): SimpleNamespace( config=SimpleNamespace( generation_config=SimpleNamespace(), max_model_len=MAX_CTX - ) + ), + io_processor=SimpleNamespace(encode_prompt=encode), ), ) monkeypatch.setattr(api_server, "generate_async", fake_generate_async) @@ -195,7 +204,7 @@ def test_with_handoff_the_prompt_does_not_depend_on_max_tokens( prefill = self._sent_prompt(server, endpoint, 1, HANDOFF) decode = self._sent_prompt(server, endpoint, 4096, HANDOFF) - assert prefill == decode == f"<{MAX_CTX - MAX_CTX // 8}>" + assert len(prefill) == len(decode) == MAX_CTX - MAX_CTX // 8 @pytest.mark.parametrize("endpoint", ["messages", "responses"]) def test_without_handoff_truncation_still_follows_max_tokens( @@ -203,8 +212,61 @@ def test_without_handoff_truncation_still_follows_max_tokens( ): server.template = "" - assert self._sent_prompt(server, endpoint, 1, None) == "" + # max_tokens=1 leaves room for the whole prompt; 4096 does not. + assert len(self._sent_prompt(server, endpoint, 1, None)) == LONG_PROMPT_TOKENS assert ( - self._sent_prompt(server, endpoint, 4096, None) - == f"<{MAX_CTX - MAX_CTX // 8}>" + len(self._sent_prompt(server, endpoint, 4096, None)) + == MAX_CTX - MAX_CTX // 8 + ) + + @pytest.mark.parametrize("endpoint", ["messages", "responses"]) + def test_input_tokens_are_the_engine_count_of_what_it_ran(self, server, endpoint): + server.template = "" + if endpoint == "messages": + usage = _messages(max_tokens=4096)["usage"] + else: + body = {"model": "m", "input": "hi", "max_output_tokens": 4096} + usage = _responses(body).json()["usage"] + assert usage["input_tokens"] == MAX_CTX - MAX_CTX // 8 + + +class TestStructuredResponsesCutFirst: + """A structured /v1/responses request builds its reasoning channel into the + sampling params before the engine call; both are seeded from the text the + engine runs.""" + + def test_over_the_cap_the_engine_gets_the_kept_ids(self, server, monkeypatch): + server.template = "" + configured = [] + seeds = [] + + async def configure(_sampling, schema, channel, **_kw): + configured.append(schema) + + monkeypatch.setattr(api_server, "_configure_structured_output", configure) + monkeypatch.setattr( + api_server, + "prompt_starts_in_reasoning", + lambda text: seeds.append(text) or False, + ) + # The fake engine answers "hello", which is not JSON. + monkeypatch.setattr( + api_server, "validate_responses_output", lambda *a, **kw: None ) + response = _responses( + { + "model": "m", + "input": "hi", + "max_output_tokens": 4096, + "text": {"format": {"type": "json_object"}}, + } + ) + + assert response.status_code == 200, response.text + sent, _kwargs = server.generate[0] + kept = MAX_CTX - MAX_CTX // 8 + assert configured == ['{"type": "object"}'] + # The engine is handed the kept ids, and does not tokenize again. + assert sent == [0] * kept + # The channel was seeded from the cut text, before the engine call. + assert seeds[0] == f"<{kept}>" diff --git a/tests/test_distributed_store.py b/tests/test_distributed_store.py index 6910fd527e..a82b7c3365 100644 --- a/tests/test_distributed_store.py +++ b/tests/test_distributed_store.py @@ -31,7 +31,11 @@ def manager(monkeypatch): monkeypatch.delenv(name, raising=False) manager = CoreManager.__new__(CoreManager) manager._init_shared_state( - SimpleNamespace(dp_load_balance="round_robin"), + SimpleNamespace( + dp_load_balance="round_robin", + keep_finished_outputs=False, + asyncio_mode=False, + ), label="test", local_engine_count=0, ) diff --git a/tests/test_dp_backpressure.py b/tests/test_dp_backpressure.py index 8abca2f278..86f25da871 100644 --- a/tests/test_dp_backpressure.py +++ b/tests/test_dp_backpressure.py @@ -121,7 +121,12 @@ def finish_one_on_rank(rank): def _init_with_strategy(strategy): mgr = CoreManager.__new__(CoreManager) mgr._init_shared_state( - SimpleNamespace(dp_load_balance=strategy, max_num_seqs=1536), + SimpleNamespace( + dp_load_balance=strategy, + max_num_seqs=1536, + keep_finished_outputs=False, + asyncio_mode=False, + ), label="test", local_engine_count=2, ) @@ -309,3 +314,23 @@ def release_loop(): assert pending == [] assert sorted(sent) == sorted(seq.id for seq in seqs) assert len(sent) == len(set(sent)) + + +def test_a_stalled_run_ends_pending_seqs_that_wait_on_a_callback(): + # The stalled-run error goes to outputs_queue, which a callback-driven + # client never reads: such a pending seq is ended through its callback. + mgr = make_mgr(n_ranks=2, depth=1) + mgr.outputs_queue = queue.Queue() + delivered = {} + seqs = [Seq() for _ in range(4)] + seqs[-1].stream_callback = lambda out: delivered.update( + {out.request_id: out.finish_reason} + ) + mgr.add_request(seqs) + mgr.fail_ranks.update({0, 1}) + for sid in list(mgr._seq_load): + mgr._release_seq_load(sid) + mgr._bp_dispatch_or_refill(raise_on_error=False) + assert delivered == {seqs[-1].id: "aborted"} + assert all(seq.id != seqs[-1].id for seq in mgr._bp_pending) + assert isinstance(mgr.outputs_queue.get_nowait(), RuntimeError) diff --git a/tests/test_engine_core_mgr_shared_state.py b/tests/test_engine_core_mgr_shared_state.py index 1ad39a9d21..2cf00bfe1f 100644 --- a/tests/test_engine_core_mgr_shared_state.py +++ b/tests/test_engine_core_mgr_shared_state.py @@ -76,6 +76,8 @@ 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", False), + asyncio_mode=kwargs.get("asyncio_mode", False), ), label=cls.__name__, local_engine_count=kwargs.get("local_engine_count", 2), @@ -130,6 +132,37 @@ 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): + # Off by default: online front ends take every output through the + # callbacks, and a queued Sequence would stay for the process's life. + # asyncio_mode reads the queue, so it keeps them too. + for keep, asyncio_mode in ((True, False), (False, True), (False, False)): + mgr = _bare_manager( + CoreManager, keep_finished_outputs=keep, asyncio_mode=asyncio_mode + ) + 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 or asyncio_mode else 0 + ) + finally: + sender.close(linger=0) + def _terminate(self): for mgr in self.managers: mgr.ctx.term() diff --git a/tests/test_llm_engine_add_request_fanout.py b/tests/test_llm_engine_add_request_fanout.py index cbbb58f66d..8de7db4a7d 100644 --- a/tests/test_llm_engine_add_request_fanout.py +++ b/tests/test_llm_engine_add_request_fanout.py @@ -167,3 +167,19 @@ def test_preprocess_still_refuses_n_greater_than_one(): io.preprocess("hello", SamplingParams(n=2)) assert io.preprocess("hello", SamplingParams(n=1)).choice_index == 0 + + +def test_offline_add_turns_on_the_finished_outputs_queue_before_submitting(): + # `step` / `generate` read outputs_queue, so the offline entry point turns + # it on (Config.keep_finished_outputs defaults to off for online servers) + # before any of its sequences can finish. + engine, _io, _submitted = _engine() + seen = [] + engine.core_mgr.keep_finished_outputs = False + engine.core_mgr.add_request = lambda seqs: seen.append( + engine.core_mgr.keep_finished_outputs + ) + + LLMEngine.add_request(engine, ["hello"], SamplingParams()) + + assert seen == [True]