diff --git a/components/src/dynamo/sglang/request_handlers/handler_base.py b/components/src/dynamo/sglang/request_handlers/handler_base.py index 82773910a948..f929073223a2 100644 --- a/components/src/dynamo/sglang/request_handlers/handler_base.py +++ b/components/src/dynamo/sglang/request_handlers/handler_base.py @@ -1134,8 +1134,22 @@ def _get_bootstrap_info(engine: sgl.Engine) -> Tuple[str, int]: return bootstrap_host, bootstrap_port + def _abort_requests(self, request_ids: set[str], context: Context) -> None: + if not request_ids: + return + if not getattr(self.engine, "tokenizer_manager", None): + logging.error("SGLang tokenizer_manager not found for abort requests") + return + for request_id in request_ids: + self.engine.tokenizer_manager.abort_request(rid=request_id, abort_all=False) + request_ids.clear() + logging.info("Aborted Request ID: %s", context.id()) + async def _handle_cancellation( - self, request_id_future: asyncio.Future, context: Context + self, + request_id_future: asyncio.Future, + context: Context, + request_ids: set[str] | None = None, ): """Background task to handle cancellation and shutdown by monitoring both signals. @@ -1143,6 +1157,7 @@ async def _handle_cancellation( request_id_future: Future that will be set with the SGLang request ID when the first response arrives. context: Context object for cancellation handling. + request_ids: Live set of unfinished parallel-choice IDs, when available. Raises: EngineShutdown: If shutdown event was triggered. @@ -1188,22 +1203,9 @@ async def _handle_cancellation( f"Cancellation or shutdown signal received for SGLang Request ID {sglang_request_id}, Context: {context.id()}" ) - # Call abort_request on the tokenizer_manager through the engine - if ( - hasattr(self.engine, "tokenizer_manager") - and self.engine.tokenizer_manager - ): - logging.info( - f"Calling SGLang abort_request for Request ID {sglang_request_id}" - ) - self.engine.tokenizer_manager.abort_request( - rid=sglang_request_id, abort_all=False - ) - logging.info(f"Aborted Request ID: {context.id()}") - else: - logging.error( - f"SGLang tokenizer_manager not found for abort request: {context.id()}" - ) + self._abort_requests( + request_ids if request_ids is not None else {sglang_request_id}, context + ) # Check which event triggered and raise EngineShutdown if shutdown if shutdown_task and shutdown_task in done: @@ -1233,7 +1235,10 @@ async def _handle_cancellation( @asynccontextmanager async def _cancellation_monitor( - self, request_id_future: asyncio.Future, context: Context + self, + request_id_future: asyncio.Future, + context: Context, + request_ids: set[str] | None = None, ) -> AsyncGenerator[asyncio.Task, None]: """ Context manager for monitoring request cancellation and shutdown. @@ -1246,6 +1251,7 @@ async def _cancellation_monitor( request_id_future: Future that will be set with the SGLang request ID when the first response arrives. context: Context object for cancellation handling + request_ids: Live set of unfinished parallel-choice IDs to abort on exit. Yields: asyncio.Task: The cancellation monitoring task being managed @@ -1254,12 +1260,15 @@ async def _cancellation_monitor( # Start the cancellation monitoring task cancellation_task = asyncio.create_task( - self._handle_cancellation(request_id_future, context) + self._handle_cancellation(request_id_future, context, request_ids) ) try: yield cancellation_task finally: + # Stream closure can beat the monitor's cancellation wake-up. + if request_ids is not None: + self._abort_requests(request_ids, context) # Clean up the background cancellation task request_id = "unknown" if request_id_future.done() and not request_id_future.cancelled(): diff --git a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py index b50566333185..7287c186fb34 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -360,12 +360,25 @@ def _build_sampling_params(self, request: Dict[str, Any]) -> Dict[str, Any]: _plain = stop_conditions.get("stop_token_ids") or [] _merged = list(set(_hidden).union(_plain)) stop_token_ids = _merged if _merged else None + # A tokenizer-free SGLang rejects positive min_new_tokens. Let + # Dynamo's decoder enforce the floor and keep the engine running so + # it cannot stop before Dynamo reaches it. + min_tokens = stop_conditions.get("min_tokens") + hold_engine_open = bool( + min_tokens + and min_tokens > 0 + and self.config.server_args.skip_tokenizer_init + ) + if hold_engine_open: + stop_token_ids = None param_mapping = { "n": sampling_opts.get("n"), "max_new_tokens": stop_conditions.get("max_tokens"), - "min_new_tokens": stop_conditions.get("min_tokens"), - "ignore_eos": stop_conditions.get("ignore_eos"), + "min_new_tokens": None if hold_engine_open else min_tokens, + "ignore_eos": True + if hold_engine_open + else stop_conditions.get("ignore_eos"), "stop_token_ids": stop_token_ids, **_sampling_option_params(sampling_opts), **self._get_guided_decoding_params( @@ -691,16 +704,23 @@ async def _process_token_stream( """ # Use Future pattern for request ID - will be set when first response arrives request_id_future: asyncio.Future[str] = asyncio.Future() + request_ids: set[str] = set() first_output_seen = False - async with self._cancellation_monitor(request_id_future, context): + async with self._cancellation_monitor(request_id_future, context, request_ids): async for res in stream_source: meta_info = res.get("meta_info", {}) - # Extract SGLang request ID from the first response and set the future - if not request_id_future.done(): - sglang_request_id = meta_info.get("id") - if sglang_request_id: + sglang_request_id = meta_info.get("id") + if sglang_request_id: + request_ids.add(sglang_request_id) + if not request_id_future.done(): request_id_future.set_result(sglang_request_id) logging.debug(f"New SGLang Request ID: {sglang_request_id}") + if meta_info.get("finish_reason"): + request_ids.discard(sglang_request_id) + if context.is_stopped(): + # A choice's first chunk can arrive after the monitor fired. + self._abort_requests(request_ids, context) + continue # Check cancellation before yielding to allow proper cleanup. # This lets SGLang proceed to the second token generation, which will @@ -829,16 +849,22 @@ async def _process_text_stream( # Use Future pattern for request ID - will be set when first response arrives request_id_future: asyncio.Future[str] = asyncio.Future() + request_ids: set[str] = set() first_output_seen = False - async with self._cancellation_monitor(request_id_future, context): + async with self._cancellation_monitor(request_id_future, context, request_ids): async for res in stream_source: meta_info = res.get("meta_info", {}) - # Extract SGLang request ID from the first response and set the future - if not request_id_future.done(): - sglang_request_id = meta_info.get("id") - if sglang_request_id: + sglang_request_id = meta_info.get("id") + if sglang_request_id: + request_ids.add(sglang_request_id) + if not request_id_future.done(): request_id_future.set_result(sglang_request_id) logging.debug(f"New SGLang Request ID: {sglang_request_id}") + if meta_info.get("finish_reason"): + request_ids.discard(sglang_request_id) + if context.is_stopped(): + self._abort_requests(request_ids, context) + continue # Check cancellation before yielding to allow proper cleanup. # This lets SGLang proceed to the second token generation, which will diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py index f38e31230bff..380896a8c57d 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -2,8 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio -from contextlib import asynccontextmanager +from contextlib import aclosing, asynccontextmanager from types import SimpleNamespace +from unittest.mock import Mock, call import pytest @@ -239,12 +240,20 @@ def test_openai_stop_sampling_params_maps_token_id_stop_array(): } -def _new_decode_handler(*, use_sglang_tokenizer: bool = False, enable_rl: bool = False): +def _new_decode_handler( + *, + use_sglang_tokenizer: bool = False, + skip_tokenizer_init: bool = False, + enable_rl: bool = False, +): handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler) handler.shutdown_event = None handler.use_sglang_tokenizer = use_sglang_tokenizer handler.config = SimpleNamespace( - server_args=SimpleNamespace(served_model_name="test-model"), + server_args=SimpleNamespace( + served_model_name="test-model", + skip_tokenizer_init=skip_tokenizer_init, + ), dynamo_args=SimpleNamespace(enable_rl=enable_rl), ) handler._first_token_source = None @@ -257,6 +266,68 @@ async def no_cancellation_monitor(*args, **kwargs): return handler +@pytest.mark.asyncio +@pytest.mark.parametrize( + "processor_name", ["_process_token_stream", "_process_text_stream"] +) +@pytest.mark.parametrize("cancel_from_client", [False, True], ids=["close", "cancel"]) +async def test_stream_cancellation_aborts_unfinished_choices( + processor_name, cancel_from_client, caplog +): + handler = _new_decode_handler() + # Exercise the real monitor, including closure before its task can run. + del handler._cancellation_monitor + cancellation = asyncio.get_running_loop().create_future() + first_abort = asyncio.Event() + abort_request = Mock(side_effect=lambda **_kwargs: first_abort.set()) + handler.engine = SimpleNamespace( + tokenizer_manager=SimpleNamespace(abort_request=abort_request) + ) + context = SimpleNamespace( + id=lambda: "context-request-id", + is_stopped=cancellation.done, + async_killed_or_stopped=lambda: cancellation, + notify_first_token=lambda: None, + ) + caplog.set_level("INFO") + + async def stream(): + for index in range(4): + if index == 3: + # A new choice can first appear after the monitor has aborted + # the choices it already knows about. + await first_abort.wait() + yield { + "index": index, + "text": "hello", + "output_ids": [11], + "meta_info": { + "id": f"choice-{index}", + "finish_reason": {"type": "length"} if index == 2 else None, + }, + } + + async with aclosing(getattr(handler, processor_name)(stream(), context)) as output: + for _ in range(3): + await anext(output) + abort_request.assert_not_called() + if cancel_from_client: + cancellation.set_result(None) + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(anext(output), timeout=5) + + expected_ids = ["choice-0", "choice-1"] + if cancel_from_client: + expected_ids.append("choice-3") + abort_request.assert_has_calls( + [call(rid=request_id, abort_all=False) for request_id in expected_ids], + any_order=True, + ) + assert abort_request.call_count == len(expected_ids) + # Preserve the context-ID log consumed by the GPU cancellation tests. + assert "Aborted Request ID: context-request-id" in caplog.messages + + @pytest.mark.asyncio @pytest.mark.parametrize( "processor_name", ["_process_token_stream", "_process_text_stream"] @@ -265,7 +336,7 @@ async def test_shutdown_abort_chunk_raises_engine_shutdown(processor_name): handler = _new_decode_handler() handler.shutdown_event = asyncio.Event() handler.shutdown_event.set() - context = SimpleNamespace(id=lambda: "request-id") + context = SimpleNamespace(id=lambda: "request-id", is_stopped=lambda: False) async def stream(): yield { @@ -747,23 +818,55 @@ def test_build_sampling_params_maps_guided_decoding_to_regex(): assert sampling_params["regex"] == r"\{\}" -def test_build_sampling_params_maps_min_tokens_for_token_requests(): - handler = _new_decode_handler(use_sglang_tokenizer=False) +def test_build_sampling_params_holds_tokenizer_free_engine_open_for_min_tokens(): + handler = _new_decode_handler( + use_sglang_tokenizer=False, + skip_tokenizer_init=True, + ) sampling_params = handler._build_sampling_params( { "sampling_options": {}, "stop_conditions": { "max_tokens": 64, - "min_tokens": 64, - "ignore_eos": True, + "min_tokens": 8, + "stop_token_ids": [7], + "stop_token_ids_hidden": [151643], }, } ) - assert sampling_params["min_new_tokens"] == 64 - assert sampling_params["max_new_tokens"] == 64 + assert "min_new_tokens" not in sampling_params + assert "stop_token_ids" not in sampling_params assert sampling_params["ignore_eos"] is True + assert sampling_params["max_new_tokens"] == 64 + + from sglang.srt.sampling.sampling_params import SamplingParams + + SamplingParams(**sampling_params).normalize(tokenizer=None) + + +def test_build_sampling_params_forwards_min_tokens_when_engine_has_tokenizer(): + handler = _new_decode_handler( + use_sglang_tokenizer=False, + skip_tokenizer_init=False, + ) + + sampling_params = handler._build_sampling_params( + { + "sampling_options": {}, + "stop_conditions": { + "max_tokens": 64, + "min_tokens": 8, + "stop_token_ids": [7], + }, + } + ) + + assert sampling_params["min_new_tokens"] == 8 + assert sampling_params["stop_token_ids"] == [7] + assert "ignore_eos" not in sampling_params + assert sampling_params["max_new_tokens"] == 64 def test_build_sampling_params_maps_min_tokens_for_sglang_tokenizer_requests(): @@ -778,7 +881,10 @@ def test_build_sampling_params_maps_min_tokens_for_sglang_tokenizer_requests(): def test_build_sampling_params_omits_min_new_tokens_when_min_tokens_absent(): - handler = _new_decode_handler(use_sglang_tokenizer=False) + handler = _new_decode_handler( + use_sglang_tokenizer=False, + skip_tokenizer_init=True, + ) sampling_params = handler._build_sampling_params( {"sampling_options": {}, "stop_conditions": {"max_tokens": 8}} @@ -788,7 +894,10 @@ def test_build_sampling_params_omits_min_new_tokens_when_min_tokens_absent(): def test_build_sampling_params_forwards_explicit_zero_min_tokens(): - handler = _new_decode_handler(use_sglang_tokenizer=False) + handler = _new_decode_handler( + use_sglang_tokenizer=False, + skip_tokenizer_init=True, + ) sampling_params = handler._build_sampling_params( {"sampling_options": {}, "stop_conditions": {"max_tokens": 8, "min_tokens": 0}} diff --git a/container/deps/requirements.test.txt b/container/deps/requirements.test.txt index fe0de26787c3..852b6a590358 100644 --- a/container/deps/requirements.test.txt +++ b/container/deps/requirements.test.txt @@ -38,6 +38,8 @@ mistral-common>=1.11.2 mypy==1.18.2 # For NATS object store verification in router tests nats-py==2.12.0 +# Router benchmark tests import graph synthesis during whole-tree collection. +networkx openai==2.32.0 # for Fault Tolerance migration tests pandas-stubs>=1.3.0 psutil<=7.0.0 # System package, may vary by platform (was >=5.0.0) diff --git a/deploy/operator/docs/footer.md b/deploy/operator/docs/footer.md index 9c762224013e..4f81bef66b51 100644 --- a/deploy/operator/docs/footer.md +++ b/deploy/operator/docs/footer.md @@ -344,7 +344,6 @@ For users who want to understand the implementation details or contribute to the - **Checkpoint / Restore**: - [`internal/checkpoint/podspec.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/checkpoint/podspec.go) - Checkpoint env var injection and volume setup - [`internal/checkpoint/resolve.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/checkpoint/resolve.go) - Checkpoint resolution logic - - [`internal/checkpoint/resource.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/checkpoint/resource.go) - Checkpoint resource management - **Constants & Annotations**: [`internal/consts/consts.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/consts/consts.go) - Defines annotation keys and other constants ## Notes diff --git a/deploy/operator/internal/controller/testdata/dgd/components/output.yaml b/deploy/operator/internal/controller/testdata/dgd/components/output.yaml index 5b023b553dfb..5bbdf33f152f 100644 --- a/deploy/operator/internal/controller/testdata/dgd/components/output.yaml +++ b/deploy/operator/internal/controller/testdata/dgd/components/output.yaml @@ -12,7 +12,7 @@ spec: podTemplate: metadata: annotations: - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-worker-hash: $var:worker-hash spec: @@ -35,7 +35,7 @@ spec: podTemplate: metadata: annotations: - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* spec: containers: - image: registry.example/dynamo-frontend:1.4.0-test diff --git a/deploy/operator/internal/controller/testdata/dgd/grove/output.yaml b/deploy/operator/internal/controller/testdata/dgd/grove/output.yaml index 67fa19a2dfda..77e90a20b35f 100644 --- a/deploy/operator/internal/controller/testdata/dgd/grove/output.yaml +++ b/deploy/operator/internal/controller/testdata/dgd/grove/output.yaml @@ -13,7 +13,7 @@ spec: cliques: - annotations: nvidia.com/dynamo-discovery-backend: kubernetes - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-component: frontend nvidia.com/dynamo-component-type: frontend @@ -102,7 +102,7 @@ spec: roleName: frontend - annotations: nvidia.com/dynamo-discovery-backend: kubernetes - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-component: decode nvidia.com/dynamo-component-type: decode @@ -215,7 +215,7 @@ spec: roleName: decode-ldr - annotations: nvidia.com/dynamo-discovery-backend: kubernetes - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-component: decode nvidia.com/dynamo-component-type: decode diff --git a/deploy/operator/internal/controller/testdata/dgdr/deployment/output.yaml b/deploy/operator/internal/controller/testdata/dgdr/deployment/output.yaml index 16f41d151ce1..9a59072f28bc 100644 --- a/deploy/operator/internal/controller/testdata/dgdr/deployment/output.yaml +++ b/deploy/operator/internal/controller/testdata/dgdr/deployment/output.yaml @@ -121,7 +121,7 @@ spec: podTemplate: metadata: annotations: - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-worker-hash: "$var:worker-hash" spec: @@ -632,7 +632,7 @@ spec: podTemplate: metadata: annotations: - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* spec: containers: - args: @@ -658,7 +658,7 @@ spec: podTemplate: metadata: annotations: - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-worker-hash: "$var:worker-hash" spec: diff --git a/deploy/operator/internal/controller/testdata/dgdr/grove/output.yaml b/deploy/operator/internal/controller/testdata/dgdr/grove/output.yaml index 2d4655c400f2..80a10e3bc305 100644 --- a/deploy/operator/internal/controller/testdata/dgdr/grove/output.yaml +++ b/deploy/operator/internal/controller/testdata/dgdr/grove/output.yaml @@ -126,7 +126,7 @@ spec: cliques: - annotations: nvidia.com/dynamo-discovery-backend: kubernetes - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-component: Frontend nvidia.com/dynamo-component-type: frontend @@ -215,7 +215,7 @@ spec: roleName: frontend - annotations: nvidia.com/dynamo-discovery-backend: kubernetes - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-component: decode nvidia.com/dynamo-component-type: decode @@ -355,7 +355,7 @@ spec: roleName: decode - annotations: nvidia.com/dynamo-discovery-backend: kubernetes - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-component: prefill nvidia.com/dynamo-component-type: prefill diff --git a/deploy/operator/internal/controller/testdata/dgdr/lws/output.yaml b/deploy/operator/internal/controller/testdata/dgdr/lws/output.yaml index 0f3589e29a6e..970063993caa 100644 --- a/deploy/operator/internal/controller/testdata/dgdr/lws/output.yaml +++ b/deploy/operator/internal/controller/testdata/dgdr/lws/output.yaml @@ -125,7 +125,7 @@ spec: podTemplate: metadata: annotations: - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-worker-hash: $var:worker-hash spec: @@ -755,7 +755,7 @@ spec: podTemplate: metadata: annotations: - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* spec: containers: - args: @@ -781,7 +781,7 @@ spec: podTemplate: metadata: annotations: - nvidia.com/dynamo-operator-origin-version: $pattern:1\.4\..* + nvidia.com/dynamo-operator-origin-version: $pattern:1\.5\..* labels: nvidia.com/dynamo-worker-hash: $var:worker-hash spec: diff --git a/docs/fern/pages/recipes/model-recipes/kimi-k3.mdx b/docs/fern/pages/recipes/model-recipes/kimi-k3.mdx index d539c84b3357..d77c5eea465a 100644 --- a/docs/fern/pages/recipes/model-recipes/kimi-k3.mdx +++ b/docs/fern/pages/recipes/model-recipes/kimi-k3.mdx @@ -230,7 +230,7 @@ All four targets serve Kimi-K3 on vLLM with FP8 KV cache, KV-aware routing, pref - Source README: [recipes/kimi-k3/README.md](https://github.com/ai-dynamo/dynamo/blob/main/recipes/kimi-k3/README.md) - Aggregated GB200: [deploy.yaml](https://github.com/ai-dynamo/dynamo/blob/main/recipes/kimi-k3/vllm/agg-gb200-agentic/deploy.yaml) -- Disaggregated GB200: [deploy.yaml](https://github.com/ai-dynamo/dynamo/blob/main/recipes/kimi-k3/vllm/disagg-gb200-agentic/deploy.yaml) +- Disaggregated GB200: [deploy.yaml](https://github.com/ai-dynamo/dynamo/blob/release/1.5.0/recipes/kimi-k3/vllm/disagg-gb200-agentic/deploy.yaml) - Aggregated GB300: [deploy.yaml](https://github.com/ai-dynamo/dynamo/blob/main/recipes/kimi-k3/vllm/agg-gb300-agentic/deploy.yaml) - Disaggregated GB300: [deploy.yaml](https://github.com/ai-dynamo/dynamo/blob/main/recipes/kimi-k3/vllm/disagg-gb300-agentic/deploy.yaml) - Setup assets: [model-cache.yaml](https://github.com/ai-dynamo/dynamo/blob/main/recipes/kimi-k3/model-cache/model-cache.yaml) and [model-download.yaml](https://github.com/ai-dynamo/dynamo/blob/main/recipes/kimi-k3/model-cache/model-download.yaml) diff --git a/docs/fern/pages/reference/kubernetes-api/additional-resources/api-reference-k8s.md b/docs/fern/pages/reference/kubernetes-api/additional-resources/api-reference-k8s.md index 968152b277e3..ae5a67356d91 100644 --- a/docs/fern/pages/reference/kubernetes-api/additional-resources/api-reference-k8s.md +++ b/docs/fern/pages/reference/kubernetes-api/additional-resources/api-reference-k8s.md @@ -3869,7 +3869,6 @@ For users who want to understand the implementation details or contribute to the - **Checkpoint / Restore**: - [`internal/checkpoint/podspec.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/checkpoint/podspec.go) - Checkpoint env var injection and volume setup - [`internal/checkpoint/resolve.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/checkpoint/resolve.go) - Checkpoint resolution logic - - [`internal/checkpoint/resource.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/checkpoint/resource.go) - Checkpoint resource management - **Constants & Annotations**: [`internal/consts/consts.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/consts/consts.go) - Defines annotation keys and other constants ## Notes diff --git a/docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx b/docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx index 2d027972aeb6..9ed97d1fa4ae 100644 --- a/docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx +++ b/docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx @@ -5089,7 +5089,6 @@ For users who want to understand the implementation details or contribute to the - **Checkpoint / Restore**: - [`internal/checkpoint/podspec.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/checkpoint/podspec.go) - Checkpoint env var injection and volume setup - [`internal/checkpoint/resolve.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/checkpoint/resolve.go) - Checkpoint resolution logic - - [`internal/checkpoint/resource.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/checkpoint/resource.go) - Checkpoint resource management - **Constants & Annotations**: [`internal/consts/consts.go`](https://github.com/ai-dynamo/dynamo/blob/main/deploy/operator/internal/consts/consts.go) - Defines annotation keys and other constants ### Notes diff --git a/docs/fern/pages/use-cases/reinforcement-learning/weight-updates.md b/docs/fern/pages/use-cases/reinforcement-learning/weight-updates.md index e6e43f253dab..341d1ba791f1 100644 --- a/docs/fern/pages/use-cases/reinforcement-learning/weight-updates.md +++ b/docs/fern/pages/use-cases/reinforcement-learning/weight-updates.md @@ -17,7 +17,7 @@ ModelExpress can stage a version from three sources: - **Artifact to inference:** an inference worker prepares a canonical checkpoint from S3, including the current exact-base XOR delta format, before engine installation. - **Inference to inference:** after an inference worker applies a version, a rank-compatible worker can pull that version from it instead of returning to the trainer. -All three refit paths are Experimental. The [ModelExpress RL refit package](https://github.com/ai-dynamo/modelexpress/tree/main/modelexpress_client/python/modelexpress_rl) contains the current source strategies. The [Dynamo vLLM refit example](https://github.com/ai-dynamo/modelexpress/tree/main/examples/rl/dynamo_vllm_refit) validates the full-weight trainer-to-inference path only; it does not qualify the S3 delta path or every backend and topology. ModelExpress startup loading remains a separate boot and scale-out workflow. +All three refit paths are Experimental. The [ModelExpress RL refit package](https://github.com/ai-dynamo/modelexpress/tree/main/modelexpress_client/python/modelexpress_rl) contains the current source strategies. The [Dynamo vLLM reshard refit example](https://github.com/ai-dynamo/modelexpress/tree/main/examples/rl/dynamo_vllm_reshard_refit) validates the full-weight trainer-to-inference path, and the [Dynamo vLLM S3 delta refit example](https://github.com/ai-dynamo/modelexpress/tree/main/examples/rl/dynamo_vllm_s3_delta_refit) exercises the artifact-to-inference path with a full snapshot followed by incremental deltas. Neither qualifies every backend and topology. ModelExpress startup loading remains a separate boot and scale-out workflow. > [!NOTE] > ModelExpress startup loading and live RL refit have different interfaces. vLLM 0.23 and later use the native `--load-format modelexpress`; `mx` remains a backward-compatible alias. Earlier plugin-based images generally use `mx`, while legacy split-loader images can expose `mx-source` and `mx-target`. Match the loader to the installed vLLM, ModelExpress package, and runtime image instead of inferring it from the refit API. See [Dynamo deprecations](../../reference/general/releases/deprecations.mdx) for the Dynamo-owned loader migration. diff --git a/examples/backends/vllm/launch/agg_lmcache_mp.sh b/examples/backends/vllm/launch/agg_lmcache_mp.sh index 763ed36157ca..1f9cbdfe7ae3 100755 --- a/examples/backends/vllm/launch/agg_lmcache_mp.sh +++ b/examples/backends/vllm/launch/agg_lmcache_mp.sh @@ -16,6 +16,12 @@ MODEL="Qwen/Qwen3-0.6B" # ---- Tunable (override via env vars) ---- MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}" MAX_CONCURRENT_SEQS="${MAX_CONCURRENT_SEQS:-2}" +# LMCache 0.5.4 can need more than 60 seconds to initialize on a cold host. +LMCACHE_STARTUP_TIMEOUT_SECONDS="${LMCACHE_STARTUP_TIMEOUT_SECONDS:-180}" +if [[ ! "$LMCACHE_STARTUP_TIMEOUT_SECONDS" =~ ^[1-9][0-9]*$ ]]; then + echo "LMCACHE_STARTUP_TIMEOUT_SECONDS must be a positive integer" >&2 + exit 2 +fi # Default KV cache cap from profiling (2x safety over min=560 MiB); ~3.8 GiB peak VRAM # Profiler/test framework overrides via env @@ -32,16 +38,30 @@ print_launch_banner "Launching Aggregated Serving (1 GPU) + LMCache" "$MODEL" "$ lmcache server \ --l1-size-gb "${LMCACHE_L1_SIZE_GB:-16}" --eviction-policy LRU \ --port "$LMCACHE_PORT" --http-port "$LMCACHE_HTTP_PORT" & +LMCACHE_PID=$! # Wait until the server's HTTP admin endpoint is healthy before launching workers. -for _ in $(seq 1 60); do - curl -sf "http://localhost:$LMCACHE_HTTP_PORT/healthcheck" >/dev/null 2>&1 && break +lmcache_ready=false +for ((attempt = 0; attempt < LMCACHE_STARTUP_TIMEOUT_SECONDS; attempt++)); do + if curl -sf "http://localhost:$LMCACHE_HTTP_PORT/healthcheck" >/dev/null 2>&1; then + lmcache_ready=true + break + fi + if ! kill -0 "$LMCACHE_PID" 2>/dev/null; then + lmcache_rc=0 + wait "$LMCACHE_PID" || lmcache_rc=$? + echo "lmcache server exited with code $lmcache_rc before becoming healthy" >&2 + if ((lmcache_rc == 0)); then + lmcache_rc=1 + fi + exit "$lmcache_rc" + fi sleep 1 done -curl -sf "http://localhost:$LMCACHE_HTTP_PORT/healthcheck" >/dev/null 2>&1 || { - echo "lmcache server failed to become healthy on :$LMCACHE_HTTP_PORT" >&2 +if [[ "$lmcache_ready" != true ]]; then + echo "lmcache server failed to become healthy on :$LMCACHE_HTTP_PORT after ${LMCACHE_STARTUP_TIMEOUT_SECONDS}s" >&2 exit 1 -} +fi python -m dynamo.frontend & diff --git a/lib/llm/benches/request_trace_finish_metadata.rs b/lib/llm/benches/request_trace_finish_metadata.rs index 582f98da7c82..3afc53e4f5dc 100644 --- a/lib/llm/benches/request_trace_finish_metadata.rs +++ b/lib/llm/benches/request_trace_finish_metadata.rs @@ -115,6 +115,8 @@ fn backend_outputs(count: usize) -> Vec { worker_trace_link: None, engine_data: None, routing_data: None, + encoder_result: None, + jailed_text: None, }) .collect() } diff --git a/lib/llm/benches/tokenizer_simple.rs b/lib/llm/benches/tokenizer_simple.rs index ebb758d497ba..a9dbe9d1b729 100644 --- a/lib/llm/benches/tokenizer_simple.rs +++ b/lib/llm/benches/tokenizer_simple.rs @@ -62,7 +62,7 @@ pub fn decode(c: &mut Criterion) { let tokenizer: Arc = Arc::new(HuggingFaceTokenizer::from_file(TEST_TOKENIZER).unwrap()); let ds = DecodeStream::new(tokenizer, &[], false); - Decoder::new(ds, StopConditions::default(), false, None) + Decoder::new(ds, StopConditions::default(), false, None, None) }, |mut decoder| { for tok in black_box(TEST_TOKS) { @@ -86,7 +86,7 @@ pub fn decode_big(c: &mut Criterion) { let tokenizer: Arc = Arc::new(HuggingFaceTokenizer::from_file(TEST_TOKENIZER).unwrap()); let ds = DecodeStream::new(tokenizer, &[], false); - Decoder::new(ds, StopConditions::default(), false, None) + Decoder::new(ds, StopConditions::default(), false, None, None) }, |mut decoder| { for tok in black_box(&BIG_TEST_TOKS) { @@ -127,7 +127,7 @@ pub fn tiktoken_decode(c: &mut Criterion) { let tokenizer: Arc = Arc::new(TikTokenTokenizer::from_file_auto(TEST_TIKTOKEN).unwrap()); let ds = DecodeStream::new(tokenizer, &[], false); - Decoder::new(ds, StopConditions::default(), false, None) + Decoder::new(ds, StopConditions::default(), false, None, None) }, |mut decoder| { for tok in black_box(&toks) { diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs index 2d7da8b3c5b2..ab67019c8661 100644 --- a/lib/llm/src/backend.rs +++ b/lib/llm/src/backend.rs @@ -76,6 +76,17 @@ struct DecoderUnfoldState { /// (e.g. SGLang with --skip-tokenizer-init forcibly drops it). tokenizer: Tokenizer, skip_special_tokens: bool, + /// Text flushed from a choice's decoder because the underlying engine stream ended + /// without ever sending that choice a terminal `finish_reason`, queued here so each + /// flushed choice can be emitted as its own synthetic final chunk. + pending_flush: Vec<(u32, String)>, + /// Set once `stream` has yielded `None`, so it is never polled again -- a `Stream` is + /// not guaranteed to be safely pollable past its first `None`. + stream_ended: bool, + /// Whether this request could still be migrated (see `DecoderParams::migration_possible`). + /// When `false`, `jailed_text` is left unset on every chunk rather than snapshotting the + /// decoder's withheld state for a checkpoint nothing can ever consume. + migration_possible: bool, } fn fill_missing_top_logprob_text( @@ -105,6 +116,14 @@ struct DecoderParams { include_stop_str_in_output: bool, tracker: Option>, n: u32, + // Withheld hidden-stop-sequence prefix carried over from a migrated attempt's last + // known-good chunk (see `PreprocessedRequest::jail_seed`). `None` on a first attempt. + jail_seed: Option, + // Whether this request could still be migrated to another worker (i.e. a `RetryManager` + // sits in front of this `Backend` and has retries configured). When `false`, no chunk + // this Backend emits can ever be reseeded into a retry, so there is no point snapshotting + // the decoder's withheld state onto every chunk via `jailed_text`. + migration_possible: bool, } impl DecoderParams { @@ -124,6 +143,8 @@ impl DecoderParams { .unwrap_or(false), tracker: request.tracker.clone(), n: request.sampling_options.n.unwrap_or(1) as u32, + jail_seed: request.jail_seed.clone(), + migration_possible: request.migration_state.is_some(), } } } @@ -168,6 +189,10 @@ impl Backend { params.stop_conditions.clone(), params.include_stop_str_in_output, params.tracker.clone(), + // Every choice starts from the same carried-over withheld prefix. This + // only matters for `n == 1` migration retries in practice; a fresh + // decoder normally starts unseeded (`None`). + params.jail_seed.clone(), ); decoders.insert(idx, decoder); } @@ -180,6 +205,9 @@ impl Backend { finished: false, tokenizer: tokenizer.clone(), skip_special_tokens: params.skip_special_tokens, + pending_flush: Vec::new(), + stream_ended: false, + migration_possible: params.migration_possible, }) } } @@ -211,13 +239,52 @@ impl return None; } - match state.stream.next().await { + // `stream` already yielded `None` once; do not poll it again (unspecified + // behavior for most `Stream` impls). Only drain the flush queue from here on. + if state.stream_ended { + return state.pending_flush.pop().map(|(idx, flushed)| { + let output = Annotated::from_data(LLMEngineOutput { + index: Some(idx), + text: Some(flushed), + finish_reason: Some(FinishReason::Stop), + ..Default::default() + }); + (output, state) + }); + } + + let output = loop { + let output = state.stream.next().await; + let choice_already_finished = output.as_ref().is_some_and(|output| { + !output.is_event() + && output.data.as_ref().is_some_and(|data| { + state + .finished_choices + .contains(&data.index.unwrap_or_default()) + }) + }); + if !choice_already_finished { + break output; + } + }; + + match output { Some(output) => { // move to state.process_output // handle any error conditions / unwraps here // events are pass thru if output.is_event() || output.data.is_none() { + // A genuine stream-level error ends generation for every choice at + // once -- `Annotated::from_err` carries no per-choice `index` the + // way `LLMEngineOutput` does, so there is no narrower target than + // "all of them". Mark them finished so the bare-EOF flush loop below + // does not synthesize a spurious successful `Stop` for any choice + // after a real error, the same hardening already applied to decode + // errors and the decoded-text bypass path. + if output.is_error() { + state.finished_choices.extend(state.decoders.keys().copied()); + } return Some((output, state)); } @@ -253,20 +320,57 @@ impl && !state.validate_engine_decode { // Text already decoded; track finish for this choice - if data.finish_reason.is_some() { - let choice_idx = data.index.unwrap_or(0); + let choice_idx = data.index.unwrap_or(0); + let has_finish = data.finish_reason.is_some(); + if has_finish { state.finished_choices.insert(choice_idx); } + // This path bypasses the decoder entirely (the engine pre-decoded + // its own text for this chunk), but the decoder for this choice is + // still alive across chunks and may be holding text withheld by a + // *previous* chunk that did go through it. Losing track of that here + // would either drop it silently (on a terminal chunk) or make a + // migration checkpoint captured from this chunk (see `jail_seed`) + // look falsely resolved, so: + if let Some(decoder) = state.decoders.get_mut(&choice_idx) { + // On a terminal chunk, flush it and put it *before* this chunk's + // own text -- it is strictly older -- rather than appending, + // which would reorder output (e.g. "there" + withheld "o" must + // come out as "othere", not "thereo"). + if has_finish + && let Some(flushed) = decoder.flush_jailed() + && let Some(data) = &mut output.data + { + let newer = data.text.take().unwrap_or_default(); + data.text = Some(flushed + &newer); + } + // Mirror the decoder's current withheld state on every chunk + // (terminal or not), even though this chunk didn't change it, so + // a checkpoint captured here reflects reality instead of always + // reading as resolved. Skipped entirely when migration can't + // happen -- nothing will ever read this checkpoint, so there is + // no reason to allocate a snapshot of it. + if state.migration_possible + && let Some(data) = &mut output.data + { + data.jailed_text = decoder.peek_jailed(); + } + } return Some((output, state)); } let data = output.data.as_ref().unwrap(); let choice_idx = data.index.unwrap_or(0); + // Snapshot the choice count before borrowing a specific decoder mutably + // below: that borrow stays alive until this choice's `peek_jailed()` call + // near the end of this arm, so `state.decoders` can't be read again + // (even just its length) in between. The count itself never changes + // after `Backend::decoder()` creates the map, so this is safe to cache. + let decoders_count = state.decoders.len(); let Some(decoder) = state.decoders.get_mut(&choice_idx) else { tracing::error!( - "engine emitted choice index {choice_idx}, but only {} choices were requested", - state.decoders.len() + "engine emitted choice index {choice_idx}, but only {decoders_count} choices were requested" ); let mut output = output; if let Some(data) = &mut output.data { @@ -277,12 +381,12 @@ impl return Some((output, state)); }; - let result = match decoder.process_token_ids(&data.token_ids) { + let mut result = match decoder.process_token_ids(&data.token_ids) { Ok(result) => result, Err(e) => { tracing::error!("Failed to process token_ids for choice {choice_idx}: {e}"); state.finished_choices.insert(choice_idx); - if state.finished_choices.len() >= state.decoders.len() { + if state.finished_choices.len() >= decoders_count { state.stream.context().stop_generating(); state.finished = true; } @@ -295,6 +399,19 @@ impl } }; + // The engine can report its own completion (e.g. it hit `max_tokens`) + // without our decoder ever detecting a local stop condition. Any text + // still withheld as a partial hidden-stop-sequence match can never + // complete at that point, so flush it now rather than silently dropping + // it -- this is the last chance before the decoder for this choice is + // discarded. + if result.stop_trigger.is_none() + && data.finish_reason.is_some() + && let Some(flushed) = decoder.flush_jailed() + { + result.text.get_or_insert_with(String::new).push_str(&flushed); + } + // NOTE: the `finish_reason` is computed from the generated `token_ids` alone. // The `data` field can have a `finish_reason` set, coming from the underlying // LLM inference `Engine`, and empty `token_ids`. See comment below for more details. @@ -343,7 +460,7 @@ impl // Once all expected choices are finished, stop the upstream generator. if finish_reason.is_some() && data.finish_reason.is_none() { state.finished_choices.insert(choice_idx); - if state.finished_choices.len() >= state.decoders.len() { + if state.finished_choices.len() >= decoders_count { state.stream.context().stop_generating(); state.finished = true; } @@ -384,14 +501,67 @@ impl data.stop_reason = stop_reason.or(data.stop_reason); } data.text = text; + + // A tokenizer-free engine can return tokens past a local + // stop in the same chunk. Keep its per-token fields aligned + // with the prefix the decoder consumed. + let consumed = tokens.len(); + if consumed < data.token_ids.len() { + data.token_ids.truncate(consumed); + if let Some(log_probs) = data.log_probs.as_mut() { + log_probs.truncate(consumed); + } + if let Some(top_logprobs) = data.top_logprobs.as_mut() { + top_logprobs.truncate(consumed); + } + } data.tokens = Some(tokens); + // Snapshot of whatever this choice's decoder is still withholding as a + // possible hidden-stop-sequence prefix after this step -- `None` once + // resolved (matched, ruled out, or flushed). Carried so a migration + // retry's fresh decoder can be reseeded from the last known-good chunk + // instead of silently losing it (see `PreprocessedRequest::jail_seed`). + // Skipped when migration can't happen: nothing will ever read this + // checkpoint, so there is no reason to allocate a snapshot of it. + if state.migration_possible { + data.jailed_text = decoder.peek_jailed(); + } output.data = Some(data); Some((output, state)) } - None => None, + None => { + // The engine's stream ended without ever sending a terminal + // `finish_reason` to some choice -- no more tokens are coming for any + // of them, so flush whatever each decoder still holds back as an + // incomplete hidden-stop-sequence match before the decoders are + // dropped, and drain the flushed choices one synthetic chunk at a time. + state.stream_ended = true; + for (idx, decoder) in state.decoders.iter_mut() { + // A choice that already finished -- successfully or with an error -- + // was already given its chance to flush (see the decoder-driven and + // decoded-text-passthrough paths above). Synthesizing another chunk + // for it here would, at best, be redundant and, at worst, turn a + // choice that ended in an error into a spurious extra `Stop`. + if state.finished_choices.contains(idx) { + continue; + } + if let Some(flushed) = decoder.flush_jailed() { + state.pending_flush.push((*idx, flushed)); + } + } + state.pending_flush.pop().map(|(idx, flushed)| { + let output = Annotated::from_data(LLMEngineOutput { + index: Some(idx), + text: Some(flushed), + finish_reason: Some(FinishReason::Stop), + ..Default::default() + }); + (output, state) + }) + } } }) .fuse(); @@ -417,6 +587,7 @@ impl worker_trace_link: data.worker_trace_link, engine_data: data.engine_data, routing_data: data.routing_data, + jailed_text: data.jailed_text, }) }) }); @@ -496,6 +667,11 @@ pub struct Decoder { // the number of bytes currently jailed jailed_bytes: usize, + + // Scratch buffers reused across `longest_hidden_prefix_suffix` calls (one per decoded + // token, for the lifetime of the request) instead of allocating fresh ones every step. + kmp_scratch: Vec, + kmp_pi_scratch: Vec, } #[allow(dead_code)] @@ -510,21 +686,49 @@ pub enum StopTrigger { } pub struct StepResult { + /// This step's own decoded text, independent of any stop-sequence withholding. + /// `SeqResult.tokens[i]` must equal the text of `token_ids[i]` so that consumers + /// zipping `tokens` with per-token logprobs (e.g. `create_logprobs`) stay aligned -- + /// logprobs describe what the model generated and must not change because a hidden + /// stop sequence also happens to be in flight. pub token: Option, + /// Text to append to the caller-visible `text`/content this step. May be `None` on a + /// step whose text is being withheld as a possible hidden-stop-sequence prefix, and may + /// carry more than one prior step's worth of text on the step that releases a withheld + /// backlog. + pub released_text: Option, pub stop_trigger: Option, } impl StepResult { + /// No stop-sequence withholding in effect: the token's own text is released as-is. fn ok(token: Option) -> Self { + Self { + token: token.clone(), + released_text: token, + stop_trigger: None, + } + } + + /// `token` and `released_text` differ, e.g. because `released_text` also carries a + /// previously withheld backlog, or because a matched stop sequence is excluded from + /// `released_text` but not from `token`. + fn ok_split(token: Option, released_text: Option) -> Self { Self { token, + released_text, stop_trigger: None, } } - fn with_stop_trigger(token: Option, stop_trigger: StopTrigger) -> Self { + fn with_stop_trigger( + token: Option, + released_text: Option, + stop_trigger: StopTrigger, + ) -> Self { Self { token, + released_text, stop_trigger: Some(stop_trigger), } } @@ -544,6 +748,9 @@ impl Decoder { stop_condition: StopConditions, include_stop_str_in_output: bool, tracker: Option>, + // Withheld hidden-stop-sequence prefix to resume from, e.g. after a migration + // retry (see `PreprocessedRequest::jail_seed`). `None` starts unseeded, as before. + jail_seed: Option, ) -> Self { let user_stop_ids: HashSet = stop_condition .stop_token_ids @@ -582,6 +789,13 @@ impl Decoder { .max() .unwrap_or(0); + // The entire seed is, by construction, text a prior attempt withheld as a partial + // hidden-stop-sequence match that had not yet resolved -- treat all of it as still + // jailed so this attempt can either complete the match or release it exactly as the + // original attempt would have, rather than leaking it or re-checking it as new text. + let jail = jail_seed.unwrap_or_default(); + let jailed_bytes = jail.len(); + Self { decode_stream, tracker, @@ -592,9 +806,11 @@ impl Decoder { visible_stop_sequences, min_tokens: stop_condition.min_tokens.unwrap_or(0), generated_tokens: 0, - jail: String::new(), + jail, jail_max_bytes, - jailed_bytes: 0, + jailed_bytes, + kmp_scratch: Vec::new(), + kmp_pi_scratch: Vec::new(), } } @@ -605,6 +821,7 @@ impl Decoder { /// In the future, this method may kick off async cpu/tokio tasks and or async cuda tasks to /// handle logits post-processing and/or other tasks. pub fn step(&mut self, token_id: TokenIdType) -> Result { + let below_min_tokens = self.generated_tokens < self.min_tokens; // increment the generated tokens self.generated_tokens += 1; @@ -620,14 +837,27 @@ impl Decoder { } // stop conditions to not apply until the minimum number of tokens have been generated - if self.generated_tokens < self.min_tokens { + if below_min_tokens { return Ok(StepResult::ok(token)); } // Check token stops. Visible token IDs are included in output. if self.visible_stop_ids.contains(&token_id) { + // Any text still withheld as a partial hidden-stop-sequence match can never + // complete now, so it goes out ahead of this (included) token's own text -- + // otherwise it would be silently dropped and, if it were released later, would + // come out of order. + let released = match (self.flush_jailed(), &token) { + (Some(mut flushed), Some(t)) => { + flushed.push_str(t); + Some(flushed) + } + (Some(flushed), None) => Some(flushed), + (None, t) => t.clone(), + }; return Ok(StepResult::with_stop_trigger( token, + released, StopTrigger::VisibleStopTokenDetected(token_id), )); } @@ -640,35 +870,48 @@ impl Decoder { } else { StopTrigger::HiddenStopTokenDetected(token_id) }; - return Ok(StepResult::with_stop_trigger(None, trigger)); + // This token's own text stays hidden, as before. But any text still withheld + // from an earlier, never-completed hidden-stop-sequence prefix match is not + // part of *this* stop and must still reach the caller. + return Ok(StepResult::with_stop_trigger( + None, + self.flush_jailed(), + trigger, + )); } // check stop sequences - the jail will always hold at least the largest stop sequence // if jail_max_bytes is 0, then there are no stop sequences if self.jail_max_bytes > 0 - && let Some(token) = &token + && let Some(token_text) = &token { - let pre_append = self.jail.len(); - self.jail.push_str(token); + // bytes at the tail of `jail` withheld by a previous step because they could + // still grow into a complete hidden stop sequence; everything before this point + // has already been released to the caller. + let release_start = self.jail.len() - self.jailed_bytes; + self.jail.push_str(token_text); // Check hidden stop sequences first (excluded from output) for seq in &self.hidden_stop_sequences { if let Some(offset) = galil_seiferas::gs_find(self.jail.as_bytes(), seq.as_bytes()) { - // return only new bytes after pre_append .. offset (excluding stop sequence) + // return only new bytes after release_start .. offset (excluding stop sequence) // example: seq = "ox", token = "boxes", return "b" - // note: this changes when we start jailing tokens for partial matches - // on the suffix of the jail with prefixes of the stop sequences // - // we might have returned a partial match, if so, then offset < pre_append - // in that case, we return the empty string - let partial_token = if offset >= pre_append { - self.jail[pre_append..offset].to_string() - } else { - "".to_string() - }; + // we might have returned a partial match, if so, then offset < release_start + // in that case, we return no text + let partial_token = (offset >= release_start) + .then(|| self.jail[release_start..offset].to_string()) + .filter(|s| !s.is_empty()); + self.jailed_bytes = 0; + // `token` (this step's own raw decoded text) is reported unchanged so + // that SeqResult.tokens[i] keeps describing token_ids[i] for logprobs; + // only the caller-visible `released_text` excludes the matched sequence. + // `token_text`'s last use was the `push_str` above, so `token` itself + // (not yet borrowed at this point) can move here instead of cloning. return Ok(StepResult::with_stop_trigger( - Some(partial_token), + token, + partial_token, StopTrigger::HiddenStopSequenceDetected(seq.to_string()), )); } @@ -679,27 +922,125 @@ impl Decoder { if let Some(offset) = galil_seiferas::gs_find(self.jail.as_bytes(), seq.as_bytes()) { // For visible stop sequences, include the stop string in the output - // Return all text from pre_append up to and including the stop sequence + // Return all text from release_start up to and including the stop sequence let stop_end = offset + seq.len(); - let token_with_stop = if stop_end > pre_append { - self.jail[pre_append..stop_end].to_string() - } else { - // Stop sequence was entirely in previously returned text - "".to_string() - }; + let token_with_stop = (stop_end > release_start) + .then(|| self.jail[release_start..stop_end].to_string()) + .filter(|s| !s.is_empty()); + self.jailed_bytes = 0; + // Same reasoning as the hidden-sequence branch above: `token` can move + // here instead of cloning. return Ok(StepResult::with_stop_trigger( - Some(token_with_stop), + token, + token_with_stop, StopTrigger::VisibleStopSequenceDetected(seq.to_string()), )); } } + // No complete match. Withhold the longest tail of `jail` that is still a viable + // prefix of some hidden stop sequence -- a later token could complete it into a + // full match. Everything else since the last release is now safe to hand back: + // it cannot be part of a hidden stop sequence, complete or partial. + self.jailed_bytes = self.longest_hidden_prefix_suffix(); + let release_end = self.jail.len() - self.jailed_bytes; + let released = (release_end > release_start) + .then(|| self.jail[release_start..release_end].to_string()); + Self::maybe_drain_to_max_bytes(&mut self.jail, self.jail_max_bytes); + // `token` is this step's own raw decoded text (for logprobs); `released` is + // what, if anything, newly clears the jail this step (for `text`/content). + return Ok(StepResult::ok_split(token, released)); } Ok(StepResult::ok(token)) } + /// Releases any text still withheld as a partial hidden-stop-sequence match. Call this + /// once it is known no further tokens are coming (e.g. the engine reports completion + /// without `Decoder` having detected a local stop condition) so a stop sequence that + /// never completed is not silently dropped from the output. + pub fn flush_jailed(&mut self) -> Option { + let flushed = self.jailed_string(); + self.jailed_bytes = 0; + flushed + } + + /// Non-consuming look at whatever is currently withheld as a possible hidden-stop- + /// sequence prefix, without releasing it. Unlike [`Self::flush_jailed`], this does not + /// end the withholding -- it exists so a caller (e.g. the streaming pipeline) can + /// snapshot the in-flight jail state onto each chunk, so it can be recovered and used to + /// reseed a fresh `Decoder` (via `jail_seed` in [`Self::new`]) if this attempt is + /// abandoned partway through, e.g. by a migration retry. + pub(crate) fn peek_jailed(&self) -> Option { + self.jailed_string() + } + + /// Returns the length, in bytes, of the longest suffix of `self.jail` that is also a + /// strict prefix of some hidden stop sequence -- text that might still grow into a + /// complete hidden stop sequence and so must not be released to the caller yet. Only + /// considers byte offsets that land on a `jail` char boundary, so the result is always + /// safe to slice with. Reuses `self.kmp_scratch`/`self.kmp_pi_scratch` across calls + /// (one per decoded token) instead of allocating fresh buffers every step. + fn longest_hidden_prefix_suffix(&mut self) -> usize { + let jail_bytes = self.jail.as_bytes(); + let mut best = 0; + for seq in &self.hidden_stop_sequences { + let seq_bytes = seq.as_bytes(); + // A full-length match would already have been caught as a complete stop; + // only strictly shorter prefixes are candidates here. Sequences that cannot + // possibly beat the current best are skipped outright. + let max_k = seq_bytes.len().saturating_sub(1).min(jail_bytes.len()); + if max_k <= best { + continue; + } + let pattern = &seq_bytes[..max_k]; + let tail_len = pattern.len().min(jail_bytes.len()); + let tail = &jail_bytes[jail_bytes.len() - tail_len..]; + + // Longest suffix of `tail` that is also a prefix of `pattern`, found in + // O(pattern.len() + tail.len()) via the standard KMP-border trick: build the + // prefix (failure) function of `pattern + sep + tail` (`sep` = 0xFF, which + // cannot occur in valid UTF-8, so no border can bridge across it) and read the + // border length back from its last entry. Borders of `pattern + sep + tail` + // strictly decrease along the classic `pi[k - 1]` chain, which is walked here + // only as far as needed to find one that also lands on a `jail` char boundary. + self.kmp_scratch.clear(); + self.kmp_scratch.extend_from_slice(pattern); + self.kmp_scratch.push(0xFF); + self.kmp_scratch.extend_from_slice(tail); + Self::kmp_prefix_function_into(&self.kmp_scratch, &mut self.kmp_pi_scratch); + + let mut k = *self.kmp_pi_scratch.last().unwrap_or(&0); + while k > best { + if self.jail.is_char_boundary(jail_bytes.len() - k) { + best = k; + break; + } + k = self.kmp_pi_scratch[k - 1]; + } + } + best + } + + /// Standard KMP prefix (failure) function, written into `pi` (cleared and reused + /// rather than allocated fresh every call): `pi[i]` is the length of the longest + /// proper prefix of `s[..=i]` that is also a suffix of it. + fn kmp_prefix_function_into(s: &[u8], pi: &mut Vec) { + pi.clear(); + pi.resize(s.len(), 0); + let mut k = 0usize; + for i in 1..s.len() { + while k > 0 && s[i] != s[k] { + k = pi[k - 1]; + } + if s[i] == s[k] { + k += 1; + } + pi[i] = k; + } + } + pub fn process_token_ids(&mut self, token_ids: &[TokenIdType]) -> Result { let mut text: Option = None; let mut tokens = Vec::with_capacity(token_ids.len()); @@ -707,13 +1048,17 @@ impl Decoder { for token_id in token_ids { let StepResult { token, + released_text, stop_trigger, } = self.step(*token_id)?; - // Always include token text (for visible stops, the stop string is already in the token) - if let Some(token) = &token { + // `text` accumulates the caller-visible content, which can lag behind and later + // release more than one step's worth at once. `tokens[i]` always reports + // token_ids[i]'s own decoded text, independent of that withholding, so per-token + // consumers (logprobs) stay aligned with token_ids. + if let Some(released_text) = &released_text { text.get_or_insert_with(|| String::with_capacity(token_ids.len())) - .push_str(token); + .push_str(released_text); } tokens.push(token); @@ -833,6 +1178,7 @@ mod tests { [102] => " Okay", [103] => "<|user|>", [101, 103] => "Okay<|user|>", + [101, 103, 102] => "Okay<|user|> Okay", _ => anyhow::bail!("unexpected token IDs: {token_ids:?}"), }; Ok(traits::DecodeResult::Complete(token.to_string())) @@ -855,20 +1201,26 @@ mod tests { &self, request: SingleIn, ) -> Result>, Error> { - let output = LLMEngineOutput { - // SGLang's output_ids_through_stop() includes the matched stop - // position. The frontend must hide its text without mutating - // this raw token sequence. - token_ids: vec![101, 103], - finish_reason: Some(FinishReason::Stop), - index: Some(0), - ..Default::default() + let output = |index, token_ids: Vec| { + let len = token_ids.len(); + Annotated::from_data(LLMEngineOutput { + token_ids, + log_probs: Some(vec![-0.1; len]), + top_logprobs: Some(vec![vec![]; len]), + index: Some(index), + ..Default::default() + }) }; + let outputs = vec![ + // Choice 0 stops inside this chunk; its suffix must be removed. + output(0, vec![101, 103, 102]), + // It keeps generating while choice 1 is unfinished; drop this chunk. + output(0, vec![102]), + output(1, vec![101, 103]), + ]; Ok(ResponseStream::new( - Box::pin(futures::stream::once(async move { - Annotated::from_data(output) - })), + Box::pin(futures::stream::iter(outputs)), request.context(), )) } @@ -1015,7 +1367,7 @@ mod tests { } #[tokio::test] - async fn test_sglang_hidden_stop_text_is_suppressed_without_mutating_tito_ids() { + async fn test_sglang_local_stop_discards_all_suffix_tokens() { use crate::protocols::common::extensions::NvExt; let tokenizer: Arc = Arc::new(CandidateDecoder); @@ -1027,27 +1379,45 @@ mod tests { stop_token_ids_hidden: Some(vec![103]), ..Default::default() }) - .sampling_options(SamplingOptions::default()) + .sampling_options(SamplingOptions { + n: Some(2), + ..Default::default() + }) .output_options(OutputOptions::default()) .build() .expect("valid preprocessed request"); let engine: ServerStreamingEngine> = Arc::new(SyntheticSglangStopEngine); - let mut stream = Operator::generate(backend.as_ref(), SingleIn::new(request), engine) + let stream = Operator::generate(backend.as_ref(), SingleIn::new(request), engine) .await .expect("backend generation succeeds"); - let output = stream - .next() - .await - .expect("backend emits a response") + let mut outputs = stream.collect::>().await; + assert_eq!( + outputs.len(), + 2, + "chunks after a local stop must be dropped" + ); + let output = outputs + .remove(0) + .data + .expect("response contains backend output"); + let second = outputs + .remove(0) .data .expect("response contains backend output"); + assert_eq!(output.index, Some(0)); assert_eq!(output.text.as_deref(), Some("Okay")); assert_eq!(output.tokens, vec![Some("Okay".to_string()), None]); assert_eq!(output.token_ids, vec![101, 103]); + assert_eq!(output.log_probs, Some(vec![-0.1; 2])); + assert_eq!(output.top_logprobs.as_ref().map(Vec::len), Some(2)); assert_eq!(output.finish_reason, Some(FinishReason::Stop)); + assert_eq!(second.index, Some(1)); + assert_eq!(second.text.as_deref(), Some("Okay")); + assert_eq!(second.token_ids, vec![101, 103]); + assert_eq!(second.finish_reason, Some(FinishReason::Stop)); let nvext = NvExt::builder() .extra_fields(vec!["completion_token_ids".to_string()]) @@ -1082,7 +1452,7 @@ mod tests { let decode_stream = crate::tokenizers::DecodeStream::new(tokenizer, &[], false); let stop_conditions = StopConditions::default(); - let mut decoder = Decoder::new(decode_stream, stop_conditions, false, None); + let mut decoder = Decoder::new(decode_stream, stop_conditions, false, None, None); let result = decoder.process_token_ids(&[42]); assert!( @@ -1105,7 +1475,7 @@ mod tests { let decode_stream = crate::tokenizers::DecodeStream::new(tokenizer, &[], false); let stop_conditions = StopConditions::default(); - let mut decoder = Decoder::new(decode_stream, stop_conditions, false, None); + let mut decoder = Decoder::new(decode_stream, stop_conditions, false, None, None); let result = decoder.process_token_ids(&[42]); let err = result.err().expect("should be Err"); @@ -1126,4 +1496,501 @@ mod tests { other => panic!("Expected FinishReason::Error, got: {:?}", other), } } + + /// A tokenizer whose per-token fragments are chosen so `1` then `2` leaves a partial + /// hidden-stop-sequence match ("STOP", a strict prefix of the configured stop + /// sequence "STOPPED") jailed, and `99` fails to decode -- used to exercise + /// EOF-time flushing and error/EOF interaction at the `Backend` (not just `Decoder`) + /// level, across multiple choices. + struct JailingTokenizer; + + impl traits::Encoder for JailingTokenizer { + fn encode(&self, _input: &str) -> anyhow::Result { + Ok(crate::tokenizers::Encoding::Sp(vec![])) + } + fn encode_batch( + &self, + _inputs: &[&str], + ) -> anyhow::Result> { + Ok(vec![]) + } + } + + impl traits::Decoder for JailingTokenizer { + // `DecodeStream::step` (the real, pinned `dynamo-tokenizers` implementation) does + // not call `decode` with a single fresh id at a time: it calls it once with the + // previously-read slice and once with that slice plus the newest id, then diffs the + // two to find the newly emitted text. Both calls must therefore accept *any* + // contextual slice of already-seen ids, not just a single one, so this decodes + // compositionally (concatenating each id's own fragment) instead of matching whole + // slices -- except token 99, which fails deliberately wherever it appears, to give + // tests a controlled decode-error boundary. + fn decode( + &self, + token_ids: &[TokenIdType], + _skip_special_tokens: bool, + ) -> anyhow::Result { + if token_ids.contains(&99) { + anyhow::bail!("simulated decode failure"); + } + let mut text = String::new(); + for &id in token_ids { + match id { + 1 => text.push_str("abc"), + 2 => text.push_str("STOP"), + other => anyhow::bail!("unexpected token id: {other}"), + } + } + Ok(traits::DecodeResult::Complete(text)) + } + } + + impl traits::Tokenizer for JailingTokenizer {} + + fn jailing_request(n: u8) -> PreprocessedRequest { + PreprocessedRequest::builder() + .model("test-model".to_string()) + .token_ids(vec![]) + .stop_conditions(StopConditions { + stop: Some(vec!["STOPPED".to_string()]), + ..Default::default() + }) + .sampling_options(SamplingOptions { + n: Some(n), + ..Default::default() + }) + .output_options(OutputOptions::default()) + .build() + .expect("valid preprocessed request") + } + + /// Wraps an inner stream and flips a shared flag if it is ever polled again after + /// already returning `None` once -- used to verify the backend's `stream_ended` guard + /// actually prevents re-polling `stream`, which most `Stream` impls do not guarantee + /// is safe. + struct EndGuardStream { + inner: std::pin::Pin> + Send>>, + ended: bool, + overpolled: Arc, + } + + impl futures::Stream for EndGuardStream { + type Item = Annotated; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + if self.ended { + self.overpolled + .store(true, std::sync::atomic::Ordering::SeqCst); + } + let poll = self.inner.as_mut().poll_next(cx); + if matches!(poll, std::task::Poll::Ready(None)) { + self.ended = true; + } + poll + } + } + + struct BareEofMultiChoiceEngine { + overpolled: Arc, + } + + #[async_trait] + impl AsyncEngine, ManyOut>, Error> + for BareEofMultiChoiceEngine + { + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + // Both choices withhold "STOP" as a never-completed partial match against + // "STOPPED", then the stream ends with no `finish_reason` ever sent for + // either -- a "bare EOF", as from an engine that simply drops the connection. + let chunks = vec![ + Annotated::from_data(LLMEngineOutput { + token_ids: vec![1], + index: Some(0), + ..Default::default() + }), + Annotated::from_data(LLMEngineOutput { + token_ids: vec![2], + index: Some(0), + ..Default::default() + }), + Annotated::from_data(LLMEngineOutput { + token_ids: vec![1], + index: Some(1), + ..Default::default() + }), + Annotated::from_data(LLMEngineOutput { + token_ids: vec![2], + index: Some(1), + ..Default::default() + }), + ]; + let guarded = EndGuardStream { + inner: Box::pin(futures::stream::iter(chunks)), + ended: false, + overpolled: self.overpolled.clone(), + }; + Ok(ResponseStream::new(Box::pin(guarded), request.context())) + } + } + + #[tokio::test] + async fn backend_flushes_multiple_jailed_choices_on_bare_eof_without_repolling() { + let tokenizer: Arc = Arc::new(JailingTokenizer); + let backend = Backend::from_tokenizer(Tokenizer::from(tokenizer)); + let request = jailing_request(2); + let overpolled = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let engine: ServerStreamingEngine> = + Arc::new(BareEofMultiChoiceEngine { + overpolled: overpolled.clone(), + }); + + let stream = Operator::generate(backend.as_ref(), SingleIn::new(request), engine) + .await + .expect("backend generation succeeds"); + let outputs: Vec<_> = stream.collect().await; + + // 2 normal chunks per choice (releasing "abc", withholding "STOP") + one + // synthetic EOF-flush chunk per choice releasing the withheld "STOP". + assert_eq!(outputs.len(), 6, "unexpected output count: {outputs:?}"); + + let mut flushed_by_index = std::collections::HashMap::new(); + for output in &outputs { + let data = output.data.as_ref().expect("every chunk carries data"); + if data.finish_reason.is_some() { + flushed_by_index.insert(data.index, data.text.clone()); + } + } + assert_eq!( + flushed_by_index.get(&Some(0)).cloned().flatten().as_deref(), + Some("STOP"), + "choice 0's withheld text must be flushed on bare EOF" + ); + assert_eq!( + flushed_by_index.get(&Some(1)).cloned().flatten().as_deref(), + Some("STOP"), + "choice 1's withheld text must be flushed on bare EOF" + ); + + assert!( + !overpolled.load(std::sync::atomic::Ordering::SeqCst), + "backend must not poll the underlying stream again after it yields None" + ); + } + + struct ErrorThenBareEofEngine; + + #[async_trait] + impl AsyncEngine, ManyOut>, Error> + for ErrorThenBareEofEngine + { + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + // Choice 0 withholds "STOP" (same as choice 1) and *then* hits a decode + // error -- so it still has jailed text outstanding at the moment it + // terminates via the error path rather than a real stop. Choice 1 withholds + // "STOP" too but never gets a `finish_reason` -- the stream just ends (bare + // EOF). Without excluding already-finished choices from the EOF flush, choice + // 0's leftover jailed text would be flushed into a second, spurious `Stop` + // chunk after its error. + let chunks = vec![ + Annotated::from_data(LLMEngineOutput { + token_ids: vec![1], + index: Some(0), + ..Default::default() + }), + Annotated::from_data(LLMEngineOutput { + token_ids: vec![2], + index: Some(0), + ..Default::default() + }), + Annotated::from_data(LLMEngineOutput { + token_ids: vec![99], + index: Some(0), + ..Default::default() + }), + Annotated::from_data(LLMEngineOutput { + token_ids: vec![1], + index: Some(1), + ..Default::default() + }), + Annotated::from_data(LLMEngineOutput { + token_ids: vec![2], + index: Some(1), + ..Default::default() + }), + ]; + Ok(ResponseStream::new( + Box::pin(futures::stream::iter(chunks)), + request.context(), + )) + } + } + + /// A choice that already ended in a decode error must not be given a second, + /// synthetic `Stop` chunk by the bare-EOF flush loop -- that would turn a genuine + /// error into what looks like a normal completion downstream. + #[tokio::test] + async fn backend_does_not_synthesize_stop_after_choice_decode_error() { + let tokenizer: Arc = Arc::new(JailingTokenizer); + let backend = Backend::from_tokenizer(Tokenizer::from(tokenizer)); + let request = jailing_request(2); + let engine: ServerStreamingEngine> = + Arc::new(ErrorThenBareEofEngine); + + let stream = Operator::generate(backend.as_ref(), SingleIn::new(request), engine) + .await + .expect("backend generation succeeds"); + let outputs: Vec<_> = stream.collect().await; + + let choice0_finishes: Vec<_> = outputs + .iter() + .filter_map(|o| o.data.as_ref()) + .filter(|d| d.index == Some(0) && d.finish_reason.is_some()) + .collect(); + assert_eq!( + choice0_finishes.len(), + 1, + "choice 0 must finish exactly once (its decode error), not again via EOF flush: {choice0_finishes:?}" + ); + assert!( + matches!( + choice0_finishes[0].finish_reason, + Some(FinishReason::Error(_)) + ), + "choice 0's only finish must be its decode error, got: {:?}", + choice0_finishes[0].finish_reason + ); + + let choice1_flush = outputs + .iter() + .filter_map(|o| o.data.as_ref()) + .find(|d| d.index == Some(1) && d.finish_reason.is_some()); + assert_eq!( + choice1_flush.and_then(|d| d.text.as_deref()), + Some("STOP"), + "choice 1's withheld text must still be flushed on the same bare EOF" + ); + } + + /// A tokenizer whose only token (`1`) decodes to `"o"` -- paired with hidden stop `"ozzy"`, + /// this lets a single decoder-driven chunk withhold "o" before later chunks switch to the + /// engine-pre-decoded-text fast path, which never calls back into this tokenizer. + struct SingleOTokenizer; + + impl traits::Encoder for SingleOTokenizer { + fn encode(&self, _input: &str) -> anyhow::Result { + Ok(crate::tokenizers::Encoding::Sp(vec![])) + } + fn encode_batch( + &self, + _inputs: &[&str], + ) -> anyhow::Result> { + Ok(vec![]) + } + } + + impl traits::Decoder for SingleOTokenizer { + fn decode( + &self, + token_ids: &[TokenIdType], + _skip_special_tokens: bool, + ) -> anyhow::Result { + Ok(traits::DecodeResult::Complete( + token_ids.iter().map(|_| "o").collect(), + )) + } + } + + impl traits::Tokenizer for SingleOTokenizer {} + + struct MixedBypassEngine; + + #[async_trait] + impl AsyncEngine, ManyOut>, Error> + for MixedBypassEngine + { + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + let chunks = vec![ + // Token-only: goes through the decoder, which withholds "o" as a viable + // prefix of the hidden stop "ozzy". + Annotated::from_data(LLMEngineOutput { + token_ids: vec![1], + index: Some(0), + ..Default::default() + }), + // Engine-pre-decoded, non-terminal: takes the fast bypass path entirely + // (never touches the decoder to produce this text), but must not make the + // withheld "o" look resolved. + Annotated::from_data(LLMEngineOutput { + token_ids: vec![], + index: Some(0), + text: Some(String::new()), + ..Default::default() + }), + // Engine-pre-decoded, terminal: the withheld "o" must be flushed ahead of + // this chunk's own (newer) text, not appended after it. + Annotated::from_data(LLMEngineOutput { + token_ids: vec![], + index: Some(0), + text: Some("there".to_string()), + finish_reason: Some(FinishReason::Length), + ..Default::default() + }), + ]; + Ok(ResponseStream::new( + Box::pin(futures::stream::iter(chunks)), + request.context(), + )) + } + } + + /// Regression test for two related bugs in the engine-pre-decoded-text fast path: it + /// must not silently erase a checkpoint of the decoder's still-withheld text just + /// because a given chunk skipped the decoder (P2, `jailed_text` erasure), and when it + /// does flush withheld text into a terminal chunk, that older text must come first, not + /// be appended after the chunk's own (newer) text (P2, terminal ordering). + #[tokio::test] + async fn backend_bypass_path_preserves_and_orders_withheld_text() { + let tokenizer: Arc = Arc::new(SingleOTokenizer); + let backend = Backend::from_tokenizer(Tokenizer::from(tokenizer)); + let request = PreprocessedRequest::builder() + .model("test-model".to_string()) + .token_ids(vec![]) + .stop_conditions(StopConditions { + stop: Some(vec!["ozzy".to_string()]), + ..Default::default() + }) + .sampling_options(SamplingOptions::default()) + .output_options(OutputOptions::default()) + // `jailed_text` is only populated when migration is possible (see + // `DecoderParams::migration_possible`); this test asserts on it. + .migration_state(Some( + crate::protocols::common::preprocessor::MigrationState::default(), + )) + .build() + .expect("valid preprocessed request"); + let engine: ServerStreamingEngine> = + Arc::new(MixedBypassEngine); + + let stream = Operator::generate(backend.as_ref(), SingleIn::new(request), engine) + .await + .expect("backend generation succeeds"); + let outputs: Vec<_> = stream.collect().await; + assert_eq!(outputs.len(), 3, "unexpected output count: {outputs:?}"); + + let data: Vec<_> = outputs + .iter() + .map(|o| o.data.as_ref().expect("every chunk carries data")) + .collect(); + + assert_eq!( + data[0].jailed_text.as_deref(), + Some("o"), + "the decoder-driven chunk must publish what it withheld" + ); + + // The non-terminal bypass chunk did not touch the decoder, so the checkpoint must + // still reflect the still-withheld "o" -- not look resolved just because this + // particular chunk skipped the decoder. + assert_eq!( + data[1].jailed_text.as_deref(), + Some("o"), + "a non-terminal bypass chunk must not erase the withheld-text checkpoint" + ); + assert_eq!(data[1].text.as_deref(), Some("")); + + // The terminal bypass chunk flushes the withheld "o" ahead of its own "there", + // producing "othere" -- not "thereo" -- and resolves the checkpoint. + assert_eq!( + data[2].text.as_deref(), + Some("othere"), + "withheld text must be flushed before, not after, this chunk's own newer text" + ); + assert_eq!(data[2].finish_reason, Some(FinishReason::Length)); + assert_eq!(data[2].jailed_text, None); + } + + struct AnnotatedErrorThenBareEofEngine; + + #[async_trait] + impl AsyncEngine, ManyOut>, Error> + for AnnotatedErrorThenBareEofEngine + { + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + // Withholds "STOP" (same jailing setup as the other Backend-level tests), then + // the stream reports a request-level error via `Annotated::from_error` -- + // unlike a decode failure, this carries no `data` and so no per-choice + // `index` -- before ending. + let chunks = vec![ + Annotated::from_data(LLMEngineOutput { + token_ids: vec![1], + index: Some(0), + ..Default::default() + }), + Annotated::from_data(LLMEngineOutput { + token_ids: vec![2], + index: Some(0), + ..Default::default() + }), + Annotated::from_error("engine reported a fatal error"), + ]; + Ok(ResponseStream::new( + Box::pin(futures::stream::iter(chunks)), + request.context(), + )) + } + } + + /// Regression test for the `Annotated`-error counterpart of + /// `backend_does_not_synthesize_stop_after_choice_decode_error`: a request-level error + /// annotation (no `data`, hence no per-choice `index`) must also mark every choice + /// finished, not just choices that fail through a decode error, so the bare-EOF flush + /// loop does not turn it into a spurious successful `Stop` afterward. + #[tokio::test] + async fn backend_does_not_synthesize_stop_after_annotated_error() { + let tokenizer: Arc = Arc::new(JailingTokenizer); + let backend = Backend::from_tokenizer(Tokenizer::from(tokenizer)); + let request = jailing_request(1); + let engine: ServerStreamingEngine> = + Arc::new(AnnotatedErrorThenBareEofEngine); + + let stream = Operator::generate(backend.as_ref(), SingleIn::new(request), engine) + .await + .expect("backend generation succeeds"); + let outputs: Vec<_> = stream.collect().await; + + let finishes: Vec<_> = outputs.iter().filter(|o| o.is_error()).collect(); + assert_eq!( + finishes.len(), + 1, + "the error must be reported exactly once, not again via EOF flush: {outputs:?}" + ); + + let synthetic_stops: Vec<_> = outputs + .iter() + .filter(|o| { + o.data + .as_ref() + .is_some_and(|d| d.finish_reason == Some(FinishReason::Stop)) + }) + .collect(); + assert!( + synthetic_stops.is_empty(), + "no choice may receive a synthetic successful Stop after a request-level error: {outputs:?}" + ); + } } diff --git a/lib/llm/src/migration.rs b/lib/llm/src/migration.rs index f3ce4154e494..3e609026c499 100644 --- a/lib/llm/src/migration.rs +++ b/lib/llm/src/migration.rs @@ -39,10 +39,13 @@ use dynamo_runtime::protocols::annotated::Annotated; /// Accessors the migration RetryManager needs from a response chunk. /// `token_ids` lets it replay already-delivered tokens; `worker_trace_link` /// lets it stamp the failed worker's span onto the next attempt's -/// `migration_link`. +/// `migration_link`; `jailed_text` lets it carry forward whatever the Backend's +/// decoder is still withholding as a possible hidden-stop-sequence prefix, so a +/// retried attempt's fresh decoder can be reseeded instead of silently losing it. pub(crate) trait HasTokenIds { fn token_ids(&self) -> &[TokenIdType]; fn worker_trace_link(&self) -> Option<&crate::protocols::common::preprocessor::TraceLink>; + fn jailed_text(&self) -> Option<&str>; } impl HasTokenIds for BackendOutput { @@ -52,6 +55,9 @@ impl HasTokenIds for BackendOutput { fn worker_trace_link(&self) -> Option<&crate::protocols::common::preprocessor::TraceLink> { self.worker_trace_link.as_ref() } + fn jailed_text(&self) -> Option<&str> { + self.jailed_text.as_deref() + } } impl HasTokenIds for LLMEngineOutput { @@ -61,6 +67,9 @@ impl HasTokenIds for LLMEngineOutput { fn worker_trace_link(&self) -> Option<&crate::protocols::common::preprocessor::TraceLink> { self.worker_trace_link.as_ref() } + fn jailed_text(&self) -> Option<&str> { + self.jailed_text.as_deref() + } } /// Error types that make a request non-migratable wherever they appear in the @@ -325,8 +334,9 @@ where // embedding prompts until a retry can represent an embedding-based continuation. // TODO: Define a replay-capability contract for attempt-local decoder and sampler state. - // Stop-string matching, generated-token penalties, and thinking-token budgets are not - // currently checkpointed across workers. + // A withheld hidden-stop-sequence prefix is now checkpointed across workers (see + // `jail_seed` / `track_response`), but generated-token penalties and thinking-token + // budgets are not. // Disable migration for structured-output (guided-decoding) requests. // Inference backends initialize the guided-decoding FSM (finite state machine) fresh @@ -655,6 +665,12 @@ where if let Some(link) = llm_engine_output.worker_trace_link() { self.last_worker_link = Some(link.clone()); } + // Snapshot whatever the Backend's decoder is currently withholding as a possible + // hidden-stop-sequence prefix, so a future retry's fresh decoder can be reseeded + // from it (`jail_seed`) instead of the withheld text simply vanishing. Overwritten + // on every chunk -- `None` once the decoder resolves it one way or the other -- so + // this always reflects the last known-good chunk's state, never a stale one. + self.request.jail_seed = llm_engine_output.jailed_text().map(str::to_string); let output_len = u32::try_from(token_ids.len()).unwrap_or(u32::MAX); if self.exceed_max_seq_len(output_len) { return; @@ -957,6 +973,7 @@ mod tests { completion_usage: None, engine_data: None, routing_data: None, + jailed_text: None, }) } @@ -2407,6 +2424,109 @@ mod tests { ); } + /// Regression test for the migration-discards-withheld-text bug: a chunk delivered + /// before a migratable error carries `jailed_text` (whatever the `Backend` decoder was + /// withholding as a possible hidden-stop-sequence prefix), and the retried attempt's + /// request must be reseeded from it via `jail_seed` rather than starting the new + /// decoder unseeded and silently losing that withheld text. + #[tokio::test] + async fn test_retry_manager_carries_jail_seed_across_migration() { + dynamo_runtime::logging::init(); + let context_id = uuid::Uuid::new_v4().to_string(); + + struct JailSeedMockEngine { + calls: Arc, + context_id: String, + } + + #[async_trait] + impl + AsyncEngine< + SingleIn, + ManyOut>, + anyhow::Error, + > for JailSeedMockEngine + { + async fn generate( + &self, + request: SingleIn, + ) -> Result>> { + let call = self.calls.fetch_add(1, Ordering::SeqCst); + let (preprocessed_request, _context) = request.transfer(()); + + if call == 0 { + // First attempt: deliver one good chunk that leaves "STOP" withheld as + // a partial hidden-stop-sequence match, then disconnect mid-stream. + assert_eq!( + preprocessed_request.jail_seed, None, + "a first attempt must not start pre-seeded" + ); + let responses = stream::iter(vec![ + Annotated::from_data(BackendOutput { + jailed_text: Some("STOP".to_string()), + ..create_mock_output(10).data.unwrap() + }), + Annotated::from_err( + DynamoError::builder() + .error_type(ErrorType::Disconnected) + .message("worker disconnected mid-stream") + .build(), + ), + ]); + let ctx = Arc::new(Controller::new(self.context_id.clone())); + Ok(ResponseStream::new(Box::pin(responses), ctx)) + } else { + // Retry attempt: the withheld text from the abandoned attempt's last + // known-good chunk must have been carried onto this request. + assert_eq!( + preprocessed_request.jail_seed.as_deref(), + Some("STOP"), + "retry request must be reseeded with the withheld jail text" + ); + let responses = stream::iter(vec![Annotated::from_data(BackendOutput { + jailed_text: None, + ..create_mock_output(11).data.unwrap() + })]); + let ctx = Arc::new(Controller::new(self.context_id.clone())); + Ok(ResponseStream::new(Box::pin(responses), ctx)) + } + } + } + + let request = create_mock_request(5); + let next_generate: ServerStreamingEngine> = + Arc::new(JailSeedMockEngine { + calls: Arc::new(AtomicU32::new(0)), + context_id: context_id.clone(), + }); + + let ctx = Arc::new(Controller::new(context_id.clone())); + let metrics = Arc::new(Metrics::new()); + let mut retry_manager = RetryManager::build( + ctx, + BTreeMap::new(), + request, + next_generate, + 1, + None, + Arc::new(TEST_MODEL.to_string()), + metrics, + None, + ) + .await + .expect("Failed to build RetryManager"); + + let mut responses = Vec::new(); + while let Some(r) = retry_manager.next().await { + responses.push(r); + } + + // One good chunk from the first attempt, one from the retry -- the in-stream + // error itself is consumed internally to drive the migration, not surfaced. + assert_eq!(responses.len(), 2); + assert!(responses.iter().all(|r| r.error.is_none())); + } + #[tokio::test] async fn test_retry_manager_cancellation_during_migration_skips_retry_dispatch() { dynamo_runtime::logging::init(); @@ -2839,4 +2959,207 @@ mod tests { "retry scheduled must name the upcoming attempt" ); } + + // --- Real-Backend migration integration tests ------------------------------------- + // + // `create_mock_output` above fabricates already-decoded `BackendOutput` text directly, + // so `test_retry_manager_carries_jail_seed_across_migration` only proves the + // `jail_seed` field gets copied between hand-built mocks -- removing the real `Backend` + // seed consumption entirely would not fail it. The tests below instead run raw + // (undetokenized) token ids through a real `crate::backend::Backend` wrapping a real + // `Decoder`, so a migration retry re-creates an actual fresh decoder the way the + // production pipeline does, and prove the checkpoint is consumed correctly by it. + + /// Token 1 decodes to "o", 2 to "there", 3 to "zzy" -- letters chosen so a hidden stop + /// of "ozzy" can complete across a migration boundary (token 1 on the first attempt, + /// the remainder on the retry), and so an ordinary continuation ("o" then "there") is + /// legible as "othere". + struct LetterTokenizer; + + impl crate::tokenizers::traits::Encoder for LetterTokenizer { + fn encode(&self, _input: &str) -> anyhow::Result { + Ok(crate::tokenizers::Encoding::Sp(vec![])) + } + fn encode_batch( + &self, + _inputs: &[&str], + ) -> anyhow::Result> { + Ok(vec![]) + } + } + + impl crate::tokenizers::traits::Decoder for LetterTokenizer { + fn decode( + &self, + token_ids: &[TokenIdType], + _skip_special_tokens: bool, + ) -> anyhow::Result { + let text: String = token_ids + .iter() + .map(|&id| match id { + 1 => "o", + 2 => "there", + 3 => "zzy", + other => panic!("unexpected token id in LetterTokenizer: {other}"), + }) + .collect(); + Ok(crate::tokenizers::traits::DecodeResult::Complete(text)) + } + } + + impl crate::tokenizers::traits::Tokenizer for LetterTokenizer {} + + fn letter_backend() -> Arc { + let tokenizer: Arc = Arc::new(LetterTokenizer); + crate::backend::Backend::from_tokenizer(crate::tokenizers::Tokenizer::from(tokenizer)) + } + + fn jail_request(stop: Option>) -> PreprocessedRequest { + PreprocessedRequest::builder() + .model(TEST_MODEL.to_string()) + .token_ids(vec![]) + .stop_conditions(StopConditions { + stop, + ..Default::default() + }) + .sampling_options(SamplingOptions::default()) + .output_options(OutputOptions::default()) + .build() + .expect("valid preprocessed request") + } + + /// Emits raw token ids for a real `Backend`/`Decoder` to detokenize: `first_attempt_tokens` + /// then a migratable disconnect on the first call, `retry_tokens` to completion on the + /// second. + struct RawTokenMigrationEngine { + calls: Arc, + first_attempt_tokens: Vec, + retry_tokens: Vec, + context_id: String, + } + + #[async_trait] + impl AsyncEngine, ManyOut>, Error> + for RawTokenMigrationEngine + { + async fn generate( + &self, + _request: SingleIn, + ) -> Result>> { + let call = self.calls.fetch_add(1, Ordering::SeqCst); + let ctx = Arc::new(Controller::new(self.context_id.clone())); + let tokens = if call == 0 { + &self.first_attempt_tokens + } else { + &self.retry_tokens + }; + let mut chunks: Vec<_> = tokens + .iter() + .map(|&id| { + Annotated::from_data(LLMEngineOutput { + token_ids: vec![id], + index: Some(0), + ..Default::default() + }) + }) + .collect(); + if call == 0 { + chunks.push(Annotated::from_err( + DynamoError::builder() + .error_type(ErrorType::Disconnected) + .message("worker disconnected mid-stream") + .build(), + )); + } + Ok(ResponseStream::new(Box::pin(stream::iter(chunks)), ctx)) + } + } + + /// Wraps a raw-token engine with a real `Backend`, so `RetryManager`'s `next_generate` + /// re-creates an actual fresh `Decoder` on every retry, exactly as the production + /// pipeline does (migration sits outside `Backend` from the response's perspective; see + /// `lib/llm/src/entrypoint/input/common.rs`). + struct BackendWrappedEngine { + backend: Arc, + raw_engine: ServerStreamingEngine>, + } + + #[async_trait] + impl AsyncEngine, ManyOut>, Error> + for BackendWrappedEngine + { + async fn generate( + &self, + request: SingleIn, + ) -> Result>> { + Operator::generate(self.backend.as_ref(), request, self.raw_engine.clone()).await + } + } + + /// Drives a `RetryManager` wrapping a real `Backend` over a scripted raw-token engine + /// and returns the concatenation of every response's visible text. + async fn run_raw_token_migration( + stop: Option>, + first_attempt_tokens: Vec, + retry_tokens: Vec, + ) -> String { + let context_id = uuid::Uuid::new_v4().to_string(); + let raw_engine: ServerStreamingEngine> = + Arc::new(RawTokenMigrationEngine { + calls: Arc::new(AtomicU32::new(0)), + first_attempt_tokens, + retry_tokens, + context_id: context_id.clone(), + }); + let next_generate: ServerStreamingEngine> = + Arc::new(BackendWrappedEngine { + backend: letter_backend(), + raw_engine, + }); + + let ctx = Arc::new(Controller::new(context_id.clone())); + let metrics = Arc::new(Metrics::new()); + let mut retry_manager = RetryManager::build( + ctx, + BTreeMap::new(), + jail_request(stop), + next_generate, + 1, + None, + Arc::new(TEST_MODEL.to_string()), + metrics, + None, + ) + .await + .expect("Failed to build RetryManager"); + + let mut text = String::new(); + while let Some(response) = retry_manager.next().await { + if let Some(t) = response.data.and_then(|data| data.text) { + text.push_str(&t); + } + } + text + } + + /// End-to-end regression for the migration-discards-withheld-text bug, through a real + /// `Backend`/`Decoder`: a hidden stop "ozzy" withholds "o" on the first attempt, the + /// worker disconnects, and the retried attempt's fresh decoder must be reseeded from + /// the checkpoint so "o" plus the retry's "there" reaches the caller as "othere" -- + /// not "there" alone. + #[tokio::test] + async fn migration_preserves_withheld_text_through_real_backend_retry() { + let text = run_raw_token_migration(Some(vec!["ozzy".to_string()]), vec![1], vec![2]).await; + assert_eq!(text, "othere"); + } + + /// Same setup, but the retry's tokens complete the hidden stop instead of abandoning + /// it: "o" (withheld, first attempt) plus "zzy" (retry) makes "ozzy", which must stay + /// fully hidden -- proving the checkpoint doesn't just prevent loss, it still + /// participates correctly in stop-sequence matching across the migration boundary. + #[tokio::test] + async fn migration_hides_stop_sequence_completed_across_real_backend_retry() { + let text = run_raw_token_migration(Some(vec!["ozzy".to_string()]), vec![1], vec![3]).await; + assert_eq!(text, "", "the completed hidden stop must not leak any text"); + } } diff --git a/lib/llm/src/protocols/common/llm_backend.rs b/lib/llm/src/protocols/common/llm_backend.rs index 3011fb94ecc8..716fc1c28e5e 100644 --- a/lib/llm/src/protocols/common/llm_backend.rs +++ b/lib/llm/src/protocols/common/llm_backend.rs @@ -149,6 +149,13 @@ pub struct BackendOutput { /// consumed by the frontend and not surfaced to clients. See [`RoutingData`](crate::protocols::common::timing::RoutingData). #[serde(default, skip_serializing_if = "Option::is_none")] pub routing_data: Option, + + /// Text the `Backend` decoder is currently withholding as a possible (but + /// unresolved) prefix of a hidden stop sequence, as of this chunk. Frontend-only + /// (Dynamo-internal): consumed by the migration `RetryManager` so a retry's fresh + /// decoder can resume with the same pending text instead of silently dropping it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub jailed_text: Option, } /// The LLM engine and backnd with manage it's own state, specifically translating how a @@ -235,6 +242,14 @@ pub struct LLMEngineOutput { /// Dynamo-internal; consumed by the frontend. See [`RoutingData`](crate::protocols::common::timing::RoutingData). #[serde(default, skip_serializing_if = "Option::is_none")] pub routing_data: Option, + + /// Text the `Backend` decoder is currently withholding as a possible (but + /// unresolved) prefix of a hidden stop sequence, as of this chunk. Engines never + /// set this; the frontend's `Backend` stage populates it locally (it is never + /// sent by a worker) so the migration `RetryManager` can seed a retry's fresh + /// decoder with it instead of silently dropping the withheld text. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub jailed_text: Option, } impl LLMEngineOutput { @@ -258,6 +273,7 @@ impl LLMEngineOutput { completion_usage: None, engine_data: None, routing_data: None, + jailed_text: None, } } @@ -281,6 +297,7 @@ impl LLMEngineOutput { completion_usage: None, engine_data: None, routing_data: None, + jailed_text: None, } } @@ -304,6 +321,7 @@ impl LLMEngineOutput { completion_usage: None, engine_data: None, routing_data: None, + jailed_text: None, } } @@ -327,6 +345,7 @@ impl LLMEngineOutput { completion_usage: None, engine_data: None, routing_data: None, + jailed_text: None, } } @@ -364,6 +383,7 @@ impl LLMEngineOutput { completion_usage: None, engine_data: None, routing_data: None, + jailed_text: None, } } } diff --git a/lib/llm/src/protocols/common/preprocessor.rs b/lib/llm/src/protocols/common/preprocessor.rs index b42a09dfe72e..841653c5529a 100644 --- a/lib/llm/src/protocols/common/preprocessor.rs +++ b/lib/llm/src/protocols/common/preprocessor.rs @@ -360,6 +360,20 @@ pub struct PreprocessedRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub migration_link: Option, + /// Text withheld by the previous attempt's decoder as a possible (but + /// unresolved) prefix of a hidden stop sequence, carried into a migration + /// retry so the new attempt's decoder does not silently drop it and can + /// still complete the match if the continuation supplies the rest of the + /// sequence. Set by the migration `RetryManager` (in-process, on its own + /// in-memory `PreprocessedRequest`) from the last successfully processed + /// response before a retry, and consumed once by `Backend` -- also + /// in-process, one hop later in the same pipeline -- when seeding the + /// retry's decoder. `#[serde(skip)]` keeps it that way: it never needs to, + /// and must not, reach a remote worker over the wire. + #[builder(default)] + #[serde(skip)] + pub(crate) jail_seed: Option, + /// Bootstrap info for disaggregated serving #[builder(default)] #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index e20a5bca4340..374ba8c4f471 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -457,6 +457,7 @@ mod tests { })), encoder_result: None, routing_data: None, + jailed_text: None, } } @@ -776,6 +777,7 @@ mod tests { "prefill_compute_time_ms": 45.6 })), routing_data: None, + jailed_text: None, } } @@ -1008,6 +1010,7 @@ mod tests { worker_trace_link: None, engine_data: None, // engine didn't provide any data routing_data: None, + jailed_text: None, }; let response = generator diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index 9eb9d5672bbe..89cb1aa274d6 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -354,6 +354,7 @@ mod tests { })), encoder_result: None, routing_data: None, + jailed_text: None, } } @@ -412,6 +413,7 @@ mod tests { "prefill_compute_time_ms": 45.6 })), routing_data: None, + jailed_text: None, } } @@ -726,6 +728,7 @@ mod tests { worker_trace_link: None, engine_data: None, // engine didn't provide any data routing_data: None, + jailed_text: None, }; let response = generator diff --git a/lib/llm/tests/chat_template_render_errors.rs b/lib/llm/tests/chat_template_render_errors.rs index 6ef109f56fe3..defeb2e174c0 100644 --- a/lib/llm/tests/chat_template_render_errors.rs +++ b/lib/llm/tests/chat_template_render_errors.rs @@ -76,6 +76,7 @@ impl AsyncEngine, ManyOut worker_trace_link: None, engine_data: None, routing_data: None, + jailed_text: None, }; Ok(ResponseStream::new( diff --git a/lib/llm/tests/http-service.rs b/lib/llm/tests/http-service.rs index c64acf8e2f7c..4f8ae9f41793 100644 --- a/lib/llm/tests/http-service.rs +++ b/lib/llm/tests/http-service.rs @@ -2457,6 +2457,7 @@ mod zero_top_logprobs { worker_trace_link: None, engine_data: None, routing_data: None, + jailed_text: None, }) .expect("backend output conversion failed") }) diff --git a/lib/llm/tests/test_stop_behavior.rs b/lib/llm/tests/test_stop_behavior.rs index e61e2612eec3..47825ceff85b 100644 --- a/lib/llm/tests/test_stop_behavior.rs +++ b/lib/llm/tests/test_stop_behavior.rs @@ -13,6 +13,21 @@ const STOP: u32 = 2; const THERE: u32 = 3; const EOS: u32 = 99; +// Tokens whose decoded fragments only form a complete stop string once several of them +// are concatenated -- used to exercise a stop sequence split across multiple decode steps. +const DELTA: u32 = 10; +const EPSILON: u32 = 11; +const ZETA: u32 = 12; +const ETA: u32 = 13; +const THETA: u32 = 14; +const OH: u32 = 15; +const OTHER: u32 = 16; + +// Single-character tokens used to build a long, self-similar (periodic) run for the +// prefix-matching regression test below. +const A: u32 = 20; +const B: u32 = 21; + struct TestTokenizer; impl tokenizer_traits::Encoder for TestTokenizer { @@ -34,6 +49,15 @@ impl tokenizer_traits::Decoder for TestTokenizer { STOP => Some("STOP"), THERE => Some("there"), EOS => Some(""), + DELTA => Some(" delta"), + EPSILON => Some(" epsilon"), + ZETA => Some(" zeta"), + ETA => Some(" eta"), + THETA => Some(" theta"), + OH => Some("o"), + OTHER => Some("there"), + A => Some("a"), + B => Some("b"), _ => Some("?"), }) .collect(); @@ -59,7 +83,7 @@ fn make_decoder( stop: stop_sequences.map(|v| v.into_iter().map(String::from).collect()), ..Default::default() }; - Decoder::new(decode_stream, stop_conditions, include_stop_str, None) + Decoder::new(decode_stream, stop_conditions, include_stop_str, None, None) } #[test] @@ -118,11 +142,17 @@ fn trailing_tokens_ignored_after_stop() { #[test] fn min_tokens_delays_stop() { - let mut decoder = make_decoder(None, Some(3), Some(vec![EOS]), None, false); + let mut decoder = make_decoder(None, Some(2), Some(vec![EOS]), None, false); let result = decoder.process_token_ids(&[HI, EOS]).unwrap(); assert_eq!(result.text.as_deref(), Some("hi")); assert!(result.stop_trigger.is_none()); + + let result = decoder.process_token_ids(&[EOS]).unwrap(); + assert!(matches!( + result.stop_trigger, + Some(StopTrigger::HiddenStopTokenDetected(id)) if id == EOS + )); } #[test] @@ -146,7 +176,7 @@ fn user_stop_token_reports_distinct_trigger() { stop_token_ids_hidden: Some(vec![EOS]), ..Default::default() }; - let mut decoder = Decoder::new(decode_stream, stop_conditions, false, None); + let mut decoder = Decoder::new(decode_stream, stop_conditions, false, None, None); let result = decoder.process_token_ids(&[HI, STOP]).unwrap(); assert_eq!(result.text.as_deref(), Some("hi")); @@ -155,3 +185,168 @@ fn user_stop_token_reports_distinct_trigger() { Some(StopTrigger::UserStopTokenDetected(id)) if id == STOP )); } + +/// A hidden stop sequence split across several decode fragments must not leak any +/// fragment before the full sequence is recognized. +#[test] +fn hidden_stop_sequence_split_across_tokens_is_not_leaked() { + let mut decoder = make_decoder(None, None, None, Some(vec![" zeta eta theta"]), false); + let result = decoder + .process_token_ids(&[DELTA, EPSILON, ZETA, ETA, THETA]) + .unwrap(); + + assert_eq!(result.text.as_deref(), Some(" delta epsilon")); + assert!(matches!( + result.stop_trigger, + Some(StopTrigger::HiddenStopSequenceDetected(ref s)) if s == " zeta eta theta" + )); + // Withholding for `text` must not disturb `tokens[i]`: it always reports token_ids[i]'s + // own decoded text, matched-hidden-sequence tokens included, so logprob consumers that + // zip `tokens` with `token_ids` stay aligned. + assert_eq!( + result.tokens, + vec![ + Some(" delta".to_string()), + Some(" epsilon".to_string()), + Some(" zeta".to_string()), + Some(" eta".to_string()), + Some(" theta".to_string()), + ] + ); +} + +/// A withheld candidate prefix that turns out not to be part of the stop sequence must be +/// released once it can no longer complete, rather than being lost forever. +#[test] +fn withheld_prefix_is_released_once_it_cannot_complete() { + let mut decoder = make_decoder(None, None, None, Some(vec!["ozzy"]), false); + let result = decoder.process_token_ids(&[OH, OTHER]).unwrap(); + + assert_eq!(result.text.as_deref(), Some("othere")); + assert!(result.stop_trigger.is_none()); + // Regression check: withholding used to bunch both tokens' text onto the releasing + // step, pairing (OH, OTHER) with ("", "othere") instead of their own ("o", "there"). + assert_eq!( + result.tokens, + vec![Some("o".to_string()), Some("there".to_string())] + ); +} + +/// A stop sequence that never completes must still be flushed when generation ends for a +/// reason our decoder did not itself detect (e.g. the engine's own `max_tokens` limit) -- +/// otherwise a partial match silently swallows real output forever. +#[test] +fn flush_jailed_releases_incomplete_partial_match() { + let mut decoder = make_decoder(None, None, None, Some(vec![" zeta eta theta"]), false); + let result = decoder.process_token_ids(&[DELTA, ZETA, ETA]).unwrap(); + + assert_eq!(result.text.as_deref(), Some(" delta")); + assert!(result.stop_trigger.is_none()); + + let flushed = decoder.flush_jailed(); + assert_eq!(flushed.as_deref(), Some(" zeta eta")); + assert_eq!(decoder.flush_jailed(), None); +} + +/// A hidden stop *token* ending the stream is a different stop from any hidden stop +/// *sequence* prefix still being withheld -- that withheld text never completed, so it must +/// still reach the caller even though this particular stop hides its own token's text. +#[test] +fn hidden_stop_token_flushes_prior_jailed_prefix() { + // "hi" (from HI) is a genuine, never-completed prefix of the hidden sequence "hiya". + let mut decoder = make_decoder(None, None, Some(vec![EOS]), Some(vec!["hiya"]), false); + let result = decoder.process_token_ids(&[HI, EOS]).unwrap(); + + assert_eq!(result.text.as_deref(), Some("hi")); + assert!(matches!( + result.stop_trigger, + Some(StopTrigger::HiddenStopTokenDetected(id)) if id == EOS + )); + assert_eq!(result.tokens, vec![Some("hi".to_string()), None]); +} + +/// Same as above but the terminating stop is a *visible* token: the flushed backlog must +/// come out before this token's own (included) text, not after or in place of it. +#[test] +fn visible_stop_token_flushes_and_orders_prior_jailed_prefix() { + let tokenizer: Arc = Arc::new(TestTokenizer); + let decode_stream = tokenizers::DecodeStream::new(tokenizer, &[], false); + let stop_conditions = StopConditions { + stop_token_ids_visible: Some(vec![STOP]), + stop: Some(vec!["hiya".to_string()]), + ..Default::default() + }; + let mut decoder = Decoder::new(decode_stream, stop_conditions, false, None, None); + let result = decoder.process_token_ids(&[HI, STOP]).unwrap(); + + assert_eq!(result.text.as_deref(), Some("hiSTOP")); + assert!(matches!( + result.stop_trigger, + Some(StopTrigger::VisibleStopTokenDetected(id)) if id == STOP + )); + assert_eq!( + result.tokens, + vec![Some("hi".to_string()), Some("STOP".to_string())] + ); +} + +/// Pins down exactly *when* a withheld candidate prefix is released, using +/// `Decoder::step`'s `released_text` directly rather than the aggregate `text` from +/// `process_token_ids`: the first token's text must stay withheld on its own step (not +/// released early, in case it still completes into the stop sequence) and only come out +/// once the following step proves it cannot. +#[test] +fn released_text_is_withheld_exactly_until_prefix_is_ruled_out() { + let mut decoder = make_decoder(None, None, None, Some(vec!["ozzy"]), false); + + let first = decoder.step(OH).unwrap(); + assert_eq!(first.token.as_deref(), Some("o")); + assert_eq!( + first.released_text, None, + "\"o\" is still a viable prefix of \"ozzy\" and must not be released yet" + ); + assert!(first.stop_trigger.is_none()); + + let second = decoder.step(OTHER).unwrap(); + assert_eq!(second.token.as_deref(), Some("there")); + assert_eq!( + second.released_text.as_deref(), + Some("othere"), + "once \"o\" can no longer complete, it must be released together with this step's own text" + ); + assert!(second.stop_trigger.is_none()); +} + +/// Regression test for a self-similar (periodic) run of withheld candidate bytes -- the +/// shape of input that made the previous byte-by-byte prefix scan quadratic (each +/// candidate length re-scans from scratch instead of reusing prior work). This asserts +/// correctness, which a unit test can pin down, not the algorithm's complexity, which it +/// cannot: a sliding match window must still land on the right byte offset when an extra +/// repeated character precedes the real match, or either a real stop is missed or content +/// is dropped/leaked. +#[test] +fn hidden_stop_sequence_survives_self_similar_prefix_run() { + let mut decoder = make_decoder(None, None, None, Some(vec!["aaaab"]), false); + // "aaaaa" (five 'a's) grows the withheld tail beyond the stop sequence's own prefix + // length one byte at a time, forcing the matcher to keep re-deriving the longest + // still-viable prefix length as the window slides, before the final 'b' completes it. + let result = decoder + .process_token_ids(&[A, A, A, A, A, B]) + .expect("decode succeeds"); + + assert_eq!( + result.text.as_deref(), + Some("a"), + "only the one 'a' that fell out of the sliding window may be released; \ + the rest belongs to the matched stop sequence and must stay hidden" + ); + assert!(matches!( + result.stop_trigger, + Some(StopTrigger::HiddenStopSequenceDetected(ref s)) if s == "aaaab" + )); + assert_eq!( + result.tokens.len(), + 6, + "one token report per input token id" + ); +} diff --git a/lib/llm/tests/test_streaming_usage.rs b/lib/llm/tests/test_streaming_usage.rs index d64ce77704a1..4025eaa56a21 100644 --- a/lib/llm/tests/test_streaming_usage.rs +++ b/lib/llm/tests/test_streaming_usage.rs @@ -116,6 +116,7 @@ fn build_backend_outputs_with_cached_tokens(cached_tokens: Option) -> Vec) -> Vec) -> Vec BackendOutput { worker_trace_link: None, engine_data: None, routing_data: None, + jailed_text: None, } } @@ -312,6 +313,7 @@ async fn test_streaming_named_tool_buffers_until_finish() { worker_trace_link: None, engine_data: None, routing_data: None, + jailed_text: None, }; let response = generator @@ -383,6 +385,7 @@ async fn test_streaming_required_tool_parallel() { worker_trace_link: None, engine_data: None, routing_data: None, + jailed_text: None, }; let response = generator @@ -456,6 +459,7 @@ fn test_no_tool_choice_outputs_normal_text() { worker_trace_link: None, engine_data: None, routing_data: None, + jailed_text: None, }; let response = generator diff --git a/lib/llm/tests/tool_choice_finish_reasons.rs b/lib/llm/tests/tool_choice_finish_reasons.rs index e3a495df55e5..dbd521772dc8 100644 --- a/lib/llm/tests/tool_choice_finish_reasons.rs +++ b/lib/llm/tests/tool_choice_finish_reasons.rs @@ -56,6 +56,7 @@ fn build_backend_output_with_finish(text: &str, finish: common::FinishReason) -> worker_trace_link: None, engine_data: None, routing_data: None, + jailed_text: None, } }