diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 0d5194b84da4..37638327f054 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -464,6 +464,7 @@ areas: github_team: '@ai-dynamo/dynamo-rl-codeowners' path_globs: - components/src/dynamo/common/rl/ + - examples/rl/ - lib/rl/ - tests/rl/ - label: performance diff --git a/CODEOWNERS b/CODEOWNERS index 6334cab97e82..0ff6f6ebc2df 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -392,6 +392,7 @@ # === rl (@ai-dynamo/dynamo-rl-codeowners) === /lib/rl/ @ai-dynamo/dynamo-rl-codeowners +/examples/rl/ @ai-dynamo/dynamo-rl-codeowners # === performance (@ai-dynamo/dynamo-performance-codeowners) === /recipes/ @ai-dynamo/dynamo-performance-codeowners diff --git a/components/src/dynamo/common/backend/logprobs.py b/components/src/dynamo/common/backend/logprobs.py index e032312e845d..37257786002f 100644 --- a/components/src/dynamo/common/backend/logprobs.py +++ b/components/src/dynamo/common/backend/logprobs.py @@ -314,52 +314,42 @@ def _parse(name: str, value: Any) -> Optional[int]: def extract_from_sglang_meta( meta_info: dict[str, Any], *, - num_output_tokens_in_chunk: Optional[int] = None, return_tokens_as_token_ids: bool = False, ) -> tuple[Optional[list[float]], Optional[list[list[dict[str, Any]]]]]: """Extract logprobs from SGLang's ``meta_info`` dict. - The pinned SGLang release and its N-1 predecessor align - ``output_token_logprobs`` and ``output_top_logprobs`` with each incremental - ``output_ids`` chunk. When provided, ``num_output_tokens_in_chunk`` keeps - malformed trailing metadata out of Dynamo's response. + Dynamo enables SGLang's incremental streaming output, so + ``output_token_logprobs`` and ``output_top_logprobs`` already align with + the current disjoint ``output_ids`` chunk and can be forwarded directly. """ output_token_logprobs = meta_info.get("output_token_logprobs") if not output_token_logprobs: return None, None - new_logprobs = output_token_logprobs - if num_output_tokens_in_chunk is not None: - new_logprobs = new_logprobs[:num_output_tokens_in_chunk] - if not new_logprobs: - return None, None - - log_probs = [float(entry[0]) for entry in new_logprobs] + log_probs = [float(entry[0]) for entry in output_token_logprobs] top_logprobs: Optional[list[list[dict[str, Any]]]] = None output_top = meta_info.get("output_top_logprobs") if output_top: - new_top = output_top[: len(new_logprobs)] - if new_top: - top_logprobs = [] - for position_entries in new_top: - if position_entries is None: - top_logprobs.append([]) - continue - position_list: list[dict[str, Any]] = [] - for rank_idx, entry in enumerate(position_entries): - tok_id = entry[1] - token_str = ( - f"token_id:{tok_id}" if return_tokens_as_token_ids else entry[2] - ) - position_list.append( - { - "rank": rank_idx + 1, - "token_id": tok_id, - "token": token_str, - "logprob": float(entry[0]), - } - ) - top_logprobs.append(position_list) + top_logprobs = [] + for position_entries in output_top: + if position_entries is None: + top_logprobs.append([]) + continue + position_list: list[dict[str, Any]] = [] + for rank_idx, entry in enumerate(position_entries): + tok_id = entry[1] + token_str = ( + f"token_id:{tok_id}" if return_tokens_as_token_ids else entry[2] + ) + position_list.append( + { + "rank": rank_idx + 1, + "token_id": tok_id, + "token": token_str, + "logprob": float(entry[0]), + } + ) + top_logprobs.append(position_list) return log_probs, top_logprobs diff --git a/components/src/dynamo/common/backend/tests/test_logprobs.py b/components/src/dynamo/common/backend/tests/test_logprobs.py index 7bdc3437f78a..fc5fe89919d8 100644 --- a/components/src/dynamo/common/backend/tests/test_logprobs.py +++ b/components/src/dynamo/common/backend/tests/test_logprobs.py @@ -405,9 +405,7 @@ def test_sglang_extract_supports_incremental_streaming_metadata(): "output_token_logprobs": [(-0.2, 2, "b")], "output_top_logprobs": [[(-0.2, 2, "b"), (-1.2, 20, "B")]], } - log_probs, top_logprobs = extract_from_sglang_meta( - meta, num_output_tokens_in_chunk=1 - ) + log_probs, top_logprobs = extract_from_sglang_meta(meta) assert log_probs == [-0.2] assert top_logprobs == [ [ @@ -456,14 +454,6 @@ def test_sglang_extract_none_top_position_becomes_empty_list(): ] -def test_sglang_extract_clamps_metadata_to_output_chunk(): - meta = { - "output_token_logprobs": [(-0.1, 1, "a"), (-0.2, 2, "b")], - } - log_probs, _ = extract_from_sglang_meta(meta, num_output_tokens_in_chunk=1) - assert log_probs == [-0.1] - - # --------------------------------------------------------------------------- # Active handler adapters. These tests import the static methods on each # backend's handler class and compare them with the shared helper. They catch diff --git a/components/src/dynamo/frontend/sglang_prepost.py b/components/src/dynamo/frontend/sglang_prepost.py index f8ff8d0040ad..c91db2899ac6 100644 --- a/components/src/dynamo/frontend/sglang_prepost.py +++ b/components/src/dynamo/frontend/sglang_prepost.py @@ -986,6 +986,7 @@ def __init__( eos_token_ids: list[int] | None = None, prompt_token_ids: list[int] | None = None, stop_strings: set[str] | None = None, + stop_token_ids: set[int] | None = None, ) -> None: self.tokenizer = tokenizer self.tool_call_parser = tool_call_parser @@ -1007,7 +1008,11 @@ def __init__( self._pending_guided_reasoning_parts: list[str] | None = ( [] if self._is_json_array_parser and reasoning_parser is not None else None ) + # The frontend owns text shaping. Keep every stop token in the raw + # engine stream, but exclude a matched stop suffix when decoding + # user-visible text. self._eos_token_ids = set(eos_token_ids or []) + self._request_stop_token_ids = set(stop_token_ids or []) self._stop_strings = stop_strings or set() self._pending_stop_text = "" self._locally_finished = False @@ -1036,11 +1041,35 @@ def __init__( # Full text accumulator for robust finish-time re-parse. self._tool_text_parts: list[str] = [] - def _strip_trailing_eos_token_ids(self, token_ids: list[int]) -> list[int]: - if not self._eos_token_ids: + def _strip_matched_stop_token_ids( + self, token_ids: list[int], stop_reason: Any + ) -> list[int]: + """Remove only the engine-reported token-stop suffix from display IDs.""" + if not token_ids: return token_ids - while token_ids and token_ids[-1] in self._eos_token_ids: - token_ids.pop() + + known_stop_ids = self._eos_token_ids | self._request_stop_token_ids + if isinstance(stop_reason, int) and not isinstance(stop_reason, bool): + matched_ids = [stop_reason] if stop_reason in known_stop_ids else [] + elif ( + isinstance(stop_reason, list) + and stop_reason + and all( + isinstance(token_id, int) + and not isinstance(token_id, bool) + and token_id in known_stop_ids + for token_id in stop_reason + ) + ): + matched_ids = stop_reason + elif stop_reason is None and token_ids[-1] in self._eos_token_ids: + # Model EOS is intentionally omitted from the public stop_reason. + matched_ids = [token_ids[-1]] + else: + matched_ids = [] + + if matched_ids and token_ids[-len(matched_ids) :] == matched_ids: + del token_ids[-len(matched_ids) :] return token_ids def _tool_call_id(self, name: str, index: int) -> str: @@ -1364,9 +1393,12 @@ def process_output(self, engine_response: dict[str, Any]) -> dict[str, Any] | No stop_reason = engine_response.get("stop_reason") log_probs = engine_response.get("log_probs") top_logprobs = engine_response.get("top_logprobs") - if finish_reason is not None: + stop_terminated = engine_response.get( + "stop_terminated", finish_reason == "stop" + ) + if stop_terminated: raw_token_count = len(token_ids) - token_ids = self._strip_trailing_eos_token_ids(list(token_ids)) + token_ids = self._strip_matched_stop_token_ids(list(token_ids), stop_reason) retained_token_count = len(token_ids) if log_probs is not None and len(log_probs) == raw_token_count: log_probs = log_probs[:retained_token_count] diff --git a/components/src/dynamo/frontend/sglang_processor.py b/components/src/dynamo/frontend/sglang_processor.py index e5d3781e6f92..35a1ee6f0b28 100644 --- a/components/src/dynamo/frontend/sglang_processor.py +++ b/components/src/dynamo/frontend/sglang_processor.py @@ -146,6 +146,17 @@ def _request_stop_strings(request: dict[str, Any]) -> set[str]: return set() +def _request_stop_token_ids(request: dict[str, Any]) -> list[int]: + """Merge ``stop_token_ids`` with the legacy integer-valued ``stop`` form.""" + values: list[Any] = list(request.get("stop_token_ids") or []) + stop = request.get("stop") + if isinstance(stop, list) and all( + isinstance(item, int) and not isinstance(item, bool) for item in stop + ): + values.extend(stop) + return _normalize_eos_token_ids(values) + + def _tokenizer_eos_token_ids(tokenizer: Any) -> list[int]: eos_token_ids = _normalize_eos_token_ids(getattr(tokenizer, "eos_token_ids", None)) if eos_token_ids: @@ -378,13 +389,12 @@ def _build_dynamo_preproc( max_tokens = request.get("max_completion_tokens") or request.get("max_tokens") stop = request.get("stop") - stop_token_ids = request.get("stop_token_ids", []) + stop_token_ids = _request_stop_token_ids(request) if isinstance(stop, str): stop = [stop] elif isinstance(stop, list) and all( isinstance(item, int) and not isinstance(item, bool) for item in stop ): - stop_token_ids = [*stop_token_ids, *stop] stop = [] elif stop is None: stop = [] @@ -612,6 +622,7 @@ async def _generator_inner( eos_token_ids=self.eos_token_ids, prompt_token_ids=pre.prompt_token_ids, stop_strings=_request_stop_strings(request), + stop_token_ids=set(_request_stop_token_ids(request)), ) async for item in self._generate_and_stream( @@ -672,6 +683,7 @@ async def _generator_inner_pool( eos_token_ids=self.eos_token_ids, prompt_token_ids=preproc_result.prompt_token_ids, stop_strings=_request_stop_strings(request), + stop_token_ids=set(_request_stop_token_ids(request)), ) async for item in self._generate_and_stream( @@ -727,6 +739,7 @@ def flush_pending( *, finish_reason: str | None, stop_reason: Any | None, + stop_terminated: bool, engine_data: Any | None, ) -> dict[str, Any]: nonlocal pending_token_ids @@ -742,6 +755,7 @@ def flush_pending( "token_ids": pending_token_ids, "finish_reason": finish_reason, "stop_reason": stop_reason, + "stop_terminated": stop_terminated, } if pending_log_probs is not None: mapped_response["log_probs"] = pending_log_probs @@ -864,6 +878,7 @@ def flush_pending( envelope = flush_pending( finish_reason=None, stop_reason=None, + stop_terminated=False, engine_data=None, ) yield envelope @@ -872,8 +887,10 @@ def flush_pending( chunk_tokens = len(new_ids) cumulative_output_tokens += chunk_tokens - finish_reason = _map_finish_reason(engine_response.get("finish_reason")) + raw_finish_reason = engine_response.get("finish_reason") + finish_reason = _map_finish_reason(raw_finish_reason) stop_reason = engine_response.get("stop_reason") + stop_terminated = raw_finish_reason in {"eos", "stop"} if usage := engine_response.get("completion_usage"): pending_usage = usage @@ -898,6 +915,7 @@ def flush_pending( envelope = flush_pending( finish_reason=finish_reason, stop_reason=stop_reason, + stop_terminated=stop_terminated, engine_data=engine_data, ) yield envelope diff --git a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py index a9afcda2b395..7010443dc266 100644 --- a/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_sglang_processor_unit.py @@ -3777,16 +3777,161 @@ def test_completed_batches_replace_decode_context(self): assert post._decode_context_ids == [ord("a")] assert post._pending_decode_ids == [] - def test_strips_all_configured_trailing_eos_token_ids(self, tokenizer): - """Any configured EOS id is stripped from the final chunk before decode.""" + def test_strips_only_the_exact_matched_stop_suffix(self, tokenizer): + """Matched metadata, not configured membership alone, selects the suffix.""" post = SglangStreamingPostProcessor( tokenizer=tokenizer, tool_call_parser=None, reasoning_parser=None, eos_token_ids=[2, 3], + stop_token_ids={4, 5}, ) - assert post._strip_trailing_eos_token_ids([10, 3, 2]) == [10] + assert post._strip_matched_stop_token_ids([10, 3, 2], None) == [10, 3] + assert post._strip_matched_stop_token_ids([10, 4, 5], 5) == [10, 4] + assert post._strip_matched_stop_token_ids([10, 4, 5], [4, 5]) == [10] + + @pytest.mark.parametrize( + "stop_config", + [ + {"stop_token_ids": {ord("A")}}, + {"eos_token_ids": [ord("A")]}, + ], + ids=["request-stop-id", "model-eos-id"], + ) + def test_length_finish_keeps_configured_token_and_logprob(self, stop_config): + """A length finish does not turn a configured final ID into a match.""" + routed_engine = FakeRoutedEngine( + items=[ + { + "token_ids": [ord("A")], + "finish_reason": "length", + "log_probs": [-0.25], + } + ] + ) + processor = SglangProcessor( + tokenizer=self.ByteTokenizer(), + routed_engine=routed_engine, + tool_call_parser_name=None, + reasoning_parser_name=None, + eos_token_ids=None, + ) + post = SglangStreamingPostProcessor( + tokenizer=self.ByteTokenizer(), + tool_call_parser=None, + reasoning_parser=None, + **stop_config, + ) + + async def collect(): + return [ + item["data"] + async for item in processor._generate_and_stream( + "req-length", {"model": "test-model"}, {}, [], post + ) + if "data" in item + ] + + chunks = asyncio.run(collect()) + + assert len(chunks) == 1 + choice = chunks[0]["choices"][0] + assert choice["delta"]["content"] == "A" + assert choice["finish_reason"] == "length" + assert choice["logprobs"]["content"] == [ + { + "token": "A", + "logprob": -0.25, + "bytes": [65], + "top_logprobs": [], + } + ] + + def test_string_match_inside_configured_stop_token_keeps_visible_prefix( + self, tokenizer + ): + """An engine-reported string match wins over configured token membership.""" + token_ids = tokenizer.encode("alphabet") + assert len(token_ids) == 1 + post = SglangStreamingPostProcessor( + tokenizer=tokenizer, + tool_call_parser=None, + reasoning_parser=None, + stop_strings={"pha"}, + stop_token_ids=set(token_ids), + ) + + choice = post.process_output( + { + "token_ids": token_ids, + "finish_reason": "stop", + "stop_reason": "pha", + "stop_terminated": True, + } + ) + + assert choice is not None + assert choice["delta"]["content"] == "al" + assert choice["finish_reason"] == "stop" + + def test_request_stop_token_id_is_hidden_by_python_frontend(self): + """Match the Rust frontend when SGLang returns a request stop token.""" + + class RequestTokenizer(self.ByteTokenizer): + chat_template = "{{ messages }}" + eos_token_id = 0 + + def encode(self, text, *, add_special_tokens=False): + del add_special_tokens + return list(text.encode()) + + def apply_chat_template(self, messages, **kwargs): + del messages, kwargs + return [1, 2, 3] + + tokenizer = RequestTokenizer() + generated_ids = tokenizer.encode("Hello world", add_special_tokens=False) + assert len(generated_ids) >= 2 + stop_token_id = generated_ids[-1] + visible_ids = generated_ids[:-1] + routed_engine = FakeRoutedEngine( + items=[ + { + # Native SGLang includes the matched stop position. + "token_ids": generated_ids, + "finish_reason": "stop", + "stop_reason": stop_token_id, + } + ] + ) + processor = SglangProcessor( + tokenizer=tokenizer, + routed_engine=routed_engine, + tool_call_parser_name=None, + reasoning_parser_name=None, + eos_token_ids=[tokenizer.eos_token_id], + ) + request = { + "model": MODEL, + "messages": [{"role": "user", "content": "Say hello"}], + "stop_token_ids": [stop_token_id], + } + + async def collect(): + return [item async for item in processor.generator(request)] + + items = asyncio.run(collect()) + content = "".join( + item["data"]["choices"][0]["delta"].get("content", "") + for item in items + if "data" in item + ) + + assert routed_engine.requests[0]["stop_conditions"]["stop_token_ids"] == [ + stop_token_id + ] + assert content == tokenizer.decode(visible_ids, skip_special_tokens=True) # --------------------------------------------------------------------------- diff --git a/components/src/dynamo/sglang/pause.py b/components/src/dynamo/sglang/pause.py deleted file mode 100644 index e75f6d7fcee6..000000000000 --- a/components/src/dynamo/sglang/pause.py +++ /dev/null @@ -1,77 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import logging -from typing import Any - -from sglang.srt.managers.io_struct import ( - ContinueGenerationReqInput, - PauseGenerationReqInput, - ReleaseMemoryOccupationReqInput, - ResumeMemoryOccupationReqInput, -) - -logger = logging.getLogger(__name__) - - -class SGLangEnginePauseController: - def __init__(self, engine: Any): - self._engine = engine - self._is_paused = False - self._generation_paused = False - - @property - def is_paused(self) -> bool: - return self._is_paused - - @property - def needs_resume_recovery(self) -> bool: - return self._generation_paused - - async def pause(self, tags: list[str] | None = None) -> bool: - if self._is_paused or self._generation_paused: - return False - - await self._engine.tokenizer_manager.pause_generation(PauseGenerationReqInput()) - self._generation_paused = True - try: - await self._engine.tokenizer_manager.release_memory_occupation( - ReleaseMemoryOccupationReqInput(tags=tags), - None, - ) - except Exception: - try: - await self._engine.tokenizer_manager.continue_generation( - ContinueGenerationReqInput() - ) - self._generation_paused = False - except Exception: - logger.exception( - "failed to resume generation after memory release failed" - ) - raise - - self._is_paused = True - return True - - async def resume(self, tags: list[str] | None = None) -> bool: - if not self._is_paused and not self._generation_paused: - return False - - if self._is_paused: - await self._engine.tokenizer_manager.resume_memory_occupation( - ResumeMemoryOccupationReqInput(tags=tags), - None, - ) - if self._generation_paused: - await self._engine.tokenizer_manager.continue_generation( - ContinueGenerationReqInput() - ) - self._generation_paused = False - return True - - def mark_resumed(self) -> None: - self._is_paused = False - self._generation_paused = False diff --git a/components/src/dynamo/sglang/request_handlers/handler_base.py b/components/src/dynamo/sglang/request_handlers/handler_base.py index 884205a8b23b..82773910a948 100644 --- a/components/src/dynamo/sglang/request_handlers/handler_base.py +++ b/components/src/dynamo/sglang/request_handlers/handler_base.py @@ -48,12 +48,19 @@ from dynamo.sglang.args import Config from dynamo.sglang.capacity import kv_event_block_size from dynamo.sglang.engine_routes import resolve_configured_engine_routes -from dynamo.sglang.pause import SGLangEnginePauseController from dynamo.sglang.publisher import DynamoSglangPublisher logger = logging.getLogger(__name__) +_DEFAULT_NATIVE_ENGINE_ROUTES = ( + ("pause_generation", "pause_generation:tm"), + ("continue_generation", "continue_generation:tm"), + ("release_memory_occupation", "release_memory_occupation:tm"), + ("resume_memory_occupation", "resume_memory_occupation:tm"), +) + + RequestT = TypeVar("RequestT") ResponseT = TypeVar("ResponseT") @@ -604,10 +611,7 @@ def __init__( # have an sgl.Engine. self.input_param_manager = InputParamManager(None) self._engine_supports_priority = False - self._pause_controller = ( - SGLangEnginePauseController(engine) if engine is not None else None - ) - self._pause_lock = asyncio.Lock() + self._engine_route_lock = asyncio.Lock() # Serializes elastic-EP scaling: SGLang tracks a single in-flight scale # phase, so concurrent scale_elastic_ep calls must not overlap. @@ -626,122 +630,6 @@ def _priority_kwargs(self, priority: Any) -> Dict[str, Any]: return {"priority": normalized} return {} - async def release_memory_occupation(self, body: dict) -> dict: - """Release GPU memory occupation and unregister from discovery. - - Args: - body: Optional dict with "tags" to target specific memory regions. - - Order of operations: - 1. Unregister from discovery - stop accepting new requests - 2. Pause generation - drain in-flight requests - 3. Release memory - safe now that no requests are active - """ - if self._pause_controller is None: - return { - "status": "error", - "message": "memory control not supported on this worker", - } - - body = body or {} - tags = body.get("tags") - async with self._pause_lock: - if self._pause_controller.is_paused: - return { - "status": "ok", - "message": "Memory already released", - } - if self._pause_controller.needs_resume_recovery: - return { - "status": "error", - "message": "resume_memory_occupation required before retrying release", - } - - unregistered = False - try: - # Stop new requests and drain in-flight work before releasing memory. - if self.generate_endpoint is not None: - await self.generate_endpoint.unregister_endpoint_instance() - unregistered = True - - await self._pause_controller.pause(tags) - - return { - "status": "ok", - "message": ( - f"Memory released for tags: {tags}" - if tags is not None - else "Memory released" - ), - } - except Exception as e: - logging.error(f"Failed to release memory occupation: {e}") - # If pause rolled back cleanly the engine is serving-safe again, - # but discovery still shows us unregistered and resume will - # early-return. Re-register so the worker rejoins the routing pool. - if ( - unregistered - and not self._pause_controller.is_paused - and not self._pause_controller.needs_resume_recovery - and self.generate_endpoint is not None - ): - try: - await self.generate_endpoint.register_endpoint_instance() - logging.info( - "Re-registered endpoint after failed memory release rollback" - ) - except Exception as reg_err: - logging.error( - f"Failed to re-register endpoint after release failure: {reg_err}" - ) - return {"status": "error", "message": str(e)} - - async def resume_memory_occupation(self, body: dict) -> dict: - """Resume GPU memory occupation and re-register to discovery. - - Args: - body: Optional dict with "tags" to target specific memory regions. - - Order of operations: - 1. Resume memory - restore GPU allocations - 2. Continue generation - ready to serve requests - 3. Re-register to discovery - allow frontend to route here - """ - if self._pause_controller is None: - return { - "status": "error", - "message": "memory control not supported on this worker", - } - - body = body or {} - tags = body.get("tags") - async with self._pause_lock: - needs_recovery = self._pause_controller.needs_resume_recovery - if not self._pause_controller.is_paused and not needs_recovery: - return { - "status": "ok", - "message": "Memory already resumed", - } - - try: - await self._pause_controller.resume(tags) - - if self.generate_endpoint is not None: - await self.generate_endpoint.register_endpoint_instance() - self._pause_controller.mark_resumed() - - return { - "status": "ok", - "message": ( - f"Memory resumed for tags: {tags}" - if tags is not None - else "Memory resumed" - ), - } - except Exception as e: - logging.error(f"Failed to resume memory occupation: {e}") - return {"status": "error", "message": str(e)} - async def clear_kv_blocks(self, request: Optional[Dict[str, Any]] = None): """Flush SGLang's local cache when no requests are active.""" tokenizer_manager = ( @@ -757,7 +645,7 @@ async def clear_kv_blocks(self, request: Optional[Dict[str, Any]] = None): return try: - async with self._pause_lock: + async with self._engine_route_lock: if getattr(tokenizer_manager, "rid_to_state", None): yield { "status": "error", @@ -975,6 +863,41 @@ async def is_scaling_elastic_ep(self, body: dict) -> dict: return backend_error return dict(self.engine.tokenizer_manager.get_elastic_ep_state()) + async def _sync_discovery_with_sglang_pause_state(self) -> None: + """Make discovery match SGLang's authoritative generation pause state.""" + tokenizer_manager = ( + getattr(self.engine, "tokenizer_manager", None) + if self.engine is not None + else None + ) + if ( + self.generate_endpoint is None + or tokenizer_manager is None + or not hasattr(tokenizer_manager, "is_pause") + ): + return + + if tokenizer_manager.is_pause: + await self.generate_endpoint.unregister_endpoint_instance() + else: + await self.generate_endpoint.register_endpoint_instance() + + async def _invoke_engine_route(self, route_handler, body: dict) -> dict: + """Invoke one engine route and then synchronize worker discovery.""" + async with self._engine_route_lock: + try: + return await route_handler(body) + finally: + await self._sync_discovery_with_sglang_pause_state() + + def _wrap_engine_route(self, route_handler): + """Add discovery synchronization to an SGLang route.""" + + async def synchronized_handler(body: dict) -> dict: + return await self._invoke_engine_route(route_handler, body) + + return synchronized_handler + def register_engine_routes(self, runtime: DistributedRuntime) -> None: """Register all engine routes for this handler. @@ -988,8 +911,6 @@ def register_engine_routes(self, runtime: DistributedRuntime) -> None: built_in_routes = { "control/start_profile": self.start_profile, "control/stop_profile": self.stop_profile, - "control/release_memory_occupation": self.release_memory_occupation, - "control/resume_memory_occupation": self.resume_memory_occupation, "control/update_weights_from_disk": self.update_weights_from_disk, "control/update_weights_from_tensor": self.update_weights_from_tensor, "control/update_weights_from_distributed": ( @@ -1013,11 +934,23 @@ def register_engine_routes(self, runtime: DistributedRuntime) -> None: "with a built-in route" ) + configured_paths = {path for path, _ in configured_routes} + # Expose the SGLang lifecycle methods without requiring --engine-route. + # An explicit route with the same public path overrides the default. + default_native_routes = resolve_configured_engine_routes( + self.engine, + [ + descriptor + for path, descriptor in _DEFAULT_NATIVE_ENGINE_ROUTES + if path not in configured_paths + ], + ) + register_model_taint_route(runtime, self.generate_endpoint) for path, handler in built_in_routes.items(): runtime.register_engine_route(path, handler) - for path, configured_handler in configured_routes: - runtime.register_engine_route(path, configured_handler) + for path, route_handler in [*default_native_routes, *configured_routes]: + runtime.register_engine_route(path, self._wrap_engine_route(route_handler)) @abstractmethod def generate(self, request: RequestT, context: Context) -> AsyncIterator[ResponseT]: 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 e7b18053d97b..b50566333185 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/decode_handler.py @@ -18,10 +18,7 @@ from dynamo.common.metadata_upload import MetadataUploader from dynamo.common.multimodal.image_loader import ImageLoader from dynamo.common.multimodal.video_loader import VideoLoader -from dynamo.common.utils.engine_response import ( - normalize_finish_reason, - trailing_stop_prefix_len, -) +from dynamo.common.utils.engine_response import normalize_finish_reason from dynamo.llm import HttpError from dynamo.llm.exceptions import EngineShutdown from dynamo.sglang._compat import ( @@ -165,77 +162,6 @@ def _user_stop_token_ids(request: Dict[str, Any]) -> set[int]: } -def _suppressed_stop_token_ids(request: Dict[str, Any]) -> set[int]: - stop_conditions = request.get("stop_conditions") or {} - if not isinstance(stop_conditions, dict): - return set() - - values = [ - *(stop_conditions.get("stop_token_ids") or []), - *(stop_conditions.get("stop_token_ids_hidden") or []), - ] - return { - token_id - for token_id in values - if isinstance(token_id, int) and not isinstance(token_id, bool) - } - - -def _stop_strings(request: Dict[str, Any]) -> set[str]: - stop_conditions = request.get("stop_conditions") - stop = ( - stop_conditions.get("stop") - if isinstance(stop_conditions, dict) - else request.get("stop") - ) - if isinstance(stop, str): - return {stop} - if isinstance(stop, list): - return {item for item in stop if isinstance(item, str)} - return set() - - -def _remove_suppressed_stop_tokens( - output_ids: list[int], - matched: Any, - suppressed_stop_token_ids: set[int] | None, -) -> tuple[list[int], int]: - if not output_ids or not suppressed_stop_token_ids: - return output_ids, 0 - - if isinstance(matched, int) and not isinstance(matched, bool): - matched_ids = [matched] - elif isinstance(matched, list) and all( - isinstance(item, int) and not isinstance(item, bool) for item in matched - ): - matched_ids = matched - else: - return output_ids, 0 - - if ( - matched_ids - and all(token_id in suppressed_stop_token_ids for token_id in matched_ids) - and len(output_ids) >= len(matched_ids) - and output_ids[-len(matched_ids) :] == matched_ids - ): - return output_ids[: -len(matched_ids)], len(matched_ids) - - return output_ids, 0 - - -def _split_trailing_suppressed_stop_tokens( - output_ids: list[int], - suppressed_stop_token_ids: set[int] | None, -) -> tuple[list[int], list[int]]: - if not output_ids or not suppressed_stop_token_ids: - return output_ids, [] - - split_at = len(output_ids) - while split_at > 0 and output_ids[split_at - 1] in suppressed_stop_token_ids: - split_at -= 1 - return output_ids[:split_at], output_ids[split_at:] - - def _openai_stop_sampling_params(request: Dict[str, Any]) -> Dict[str, Any]: stop = request.get("stop") if isinstance(stop, str): @@ -434,6 +360,7 @@ 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 + param_mapping = { "n": sampling_opts.get("n"), "max_new_tokens": stop_conditions.get("max_tokens"), @@ -473,12 +400,11 @@ def _build_logprob_kwargs(request: Dict[str, Any]) -> Dict[str, Any]: @staticmethod def _extract_logprobs( meta_info: Dict[str, Any], - num_output_tokens_in_chunk: Optional[int] = None, + *, return_tokens_as_token_ids: bool = False, ) -> tuple: return _shared_logprobs.extract_from_sglang_meta( meta_info, - num_output_tokens_in_chunk=num_output_tokens_in_chunk, return_tokens_as_token_ids=return_tokens_as_token_ids, ) @@ -567,7 +493,6 @@ async def generate( output_options.get("return_tokens_as_token_ids") ) user_stop_token_ids = _user_stop_token_ids(request) - suppressed_stop_token_ids = _suppressed_stop_token_ids(request) lora_path = self._resolve_lora(request) if lora_path: @@ -624,7 +549,6 @@ async def generate( context, return_tokens_as_token_ids, user_stop_token_ids=user_stop_token_ids, - suppressed_stop_token_ids=suppressed_stop_token_ids, metadata_uploader=metadata_uploader, ): yield out @@ -708,7 +632,6 @@ async def generate( context, return_tokens_as_token_ids, user_stop_token_ids=user_stop_token_ids, - suppressed_stop_token_ids=suppressed_stop_token_ids, metadata_uploader=metadata_uploader, ): yield out @@ -752,7 +675,6 @@ async def _process_token_stream( context: Context, return_tokens_as_token_ids: bool = False, user_stop_token_ids: set[int] | None = None, - suppressed_stop_token_ids: set[int] | None = None, metadata_uploader: MetadataUploader | None = None, ) -> AsyncGenerator[Dict[str, Any], None]: """Process token-based stream output. @@ -770,10 +692,6 @@ 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() first_output_seen = False - pending_stop_tokens_per_choice: dict[int, list[int]] = {} - pending_log_probs_per_choice: dict[int, list[Any]] = {} - pending_top_logprobs_per_choice: dict[int, list[Any]] = {} - completion_tokens_per_choice: dict[int, int] = {} async with self._cancellation_monitor(request_id_future, context): async for res in stream_source: meta_info = res.get("meta_info", {}) @@ -814,60 +732,7 @@ async def _process_token_stream( out["stop_reason"] = stop_reason # With stream_output=True, output_ids contains only new tokens (disjoint) - raw_output_ids = list(res.get("output_ids", [])) - if raw_output_ids and not first_output_seen: - first_output_seen = True - context.notify_first_token() - log_probs = None - top_logprobs = None - if metadata_uploader is None: - # Extract logprobs for new tokens if available. - log_probs, top_logprobs = self._extract_logprobs( - meta_info, - num_output_tokens_in_chunk=len(raw_output_ids), - return_tokens_as_token_ids=return_tokens_as_token_ids, - ) - - output_ids = pending_stop_tokens_per_choice.pop(output_idx, []) - output_ids.extend(raw_output_ids) - if log_probs is not None: - log_probs = pending_log_probs_per_choice.pop(output_idx, []) + list( - log_probs - ) - if top_logprobs is not None: - top_logprobs = pending_top_logprobs_per_choice.pop( - output_idx, [] - ) + list(top_logprobs) - - matched = finish_reason.get("matched") if finish_reason else None - output_ids, removed_count = _remove_suppressed_stop_tokens( - output_ids, matched, suppressed_stop_token_ids - ) - if removed_count: - if log_probs is not None: - log_probs = log_probs[:-removed_count] - if top_logprobs is not None: - top_logprobs = top_logprobs[:-removed_count] - elif not finish_reason: - ( - output_ids, - pending_stop_tokens, - ) = _split_trailing_suppressed_stop_tokens( - output_ids, suppressed_stop_token_ids - ) - if pending_stop_tokens: - pending_stop_tokens_per_choice[output_idx] = pending_stop_tokens - if log_probs is not None: - pending_log_probs_per_choice[output_idx] = log_probs[ - -len(pending_stop_tokens) : - ] - log_probs = log_probs[: -len(pending_stop_tokens)] - if top_logprobs is not None: - pending_top_logprobs_per_choice[output_idx] = top_logprobs[ - -len(pending_stop_tokens) : - ] - top_logprobs = top_logprobs[: -len(pending_stop_tokens)] - + output_ids = res.get("output_ids", []) # Empty, non-final chunks can happen during scheduler idle ticks. # Keep waiting for the next chunk unless cancellation was requested. if not output_ids and not finish_reason: @@ -875,13 +740,21 @@ async def _process_token_stream( break continue + if output_ids and not first_output_seen: + first_output_seen = True + context.notify_first_token() + # Pass through disjoint token segments directly out["token_ids"] = output_ids - - if log_probs is not None: - out["log_probs"] = log_probs - if top_logprobs is not None: - out["top_logprobs"] = top_logprobs + if metadata_uploader is None: + log_probs, top_logprobs = self._extract_logprobs( + meta_info, + return_tokens_as_token_ids=return_tokens_as_token_ids, + ) + if log_probs is not None: + out["log_probs"] = log_probs + if top_logprobs is not None: + out["top_logprobs"] = top_logprobs engine_data: dict[str, Any] = dict(res.get("engine_data") or {}) routed_experts = meta_info.get("routed_experts") @@ -905,14 +778,10 @@ async def _process_token_stream( if cached_tokens is not None and cached_tokens > 0: prefill_prompt_tokens_details = {"cached_tokens": cached_tokens} if input_tokens is not None and completion_tokens is not None: - completion_tokens_per_choice[output_idx] = completion_tokens - request_completion_tokens = sum( - completion_tokens_per_choice.values() - ) completion_usage = { "prompt_tokens": input_tokens, - "completion_tokens": request_completion_tokens, - "total_tokens": input_tokens + request_completion_tokens, + "completion_tokens": completion_tokens, + "total_tokens": input_tokens + completion_tokens, } if prefill_prompt_tokens_details is not None: completion_usage[ @@ -957,11 +826,6 @@ async def _process_text_stream( OpenAI-formatted chat completion chunk dicts. """ request = request or {} - stop_strings = _stop_strings(request) - # SGLang text chunks are cumulative per choice. Keep independent text - # offsets so interleaved n>1 choices do not compute deltas from each - # other's previous text. - text_counts_per_choice: dict[int, int] = {} # Use Future pattern for request ID - will be set when first response arrives request_id_future: asyncio.Future[str] = asyncio.Future() @@ -984,7 +848,10 @@ async def _process_text_stream( # Same defaulting as token mode: non-n chunks are choice 0. index = res.get("index") or 0 - text = res.get("text", "") + # Dynamo forces incremental_streaming_output=True, so SGLang + # has already produced the client-facing disjoint text delta. + # Its default detokenizer also buffers and trims stop markers. + delta = res.get("text", "") finish_reason = meta_info["finish_reason"] if finish_reason: @@ -1001,20 +868,6 @@ async def _process_text_stream( finish_reason_type = normalize_finish_reason(finish_reason["type"]) else: finish_reason_type = None - next_count = len(text) - count = text_counts_per_choice.get(index, 0) - visible_text_len = next_count - matched = finish_reason.get("matched") if finish_reason else None - if ( - isinstance(matched, str) - and matched in stop_strings - and text.endswith(matched) - ): - visible_text_len = len(text) - len(matched) - elif not finish_reason: - visible_text_len -= trailing_stop_prefix_len(text, stop_strings) - - delta = text[count:visible_text_len] if visible_text_len > count else "" if res.get("output_ids") and not first_output_seen: first_output_seen = True context.notify_first_token() @@ -1063,4 +916,3 @@ async def _process_text_stream( response["nvext"] = response_nvext if not context.is_stopped(): yield response - text_counts_per_choice[index] = visible_text_len diff --git a/components/src/dynamo/sglang/snapshot.py b/components/src/dynamo/sglang/snapshot.py index 8279a8fa28ed..5e47fcbb1d4b 100644 --- a/components/src/dynamo/sglang/snapshot.py +++ b/components/src/dynamo/sglang/snapshot.py @@ -11,6 +11,12 @@ from typing import Any import sglang as sgl +from sglang.srt.managers.io_struct import ( + ContinueGenerationReqInput, + PauseGenerationReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, +) from dynamo.common.snapshot.lifecycle import ( EngineSnapshotController, @@ -19,11 +25,31 @@ ) from dynamo.sglang._compat import override_server_args, resolved_server_args -from .pause import SGLangEnginePauseController - logger = logging.getLogger(__name__) +class _SGLangSnapshotLifecycle: + """Adapt SGLang's native controls to the snapshot lifecycle interface.""" + + def __init__(self, engine: sgl.Engine): + self._tokenizer_manager = engine.tokenizer_manager + + async def pause(self) -> None: + await self._tokenizer_manager.pause_generation(PauseGenerationReqInput()) + await self._tokenizer_manager.release_memory_occupation( + ReleaseMemoryOccupationReqInput(), None + ) + + async def resume(self) -> None: + await self._tokenizer_manager.resume_memory_occupation( + ResumeMemoryOccupationReqInput(), None + ) + await self._tokenizer_manager.continue_generation(ContinueGenerationReqInput()) + + def mark_resumed(self) -> None: + pass + + async def warmup_engine(engine: sgl.Engine, server_args: Any) -> None: """Warm up the direct SGLang Engine generation path before snapshot capture.""" DUMMY_WARMUP_MAX_NEW_TOKENS = 8 @@ -171,7 +197,7 @@ async def prepare_snapshot_engine( snapshot_controller = EngineSnapshotController( engine=engine, - pause_controller=SGLangEnginePauseController(engine), + pause_controller=_SGLangSnapshotLifecycle(engine), snapshot_config=snapshot_config, ) if not await snapshot_controller.wait_for_restore(): 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 7c817caab3b2..f38e31230bff 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_handler.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio -import json from contextlib import asynccontextmanager from types import SimpleNamespace @@ -26,8 +25,6 @@ _extract_sglang_stop_reason, _nvext_extra_field_requested, _openai_stop_sampling_params, - _remove_suppressed_stop_tokens, - _suppressed_stop_token_ids, _user_stop_token_ids, ) from dynamo.sglang.request_handlers.llm.mm_disagg_utils import ( @@ -36,10 +33,7 @@ raise_if_unextracted_multimodal, ) from dynamo.sglang.request_handlers.llm.prefill_handler import PrefillWorkerHandler -from dynamo.sglang.request_handlers.multimodal.worker_handler import ( - SglangUtils, - StreamProcessor, -) +from dynamo.sglang.request_handlers.multimodal.worker_handler import SglangUtils pytestmark = [ pytest.mark.unit, @@ -210,7 +204,6 @@ def test_user_stop_token_ids_ignores_hidden_ids(): "stop_conditions": { "stop_token_ids": [576], "stop_token_ids_hidden": [128001], - "stop_token_ids_visible": [128009], } } ) == {576} @@ -229,33 +222,6 @@ def test_user_stop_token_ids_treats_token_id_display_as_string_stop(): assert _user_stop_token_ids({"stop": ["token_id:576"]}) == set() -def test_suppressed_stop_token_ids_includes_plain_and_hidden_stop_token_ids(): - assert _suppressed_stop_token_ids( - { - "stop_conditions": { - "stop_token_ids": [576], - "stop_token_ids_hidden": [128001], - "stop_token_ids_visible": [128009], - } - } - ) == {576, 128001} - - -def test_remove_suppressed_stop_tokens_removes_matched_tail_sequence(): - assert _remove_suppressed_stop_tokens( - [101, 128001, 128009], [128001, 128009], {128001, 128009} - ) == ([101], 2) - - -def test_remove_suppressed_stop_tokens_keeps_visible_or_non_tail_matches(): - assert _remove_suppressed_stop_tokens( - [101, 128001, 128009], [128001, 128009], {128001} - ) == ([101, 128001, 128009], 0) - assert _remove_suppressed_stop_tokens( - [128001, 128009, 101], [128001, 128009], {128001, 128009} - ) == ([128001, 128009, 101], 0) - - def test_openai_stop_sampling_params_preserves_string_stops(): assert _openai_stop_sampling_params({"stop": "END"}) == {"stop": "END"} assert _openai_stop_sampling_params({"stop": ["END"]}) == {"stop": ["END"]} @@ -716,24 +682,6 @@ def test_build_sampling_params_passes_n_for_token_requests(): assert sampling_params["max_new_tokens"] == 8 -def test_build_sampling_params_omits_string_stop_for_token_requests(): - handler = _new_decode_handler(use_sglang_tokenizer=False) - - sampling_params = handler._build_sampling_params( - { - "sampling_options": {}, - "stop_conditions": { - "max_tokens": 8, - "stop": ["<|user|>"], - "stop_token_ids_hidden": [128001], - }, - } - ) - - assert "stop" not in sampling_params - assert sampling_params["stop_token_ids"] == [128001] - - def test_build_sampling_params_forwards_repetition_controls_for_token_requests(): handler = _new_decode_handler(use_sglang_tokenizer=False) @@ -1035,7 +983,6 @@ def test_extract_logprobs_formats_top_tokens_as_token_ids(): "output_token_logprobs": [(-0.1, 101, "a")], "output_top_logprobs": [[(-0.1, 101, "a"), (-0.2, 102, "b")]], }, - 1, return_tokens_as_token_ids=True, ) @@ -1193,11 +1140,7 @@ async def test_process_token_stream_treats_completion_usage_as_optional(): ) assert chunks == [ - { - "index": 0, - "finish_reason": "stop", - "token_ids": [], - }, + {"index": 0, "finish_reason": "stop", "token_ids": []}, { "index": 1, "finish_reason": "stop", @@ -1211,48 +1154,6 @@ async def test_process_token_stream_treats_completion_usage_as_optional(): ] -@pytest.mark.asyncio -async def test_process_token_stream_reports_request_total_completion_usage(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_token_stream( - _stream( - [ - { - "index": 0, - "output_ids": [], - "meta_info": { - "id": "request-1", - "finish_reason": {"type": "stop"}, - "prompt_tokens": 2, - "completion_tokens": 1, - }, - }, - { - "index": 1, - "output_ids": [], - "meta_info": { - "id": "request-1", - "finish_reason": {"type": "stop"}, - "prompt_tokens": 2, - "completion_tokens": 2, - }, - }, - ] - ), - _Context(), - ) - ) - - assert chunks[0]["completion_usage"]["completion_tokens"] == 1 - assert chunks[1]["completion_usage"] == { - "prompt_tokens": 2, - "completion_tokens": 3, - "total_tokens": 5, - } - - @pytest.mark.asyncio async def test_process_token_stream_accepts_incremental_logprob_arrays(): handler = _new_decode_handler() @@ -1269,6 +1170,15 @@ async def test_process_token_stream_accepts_incremental_logprob_arrays(): "id": "request-1", "finish_reason": None, "output_token_logprobs": [(-0.1, 101, "")], + "output_top_logprobs": [ + [ + (-0.1, 101, "a"), + (-0.2, 201, "b"), + (-0.3, 202, "c"), + (-0.4, 203, "d"), + (-0.5, 204, "e"), + ] + ], }, "engine_data": {"native_chunk": 1}, }, @@ -1280,6 +1190,15 @@ async def test_process_token_stream_accepts_incremental_logprob_arrays(): "id": "request-1", "finish_reason": None, "output_token_logprobs": [(-0.3, 102, "")], + "output_top_logprobs": [ + [ + (-0.3, 102, "c"), + (-0.4, 301, "d"), + (-0.5, 302, "e"), + (-0.6, 303, "f"), + (-0.7, 304, "g"), + ] + ], }, "engine_data": {"native_chunk": 2}, }, @@ -1291,6 +1210,9 @@ async def test_process_token_stream_accepts_incremental_logprob_arrays(): assert [chunk["token_ids"] for chunk in chunks] == [[101], [102]] assert [chunk["log_probs"] for chunk in chunks] == [[-0.1], [-0.3]] + assert [ + [len(position) for position in chunk["top_logprobs"]] for chunk in chunks + ] == [[5], [5]] assert all("text" not in chunk for chunk in chunks) assert all("tokens" not in chunk for chunk in chunks) assert [chunk["engine_data"]["native_chunk"] for chunk in chunks] == [1, 2] @@ -1463,7 +1385,7 @@ async def upload_choice(self, choice_index, metadata): @pytest.mark.asyncio -async def test_process_text_stream_tracks_delta_per_choice_index(): +async def test_process_text_stream_forwards_incremental_text_per_choice(): handler = _new_decode_handler() chunks = await _collect( @@ -1482,12 +1404,12 @@ async def test_process_text_stream_tracks_delta_per_choice_index(): }, { "index": 0, - "text": "Hello", + "text": "llo", "meta_info": {"id": "request-1", "finish_reason": None}, }, { "index": 1, - "text": "Good", + "text": "od", "meta_info": {"id": "request-1", "finish_reason": None}, }, ] @@ -1506,6 +1428,50 @@ async def test_process_text_stream_tracks_delta_per_choice_index(): ] +@pytest.mark.asyncio +async def test_process_text_stream_forwards_sglang_trimmed_stop_response(): + """SGLang trims display text but retains matched token IDs upstream.""" + handler = _new_decode_handler() + + chunks = await _collect( + handler._process_text_stream( + _stream( + [ + { + "index": 0, + "text": " delta", + "output_ids": [9477], + "meta_info": {"id": "request-1", "finish_reason": None}, + }, + { + "index": 0, + # Vanilla SGLang's default no_stop_trim=False response: + # the text is already trimmed while output_ids still + # include the three-token matched stop string. + "text": " epsilon", + "output_ids": [31204, 1147, 1915, 38997, 18526], + "meta_info": { + "id": "request-1", + "finish_reason": { + "type": "stop", + "matched": " zeta eta theta", + }, + }, + }, + ] + ), + _Context(), + request={"stop": [" zeta eta theta"]}, + ) + ) + + assert [chunk["choices"][0]["delta"]["content"] for chunk in chunks] == [ + " delta", + " epsilon", + ] + assert chunks[-1]["choices"][0]["finish_reason"] == "stop" + + @pytest.mark.asyncio async def test_process_text_stream_notifies_for_empty_decoded_token(): handler = _new_decode_handler() @@ -1564,82 +1530,6 @@ async def test_process_text_stream_stop_reason_uses_response_nvext(): assert chunks[0]["nvext"]["stop_reason"] == "END" -@pytest.mark.asyncio -async def test_process_text_stream_buffers_split_stop_string_suffix(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_text_stream( - _stream( - [ - { - "index": 0, - "text": "Hello<|us", - "meta_info": {"id": "request-1", "finish_reason": None}, - }, - { - "index": 0, - "text": "Hello<|user|>", - "meta_info": { - "id": "request-1", - "finish_reason": { - "type": "stop", - "matched": "<|user|>", - }, - }, - }, - ] - ), - _Context(), - request={"stop": ["<|user|>"]}, - ) - ) - - assert [chunk["choices"][0]["delta"]["content"] for chunk in chunks] == [ - "Hello", - "", - ] - assert chunks[-1]["choices"][0]["finish_reason"] == "stop" - - -@pytest.mark.asyncio -async def test_process_text_stream_removes_matched_stop_string_suffix(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_text_stream( - _stream( - [ - { - "index": 0, - "text": "Hello", - "meta_info": {"id": "request-1", "finish_reason": None}, - }, - { - "index": 0, - "text": "Hello<|user|>", - "meta_info": { - "id": "request-1", - "finish_reason": { - "type": "stop", - "matched": "<|user|>", - }, - }, - }, - ] - ), - _Context(), - request={"stop": ["<|user|>"]}, - ) - ) - - assert [chunk["choices"][0]["delta"]["content"] for chunk in chunks] == [ - "Hello", - "", - ] - assert chunks[-1]["choices"][0]["finish_reason"] == "stop" - - @pytest.mark.asyncio async def test_process_text_stream_stop_reason_requires_nvext_extra_field(): handler = _new_decode_handler() @@ -1728,7 +1618,7 @@ async def test_process_text_stream_uploads_routed_experts(tmp_path): @pytest.mark.asyncio -async def test_process_token_stream_suppresses_hidden_stop_token_reason(): +async def test_process_token_stream_preserves_hidden_stop_token_but_not_reason(): handler = _new_decode_handler() chunks = await _collect( @@ -1753,344 +1643,9 @@ async def test_process_token_stream_suppresses_hidden_stop_token_reason(): ) ) + assert chunks[0]["token_ids"] == [128001] assert "stop_reason" not in chunks[0] -@pytest.mark.asyncio -async def test_process_token_stream_removes_matched_hidden_stop_token(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_token_stream( - _stream( - [ - { - "index": 0, - "output_ids": [101, 128001], - "meta_info": { - "id": "request-1", - "finish_reason": {"type": "stop", "matched": 128001}, - }, - } - ] - ), - _Context(), - user_stop_token_ids={576}, - suppressed_stop_token_ids={128001}, - ) - ) - - assert chunks == [ - { - "index": 0, - "finish_reason": "stop", - "token_ids": [101], - } - ] - - -@pytest.mark.asyncio -async def test_process_token_stream_removes_plain_stop_token_and_keeps_reason(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_token_stream( - _stream( - [ - { - "index": 0, - "output_ids": [101, 576], - "meta_info": { - "id": "request-1", - "finish_reason": {"type": "stop", "matched": 576}, - }, - } - ] - ), - _Context(), - user_stop_token_ids={576}, - suppressed_stop_token_ids={576}, - ) - ) - - assert chunks == [ - { - "index": 0, - "finish_reason": "stop", - "stop_reason": 576, - "token_ids": [101], - } - ] - - -@pytest.mark.asyncio -async def test_process_token_stream_removes_matched_hidden_stop_token_sequence(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_token_stream( - _stream( - [ - { - "index": 0, - "output_ids": [101, 128001, 128009], - "meta_info": { - "id": "request-1", - "finish_reason": { - "type": "stop", - "matched": [128001, 128009], - }, - }, - } - ] - ), - _Context(), - user_stop_token_ids={576}, - suppressed_stop_token_ids={128001, 128009}, - ) - ) - - assert chunks == [ - { - "index": 0, - "finish_reason": "stop", - "token_ids": [101], - } - ] - - -@pytest.mark.asyncio -async def test_process_token_stream_buffers_split_hidden_stop_token_sequence(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_token_stream( - _stream( - [ - { - "index": 0, - "output_ids": [101, 128001], - "meta_info": {"id": "request-1", "finish_reason": None}, - }, - { - "index": 0, - "output_ids": [128009], - "meta_info": { - "id": "request-1", - "finish_reason": { - "type": "stop", - "matched": [128001, 128009], - }, - }, - }, - ] - ), - _Context(), - user_stop_token_ids={576}, - suppressed_stop_token_ids={128001, 128009}, - ) - ) - - assert chunks == [ - {"index": 0, "token_ids": [101]}, - { - "index": 0, - "finish_reason": "stop", - "token_ids": [], - }, - ] - - -@pytest.mark.asyncio -async def test_process_token_stream_trims_logprobs_for_suppressed_stop_token(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_token_stream( - _stream( - [ - { - "index": 0, - "output_ids": [101, 128001], - "meta_info": { - "id": "request-1", - "finish_reason": {"type": "stop", "matched": 128001}, - "output_token_logprobs": [ - (-0.1, 101, "a"), - (-0.2, 128001, "<|user|>"), - ], - "output_top_logprobs": [ - [(-0.1, 101, "a")], - [(-0.2, 128001, "<|user|>")], - ], - }, - } - ] - ), - _Context(), - user_stop_token_ids={576}, - suppressed_stop_token_ids={128001}, - ) - ) - - assert chunks == [ - { - "index": 0, - "finish_reason": "stop", - "token_ids": [101], - "log_probs": [-0.1], - "top_logprobs": [ - [{"rank": 1, "token_id": 101, "token": "a", "logprob": -0.1}] - ], - } - ] - - -@pytest.mark.asyncio -async def test_process_token_stream_keeps_final_stop_when_hidden_token_removed_to_empty(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_token_stream( - _stream( - [ - { - "index": 0, - "output_ids": [128001], - "meta_info": { - "id": "request-1", - "finish_reason": {"type": "stop", "matched": 128001}, - }, - } - ] - ), - _Context(), - user_stop_token_ids={576}, - suppressed_stop_token_ids={128001}, - ) - ) - - assert chunks == [ - { - "index": 0, - "finish_reason": "stop", - "token_ids": [], - } - ] - - -@pytest.mark.asyncio -async def test_process_token_stream_keeps_visible_stop_token(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_token_stream( - _stream( - [ - { - "index": 0, - "output_ids": [101, 128001], - "meta_info": { - "id": "request-1", - "finish_reason": {"type": "stop", "matched": 128001}, - }, - } - ] - ), - _Context(), - suppressed_stop_token_ids=set(), - ) - ) - - assert chunks == [ - { - "index": 0, - "finish_reason": "stop", - "stop_reason": 128001, - "token_ids": [101, 128001], - } - ] - - -@pytest.mark.asyncio -async def test_process_token_stream_keeps_hidden_stop_token_when_match_is_not_tail(): - handler = _new_decode_handler() - - chunks = await _collect( - handler._process_token_stream( - _stream( - [ - { - "index": 0, - "output_ids": [128001, 101], - "meta_info": { - "id": "request-1", - "finish_reason": {"type": "stop", "matched": 128001}, - }, - } - ] - ), - _Context(), - user_stop_token_ids={576}, - suppressed_stop_token_ids={128001}, - ) - ) - - assert chunks == [ - { - "index": 0, - "finish_reason": "stop", - "token_ids": [128001, 101], - } - ] - - -@pytest.mark.asyncio -async def test_multimodal_stream_keeps_reading_after_one_choice_finishes(): - chunks = await _collect( - StreamProcessor.process_sglang_stream( - _stream( - [ - { - "index": 0, - "output_ids": [101], - "text": "a", - "meta_info": {"finish_reason": None}, - }, - { - "index": 1, - "output_ids": [201], - "text": "b", - "meta_info": {"finish_reason": None}, - }, - { - "index": 0, - "output_ids": [], - "text": "a", - "meta_info": {"finish_reason": {"type": "stop"}}, - }, - { - "index": 1, - "output_ids": [], - "text": "b", - "meta_info": {"finish_reason": {"type": "stop"}}, - }, - ] - ) - ) - ) - - outputs = [json.loads(chunk) for chunk in chunks] - - assert [output["index"] for output in outputs] == [0, 1, 0, 1] - assert [output["finished"] for output in outputs] == [False, False, True, True] - assert [output.get("finish_reason") for output in outputs] == [ - None, - None, - "stop", - "stop", - ] - - async def _collect(stream): return [item async for item in stream] diff --git a/components/src/dynamo/sglang/tests/test_sglang_elastic_ep.py b/components/src/dynamo/sglang/tests/test_sglang_elastic_ep.py index dd107fa39f48..933fbd6cef50 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_elastic_ep.py +++ b/components/src/dynamo/sglang/tests/test_sglang_elastic_ep.py @@ -66,6 +66,10 @@ def _make_handler( tokenizer_manager = SimpleNamespace( server_args=SimpleNamespace(elastic_ep_backend=backend), + pause_generation=lambda: None, + continue_generation=lambda: None, + release_memory_occupation=lambda: None, + resume_memory_occupation=lambda: None, get_elastic_ep_state=MagicMock( return_value=state if state is not None diff --git a/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py b/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py index 1174ee5f9f43..b6cb0e5103f0 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py +++ b/components/src/dynamo/sglang/tests/test_sglang_memory_occupation_handlers.py @@ -2,17 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio -import sys -import types from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest -from dynamo.sglang.request_handlers.handler_base import ( - BaseWorkerHandler, - SGLangEnginePauseController, -) +from dynamo.sglang.request_handlers.handler_base import BaseWorkerHandler pytestmark = [ pytest.mark.unit, @@ -22,24 +17,6 @@ ] -@pytest.fixture(autouse=True) -def _stub_sglang_io_struct(monkeypatch): - """Keep unit tests independent from CUDA-only sglang imports.""" - - io_struct = types.ModuleType("sglang.srt.managers.io_struct") - - class _Req: - def __init__(self, tags=None): - self.tags = tags - - io_struct.PauseGenerationReqInput = _Req - io_struct.ReleaseMemoryOccupationReqInput = _Req - io_struct.ResumeMemoryOccupationReqInput = _Req - io_struct.ContinueGenerationReqInput = _Req - - monkeypatch.setitem(sys.modules, "sglang.srt.managers.io_struct", io_struct) - - class _TestWorkerHandler(BaseWorkerHandler): async def generate(self, request, context): yield {} @@ -50,6 +27,7 @@ def handler(): handler = _TestWorkerHandler.__new__(_TestWorkerHandler) handler.engine = SimpleNamespace( tokenizer_manager=SimpleNamespace( + is_pause=False, pause_generation=AsyncMock(), release_memory_occupation=AsyncMock(), resume_memory_occupation=AsyncMock(), @@ -65,165 +43,103 @@ def handler(): unregister_endpoint_instance=AsyncMock(), register_endpoint_instance=AsyncMock(), ) - handler._pause_controller = SGLangEnginePauseController(handler.engine) - handler._pause_lock = asyncio.Lock() + handler._engine_route_lock = asyncio.Lock() return handler -@pytest.mark.asyncio -async def test_release_and_resume_are_idempotent(handler): - first_release = await handler.release_memory_occupation({}) - second_release = await handler.release_memory_occupation({}) - - first_resume = await handler.resume_memory_occupation({}) - second_resume = await handler.resume_memory_occupation({}) - - assert first_release["status"] == "ok" - assert second_release["status"] == "ok" - assert first_resume["status"] == "ok" - assert second_resume["status"] == "ok" - assert second_release["message"] == "Memory already released" - assert second_resume["message"] == "Memory already resumed" - - release_req = ( - handler.engine.tokenizer_manager.release_memory_occupation.await_args.args[0] - ) - resume_req = ( - handler.engine.tokenizer_manager.resume_memory_occupation.await_args.args[0] - ) - assert release_req.tags is None - assert resume_req.tags is None +def _registered_engine_routes(handler, configured_routes=None): + registered = {} - handler.engine.tokenizer_manager.pause_generation.assert_awaited_once() - handler.engine.tokenizer_manager.release_memory_occupation.assert_awaited_once() - handler.generate_endpoint.unregister_endpoint_instance.assert_awaited_once() + class Runtime: + def register_engine_route(self, path, route_handler): + registered[path] = route_handler - handler.engine.tokenizer_manager.resume_memory_occupation.assert_awaited_once() - handler.engine.tokenizer_manager.continue_generation.assert_awaited_once() - handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() + handler.config = SimpleNamespace( + dynamo_args=SimpleNamespace(engine_routes=configured_routes or []) + ) + handler.register_engine_routes(Runtime()) + return registered -@pytest.mark.asyncio -async def test_memory_occupation_handlers_forward_tags_exactly(handler): - await handler.release_memory_occupation({"tags": []}) - resume_result = await handler.resume_memory_occupation({"tags": []}) +def _make_native_manager_methods_routable(manager): + mocks = {} + for method_name in ( + "pause_generation", + "continue_generation", + "release_memory_occupation", + "resume_memory_occupation", + ): + method_mock = getattr(manager, method_name) + mocks[method_name] = method_mock - assert resume_result["status"] == "ok" - release_req = ( - handler.engine.tokenizer_manager.release_memory_occupation.await_args.args[0] - ) - resume_req = ( - handler.engine.tokenizer_manager.resume_memory_occupation.await_args.args[0] - ) - assert release_req.tags == [] - assert resume_req.tags == [] - - handler.engine.tokenizer_manager.pause_generation.reset_mock() - handler.engine.tokenizer_manager.release_memory_occupation.reset_mock() - handler.engine.tokenizer_manager.resume_memory_occupation.reset_mock() - handler.engine.tokenizer_manager.continue_generation.reset_mock() - handler.generate_endpoint.unregister_endpoint_instance.reset_mock() - handler.generate_endpoint.register_endpoint_instance.reset_mock() - - await handler.release_memory_occupation({"tags": ["weights"]}) - resume_result = await handler.resume_memory_occupation({}) - - assert resume_result["status"] == "ok" - release_req = ( - handler.engine.tokenizer_manager.release_memory_occupation.await_args.args[0] - ) - resume_req = ( - handler.engine.tokenizer_manager.resume_memory_occupation.await_args.args[0] - ) - assert release_req.tags == ["weights"] - assert resume_req.tags is None - handler.engine.tokenizer_manager.continue_generation.assert_awaited_once() - handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() + async def route_method(_method_mock=method_mock, **kwargs): + return await _method_mock(**kwargs) + + setattr(manager, method_name, route_method) + return mocks @pytest.mark.asyncio -async def test_resume_recovers_generation_pause_after_failed_release_rollback(handler): +async def test_native_memory_routes_follow_sglang_pause_state(handler): manager = handler.engine.tokenizer_manager - manager.release_memory_occupation = AsyncMock( - side_effect=RuntimeError("release failed") - ) - failed_continue = AsyncMock(side_effect=RuntimeError("continue failed")) - manager.continue_generation = failed_continue + method_mocks = _make_native_manager_methods_routable(manager) - release_result = await handler.release_memory_occupation({}) + async def pause_generation(**_kwargs): + manager.is_pause = True - assert release_result["status"] == "error" - assert handler._pause_controller.is_paused is False - assert handler._pause_controller.needs_resume_recovery is True - failed_continue.assert_awaited_once() - handler.generate_endpoint.unregister_endpoint_instance.assert_awaited_once() + async def continue_generation(**_kwargs): + manager.is_pause = False - manager.continue_generation = AsyncMock() - resume_result = await handler.resume_memory_occupation({}) + method_mocks["pause_generation"].side_effect = pause_generation + method_mocks["continue_generation"].side_effect = continue_generation + routes = _registered_engine_routes(handler) - assert resume_result["status"] == "ok" - manager.resume_memory_occupation.assert_not_awaited() - manager.continue_generation.assert_awaited_once() - handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() - assert handler._pause_controller.needs_resume_recovery is False + assert { + "pause_generation", + "continue_generation", + "release_memory_occupation", + "resume_memory_occupation", + }.issubset(routes) + await routes["pause_generation"]({}) + await routes["release_memory_occupation"]({"tags": ["weights", "kv_cache"]}) + await routes["resume_memory_occupation"]({"tags": ["weights"]}) -@pytest.mark.asyncio -async def test_release_reregisters_after_clean_pause_rollback(handler): - manager = handler.engine.tokenizer_manager - manager.release_memory_occupation = AsyncMock( - side_effect=RuntimeError("release failed") - ) + assert manager.is_pause is True + assert handler.generate_endpoint.unregister_endpoint_instance.await_count == 3 + handler.generate_endpoint.register_endpoint_instance.assert_not_awaited() - release_result = await handler.release_memory_occupation({}) + await routes["resume_memory_occupation"]({"tags": ["kv_cache"]}) - # Rollback continued generation cleanly, so the engine is serving-safe again. - assert release_result["status"] == "error" - assert handler._pause_controller.is_paused is False - assert handler._pause_controller.needs_resume_recovery is False - manager.continue_generation.assert_awaited_once() - handler.generate_endpoint.unregister_endpoint_instance.assert_awaited_once() - # The endpoint must rejoin the routing pool since resume will early-return. - handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() + assert handler.generate_endpoint.unregister_endpoint_instance.await_count == 4 + handler.generate_endpoint.register_endpoint_instance.assert_not_awaited() + await routes["continue_generation"]({}) -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("method_name", "endpoint_method"), - [ - ("release_memory_occupation", "unregister_endpoint_instance"), - ("resume_memory_occupation", "register_endpoint_instance"), - ], -) -async def test_memory_control_returns_error_without_pause_controller( - handler, - method_name, - endpoint_method, -): - handler.engine = None - handler._pause_controller = None + assert manager.is_pause is False + handler.generate_endpoint.register_endpoint_instance.assert_awaited_once() - result = await getattr(handler, method_name)({}) - assert result == { - "status": "error", - "message": "memory control not supported on this worker", - } - getattr(handler.generate_endpoint, endpoint_method).assert_not_awaited() +@pytest.mark.asyncio +async def test_configured_engine_route_overrides_default_and_syncs_discovery(handler): + custom_pause_mock = AsyncMock(return_value={"custom": True}) + async def custom_pause(): + return await custom_pause_mock() -@pytest.mark.asyncio -async def test_resume_keeps_paused_state_when_register_fails(handler): - await handler.release_memory_occupation({}) - handler.generate_endpoint.register_endpoint_instance = AsyncMock( - side_effect=RuntimeError("discovery write timeout") + handler.engine.custom_pause = custom_pause + handler.engine.tokenizer_manager.is_pause = True + method_mocks = _make_native_manager_methods_routable( + handler.engine.tokenizer_manager ) + routes = _registered_engine_routes(handler, ["pause_generation=custom_pause"]) - result = await handler.resume_memory_occupation({}) + result = await routes["pause_generation"]({}) - assert result["status"] == "error" - assert handler._pause_controller is not None - assert handler._pause_controller.is_paused is True + assert result == {"custom": True} + custom_pause_mock.assert_awaited_once_with() + method_mocks["pause_generation"].assert_not_awaited() + handler.generate_endpoint.unregister_endpoint_instance.assert_awaited_once() @pytest.mark.asyncio diff --git a/components/src/dynamo/sglang/tests/test_sglang_unit.py b/components/src/dynamo/sglang/tests/test_sglang_unit.py index 97a607c2a4eb..520a304bdbc7 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_unit.py +++ b/components/src/dynamo/sglang/tests/test_sglang_unit.py @@ -271,7 +271,14 @@ def test_engine_generate_capability_registration_gate( def test_builtin_engine_routes_include_model_taint_update(monkeypatch): handler = object.__new__(DecodeWorkerHandler) - handler.engine = SimpleNamespace() + handler.engine = SimpleNamespace( + tokenizer_manager=SimpleNamespace( + pause_generation=lambda: None, + continue_generation=lambda: None, + release_memory_occupation=lambda: None, + resume_memory_occupation=lambda: None, + ) + ) handler.generate_endpoint = object() handler.config = SimpleNamespace(dynamo_args=SimpleNamespace(engine_routes=[])) @@ -296,6 +303,10 @@ def register_engine_route(self, path, route_handler): assert {path for path, _ in registered_routes} >= { "control/start_profile", "control/stop_profile", + "pause_generation", + "continue_generation", + "release_memory_occupation", + "resume_memory_occupation", } diff --git a/components/src/dynamo/thunderagent_router/README.md b/components/src/dynamo/thunderagent_router/README.md index 44ad1f85a1f3..33b12b9745ee 100644 --- a/components/src/dynamo/thunderagent_router/README.md +++ b/components/src/dynamo/thunderagent_router/README.md @@ -58,6 +58,12 @@ chat/completions surface; if the wrapped backend also advertises that surface for the same model, the frontend may route requests directly to the backend and bypass ThunderAgent lifecycle handling such as `x-dynamo-session-final`. +For clients that require SGLang's native `/generate` API, add +`--publish-sglang-generate` to ThunderAgent. This advertises the existing +SGLang-native frontend route through ThunderAgent while the wrapped worker +continues to use `--endpoint-types none`. This option requires `--model-name` +so ThunderAgent can register the native capability for that model. + The control-loop knobs (`--pause-threshold`, `--pause-target`, `--resume-hysteresis`, `--scheduler-interval-seconds`, …) and their defaults are documented in [docs/agents/thunderagent-router.md](../../../../docs/fern/pages/use-cases/agents/thunderagent-program-scheduler.md#utilization-driven-control-loop). diff --git a/components/src/dynamo/thunderagent_router/__main__.py b/components/src/dynamo/thunderagent_router/__main__.py index 4a36bdb94e5f..84cd1ef7413d 100644 --- a/components/src/dynamo/thunderagent_router/__main__.py +++ b/components/src/dynamo/thunderagent_router/__main__.py @@ -44,6 +44,15 @@ logger = logging.getLogger(__name__) +def _publish_sglang_generate_capability( + runtime_config: ModelRuntimeConfig, +) -> None: + """Advertise native SGLang generate without importing SGLang by default.""" + from dynamo.sglang.engine_generate import SGLANG_GENERATE_CAPABILITY + + runtime_config.set_engine_specific(SGLANG_GENERATE_CAPABILITY, "true") + + def _extract_program_id(request: dict[str, Any]) -> Optional[str]: ctx = request.get("agent_context") if not isinstance(ctx, dict): @@ -451,6 +460,9 @@ async def worker(runtime: DistributedRuntime) -> None: runtime_cfg.tool_call_parser = config.tool_call_parser if config.reasoning_parser: runtime_cfg.reasoning_parser = config.reasoning_parser + if config.publish_sglang_generate: + _publish_sglang_generate_capability(runtime_cfg) + logger.info("Published SGLang engine-native generate capability") await register_model( model_input=ModelInput.Tokens, model_type=ModelType.Chat | ModelType.Completions, diff --git a/components/src/dynamo/thunderagent_router/args.py b/components/src/dynamo/thunderagent_router/args.py index 672295fbc3b6..fcafb73e608c 100644 --- a/components/src/dynamo/thunderagent_router/args.py +++ b/components/src/dynamo/thunderagent_router/args.py @@ -9,7 +9,7 @@ from typing import Optional from dynamo.common.configuration.arg_group import ArgGroup -from dynamo.common.configuration.utils import add_argument +from dynamo.common.configuration.utils import add_argument, add_negatable_bool_argument from dynamo.router.args import ( DynamoRouterArgGroup, DynamoRouterConfig, @@ -36,6 +36,7 @@ class ThunderAgentRouterConfig(DynamoRouterConfig): model_path: Optional[str] = None tool_call_parser: Optional[str] = None reasoning_parser: Optional[str] = None + publish_sglang_generate: bool = False def to_thunderagent_config(self) -> ThunderAgentConfig: return ThunderAgentConfig( @@ -69,6 +70,8 @@ def validate(self) -> None: # type: ignore[override] raise ValueError("--scheduler-interval-seconds must be > 0") if self.resume_timeout_seconds <= 0: raise ValueError("--resume-timeout-seconds must be > 0") + if self.publish_sglang_generate and not self.model_name: + raise ValueError("--publish-sglang-generate requires --model-name") class ThunderAgentArgGroup(ArgGroup): @@ -218,6 +221,15 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: "is set.", arg_type=str, ) + add_negatable_bool_argument( + g, + flag_name="--publish-sglang-generate", + env_var="DYN_THUNDERAGENT_PUBLISH_SGLANG_GENERATE", + default=False, + help="Advertise SGLang's native /generate API through the " + "ThunderAgent router. Enable only when --endpoint targets a " + "Dynamo SGLang worker that publishes native generate support.", + ) def parse_args(argv: Optional[list[str]] = None) -> ThunderAgentRouterConfig: diff --git a/components/src/dynamo/thunderagent_router/tests/test_args.py b/components/src/dynamo/thunderagent_router/tests/test_args.py new file mode 100644 index 000000000000..c699dca0466a --- /dev/null +++ b/components/src/dynamo/thunderagent_router/tests/test_args.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for ThunderAgent router CLI configuration.""" + +from __future__ import annotations + +import sys +from types import ModuleType + +import pytest + +from dynamo.thunderagent_router.__main__ import _publish_sglang_generate_capability +from dynamo.thunderagent_router.args import parse_args + +pytestmark = [pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0] + + +def test_sglang_generate_capability_is_opt_in() -> None: + base = [ + "--endpoint", + "dynamo.sglang.generate", + "--model-name", + "Qwen/Qwen3-0.6B", + ] + assert parse_args(base).publish_sglang_generate is False + assert ( + parse_args([*base, "--publish-sglang-generate"]).publish_sglang_generate is True + ) + assert ( + parse_args([*base, "--no-publish-sglang-generate"]).publish_sglang_generate + is False + ) + + +def test_sglang_generate_capability_reads_environment(monkeypatch) -> None: + monkeypatch.setenv("DYN_THUNDERAGENT_PUBLISH_SGLANG_GENERATE", "true") + assert ( + parse_args( + [ + "--endpoint", + "dynamo.sglang.generate", + "--model-name", + "Qwen/Qwen3-0.6B", + ] + ).publish_sglang_generate + is True + ) + + +def test_sglang_generate_capability_requires_model_name() -> None: + with pytest.raises( + ValueError, match="--publish-sglang-generate requires --model-name" + ): + parse_args( + [ + "--endpoint", + "dynamo.sglang.generate", + "--publish-sglang-generate", + ] + ) + + +def test_publish_sglang_generate_capability_uses_backend_contract(monkeypatch) -> None: + module = ModuleType("dynamo.sglang.engine_generate") + module.SGLANG_GENERATE_CAPABILITY = "sglang_generate" + monkeypatch.setitem(sys.modules, module.__name__, module) + + published: list[tuple[str, str]] = [] + + class RuntimeConfig: + def set_engine_specific(self, key: str, value: str) -> None: + published.append((key, value)) + + runtime_config = RuntimeConfig() + + _publish_sglang_generate_capability(runtime_config) # type: ignore[arg-type] + assert published == [("sglang_generate", "true")] diff --git a/docs/fern/index.yml b/docs/fern/index.yml index c1f46e1c0b0a..178acd4f0b67 100644 --- a/docs/fern/index.yml +++ b/docs/fern/index.yml @@ -665,6 +665,8 @@ navigation: path: pages/use-cases/reinforcement-learning/verl.md - page: NeMo RL Integration path: pages/use-cases/reinforcement-learning/nemo-rl.md + - page: Slime Integration + path: pages/use-cases/reinforcement-learning/slime.md - page: External Trainer Integration path: pages/use-cases/reinforcement-learning/implementation-guide.md - page: RL Integration Reference diff --git a/docs/fern/pages/developer-guide/additional-resources/nvidia-request-extensions-nvext.md b/docs/fern/pages/developer-guide/additional-resources/nvidia-request-extensions-nvext.md index 41fc91c15c8e..28ef0272733b 100644 --- a/docs/fern/pages/developer-guide/additional-resources/nvidia-request-extensions-nvext.md +++ b/docs/fern/pages/developer-guide/additional-resources/nvidia-request-extensions-nvext.md @@ -37,7 +37,7 @@ Include `nvext` as a top-level field alongside standard OpenAI-compatible fields | `token_data` | `u32[]` | `None` | Preprocessor | Pre-tokenized prompt tokens. When present, the frontend skips tokenization. | | `max_thinking_tokens` | `u32` | `None` | Backend | Maximum thinking tokens allowed (passed through to backends). | | `cache_salt` | `string` | `None` | Router / supported backends | Namespaces Dynamo KV routing. vLLM and TensorRT-LLM also isolate backend KV-cache reuse; see [Backend support](#backend-support). This is the recommended cache-isolation input. | -| `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`, `"engine_data"`, `"stop_reason"`, `"completion_token_ids"`, `"prompt_logprobs"`. | +| `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`, `"engine_data"`, `"stop_reason"`, `"detailed_finish_reason"`, `"prompt_token_ids"`, `"completion_token_ids"`, `"prompt_logprobs"`. | | `metadata_upload` | object | `None` | SGLang backend | Uploads final cumulative SGLang `meta_info` out of band. The object accepts one required `url` field. Requires an RL-enabled SGLang worker. | | `prefill_worker_id` | `u64` | `None` | Router | Routes the request to a specific prefill worker (disaggregated serving). | | `decode_worker_id` | `u64` | `None` | Router | Routes the request to a specific decode worker (disaggregated serving). | @@ -248,10 +248,56 @@ When the client requests response metadata via `extra_fields`, the response incl | `routed_experts` | `extra_fields: ["routed_experts"]` | Backend-specific routed expert capture payload returned by compatible vLLM and SGLang engines. | | `engine_data` | `extra_fields: ["engine_data"]` | Opaque backend-provided engine metadata. | | `stop_reason` | `extra_fields: ["stop_reason"]` | Backend-specific matched stop condition, returned under `nvext` because it is not part of the OpenAI completions schema. Dynamo currently serves this as a response-level field for single-choice requests; supporting `n > 1` will require an indexed per-choice shape. | +| `detailed_finish_reason` | `extra_fields: ["detailed_finish_reason"]` | Dynamo's internal finish reason before OpenAI conversion. Backend-specific values are normalized first. See [Detailed Finish Reason](#detailed-finish-reason). | +| `prompt_token_ids` | `extra_fields: ["prompt_token_ids"]` | Effective single-prompt token sequence used after preprocessing, including pre-tokenized input supplied through the request. Emitted on the final response. | | `completion_token_ids` | `extra_fields: ["completion_token_ids"]` | Generated token IDs. Requires a single prompt and one generated choice. | | `prompt_logprobs` | `extra_fields: ["prompt_logprobs"]` | Prompt log probabilities requested with the top-level `prompt_logprobs` field. Emitted on the final response. | | `token_ids` | Automatic (GAIE Stage 1) | Tokenized prompt for reuse in Stage 2 query-only mode. | +### Detailed Finish Reason + +The OpenAI-compatible API uses standard values in `finish_reason`. +Dynamo maps an external cancellation to `stop` to keep this API compatible. + +The `nvext.detailed_finish_reason` field contains Dynamo's internal finish reason. +Dynamo normalizes backend-specific values before it creates this field. +For example, SGLang's `abort` becomes `cancelled`. +The field supports these values: + +- `eos` +- `length` +- `stop` +- `cancelled` +- `content_filter` + +Add `detailed_finish_reason` to `nvext.extra_fields`: + +```json +{ + "nvext": { + "extra_fields": ["detailed_finish_reason"] + } +} +``` + +If SGLang aborts the request, Dynamo returns this response data: + +```json +{ + "choices": [ + { + "finish_reason": "stop" + } + ], + "nvext": { + "detailed_finish_reason": "cancelled" + } +} +``` + +This response field is available for `/v1/chat/completions` and `/v1/completions`. +Dynamo supports this field for requests with one choice. + ### Example response `nvext` ```json diff --git a/docs/fern/pages/use-cases/reinforcement-learning/implementation-guide.md b/docs/fern/pages/use-cases/reinforcement-learning/implementation-guide.md index 511096f602b1..c282cd936aca 100644 --- a/docs/fern/pages/use-cases/reinforcement-learning/implementation-guide.md +++ b/docs/fern/pages/use-cases/reinforcement-learning/implementation-guide.md @@ -37,7 +37,7 @@ curl -N http://localhost:8000/generate \ Set `DYN_SGLANG_ENABLE_GENERATE=1` on the frontend. The worker must accept token input; do not start it with `--use-sglang-tokenizer`, which selects text input and prevents the worker from advertising `/generate`. -Use an OpenAI-compatible route when the adapter needs one envelope across backends or named NVIDIA request extensions. This example bypasses frontend tokenization and asks Dynamo to return the generated token IDs and prompt log probabilities: +Use an OpenAI-compatible route when the adapter needs one envelope across backends or named NVIDIA request extensions. This example bypasses frontend tokenization and asks Dynamo to return the effective prompt token sequence, generated token IDs, and prompt log probabilities: ```bash curl http://localhost:8000/v1/completions \ @@ -51,11 +51,13 @@ curl http://localhost:8000/v1/completions \ "prompt_logprobs": 5, "nvext": { "token_data": [151644, 8948, 198], - "extra_fields": ["completion_token_ids", "prompt_logprobs"] + "extra_fields": ["prompt_token_ids", "completion_token_ids", "prompt_logprobs"] } }' ``` +`nvext.token_data` supplies pre-tokenized request input. `nvext.extra_fields` selects response metadata; `prompt_token_ids` returns the effective single-prompt token sequence used after preprocessing, including pre-tokenized input. + Treat the engine's token sequence as authoritative. Do not reconstruct training tokens from generated text. Before accepting a sample, check token and log-probability alignment, the terminal state, the selected model and tokenizer, and the framework's mask and retry rules. See [RL Integration Reference](integration-reference.md#know-what-returns-to-the-trainer) for the data Dynamo can return and what remains framework-owned. ## Discover and Administer Workers diff --git a/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md b/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md index 0448ede544b8..495b8bf74212 100644 --- a/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md +++ b/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md @@ -68,6 +68,7 @@ Before admitting a sample to training, verify: | Data | Dynamo surface | Contract boundary | |---|---|---| +| Prompt token IDs | Named `nvext.prompt_token_ids` | Returns the effective single-prompt token sequence used after preprocessing, including token arrays or `nvext.token_data` supplied through the request. | | Generated token IDs | Named `nvext.completion_token_ids`, or the native SGLang stream | Use the engine-returned sequence; exact placement depends on the selected interface. | | Selected and prompt log probabilities | Standard completion log probabilities, named `nvext.prompt_logprobs`, or native SGLang metadata | Check support and alignment on the exact backend and response path. | | Routed experts and raw engine data | Opt-in `nvext.routed_experts` or `nvext.engine_data` | Backend-specific. Prefer named fields over the raw engine payload. | @@ -206,7 +207,7 @@ Even two vLLM deployments can expose different route families because the Python |---|---|---|---| | [verl](verl.md) | Experimental | Shared Dynamo frontend with colocated vLLM rollout workers | The public recipe owns CUDA IPC updates through Ray/ZMQ control. Choose native Dynamo routing or ThunderAgent as distinct variants. | | [NeMo RL](nemo-rl.md) | Experimental | NeMo RL-managed Dynamo/vLLM fleet on Slurm and Ray | Fixed non-colocated fleet with framework-owned NCCL refit; not an external Dynamo deployment. | -| [SLIME](https://github.com/THUDM/slime) | Integration in progress | Proposed shared SGLang `/generate` path with direct engine control | No merged, maintained Dynamo recipe yet; discovery and update ownership remain unsettled. | +| [Slime](slime.md) | Experimental fixed worker set | Dynamo sidecars manage stock SGLang engines for [streaming external rollouts](https://github.com/THUDM/slime/pull/2272) | The [fixed-worker-set example](https://github.com/ai-dynamo/dynamo/blob/main/examples/README.md#integration-examples) sends generation through the Dynamo frontend. Slime uses static engine Services for control and weight updates. The example does not support elastic discovery or external-engine recovery. | | [Prime-RL](https://github.com/PrimeIntellect-ai/prime-rl) | Router available; full integration in progress | Prime-RL documents the Dynamo router as a drop-in option; a proposed Dynamo/vLLM sidecar adds worker discovery and external updates | Routing can be evaluated today, but the full adapter remains in upstream development and is not a released compatibility surface. | | [OpenRLHF](https://github.com/OpenRLHF/OpenRLHF), [Miles](https://github.com/fleet-ai/miles-fleet), [SkyRL](https://github.com/NovaSky-AI/SkyRL), and [Polar](https://github.com/NVIDIA-NeMo/ProRL-Agent-Server) | No Dynamo guide | No maintained Dynamo adapter was found in the reviewed public sources | Add a guide only after a maintained integration completes the same generation, update, failure, and ownership checks. | diff --git a/docs/fern/pages/use-cases/reinforcement-learning/overview.md b/docs/fern/pages/use-cases/reinforcement-learning/overview.md index c715987ec967..e2472e2ae996 100644 --- a/docs/fern/pages/use-cases/reinforcement-learning/overview.md +++ b/docs/fern/pages/use-cases/reinforcement-learning/overview.md @@ -13,7 +13,9 @@ Use Dynamo when rollout serving has become a distributed-systems problem: many w In production, [Cognition used Dynamo while training SWE-1.7](https://cognition.com/blog/swe-1-7) to manage inference-engine lifecycles and route inference across a multi-cluster RL system. When a replica failed, Dynamo rerouted inference to another worker and rescheduled the replica so the rollout pipeline could remain available while the latest policy state was restored. -Dynamo also connects to a growing RL framework ecosystem. [verl](https://github.com/verl-project/verl-recipe/tree/main/dynamo) publishes a Dynamo rollout recipe, [NeMo RL](https://github.com/NVIDIA-NeMo/RL/tree/main/nemo_rl/models/generation/dynamo) includes a managed Dynamo generation backend, and [Prime-RL](https://www.primeintellect.ai/blog/rl-at-1t-scale) supports the Dynamo router as a drop-in routing option. Integration work with [SLIME](integration-reference.md#framework-compatibility) extends the same serving capabilities to SGLang-based rollout stacks. Together, these paths demonstrate that Dynamo can support different training frameworks and deployment models while leaving RL semantics with the trainer. +Dynamo connects to these RL frameworks. [verl](https://github.com/verl-project/verl-recipe/tree/main/dynamo) publishes a Dynamo rollout recipe. [NeMo RL](https://github.com/NVIDIA-NeMo/RL/tree/main/nemo_rl/models/generation/dynamo) includes a managed Dynamo generation backend. [Prime-RL](https://www.primeintellect.ai/blog/rl-at-1t-scale) supports the Dynamo router as a drop-in routing option. The experimental [Slime integration](slime.md) runs stock SGLang engines with Dynamo sidecars. + +These integrations support different training frameworks and deployment models. The trainer retains ownership of RL semantics. ## Keep Rollouts Running Through Failures @@ -46,7 +48,7 @@ Start with a specific bottleneck and a baseline metric. Measure cache reuse, wor | Custom or external trainer | Experimental contract | [External Trainer Integration](implementation-guide.md) for generation, discovery, administration, and policy updates. | | verl | Experimental | [verl Integration](verl.md) for the public colocated Dynamo/vLLM recipe. | | NeMo RL | Experimental | [NeMo RL Integration](nemo-rl.md) for the managed Slurm/Ray Dynamo backend. | -| SLIME | Integration in progress | Review the current boundary in [Framework Compatibility](integration-reference.md#framework-compatibility). | +| Slime | Experimental fixed worker set | Use the [Slime Integration](slime.md) with the streaming external-rollout support from Slime. | | Prime-RL | Routing available; integration in progress | See [Prime-RL's routing overview](https://www.primeintellect.ai/blog/rl-at-1t-scale) and the current boundary in [Framework Compatibility](integration-reference.md#framework-compatibility). | Experimental guides have runnable upstream artifacts but do not make a general compatibility promise. Integrations in progress remain in the compatibility table until a maintained path lands. diff --git a/docs/fern/pages/use-cases/reinforcement-learning/slime.md b/docs/fern/pages/use-cases/reinforcement-learning/slime.md new file mode 100644 index 000000000000..342f8a4905ff --- /dev/null +++ b/docs/fern/pages/use-cases/reinforcement-learning/slime.md @@ -0,0 +1,150 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: Slime Integration +subtitle: Run Slime external rollouts through Dynamo on Kubernetes +--- + +**Experimental.** This integration connects Slime to a fixed set of SGLang engines that a DynamoGraphDeployment manages. + +Slime sends rollout generation through the Dynamo frontend. It sends control and weight-update requests to each SGLang engine. + +## Integration Shape + +| Concern | Current path | +|---|---| +| Training and rollout orchestration | Slime | +| Generation | Dynamo frontend with SGLang workers | +| Routing | Dynamo routes rollout requests to the fixed worker set | +| Policy update | Slime calls the native SGLang APIs for each worker | +| Service discovery | Fixed Kubernetes Services | +| Deployment | DynamoGraphDeployment on Kubernetes | + +Slime starts an SGLang router to track the external engines. The custom generator sends rollout requests to the Dynamo frontend instead. + +The Slime external-engine support lives in the upstream Slime repository. The Dynamo-specific deployment files live in the Dynamo example. + +| Source | Files | +|---|---| +| Slime | [`slime/rollout/`](https://github.com/THUDM/slime/tree/4c1ab40203952b3dcc8582b653f3a83f2c6e8128/slime/rollout) | +| Dynamo | [`examples/rl/slime/`](https://github.com/ai-dynamo/dynamo/blob/main/examples/README.md#integration-examples) | + +## Prerequisites + +- Install the [Dynamo Kubernetes Platform](../../kubernetes/installation/install-dynamo.md) on a GPU cluster. +- Install the NVIDIA device plugin for the `nvidia.com/gpu` resource. +- Install `kubectl` and `envsubst` on the deployment host. +- Select matching Dynamo frontend and SGLang runtime image versions. +- Allocate two GPUs for the SGLang workers. +- Allocate the additional resources that the Slime training job requires. +- Use a Slime revision that includes [THUDM/slime#2272](https://github.com/THUDM/slime/pull/2272). + +## Prepare the Source + +Clone the Slime merge commit that added streaming external rollouts: + +```bash +git clone https://github.com/THUDM/slime.git +git -C slime checkout 4c1ab40203952b3dcc8582b653f3a83f2c6e8128 +test -f slime/slime/rollout/sglang_streaming_rollout.py +``` + +Install Slime with the [upstream instructions](https://github.com/THUDM/slime/blob/4c1ab40203952b3dcc8582b653f3a83f2c6e8128/README.md). Keep the checkout path for `SLIME_HOME`. + +Clone Dynamo to get the deployment example: + +```bash +git clone https://github.com/ai-dynamo/dynamo.git +cd dynamo +``` + +## Deploy the Worker Set + +Set the Dynamo images to the same version. The deployment uses `Qwen/Qwen3-0.6B` by default. + +```bash +export KUBE_CONTEXT= +export NAMESPACE= +export DYNAMO_FRONTEND_IMAGE=nvcr.io/nvidia/ai-dynamo/dynamo-frontend: +export SGLANG_RUNTIME_IMAGE=nvcr.io/nvidia/ai-dynamo/sglang-runtime: +export DYNAMO_RUNTIME_VERSION= +export MODEL_PATH=Qwen/Qwen3-0.6B +examples/rl/slime/deploy-dynamo.sh +``` + +Each worker Pod runs SGLang and loads `dynamo.sglang.sidecar` through the SGLang `--sidecar` option. + +The manifest creates these Services: + +- `slime-sglang-rollout:8000` exposes the Dynamo frontend. +- `slime-sglang-engine-0:30000` exposes the first native SGLang API. +- `slime-sglang-engine-1:30000` exposes the second native SGLang API. + +Each engine Service selects one worker component. Keep each worker component at one replica. + +Restrict access to the engine Services. These Services expose administrative and weight-update APIs. + +## Make Sure That the Engines Are Ready + +Run these commands from the Slime environment or another Pod in the deployment namespace: + +```bash +curl --fail-with-body http://slime-sglang-engine-0:30000/health_generate +curl --fail-with-body http://slime-sglang-engine-1:30000/health_generate +curl --fail-with-body http://slime-sglang-rollout:8000/health +``` + +The first two commands show that the SGLang engines are ready. The third command shows that the Dynamo frontend is ready. + +## Run Slime + +Set the fixed engine addresses and the Dynamo frontend address. Then add the arguments that your training workload requires. + +```bash +export SLIME_HOME= +export DYNAMO_ENGINE_ADDRS="slime-sglang-engine-0:30000 slime-sglang-engine-1:30000" +export DYNAMO_ROLLOUT_URL="http://slime-sglang-rollout:8000" +examples/rl/slime/launch-slime.sh \ + --hf-checkpoint Qwen/Qwen3-0.6B \ + --prompt-data \ + --input-key prompt \ + --rm-type random \ + --num-rollout 1 \ + --rollout-batch-size 2 \ + --n-samples-per-prompt 2 \ + --rollout-max-response-len 128 \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 1 +``` + +The launcher passes the fixed worker list to Slime. It also enables incremental streaming output. + +The `dynamo_generate.generate_streaming` adapter sends generation to `DYNAMO_ROLLOUT_URL`. Slime calls each fixed engine for control operations and weight updates. + +## Make Sure That the Integration Works + +Complete these steps before you increase the workload size: + +1. Run one rollout with the selected model and weight transport. +2. Apply one policy update to both engines. +3. Run one rollout after the update. +4. Stop one worker during a test run. +5. Make sure that Slime stops or reports the worker failure. +6. Restart the Slime job after Kubernetes replaces the worker Pod. + +The integration does not recover a replaced external engine during a Slime job. + +## Current Limitations + +- The example uses a fixed worker set. +- The example does not support dynamic endpoint registration. +- Slime must restart after Kubernetes replaces an engine Pod. +- The native engine Services require network isolation. +- The example does not provide a fleet-wide policy-update transaction or automatic rollback. + +## Upstream Resources + +- [Dynamo Slime example](https://github.com/ai-dynamo/dynamo/blob/main/examples/README.md#integration-examples) +- [Slime streaming external-rollout change](https://github.com/THUDM/slime/pull/2272) +- [Slime external rollout engine guide](https://github.com/THUDM/slime/blob/4c1ab40203952b3dcc8582b653f3a83f2c6e8128/docs/en/advanced/external-rollout-engines.md) +- [Shared RL integration reference](integration-reference.md) diff --git a/examples/README.md b/examples/README.md index 620539feb215..73106589853f 100644 --- a/examples/README.md +++ b/examples/README.md @@ -58,7 +58,8 @@ Platform-specific manifests and templates for production environments. Deploymen End-to-end examples that connect Dynamo to adjacent inference services: -- **[llm-d Batch Gateway](deployments/llm-d-batch-gateway/README.md)** - Experimental OpenAI Batch lifecycle on a dedicated Dynamo worker pool +- **[llm-d Batch Gateway](deployments/llm-d-batch-gateway/README.md)** provides an experimental OpenAI Batch lifecycle on a dedicated Dynamo worker pool. +- **[Slime External Rollouts](rl/slime/README.md)** connects a fixed Slime worker set to stock SGLang engines and Dynamo sidecars. ## Runtime Examples diff --git a/examples/rl/slime/README.md b/examples/rl/slime/README.md new file mode 100644 index 000000000000..4de1ebb12cf8 --- /dev/null +++ b/examples/rl/slime/README.md @@ -0,0 +1,160 @@ + + +# Slime External Rollouts with Dynamo + +**Experimental.** This example runs Slime with two SGLang engines in a fixed +worker set. A DynamoGraphDeployment manages the engines. Each worker Pod runs +SGLang and the Dynamo sidecar in one runtime container. Slime connects through +stable Kubernetes Services. Slime uses the Dynamo frontend for incremental +streaming responses. + +This example uses the external-engine and streaming support from +[THUDM/slime#2272](https://github.com/THUDM/slime/pull/2272). Slime queries +`/server_info` on each engine. Slime registers the fixed addresses with its +SGLang router. Slime calls the native SGLang control and weight-update endpoints. +The included custom generator sends rollout requests to the Dynamo frontend. + +> [!IMPORTANT] +> Slime starts an SGLang router to track the external engines. The custom +> generator does not send rollout requests to that router. +> +> Keep each worker component at one replica. If you change the worker set, +> restart Slime. + +## Prerequisites + +- Install the Dynamo Kubernetes Platform on a GPU-capable Kubernetes cluster. +- Install the NVIDIA device plugin for the `nvidia.com/gpu` resource. +- Use a Dynamo SGLang runtime image that contains `dynamo.sglang.sidecar`. +- For a gated model, create an `hf-token-secret` secret that contains + `HF_TOKEN` in the deployment namespace. The public default model does not + require this secret. +- Use a Slime revision that contains + [THUDM/slime#2272](https://github.com/THUDM/slime/pull/2272). +- Run Slime in a location that can resolve and reach the worker Services. The + same Kubernetes namespace meets this requirement. +- Install `envsubst` on the machine that deploys the manifest. + +## Deploy the fixed worker set + +Set the Dynamo frontend image and the SGLang runtime image to the same Dynamo +version. Set `DYNAMO_RUNTIME_VERSION` to that version. + +The script uses `MODEL_PATH=Qwen/Qwen3-0.6B` by default. + +```bash +export KUBE_CONTEXT= +export NAMESPACE= +export DYNAMO_FRONTEND_IMAGE=nvcr.io/nvidia/ai-dynamo/dynamo-frontend: +export SGLANG_RUNTIME_IMAGE=nvcr.io/nvidia/ai-dynamo/sglang-runtime: +export DYNAMO_RUNTIME_VERSION= +export MODEL_PATH=Qwen/Qwen3-0.6B +examples/rl/slime/deploy-dynamo.sh +``` + +If you use an unreleased version, pin nightly tags from the same date. + +SGLang loads `dynamo.sglang.sidecar` through its `--sidecar` option. It passes +the local gRPC endpoint to the module through `--sidecar-args`. This path does +not require a separate Dynamo sidecar image. + +Each worker Pod requests one `nvidia.com/gpu` resource. The worker Pods also +tolerate the standard `nvidia.com/gpu=true:NoSchedule` GPU-node taint. + +Add only the node selectors and tolerations that your cluster administrator +assigns to your workload. + +The manifest creates three Services: + +- `slime-sglang-rollout:8000` exposes the Dynamo frontend for rollout requests. +- `slime-sglang-engine-0:30000` exposes the first native SGLang control API. +- `slime-sglang-engine-1:30000` exposes the second native SGLang control API. + +The `engine-*` names do not conflict with the operator-owned worker Services. +The operator uses those worker Services for Dynamo discovery on port 9090. + +Each engine Service selects one worker component with one replica. The Service +name stays stable after Kubernetes replaces its Pod. Slime does not recover an +external engine automatically. + +If Kubernetes restarts or replaces an engine Pod, restart the Slime job. + +The native engine Services expose administrative APIs and weight-update APIs. + +Restrict access to these Services to the training network. + +## Engine Readiness + +Run these commands from the Slime environment or another Pod in the deployment +namespace: + +```bash +curl --fail-with-body http://slime-sglang-engine-0:30000/health_generate +curl --fail-with-body http://slime-sglang-engine-1:30000/health_generate +curl --fail-with-body http://slime-sglang-rollout:8000/health +``` + +The first two commands make sure that both fixed engines are reachable. The +third command makes sure that the Dynamo frontend is ready. + +## Start Slime + +`launch-slime.sh` supplies the fixed external-engine list. It uses the streaming +generator from THUDM/slime#2272. + +Add the model, dataset, trainer, weight-update, and resource arguments that your +workload requires. + +```bash +export SLIME_HOME= +export DYNAMO_ENGINE_ADDRS="slime-sglang-engine-0:30000 slime-sglang-engine-1:30000" +export DYNAMO_ROLLOUT_URL="http://slime-sglang-rollout:8000" +examples/rl/slime/launch-slime.sh \ + --hf-checkpoint Qwen/Qwen3-0.6B \ + --prompt-data \ + --input-key prompt \ + --rm-type random \ + --num-rollout 1 \ + --rollout-batch-size 2 \ + --n-samples-per-prompt 2 \ + --rollout-max-response-len 128 \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 1 +``` + +The launcher passes these integration arguments: + +```text +--rollout-external-engine-addrs +--rollout-function-path slime.rollout.sglang_rollout.generate_rollout +--custom-generate-function-path dynamo_generate.generate_streaming +--sglang-incremental-streaming-output +``` + +The `dynamo_generate.generate_streaming` adapter copies the Slime arguments. +Then the adapter sets the generation address to `DYNAMO_ROLLOUT_URL`. The +adapter calls the streaming generator from THUDM/slime#2272. Slime calls each +fixed engine for control operations and weight updates. + +The `--sglang-incremental-streaming-output` flag must match the +`--incremental-streaming-output` configuration in `dynamo.yaml`. + +Read the Slime +[external rollout engine guide](https://github.com/THUDM/slime/blob/main/docs/en/advanced/external-rollout-engines.md) +for NCCL updates, full-checkpoint disk updates, and delta disk updates. + +## Current Boundary + +This example uses a fixed worker set. It does not support elastic discovery, +dynamic endpoint registration, or external-engine fault recovery. + +Before you use the example for training, complete these tests: + +1. Run one rollout with the selected model and weight transport. +2. Apply one policy update. +3. Run one post-update rollout. +4. Simulate one worker failure. +5. Make sure that the system recovers. diff --git a/examples/rl/slime/deploy-dynamo.sh b/examples/rl/slime/deploy-dynamo.sh new file mode 100755 index 000000000000..9cd391974244 --- /dev/null +++ b/examples/rl/slime/deploy-dynamo.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Deploy the fixed Dynamo SGLang worker set used by the Slime example. + +set -euo pipefail + +: "${KUBE_CONTEXT:?Set KUBE_CONTEXT to the target Kubernetes context}" +: "${NAMESPACE:?Set NAMESPACE to the target Kubernetes namespace}" +: "${DYNAMO_FRONTEND_IMAGE:?Set DYNAMO_FRONTEND_IMAGE to a Dynamo frontend image}" +: "${SGLANG_RUNTIME_IMAGE:?Set SGLANG_RUNTIME_IMAGE to a Dynamo SGLang runtime image}" +: "${DYNAMO_RUNTIME_VERSION:?Set DYNAMO_RUNTIME_VERSION to the matching Dynamo version}" +: "${MODEL_PATH:=Qwen/Qwen3-0.6B}" + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +export DYNAMO_FRONTEND_IMAGE SGLANG_RUNTIME_IMAGE DYNAMO_RUNTIME_VERSION MODEL_PATH +SUBSTITUTIONS='${DYNAMO_FRONTEND_IMAGE} ${SGLANG_RUNTIME_IMAGE} ${DYNAMO_RUNTIME_VERSION} ${MODEL_PATH}' + +envsubst "${SUBSTITUTIONS}" <"${SCRIPT_DIR}/dynamo.yaml" | + kubectl --context="${KUBE_CONTEXT}" --namespace="${NAMESPACE}" apply -f - diff --git a/examples/rl/slime/dynamo.yaml b/examples/rl/slime/dynamo.yaml new file mode 100644 index 000000000000..cbe2e2fd2789 --- /dev/null +++ b/examples/rl/slime/dynamo.yaml @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Each engine Service must select exactly one engine. Slime uses this fixed +# worker set for control requests. +# +# Keep each worker component at one replica. +apiVersion: v1 +kind: Service +metadata: + name: slime-sglang-rollout +spec: + ports: + - name: http + port: 8000 + targetPort: 8000 + selector: + dynamo.nvidia.com/slime-role: frontend +--- +apiVersion: v1 +kind: Service +metadata: + name: slime-sglang-engine-0 +spec: + ports: + - name: sglang-http + port: 30000 + targetPort: 30000 + selector: + dynamo.nvidia.com/slime-role: worker-0 +--- +apiVersion: v1 +kind: Service +metadata: + name: slime-sglang-engine-1 +spec: + ports: + - name: sglang-http + port: 30000 + targetPort: 30000 + selector: + dynamo.nvidia.com/slime-role: worker-1 +--- +apiVersion: nvidia.com/v1beta1 +kind: DynamoGraphDeployment +metadata: + name: slime-sglang +spec: + backendFramework: sglang + components: + - name: frontend + type: frontend + replicas: 1 + runtimeVersionOverride: ${DYNAMO_RUNTIME_VERSION} + podTemplate: + metadata: + labels: + dynamo.nvidia.com/slime-role: frontend + spec: + containers: + - name: main + image: ${DYNAMO_FRONTEND_IMAGE} + env: + - name: DYN_SGLANG_ENABLE_GENERATE + value: "1" + - name: worker-0 + type: worker + replicas: 1 + runtimeVersionOverride: ${DYNAMO_RUNTIME_VERSION} + podTemplate: + metadata: + labels: + dynamo.nvidia.com/slime-role: worker-0 + spec: + tolerations: + - key: nvidia.com/gpu + operator: Equal + value: "true" + effect: NoSchedule + containers: + - name: main + image: ${SGLANG_RUNTIME_IMAGE} + startupProbe: + httpGet: + path: /health_generate + port: 30000 + periodSeconds: 5 + timeoutSeconds: 5 + failureThreshold: 60 + readinessProbe: + httpGet: + path: /health_generate + port: 30000 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + envFrom: + - secretRef: + name: hf-token-secret + optional: true + resources: + requests: + nvidia.com/gpu: "1" + limits: + nvidia.com/gpu: "1" + command: + - python3 + - -m + - sglang.launch_server + args: + - --model-path + - ${MODEL_PATH} + - --served-model-name + - ${MODEL_PATH} + - --host + - 0.0.0.0 + - --port + - "30000" + - --grpc-port + - "30001" + - --incremental-streaming-output + - --sidecar + - dynamo.sglang.sidecar + - --sidecar-args + - '["--grpc-endpoint","http://127.0.0.1:30001"]' + - name: worker-1 + type: worker + replicas: 1 + runtimeVersionOverride: ${DYNAMO_RUNTIME_VERSION} + podTemplate: + metadata: + labels: + dynamo.nvidia.com/slime-role: worker-1 + spec: + tolerations: + - key: nvidia.com/gpu + operator: Equal + value: "true" + effect: NoSchedule + containers: + - name: main + image: ${SGLANG_RUNTIME_IMAGE} + startupProbe: + httpGet: + path: /health_generate + port: 30000 + periodSeconds: 5 + timeoutSeconds: 5 + failureThreshold: 60 + readinessProbe: + httpGet: + path: /health_generate + port: 30000 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + envFrom: + - secretRef: + name: hf-token-secret + optional: true + resources: + requests: + nvidia.com/gpu: "1" + limits: + nvidia.com/gpu: "1" + command: + - python3 + - -m + - sglang.launch_server + args: + - --model-path + - ${MODEL_PATH} + - --served-model-name + - ${MODEL_PATH} + - --host + - 0.0.0.0 + - --port + - "30000" + - --grpc-port + - "30001" + - --incremental-streaming-output + - --sidecar + - dynamo.sglang.sidecar + - --sidecar-args + - '["--grpc-endpoint","http://127.0.0.1:30001"]' diff --git a/examples/rl/slime/dynamo_generate.py b/examples/rl/slime/dynamo_generate.py new file mode 100644 index 000000000000..393c035f41d0 --- /dev/null +++ b/examples/rl/slime/dynamo_generate.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Routes Slime streaming rollout requests through the Dynamo frontend.""" + +import copy +import os +from argparse import Namespace +from urllib.parse import urlsplit + + +def _frontend_args(args: Namespace) -> Namespace: + rollout_url = os.environ.get("DYNAMO_ROLLOUT_URL") + if not rollout_url: + raise ValueError("Set DYNAMO_ROLLOUT_URL to the Dynamo frontend base URL.") + + parsed = urlsplit(rollout_url) + if parsed.scheme != "http" or parsed.hostname is None: + raise ValueError("DYNAMO_ROLLOUT_URL must use http:// and include a host.") + if parsed.path not in ("", "/") or parsed.query or parsed.fragment: + raise ValueError( + "DYNAMO_ROLLOUT_URL must not include a path, query, or fragment." + ) + + try: + port = parsed.port + except ValueError as error: + raise ValueError("DYNAMO_ROLLOUT_URL contains an invalid port.") from error + if port is None: + raise ValueError("DYNAMO_ROLLOUT_URL must include a port.") + + host = parsed.hostname + if ":" in host: + host = f"[{host}]" + + frontend_args = copy.copy(args) + frontend_args.sglang_router_ip = host + frontend_args.sglang_router_port = port + return frontend_args + + +async def generate_streaming(args: Namespace, sample, sampling_params): + from slime.rollout.sglang_streaming_rollout import ( + generate_streaming as slime_generate_streaming, + ) + + return await slime_generate_streaming(_frontend_args(args), sample, sampling_params) + + +generate_streaming.abort_mode = "request" diff --git a/examples/rl/slime/launch-slime.sh b/examples/rl/slime/launch-slime.sh new file mode 100755 index 000000000000..5d65b8fbed60 --- /dev/null +++ b/examples/rl/slime/launch-slime.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Start Slime against a fixed set of SGLang engines managed by Dynamo. + +set -euo pipefail + +: "${SLIME_HOME:?Set SLIME_HOME to the Slime checkout or installed source tree}" +: "${DYNAMO_ENGINE_ADDRS:=slime-sglang-engine-0:30000 slime-sglang-engine-1:30000}" +: "${DYNAMO_ROLLOUT_URL:=http://slime-sglang-rollout:8000}" + +EXAMPLE_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +export DYNAMO_ROLLOUT_URL +export PYTHONPATH="${EXAMPLE_DIR}${PYTHONPATH:+:${PYTHONPATH}}" + +read -r -a ENGINE_ADDRS <<<"${DYNAMO_ENGINE_ADDRS}" + +exec python3 "${SLIME_HOME}/train.py" \ + --rollout-external-engine-addrs "${ENGINE_ADDRS[@]}" \ + --rollout-function-path slime.rollout.sglang_rollout.generate_rollout \ + --custom-generate-function-path \ + dynamo_generate.generate_streaming \ + --sglang-incremental-streaming-output \ + "$@" diff --git a/lib/bindings/python/examples/hello_world/server_sglang.py b/lib/bindings/python/examples/hello_world/server_sglang.py index f70551c1bee8..5320bf949d55 100644 --- a/lib/bindings/python/examples/hello_world/server_sglang.py +++ b/lib/bindings/python/examples/hello_world/server_sglang.py @@ -61,7 +61,6 @@ async def generate(self, request): # sglang defaults this to 128 "max_new_tokens": request["stop_conditions"]["max_tokens"], } - num_output_tokens_so_far = 0 gen = await self.engine_client.async_generate( input_ids=request["token_ids"], sampling_params=sampling_params, stream=True ) @@ -69,14 +68,10 @@ async def generate(self, request): # res is a dict finish_reason = res["meta_info"]["finish_reason"] + out = {"token_ids": res["output_ids"]} if finish_reason: - # Don't forward the stop token - out = {"token_ids": [], "finish_reason": finish_reason["type"]} - else: - next_total_toks = len(res["output_ids"]) - out = {"token_ids": res["output_ids"][num_output_tokens_so_far:]} + out["finish_reason"] = finish_reason["type"] yield out - num_output_tokens_so_far = next_total_toks @dynamo_worker() @@ -102,6 +97,7 @@ async def init(runtime: DistributedRuntime, config: Config): engine_args = ServerArgs( model_path=config.model, skip_tokenizer_init=True, + incremental_streaming_output=True, ) engine_client = sglang.Engine(server_args=engine_args) diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs index ad9b3df7bd23..2d7da8b3c5b2 100644 --- a/lib/llm/src/backend.rs +++ b/lib/llm/src/backend.rs @@ -831,6 +831,8 @@ mod tests { [] => "", [101] => "Okay", [102] => " Okay", + [103] => "<|user|>", + [101, 103] => "Okay<|user|>", _ => anyhow::bail!("unexpected token IDs: {token_ids:?}"), }; Ok(traits::DecodeResult::Complete(token.to_string())) @@ -843,6 +845,35 @@ mod tests { engine_decodes_text: bool, } + struct SyntheticSglangStopEngine; + + #[async_trait] + impl AsyncEngine, ManyOut>, Error> + for SyntheticSglangStopEngine + { + async fn generate( + &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() + }; + + Ok(ResponseStream::new( + Box::pin(futures::stream::once(async move { + Annotated::from_data(output) + })), + request.context(), + )) + } + } + #[async_trait] impl AsyncEngine, ManyOut>, Error> for SyntheticSglangEngine @@ -983,6 +1014,65 @@ mod tests { assert_sglang_top_logprobs_are_decoded_in_openai_response(true).await; } + #[tokio::test] + async fn test_sglang_hidden_stop_text_is_suppressed_without_mutating_tito_ids() { + use crate::protocols::common::extensions::NvExt; + + let tokenizer: Arc = Arc::new(CandidateDecoder); + let backend = Backend::from_tokenizer(Tokenizer::from(tokenizer)); + let request = PreprocessedRequest::builder() + .model("test-model".to_string()) + .token_ids(vec![]) + .stop_conditions(StopConditions { + stop_token_ids_hidden: Some(vec![103]), + ..Default::default() + }) + .sampling_options(SamplingOptions::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) + .await + .expect("backend generation succeeds"); + let output = stream + .next() + .await + .expect("backend emits a response") + .data + .expect("response contains backend output"); + + 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.finish_reason, Some(FinishReason::Stop)); + + let nvext = NvExt::builder() + .extra_fields(vec!["completion_token_ids".to_string()]) + .build() + .expect("valid nvext selection"); + let options = DeltaGeneratorOptions::new(None, None, false, Some(&nvext)); + let mut generator = DeltaGenerator::new( + "test-model".to_string(), + options, + "test-request".to_string(), + ); + let response = generator + .choice_from_postprocessor(output) + .expect("OpenAI response conversion succeeds"); + + assert_eq!( + response + .nvext + .as_ref() + .and_then(|fields| fields.get("completion_token_ids")), + Some(&serde_json::json!([101, 103])) + ); + assert_eq!(generator.get_usage().completion_tokens, 2); + } + /// When the tokenizer's decode() returns Err, Decoder::process_token_ids() /// should propagate the error. In the backend unfold closure, this error /// gets caught and converted to FinishReason::Error. diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 1ddd0bc7103a..6921185fb828 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -4237,9 +4237,27 @@ impl OpenAIPreprocessor { } } + Self::capture_prompt_token_ids(request, tracker, &tokens_out); + Ok((tokens_out, annotations)) } + /// Retain the authoritative rendered prompt only for clients that request + /// the named response field. Ordinary requests do not clone the token list. + fn capture_prompt_token_ids( + request: &R, + tracker: Option<&RequestTracker>, + token_ids: &[TokenIdType], + ) { + let requested = request + .nvext() + .and_then(|nvext| nvext.extra_fields.as_ref()) + .is_some_and(|fields| fields.iter().any(|field| field == "prompt_token_ids")); + if requested && let Some(tracker) = tracker { + tracker.set_prompt_token_ids(token_ids.to_vec()); + } + } + fn prompt_overflow_error(token_count: usize, combined_limit: usize) -> DynamoError { DynamoError::builder() .error_type(ErrorType::InvalidArgument) @@ -8857,6 +8875,33 @@ mod tests { ); } + #[test] + fn capture_prompt_token_ids_is_exact_and_opt_in() { + let requested: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "nvext": {"extra_fields": ["prompt_token_ids"]} + })) + .unwrap(); + let tracker = RequestTracker::new(); + + OpenAIPreprocessor::capture_prompt_token_ids(&requested, Some(&tracker), &[101, 102, 103]); + assert_eq!(tracker.prompt_token_ids(), Some(&[101u32, 102, 103][..])); + + let ordinary: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}] + })) + .unwrap(); + let ordinary_tracker = RequestTracker::new(); + OpenAIPreprocessor::capture_prompt_token_ids( + &ordinary, + Some(&ordinary_tracker), + &[201, 202], + ); + assert!(ordinary_tracker.prompt_token_ids().is_none()); + } + fn chat_request_with_args( chat_template_args: Option>, ) -> NvCreateChatCompletionRequest { diff --git a/lib/llm/src/protocols/common/extensions.rs b/lib/llm/src/protocols/common/extensions.rs index d0598b30b04f..7d48a3c228a9 100644 --- a/lib/llm/src/protocols/common/extensions.rs +++ b/lib/llm/src/protocols/common/extensions.rs @@ -13,8 +13,9 @@ use crate::protocols::TokenIdType; use crate::protocols::agents::{ AgentContextHeaderValues, agent_context_header_values, session_affinity_header_value, }; +use crate::protocols::common::FinishReason; use crate::protocols::common::llm_backend::PromptLogprobs; -use crate::protocols::common::timing::TimingInfo; +use crate::protocols::common::timing::{RequestTracker, TimingInfo}; /// Request-level taint constraints carried by `nvext.routing_constraints`. #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)] @@ -666,9 +667,19 @@ pub struct NvExtResponse { #[serde(skip_serializing_if = "Option::is_none")] pub stop_reason: Option, + /// Dynamo's internal finish reason before OpenAI conversion. + /// + /// Backend-specific reasons are normalized before this value is set. + /// For example, SGLang's `abort` becomes `cancelled`. + #[serde(skip_serializing_if = "Option::is_none")] + pub detailed_finish_reason: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub completion_token_ids: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_token_ids: Option>, + #[serde(skip_serializing_if = "Option::is_none")] pub prompt_logprobs: Option, } @@ -719,10 +730,23 @@ pub struct NvExtResponseFieldSelection { pub routed_experts: bool, pub engine_data: bool, pub stop_reason: bool, + pub detailed_finish_reason: bool, pub completion_token_ids: bool, + pub prompt_token_ids: bool, pub prompt_logprobs: bool, } +/// Backend data available when building an NVExt response. +#[derive(Debug, Default)] +pub struct NvExtResponseInput<'a> { + pub tracker: Option<&'a RequestTracker>, + pub finish_reason: Option<&'a FinishReason>, + pub engine_data: Option, + pub stop_reason: Option, + pub completion_token_ids: Option<&'a [TokenIdType]>, + pub prompt_logprobs: Option, +} + impl NvExtResponseFieldSelection { pub fn from_nvext(nvext: Option<&NvExt>) -> Self { let Some(ext) = nvext else { @@ -738,7 +762,9 @@ impl NvExtResponseFieldSelection { "routed_experts" => selection.routed_experts = true, "engine_data" => selection.engine_data = true, "stop_reason" => selection.stop_reason = true, + "detailed_finish_reason" => selection.detailed_finish_reason = true, "completion_token_ids" => selection.completion_token_ids = true, + "prompt_token_ids" => selection.prompt_token_ids = true, "prompt_logprobs" => selection.prompt_logprobs = true, _ => {} } @@ -751,30 +777,26 @@ impl NvExtResponseFieldSelection { selection } - #[allow(clippy::too_many_arguments)] - pub fn build_response_nvext( - &self, - tracker: Option<&std::sync::Arc>, - finish_reason_present: bool, - engine_data_from_backend: Option, - stop_reason_from_backend: Option, - completion_token_ids_from_backend: Option<&[TokenIdType]>, - prompt_logprobs_from_backend: Option, - ) -> Option { + pub fn build_response_nvext(&self, input: NvExtResponseInput<'_>) -> Option { + let finish_reason_present = input.finish_reason.is_some(); + let worker_id = if self.worker_id { - tracker.and_then(|t| t.get_worker_info()) + input.tracker.and_then(RequestTracker::get_worker_info) } else { None }; let token_ids = if self.token_ids { - tracker.and_then(|t| t.query_token_ids().map(<[u32]>::to_vec)) + input + .tracker + .and_then(|tracker| tracker.query_token_ids().map(<[u32]>::to_vec)) } else { None }; let routed_experts = if self.routed_experts { - engine_data_from_backend + input + .engine_data .as_ref() .and_then(|data| data.get("routed_experts")) .cloned() @@ -783,31 +805,49 @@ impl NvExtResponseFieldSelection { }; let timing = if finish_reason_present && self.timing { - tracker.map(|t| t.get_timing_info()) + input.tracker.map(RequestTracker::get_timing_info) } else { None }; let engine_data = if self.engine_data { - engine_data_from_backend + input.engine_data } else { None }; let stop_reason = if self.stop_reason { - stop_reason_from_backend.and_then(|reason| serde_json::to_value(reason).ok()) + input + .stop_reason + .and_then(|reason| serde_json::to_value(reason).ok()) + } else { + None + }; + + let detailed_finish_reason = if self.detailed_finish_reason { + input.finish_reason.map(ToString::to_string) } else { None }; let completion_token_ids = if self.completion_token_ids { - completion_token_ids_from_backend.map(<[u32]>::to_vec) + input.completion_token_ids.map(<[u32]>::to_vec) + } else { + None + }; + + // Prompt IDs can be large, so emit them only once on the terminal + // chunk. The tracker stores them only for an explicit opt-in request. + let prompt_token_ids = if self.prompt_token_ids && finish_reason_present { + input + .tracker + .and_then(|tracker| tracker.prompt_token_ids().map(<[u32]>::to_vec)) } else { None }; let prompt_logprobs = if self.prompt_logprobs && finish_reason_present { - prompt_logprobs_from_backend + input.prompt_logprobs } else { None }; @@ -818,7 +858,9 @@ impl NvExtResponseFieldSelection { && timing.is_none() && engine_data.is_none() && stop_reason.is_none() + && detailed_finish_reason.is_none() && completion_token_ids.is_none() + && prompt_token_ids.is_none() && prompt_logprobs.is_none() { return None; @@ -831,7 +873,9 @@ impl NvExtResponseFieldSelection { routed_experts, engine_data, stop_reason, + detailed_finish_reason, completion_token_ids, + prompt_token_ids, prompt_logprobs, }) } @@ -1605,7 +1649,12 @@ mod tests { #[test] fn response_field_selection_respects_extra_fields() { let nvext = NvExt::builder() - .extra_fields(vec!["worker_id".to_string(), "routed_experts".to_string()]) + .extra_fields(vec![ + "worker_id".to_string(), + "routed_experts".to_string(), + "prompt_token_ids".to_string(), + "detailed_finish_reason".to_string(), + ]) .build() .unwrap(); @@ -1614,6 +1663,8 @@ mod tests { NvExtResponseFieldSelection { worker_id: true, routed_experts: true, + prompt_token_ids: true, + detailed_finish_reason: true, ..Default::default() } ); @@ -1672,6 +1723,14 @@ mod tests { tracker } + fn tracker_with_prompt_token_ids() + -> std::sync::Arc { + use crate::protocols::common::timing::RequestTracker; + let tracker = std::sync::Arc::new(RequestTracker::new()); + tracker.set_prompt_token_ids(vec![101u32, 102, 103]); + tracker + } + fn tracker_with_forwarded_worker_info() -> std::sync::Arc { use crate::protocols::common::timing::RequestTracker; @@ -1687,9 +1746,13 @@ mod tests { #[test] fn build_response_nvext_all_false_returns_none() { + let finish_reason = FinishReason::Cancelled; assert!( NvExtResponseFieldSelection::default() - .build_response_nvext(None, false, None, None, None, None) + .build_response_nvext(NvExtResponseInput { + finish_reason: Some(&finish_reason), + ..Default::default() + }) .is_none() ); } @@ -1703,7 +1766,10 @@ mod tests { let tracker = tracker_with_prefill_worker(); let out = selection - .build_response_nvext(Some(&tracker), false, None, None, None, None) + .build_response_nvext(NvExtResponseInput { + tracker: Some(&tracker), + ..Default::default() + }) .expect("worker_id should emit regardless of finish_reason"); assert!(out.worker_id.is_some()); @@ -1721,7 +1787,10 @@ mod tests { let tracker = tracker_with_forwarded_worker_info(); let out = selection - .build_response_nvext(Some(&tracker), false, None, None, None, None) + .build_response_nvext(NvExtResponseInput { + tracker: Some(&tracker), + ..Default::default() + }) .expect("forwarded worker_id should surface in nvext"); assert_eq!( @@ -1745,12 +1814,20 @@ mod tests { assert!( selection - .build_response_nvext(Some(&tracker), false, None, None, None, None) + .build_response_nvext(NvExtResponseInput { + tracker: Some(&tracker), + ..Default::default() + }) .is_none() ); + let finish_reason = FinishReason::Stop; let out = selection - .build_response_nvext(Some(&tracker), true, None, None, None, None) + .build_response_nvext(NvExtResponseInput { + tracker: Some(&tracker), + finish_reason: Some(&finish_reason), + ..Default::default() + }) .expect("timing should emit on finish"); assert!(out.timing.is_some()); } @@ -1764,7 +1841,10 @@ mod tests { let tracker = tracker_with_query_token_ids(); let out = selection - .build_response_nvext(Some(&tracker), false, None, None, None, None) + .build_response_nvext(NvExtResponseInput { + tracker: Some(&tracker), + ..Default::default() + }) .expect("token_ids should emit when present"); assert_eq!(out.token_ids, Some(vec![11u32, 22, 33])); @@ -1779,7 +1859,10 @@ mod tests { let engine_data = serde_json::json!({ "routed_experts": {"layer_0": [1, 3]} }); let out = selection - .build_response_nvext(None, false, Some(engine_data), None, None, None) + .build_response_nvext(NvExtResponseInput { + engine_data: Some(engine_data), + ..Default::default() + }) .expect("routed_experts should emit when present"); assert_eq!( @@ -1796,13 +1879,62 @@ mod tests { }; let out = selection - .build_response_nvext(None, false, None, None, Some(&[101u32, 102, 103]), None) + .build_response_nvext(NvExtResponseInput { + completion_token_ids: Some(&[101u32, 102, 103]), + ..Default::default() + }) .expect("completion_token_ids should emit when requested and present"); assert_eq!(out.completion_token_ids, Some(vec![101u32, 102, 103])); assert!(out.prompt_logprobs.is_none()); } + #[test] + fn build_response_nvext_detailed_finish_reason_pass_through() { + let selection = NvExtResponseFieldSelection { + detailed_finish_reason: true, + ..Default::default() + }; + + let finish_reason = FinishReason::Cancelled; + let out = selection + .build_response_nvext(NvExtResponseInput { + finish_reason: Some(&finish_reason), + ..Default::default() + }) + .expect("detailed_finish_reason should emit when requested and present"); + + assert_eq!(out.detailed_finish_reason.as_deref(), Some("cancelled")); + } + + #[test] + fn build_response_nvext_prompt_token_ids_final_chunk_only() { + let selection = NvExtResponseFieldSelection { + prompt_token_ids: true, + ..Default::default() + }; + let tracker = tracker_with_prompt_token_ids(); + + assert!( + selection + .build_response_nvext(NvExtResponseInput { + tracker: Some(&tracker), + ..Default::default() + }) + .is_none() + ); + + let finish_reason = FinishReason::Stop; + let out = selection + .build_response_nvext(NvExtResponseInput { + tracker: Some(&tracker), + finish_reason: Some(&finish_reason), + ..Default::default() + }) + .expect("prompt_token_ids should emit on the final chunk"); + assert_eq!(out.prompt_token_ids, Some(vec![101u32, 102, 103])); + } + #[test] fn build_response_nvext_prompt_logprobs_final_chunk_only() { let selection = NvExtResponseFieldSelection { @@ -1822,12 +1954,20 @@ mod tests { assert!( selection - .build_response_nvext(None, false, None, None, None, Some(payload.clone())) + .build_response_nvext(NvExtResponseInput { + prompt_logprobs: Some(payload.clone()), + ..Default::default() + }) .is_none() ); + let finish_reason = FinishReason::Stop; let out = selection - .build_response_nvext(None, true, None, None, None, Some(payload)) + .build_response_nvext(NvExtResponseInput { + finish_reason: Some(&finish_reason), + prompt_logprobs: Some(payload), + ..Default::default() + }) .expect("prompt_logprobs should emit on the final chunk"); let got = out.prompt_logprobs.expect("prompt_logprobs payload"); assert_eq!(got.len(), 2); diff --git a/lib/llm/src/protocols/common/timing.rs b/lib/llm/src/protocols/common/timing.rs index 79347ef32905..2fbcb3e44454 100644 --- a/lib/llm/src/protocols/common/timing.rs +++ b/lib/llm/src/protocols/common/timing.rs @@ -179,6 +179,11 @@ pub struct RequestTracker { /// re-tokenizing. Lives here rather than on `routing_data` because the preprocessor /// drains `routing_data` before the delta generator runs. First-write-wins. external_query_token_ids: OnceLock>, + + /// Frontend-rendered prompt tokens retained only for an explicit + /// `nvext.extra_fields=["prompt_token_ids"]` response request. + /// First-write-wins because one response generator owns one prompt. + prompt_token_ids: OnceLock>, } /// Data a standalone router (running the `PushRouter` bindings in its own process) @@ -238,6 +243,7 @@ impl RequestTracker { prefill_complete_time: OnceLock::new(), external_timing: OnceLock::new(), external_query_token_ids: OnceLock::new(), + prompt_token_ids: OnceLock::new(), } } @@ -633,6 +639,12 @@ impl RequestTracker { let _ = self.external_query_token_ids.set(token_ids); } + /// Retain the frontend's authoritative rendered prompt token sequence for + /// opt-in response emission. First-write-wins. + pub fn set_prompt_token_ids(&self, token_ids: Vec) { + let _ = self.prompt_token_ids.set(token_ids); + } + /// Overlay worker attribution forwarded by a standalone router (on the first chunk's /// `routing_data.worker_id`) onto this tracker so `get_worker_info`, the metrics /// annotation, and `build_response_nvext` surface it on the split-router path. @@ -657,6 +669,11 @@ impl RequestTracker { self.external_query_token_ids.get().map(Vec::as_slice) } + /// The frontend-rendered prompt token sequence, when explicitly retained. + pub fn prompt_token_ids(&self) -> Option<&[u32]> { + self.prompt_token_ids.get().map(Vec::as_slice) + } + /// Per-request timing. Starts from this process's own measurements; in the split-router /// topology, router-injected fields (`prefill_*`/`ttft_ms`/`kv_hit_rate`/`router_queue_depth`/ /// `kv_transfer_estimated_latency_ms`) fill the gaps the frontend can't measure. This @@ -892,6 +909,19 @@ mod tests { assert_eq!(tracker.query_token_ids(), Some(&[11u32, 22, 33][..])); } + #[test] + fn test_prompt_token_ids_round_trip() { + let tracker = RequestTracker::new(); + assert!(tracker.prompt_token_ids().is_none()); + + tracker.set_prompt_token_ids(vec![101, 102, 103]); + assert_eq!(tracker.prompt_token_ids(), Some(&[101u32, 102, 103][..])); + + // One request has one authoritative rendered prompt. + tracker.set_prompt_token_ids(vec![999]); + assert_eq!(tracker.prompt_token_ids(), Some(&[101u32, 102, 103][..])); + } + #[test] fn test_set_external_worker_info_round_trip() { let tracker = RequestTracker::new(); diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index 0b33231e6883..e20a5bca4340 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -6,7 +6,11 @@ use std::{collections::HashSet, sync::Arc}; use super::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse}; use crate::{ protocols::{ - common::{self, extensions::NvExtProvider, timing::RequestTracker}, + common::{ + self, + extensions::{NvExtProvider, NvExtResponseInput}, + timing::RequestTracker, + }, openai::{ convert_backend_top_logprobs, delta_common::{self, DeltaGeneratorOptions, DeltaGeneratorState}, @@ -252,7 +256,7 @@ impl crate::protocols::openai::DeltaGeneratorExt Some(dynamo_protocols::types::FinishReason::Stop), Some(common::FinishReason::Stop) => Some(dynamo_protocols::types::FinishReason::Stop), Some(common::FinishReason::Length) => { @@ -265,7 +269,7 @@ impl crate::protocols::openai::DeltaGeneratorExt { - return Err(anyhow::anyhow!(err_msg)); + return Err(anyhow::anyhow!(err_msg.clone())); } None => None, }; @@ -288,14 +292,19 @@ impl crate::protocols::openai::DeltaGeneratorExt for // Backend errors are response errors, not successful OpenAI stop reasons. // Keep completions aligned with the chat-completions delta generator. - let finish_reason = match delta.finish_reason { + let finish_reason = match delta.finish_reason.as_ref() { Some(common::FinishReason::Error(err_msg)) => { self.state.tracker_ref().record_finish(); - return Err(anyhow::anyhow!(err_msg)); + return Err(anyhow::anyhow!(err_msg.clone())); } - Some(reason) => Some(reason.into()), + Some(reason) => Some(reason.clone().into()), None => None, }; let stop_reason = delta.stop_reason.clone(); @@ -234,14 +238,19 @@ impl crate::protocols::openai::DeltaGeneratorExt for // rules stay in one place. let prompt_logprobs_payload = common::llm_backend::prompt_logprobs_from_engine_data(delta.engine_data.as_ref()); - if let Some(nvext_response) = self.state.options().response_fields.build_response_nvext( - Some(self.state.tracker_ref()), - finish_reason.is_some(), - delta.engine_data, - stop_reason, - completion_token_ids_for_nvext.as_deref(), - prompt_logprobs_payload, - ) && let Ok(nvext_json) = serde_json::to_value(&nvext_response) + if let Some(nvext_response) = + self.state + .options() + .response_fields + .build_response_nvext(NvExtResponseInput { + tracker: Some(self.state.tracker_ref()), + finish_reason: delta.finish_reason.as_ref(), + engine_data: delta.engine_data, + stop_reason, + completion_token_ids: completion_token_ids_for_nvext.as_deref(), + prompt_logprobs: prompt_logprobs_payload, + }) + && let Ok(nvext_json) = serde_json::to_value(&nvext_response) { response.nvext = Some(nvext_json); if let Some(ref info) = nvext_response.worker_id { @@ -298,7 +307,7 @@ mod tests { use super::*; use crate::protocols::common::{self, llm_backend::BackendOutput, timing::WORKER_TYPE_PREFILL}; use crate::protocols::openai::DeltaGeneratorExt; - use dynamo_protocols::types::{CompletionUsage, CreateCompletionRequestArgs, Prompt}; + use dynamo_protocols::types::{CreateCompletionRequestArgs, Prompt}; fn create_test_request() -> NvCreateCompletionRequest { let inner = CreateCompletionRequestArgs::default() @@ -360,85 +369,6 @@ mod tests { assert_eq!(response.inner.model, "test-model"); } - #[test] - fn test_completion_tokens_use_backend_usage_when_higher() { - let request = create_test_request(); - let mut generator = request.response_generator("req-backend-usage".to_string()); - - let mut backend_output = final_backend_output(); - backend_output.token_ids.clear(); - backend_output.tokens.clear(); - backend_output.completion_usage = Some(CompletionUsage { - prompt_tokens: 5, - completion_tokens: 1, - total_tokens: 6, - prompt_tokens_details: None, - completion_tokens_details: None, - }); - - generator - .choice_from_postprocessor(backend_output) - .expect("choice generation"); - - let usage = generator.get_usage(); - assert_eq!(usage.prompt_tokens, 5); - assert_eq!(usage.completion_tokens, 1); - assert_eq!(usage.total_tokens, 6); - } - - #[test] - fn test_completion_tokens_treat_backend_usage_as_request_total() { - let mut request = create_test_request(); - request.inner.n = Some(2); - let mut generator = request.response_generator("req-multi-choice-usage".to_string()); - - for index in 0..2 { - let mut backend_output = final_backend_output(); - backend_output.index = Some(index); - backend_output.token_ids.clear(); - backend_output.tokens.clear(); - backend_output.completion_usage = Some(CompletionUsage { - prompt_tokens: 5, - completion_tokens: 5, - total_tokens: 10, - prompt_tokens_details: None, - completion_tokens_details: None, - }); - generator - .choice_from_postprocessor(backend_output) - .expect("choice generation"); - } - - let usage = generator.get_usage(); - assert_eq!(usage.prompt_tokens, 5); - assert_eq!(usage.completion_tokens, 5); - assert_eq!(usage.total_tokens, 10); - } - - #[test] - fn test_completion_tokens_keep_aggregated_count_when_backend_usage_is_zero() { - let request = create_test_request(); - let mut generator = request.response_generator("req-backend-zero-usage".to_string()); - - let mut backend_output = final_backend_output(); - backend_output.completion_usage = Some(CompletionUsage { - prompt_tokens: 5, - completion_tokens: 0, - total_tokens: 5, - prompt_tokens_details: None, - completion_tokens_details: None, - }); - - generator - .choice_from_postprocessor(backend_output) - .expect("choice generation"); - - let usage = generator.get_usage(); - assert_eq!(usage.prompt_tokens, 5); - assert_eq!(usage.completion_tokens, 1); - assert_eq!(usage.total_tokens, 6); - } - fn create_test_request_with_extra_fields(fields: Vec) -> NvCreateCompletionRequest { let inner = CreateCompletionRequestArgs::default() .model("test-model") @@ -554,6 +484,26 @@ mod tests { assert_eq!(response_json["nvext"]["stop_reason"], "END"); } + #[test] + fn test_cancelled_detailed_finish_reason_preserves_openai_finish_reason() { + let request = + create_test_request_with_extra_fields(vec!["detailed_finish_reason".to_string()]); + let mut generator = request.response_generator("req-cancelled-nvext".to_string()); + let mut output = final_backend_output(); + output.finish_reason = Some(common::FinishReason::Cancelled); + + let response = generator + .choice_from_postprocessor(output) + .expect("choice generation"); + let response_json = serde_json::to_value(response).expect("serialize response"); + + assert_eq!(response_json["choices"][0]["finish_reason"], "stop"); + assert_eq!( + response_json["nvext"]["detailed_finish_reason"], + "cancelled" + ); + } + #[test] fn test_logprobs_zero_emits_chosen_token_logprob() { let mut request = create_test_request(); diff --git a/lib/llm/src/protocols/openai/delta_common.rs b/lib/llm/src/protocols/openai/delta_common.rs index 469c7123ce1a..0ee9312e3fb6 100644 --- a/lib/llm/src/protocols/openai/delta_common.rs +++ b/lib/llm/src/protocols/openai/delta_common.rs @@ -152,10 +152,6 @@ impl DeltaGeneratorState { // the embedding sequence length computed by the worker. if let Some(completion_usage) = output.completion_usage.as_ref() { self.usage.prompt_tokens = completion_usage.prompt_tokens; - self.usage.completion_tokens = self - .usage - .completion_tokens - .max(completion_usage.completion_tokens); if let Some(prompt_details) = completion_usage.prompt_tokens_details.as_ref() { self.usage.prompt_tokens_details = Some(prompt_details.clone()); diff --git a/tests/gpu_memory_service/common/runtime.py b/tests/gpu_memory_service/common/runtime.py index 62403a182c61..54774b38cb05 100644 --- a/tests/gpu_memory_service/common/runtime.py +++ b/tests/gpu_memory_service/common/runtime.py @@ -200,9 +200,23 @@ def _request_engine( payload: dict, timeout: int, action: str, + ) -> dict: + return self._post_engine( + f"control/{route}", + payload, + timeout, + action, + ) + + def _post_engine( + self, + route: str, + payload: dict, + timeout: int, + action: str, ) -> dict: response = requests.post( - f"http://localhost:{self.system_port}/engine/control/{route}", + f"http://localhost:{self.system_port}/engine/{route}", json=payload, timeout=timeout, ) @@ -398,9 +412,6 @@ def pause_payload(self) -> dict: class SGLangWithGMSProcess(GMSEngineProcess): - pause_route = "release_memory_occupation" - resume_route = "resume_memory_occupation" - def __init__( self, request, @@ -456,3 +467,32 @@ def env_updates(self) -> dict[str, str]: def pause_payload(self) -> dict: return {} + + def pause(self) -> dict: + result = self._post_engine( + "pause_generation", + {"mode": "retract"}, + 30, + "pause generation", + ) + self._post_engine( + "release_memory_occupation", + self.pause_payload(), + 30, + "release memory occupation", + ) + return result + + def resume(self, timeout: int = 30) -> dict: + self._post_engine( + "resume_memory_occupation", + self.resume_payload(), + timeout, + "resume memory occupation", + ) + return self._post_engine( + "continue_generation", + {}, + timeout, + "continue generation", + ) diff --git a/tests/test_sglang_gms_runtime_routes.py b/tests/test_sglang_gms_runtime_routes.py new file mode 100644 index 000000000000..8c5a7e05351e --- /dev/null +++ b/tests/test_sglang_gms_runtime_routes.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import Mock, call + +import pytest + +from tests.gpu_memory_service.common.runtime import SGLangWithGMSProcess + +pytestmark = [ + pytest.mark.pre_merge, + pytest.mark.gpu_0, + pytest.mark.unit, + pytest.mark.sglang, + pytest.mark.fault_tolerance, + pytest.mark.parallel, +] + + +def test_sglang_pause_resume_use_native_engine_routes(): + process = object.__new__(SGLangWithGMSProcess) + process._post_engine = Mock( + side_effect=[ + {"status": "ok"}, + None, + None, + {"status": "ok"}, + ] + ) + + assert process.pause() == {"status": "ok"} + assert process.resume(timeout=45) == {"status": "ok"} + assert process._post_engine.call_args_list == [ + call("pause_generation", {"mode": "retract"}, 30, "pause generation"), + call("release_memory_occupation", {}, 30, "release memory occupation"), + call("resume_memory_occupation", {}, 45, "resume memory occupation"), + call("continue_generation", {}, 45, "continue generation"), + ]