diff --git a/nemo_gym/openai_utils.py b/nemo_gym/openai_utils.py index 47f85caf76..8cd5685821 100644 --- a/nemo_gym/openai_utils.py +++ b/nemo_gym/openai_utils.py @@ -624,3 +624,18 @@ async def create_tokenize(self, **kwargs): await self._raise_for_status(response, request_kwargs) return await get_response_json(response) + + async def create_generate(self, **kwargs): + # SGLang's native generation endpoint. The public Chat-Completions + # response contract does not guarantee exact sampled integer token IDs, + # while /generate with return_logprob=True exposes selected-token IDs + # and logprobs. This endpoint lives at the server root, not under /v1. + base_url = self.base_url.removesuffix("/v1") + request_kwargs = dict( + url=f"{base_url}/generate", + json=kwargs, + ) + response = await self._request(method="POST", **request_kwargs) + + await self._raise_for_status(response, request_kwargs) + return await get_response_json(response) diff --git a/responses_api_models/sglang_model/README.md b/responses_api_models/sglang_model/README.md new file mode 100644 index 0000000000..59d54fb487 --- /dev/null +++ b/responses_api_models/sglang_model/README.md @@ -0,0 +1,53 @@ +# SGLang model server + +This Responses-API model server connects NeMo Gym to an SGLang server managed +outside Gym. It preserves the exact prompt IDs, sampled token IDs, and sampled +token logprobs required by token-level RL training. + +The adapter currently renders the chat template locally and calls SGLang's +native `/generate` endpoint with `return_logprob=true`. Gym's public +Chat-Completions response contract does not guarantee exact sampled integer +token IDs, so moving this transport to `/v1/chat/completions` would currently +lose a required training invariant. If that endpoint gains a stable +token-ID/logprob contract, the transport can change without changing the +session-splice or context-overflow rules below. + +For a multi-turn session, the adapter caches the token sequence and splices +each prior assistant turn's exact sampled IDs into the next prompt. It never +re-tokenizes those sampled turns. Tools and chat-template kwargs must therefore +remain fixed for the life of a session; the adapter fails loudly if they +change. If a prompt already fills `context_length`, the adapter returns a +terminal response with `finish_reason="length"` instead of truncating the +prefix. Both behaviors preserve the trainer's prefix contiguity invariant. + +The session cache is process-local. Run one Gym worker per model-server +instance, or provide sticky routing that keeps every turn of a session on the +same worker. + +## Configuration + +See `configs/sglang_model_for_training.yaml`. + +- `base_url`: the SGLang URL; either a bare server URL or one ending in `/v1`. +- `model`: a tokenizer/model identifier available to the Gym server. +- `context_length`: the SGLang server's total context limit. +- `sglang_chat_template`: optional inline copy of the server/training template. +- `sglang_chat_template_path`: optional path to the exact server chat template. +- `sglang_tool_format`: `hermes` or `qwen3_coder`. +- `trust_remote_code`: forwarded to the local tokenizer loader; defaults to + `false`. + +The leaf package pins `transformers==5.6.0`, matching NeMo RL's SGLang worker +environment. The example config leaves `context_length` mandatory (`???`) so a +mismatched server limit cannot be selected silently. + +## CPU-only tests + +From the Gym checkout: + +```bash +uv run --extra dev pytest responses_api_models/sglang_model/tests +``` + +The direct tests mock the tokenizer and SGLang HTTP client; they do not load +weights or require a GPU. diff --git a/responses_api_models/sglang_model/__init__.py b/responses_api_models/sglang_model/__init__.py new file mode 100644 index 0000000000..52a7a9daf0 --- /dev/null +++ b/responses_api_models/sglang_model/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/responses_api_models/sglang_model/_logic.py b/responses_api_models/sglang_model/_logic.py new file mode 100644 index 0000000000..19c6c93e13 --- /dev/null +++ b/responses_api_models/sglang_model/_logic.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Framework-free helpers for the SGLang model server.""" + +from typing import Any, Dict, List, Tuple + + +def _extract_output_ids( + result: Dict[str, Any], + meta: Dict[str, Any], +) -> List[int] | None: + output_ids = meta.get("output_ids") + if output_ids is None: + output_ids = result.get("output_ids") + if output_ids is None: + return None + if not isinstance(output_ids, (list, tuple)): + raise RuntimeError(f"Malformed SGLang output_ids field: expected an array, got {type(output_ids).__name__}") + return [int(token_id) for token_id in output_ids] + + +def _validate_selected_ids( + selected_ids: List[int], + output_ids: List[int] | None, +) -> None: + if output_ids is not None and selected_ids != output_ids: + raise RuntimeError( + "SGLang returned generated-token IDs that do not match output_ids: " + f"selected={selected_ids!r}, output_ids={output_ids!r}" + ) + + +def extract_generated_tokens_and_logprobs( + result: Dict[str, Any], +) -> Tuple[List[int], List[float]]: + """Return aligned generated token IDs and logprobs from ``/generate``. + + SGLang releases have emitted the selected-token data as dictionaries, + tuples, or parallel value/index arrays. Missing or malformed data is a + hard error because silently returning empty arrays would invalidate the + training loss mask. + """ + meta = result.get("meta_info") or {} + output_ids = _extract_output_ids(result, meta) + if "output_token_logprobs" in meta: + entries = meta["output_token_logprobs"] + if not isinstance(entries, (list, tuple)): + raise RuntimeError( + f"Malformed SGLang output_token_logprobs field: expected an array, got {type(entries).__name__}" + ) + if not entries: + _validate_selected_ids([], output_ids) + return [], [] + token_ids: List[int] = [] + logprobs: List[float] = [] + for entry in entries: + if isinstance(entry, dict): + token_id = entry.get("token_id", entry.get("id")) + logprob = entry.get("logprob") + elif isinstance(entry, (list, tuple)) and len(entry) >= 2: + logprob, token_id = entry[0], entry[1] + else: + raise RuntimeError(f"Malformed SGLang output_token_logprobs entry: {entry!r}") + if token_id is None or logprob is None: + raise RuntimeError(f"Malformed SGLang output_token_logprobs entry: {entry!r}") + token_ids.append(int(token_id)) + logprobs.append(float(logprob)) + _validate_selected_ids(token_ids, output_ids) + return token_ids, logprobs + + values_present = "output_token_logprobs_val" in meta or "output_token_logprobs_val" in result + values = meta.get( + "output_token_logprobs_val", + result.get("output_token_logprobs_val"), + ) + indexes = meta.get("output_token_logprobs_idx") + if indexes is None: + indexes = result.get("output_token_logprobs_idx") + + if values_present: + if not isinstance(values, (list, tuple)): + raise RuntimeError( + f"Malformed SGLang output_token_logprobs_val field: expected an array, got {type(values).__name__}" + ) + selected_ids = indexes if indexes is not None else output_ids + if not values: + normalized_ids = [int(token_id) for token_id in selected_ids] if selected_ids is not None else [] + if normalized_ids: + raise RuntimeError( + f"SGLang returned mismatched generation fields: {len(normalized_ids)} ids for 0 logprobs" + ) + _validate_selected_ids(normalized_ids, output_ids) + return [], [] + if not selected_ids or len(selected_ids) != len(values): + id_count = len(selected_ids) if selected_ids is not None else 0 + raise RuntimeError( + f"SGLang returned mismatched generation fields: {id_count} ids for {len(values)} logprobs" + ) + normalized_ids = [int(token_id) for token_id in selected_ids] + _validate_selected_ids(normalized_ids, output_ids) + return normalized_ids, [float(logprob) for logprob in values] + + result_keys = sorted(result) + meta_keys = sorted(meta) + raise RuntimeError( + "SGLang /generate returned no generated-token logprobs " + f"(result keys={result_keys}, meta_info keys={meta_keys}). " + "Ensure return_logprob=true is supported and honored." + ) diff --git a/responses_api_models/sglang_model/app.py b/responses_api_models/sglang_model/app.py new file mode 100644 index 0000000000..80f6999543 --- /dev/null +++ b/responses_api_models/sglang_model/app.py @@ -0,0 +1,517 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NeMo Gym model server backed by SGLang's native ``/generate`` endpoint.""" + +import json +import re +from copy import deepcopy +from time import time +from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple +from uuid import uuid4 + +from aiohttp.client_exceptions import ClientResponseError +from fastapi import Request +from pydantic import Field + +from nemo_gym.base_responses_api_model import Body +from nemo_gym.openai_utils import ( + NeMoGymChatCompletion, + NeMoGymChatCompletionCreateParamsNonStreaming, +) +from nemo_gym.server_utils import SESSION_ID_KEY, is_nemo_gym_fastapi_entrypoint +from responses_api_models.sglang_model._logic import extract_generated_tokens_and_logprobs +from responses_api_models.sglang_model.tool_parsers import ( + normalize_tool_call_arguments, + parse_qwen3_coder_tool_calls, +) +from responses_api_models.vllm_model.app import VLLMModel, VLLMModelConfig + + +class SGLangModelConfig(VLLMModelConfig): + """Configuration for exact-token SGLang generation.""" + + context_length: int = Field(gt=0) + trust_remote_code: bool = False + sglang_chat_template: Optional[str] = None + sglang_chat_template_path: Optional[str] = None + sglang_tool_format: Literal["hermes", "qwen3_coder"] = "hermes" + + +class SGLangModel(VLLMModel): + """Responses-API adapter that preserves exact sampled token IDs.""" + + config: SGLangModelConfig + + _SGLANG_TOOL_CALL_PATTERN: ClassVar = re.compile( + r"\s*(\{.*?\})\s*", + re.DOTALL, + ) + _SGLANG_ARGS_PATTERN: ClassVar = re.compile( + r'"arguments"\s*:\s*(.*)\}\s*$', + re.DOTALL, + ) + _SGLANG_EOS_MARKERS: ClassVar = ("<|im_end|>", "<|endoftext|>") + + def _post_init(self) -> None: + super()._post_init() + self._sglang_tokenizer: Any = None + self._sglang_chat_template: Optional[str] = None + self._sglang_session_seq: Dict[str, Dict[str, Any]] = {} + self._sglang_eos_nl_ids: Optional[List[int]] = None + + def _get_sglang_tokenizer(self) -> Any: + if self._sglang_tokenizer is None: + from transformers import AutoTokenizer + + self._sglang_tokenizer = AutoTokenizer.from_pretrained( + self.config.model, + trust_remote_code=self.config.trust_remote_code, + ) + return self._sglang_tokenizer + + def _get_sglang_chat_template(self) -> Optional[str]: + if self.config.sglang_chat_template is not None: + return self.config.sglang_chat_template + if self._sglang_chat_template is None and self.config.sglang_chat_template_path: + with open(self.config.sglang_chat_template_path) as template_file: + self._sglang_chat_template = template_file.read() + return self._sglang_chat_template + + async def chat_completions( + self, + request: Request, + body: NeMoGymChatCompletionCreateParamsNonStreaming = Body(), + ) -> NeMoGymChatCompletion: + """Generate without applying the vLLM-specific request preprocessing.""" + return await self._sglang_chat_completion( + request, + body.model_dump(exclude_unset=True), + ) + + def _full_sglang_tokenize( + self, + messages: List[Any], + tools: Any, + chat_template_kwargs: Dict[str, Any], + ) -> List[int]: + """Render and tokenize a complete prompt on a cache miss.""" + encoded = self._get_sglang_tokenizer().apply_chat_template( + normalize_tool_call_arguments(messages), + tools=tools, + chat_template=self._get_sglang_chat_template(), + add_generation_prompt=True, + tokenize=True, + **chat_template_kwargs, + ) + if isinstance(encoded, dict) or hasattr(encoded, "input_ids"): + encoded = encoded["input_ids"] + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + if encoded and isinstance(encoded[0], (list, tuple)): + encoded = encoded[0] + return [int(token_id) for token_id in encoded] + + def _sglang_eos_nl(self) -> List[int]: + if self._sglang_eos_nl_ids is None: + encoded = self._get_sglang_tokenizer()( + "<|im_end|>\n", + add_special_tokens=False, + ) + self._sglang_eos_nl_ids = [int(token_id) for token_id in encoded["input_ids"]] + return self._sglang_eos_nl_ids + + def _sglang_followup_fragment_ids( + self, + new_messages: List[Any], + chat_template_kwargs: Dict[str, Any], + ) -> Optional[List[int]]: + """Render the new messages and next assistant header as a token fragment. + + The fragment is derived by differencing two template renders against an + anchor assistant turn. Returning ``None`` asks the caller to fall back + to a complete render when the template is not splice-friendly. + """ + tokenizer = self._get_sglang_tokenizer() + chat_template = self._get_sglang_chat_template() + anchor = [{"role": "assistant", "content": "X"}] + try: + full = tokenizer.apply_chat_template( + anchor + list(new_messages), + tools=None, + chat_template=chat_template, + add_generation_prompt=True, + tokenize=False, + **chat_template_kwargs, + ) + base = tokenizer.apply_chat_template( + anchor, + tools=None, + chat_template=chat_template, + add_generation_prompt=False, + tokenize=False, + **chat_template_kwargs, + ) + except Exception: + return None + if not isinstance(full, str) or not isinstance(base, str) or not full.startswith(base): + return None + encoded = tokenizer(full[len(base) :], add_special_tokens=False) + return [int(token_id) for token_id in encoded["input_ids"]] + + @staticmethod + def _sglang_msg_sig(message: Dict[str, Any]) -> Tuple[Any, str, str]: + return ( + message.get("role"), + json.dumps(message.get("content"), sort_keys=True, default=str), + json.dumps(message.get("tool_calls"), sort_keys=True, default=str), + ) + + @classmethod + def _sglang_messages_match( + cls, + left: List[Any], + right: List[Any], + ) -> bool: + return len(left) == len(right) and all( + cls._sglang_msg_sig(left_message) == cls._sglang_msg_sig(right_message) + for left_message, right_message in zip(left, right) + ) + + def _sglang_rendering_sig( + self, + tools: Any, + chat_template_kwargs: Dict[str, Any], + ) -> Tuple[str, str, Optional[str]]: + """Identify inputs that affect the cached prompt rendering.""" + return ( + json.dumps(tools, sort_keys=True, default=str), + json.dumps(chat_template_kwargs, sort_keys=True, default=str), + self._get_sglang_chat_template(), + ) + + def _build_sglang_prompt_ids( + self, + request: Request, + messages: List[Any], + tools: Any, + chat_template_kwargs: Dict[str, Any], + ) -> Tuple[List[int], Optional[str]]: + """Build a prompt, splicing the preceding turn's exact sampled IDs.""" + try: + session_id = request.session.get(SESSION_ID_KEY) + except Exception: + session_id = None + if session_id is not None: + state = self._sglang_session_seq.get(session_id) + rendering_sig = self._sglang_rendering_sig( + tools, + chat_template_kwargs, + ) + if state is not None and state.get("rendering_sig") != rendering_sig: + raise RuntimeError( + "SGLang session tools or chat-template inputs changed after " + "sampled tokens were cached. Start a new session instead of " + "re-tokenizing the existing trajectory." + ) + if state is not None: + previous_messages = state["messages"] + previous_count = len(previous_messages) + if ( + len(messages) > previous_count + and messages[previous_count].get("role") == "assistant" + and all(message.get("role") != "assistant" for message in messages[previous_count + 1 :]) + and self._sglang_messages_match( + messages[:previous_count], + previous_messages, + ) + ): + fragment = self._sglang_followup_fragment_ids( + messages[previous_count + 1 :], + chat_template_kwargs, + ) + if fragment is not None: + return state["seq"] + fragment, session_id + return ( + self._full_sglang_tokenize( + messages, + tools, + chat_template_kwargs, + ), + session_id, + ) + + def _update_sglang_session_seq( + self, + session_id: Optional[str], + messages: List[Any], + prompt_token_ids: List[int], + generation_token_ids: List[int], + tools: Any, + chat_template_kwargs: Dict[str, Any], + ) -> None: + """Cache the exact token sequence through the generated assistant turn.""" + if session_id is None: + return + eos_newline_ids = self._sglang_eos_nl() + sequence = list(prompt_token_ids) + list(generation_token_ids) + max_overlap = min(len(sequence), len(eos_newline_ids)) + overlap = next( + ( + overlap_size + for overlap_size in range(max_overlap, 0, -1) + if sequence[-overlap_size:] == eos_newline_ids[:overlap_size] + ), + 0, + ) + sequence += eos_newline_ids[overlap:] + + self._sglang_session_seq.pop(session_id, None) + while len(self._sglang_session_seq) >= 8192: + self._sglang_session_seq.pop(next(iter(self._sglang_session_seq)), None) + self._sglang_session_seq[session_id] = { + "messages": list(messages), + "seq": sequence, + "rendering_sig": self._sglang_rendering_sig( + tools, + chat_template_kwargs, + ), + } + + def _parse_sglang_generation( + self, + text: str, + tools: Optional[List[Dict[str, Any]]] = None, + ) -> Tuple[Optional[str], str, List[Dict[str, Any]]]: + """Reconstruct reasoning, visible content, and tool calls from raw text.""" + reasoning_content: Optional[str] = None + if self.config.uses_reasoning_parser and "" in text: + reasoning_content, _, remainder = text.partition("") + else: + remainder = text + + if self.config.sglang_tool_format == "qwen3_coder": + tool_calls, content = parse_qwen3_coder_tool_calls(remainder, tools) + return reasoning_content, content, tool_calls + + tool_calls: List[Dict[str, Any]] = [] + for match in self._SGLANG_TOOL_CALL_PATTERN.finditer(remainder): + block = match.group(1) + try: + parsed = json.loads(block) + except json.JSONDecodeError: + continue + arguments_match = self._SGLANG_ARGS_PATTERN.search(block) + arguments = ( + arguments_match.group(1).strip() + if arguments_match is not None + else json.dumps(parsed.get("arguments", {})) + ) + tool_calls.append( + { + "id": f"call_{uuid4().hex}", + "type": "function", + "function": { + "name": parsed.get("name"), + "arguments": arguments, + }, + } + ) + + content = self._SGLANG_TOOL_CALL_PATTERN.sub("", remainder).strip() + return reasoning_content, content, tool_calls + + def _sglang_length_finish( + self, + prompt_token_ids: List[int], + ) -> NeMoGymChatCompletion: + """Terminate an over-context turn without truncating its prompt.""" + message: Dict[str, Any] = { + "role": "assistant", + "content": None, + "tool_calls": None, + } + if self.config.return_token_id_information: + message.update( + { + "prompt_token_ids": list(prompt_token_ids), + "generation_token_ids": [], + "generation_log_probs": [], + } + ) + return NeMoGymChatCompletion.model_validate( + { + "id": f"chtcmpl-{uuid4().hex}", + "object": "chat.completion", + "created": int(time()), + "model": self.config.model, + "choices": [ + { + "index": 0, + "finish_reason": "length", + "message": message, + "logprobs": None, + } + ], + "usage": { + "prompt_tokens": len(prompt_token_ids), + "completion_tokens": 0, + "total_tokens": len(prompt_token_ids), + }, + } + ) + + async def _sglang_chat_completion( + self, + request: Request, + body_dict: Dict[str, Any], + ) -> NeMoGymChatCompletion: + """Generate exact training tokens through SGLang's native endpoint.""" + client = self._resolve_client(request) + + messages = body_dict["messages"] + if self.config.replace_developer_role_with_system: + for message in messages: + if message.get("role") == "developer": + message["role"] = "system" + tools = body_dict.get("tools") + + chat_template_kwargs: Dict[str, Any] = {} + if self.config.chat_template_kwargs: + chat_template_kwargs = deepcopy(self.config.chat_template_kwargs) + metadata = body_dict.get("metadata") or {} + chat_template_kwargs.update( + json.loads(metadata.get("chat_template_kwargs", "{}")), + ) + + tokenizer = self._get_sglang_tokenizer() + prompt_token_ids, session_id = self._build_sglang_prompt_ids( + request, + messages, + tools, + chat_template_kwargs, + ) + + remaining_context = self.config.context_length - len(prompt_token_ids) + if remaining_context <= 0: + return self._sglang_length_finish(prompt_token_ids) + + sampling_params: Dict[str, Any] = {"spaces_between_special_tokens": False} + max_new_tokens = body_dict.get("max_completion_tokens") or body_dict.get("max_tokens") or None + if max_new_tokens is None: + max_new_tokens = remaining_context - 8 + max_new_tokens = max(1, max_new_tokens) + else: + max_new_tokens = min(max_new_tokens, remaining_context) + sampling_params["max_new_tokens"] = max_new_tokens + for key in ("temperature", "top_p", "top_k", "stop"): + if body_dict.get(key) is not None: + sampling_params[key] = body_dict[key] + + try: + result = await client.create_generate( + input_ids=prompt_token_ids, + sampling_params=sampling_params, + return_logprob=True, + ) + except ClientResponseError as error: + try: + error_body = error.response_content.decode() + except Exception: + error_body = str(error) + if any( + fragment in error_body + for fragment in ( + "context length", + "longer than", + "max_total", + "is longer", + ) + ): + return self._sglang_length_finish(prompt_token_ids) + raise + + meta_info = result.get("meta_info") or {} + generation_token_ids, generation_log_probs = extract_generated_tokens_and_logprobs( + result, + ) + self._update_sglang_session_seq( + session_id, + messages, + prompt_token_ids, + generation_token_ids, + tools, + chat_template_kwargs, + ) + + generated_text = tokenizer.decode( + generation_token_ids, + skip_special_tokens=False, + spaces_between_special_tokens=False, + ) + stripped = True + while stripped: + stripped = False + generated_text = generated_text.rstrip("\n") + for eos_marker in self._SGLANG_EOS_MARKERS: + if generated_text.endswith(eos_marker): + generated_text = generated_text[: -len(eos_marker)] + stripped = True + reasoning_content, content, tool_calls = self._parse_sglang_generation( + generated_text, + tools=tools, + ) + + finish = meta_info.get("finish_reason") + if isinstance(finish, dict): + finish = finish.get("type") + if finish == "length": + finish_reason = "length" + elif tool_calls: + finish_reason = "tool_calls" + else: + finish_reason = "stop" + + if self.config.uses_reasoning_parser and reasoning_content: + content = self._converter._wrap_reasoning_in_think_tags([reasoning_content]) + (content or "") + + message: Dict[str, Any] = { + "role": "assistant", + "content": content or None, + "tool_calls": tool_calls or None, + } + if self.config.return_token_id_information: + message.update( + { + "prompt_token_ids": prompt_token_ids, + "generation_token_ids": generation_token_ids, + "generation_log_probs": generation_log_probs, + } + ) + + return NeMoGymChatCompletion.model_validate( + { + "id": f"chtcmpl-{uuid4().hex}", + "object": "chat.completion", + "created": int(time()), + "model": self.config.model, + "choices": [ + { + "index": 0, + "finish_reason": finish_reason, + "message": message, + "logprobs": None, + } + ], + "usage": { + "prompt_tokens": len(prompt_token_ids), + "completion_tokens": len(generation_token_ids), + "total_tokens": len(prompt_token_ids) + len(generation_token_ids), + }, + } + ) + + +if __name__ == "__main__": + SGLangModel.run_webserver() +elif is_nemo_gym_fastapi_entrypoint(__file__): + app = SGLangModel.run_webserver() # noqa: F401 diff --git a/responses_api_models/sglang_model/configs/sglang_model_for_training.yaml b/responses_api_models/sglang_model/configs/sglang_model_for_training.yaml new file mode 100644 index 0000000000..2278050cc9 --- /dev/null +++ b/responses_api_models/sglang_model/configs/sglang_model_for_training.yaml @@ -0,0 +1,12 @@ +# SGLang-backed policy model server for token-level RL training. +policy_model: + responses_api_models: + sglang_model: + entrypoint: app.py + base_url: ${policy_base_url} + api_key: ${policy_api_key} + model: ${policy_model_name} + # Required: set this to the SGLang server's max total sequence length. + context_length: ??? + return_token_id_information: true + uses_reasoning_parser: true diff --git a/responses_api_models/sglang_model/pyproject.toml b/responses_api_models/sglang_model/pyproject.toml new file mode 100644 index 0000000000..811385f849 --- /dev/null +++ b/responses_api_models/sglang_model/pyproject.toml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[project] +name = "sglang-model" +version = "0.2.0rc0" +requires-python = ">=3.12" +dependencies = [ + "nemo-gym[dev]", + # Must match the transformers pin in the NeMo-RL sglang extra. This adapter + # renders the chat template and tokenizes locally, so a tokenizer or + # template delta between adapter and server shows up as a rollout + # contiguity failure rather than an import error. + "transformers==5.8.1", +] + +[build-system] +build-backend = "setuptools.build_meta" +requires = ["setuptools>=61", "setuptools-scm"] + +[tool.setuptools.packages.find] +where = [".."] +include = ["sglang_model"] + +[tool.uv.sources] +nemo-gym = { path = "../..", editable = true } diff --git a/responses_api_models/sglang_model/tests/__init__.py b/responses_api_models/sglang_model/tests/__init__.py new file mode 100644 index 0000000000..52a7a9daf0 --- /dev/null +++ b/responses_api_models/sglang_model/tests/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/responses_api_models/sglang_model/tests/test_app.py b/responses_api_models/sglang_model/tests/test_app.py new file mode 100644 index 0000000000..a4d541b8a4 --- /dev/null +++ b/responses_api_models/sglang_model/tests/test_app.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from nemo_gym.openai_utils import NeMoGymChatCompletionCreateParamsNonStreaming +from nemo_gym.server_utils import SESSION_ID_KEY, ServerClient +from responses_api_models.sglang_model.app import SGLangModel, SGLangModelConfig + + +class FakeTokenizer: + def __init__(self, full_prompt_ids: list[int], decoded: str = "answer") -> None: + self.full_prompt_ids = full_prompt_ids + self.decoded = decoded + self.decode_calls: list[dict] = [] + + def apply_chat_template( + self, + messages, + *, + tools=None, + chat_template=None, + add_generation_prompt, + tokenize, + **kwargs, + ): + if tokenize: + return list(self.full_prompt_ids) + if len(messages) == 1 and messages[0] == {"role": "assistant", "content": "X"}: + assert add_generation_prompt is False + return "ANCHOR" + assert add_generation_prompt is True + return "ANCHORFOLLOWUP" + + def __call__(self, text: str, *, add_special_tokens: bool): + assert add_special_tokens is False + if text == "<|im_end|>\n": + return {"input_ids": [90, 91]} + if text == "FOLLOWUP": + return {"input_ids": [30, 31]} + raise AssertionError(f"unexpected tokenization input: {text!r}") + + def decode(self, token_ids, *, skip_special_tokens: bool, spaces_between_special_tokens: bool): + self.decode_calls.append( + { + "token_ids": list(token_ids), + "skip_special_tokens": skip_special_tokens, + "spaces_between_special_tokens": spaces_between_special_tokens, + } + ) + return self.decoded + + +class FakeSGLangClient: + def __init__(self, result: dict) -> None: + self.result = result + self.calls: list[dict] = [] + + async def create_generate(self, **kwargs): + self.calls.append(kwargs) + return self.result + + +def make_model( + *, + context_length: int = 64, + tokenizer: FakeTokenizer | None = None, + client: FakeSGLangClient | None = None, +) -> SGLangModel: + config = SGLangModelConfig( + host="0.0.0.0", + port=8080, + entrypoint="", + name="sglang_model", + base_url="http://localhost:30000/v1", + api_key="unused", # pragma: allowlist secret + model="local-tokenizer", + context_length=context_length, + return_token_id_information=True, + uses_reasoning_parser=True, + ) + model = SGLangModel( + config=config, + server_client=MagicMock(spec=ServerClient, global_config_dict={}), + ) + if tokenizer is not None: + model._sglang_tokenizer = tokenizer + if client is not None: + model._clients = [client] + return model + + +async def test_generate_path_preserves_training_ids_reasoning_and_tools() -> None: + tokenizer = FakeTokenizer( + [1, 2, 3], + decoded=( + 'private reasoning{"name":"shell","arguments":{"command":"ls"}}<|im_end|>' + ), + ) + client = FakeSGLangClient( + { + "meta_info": { + "output_ids": [11, 12], + "output_token_logprobs": [ + {"token_id": 11, "logprob": -0.1}, + {"id": 12, "logprob": -0.2}, + ], + "finish_reason": {"type": "stop"}, + } + } + ) + model = make_model(tokenizer=tokenizer, client=client) + body = NeMoGymChatCompletionCreateParamsNonStreaming( + messages=[{"role": "user", "content": "inspect"}], + max_tokens=8, + ) + request = SimpleNamespace(session={SESSION_ID_KEY: "session-1"}) + + response = await model.chat_completions(request, body) + + assert client.calls == [ + { + "input_ids": [1, 2, 3], + "sampling_params": { + "spaces_between_special_tokens": False, + "max_new_tokens": 8, + }, + "return_logprob": True, + } + ] + choice = response.choices[0] + assert choice.finish_reason == "tool_calls" + assert choice.message.content == "private reasoning" + assert choice.message.prompt_token_ids == [1, 2, 3] + assert choice.message.generation_token_ids == [11, 12] + assert choice.message.generation_log_probs == [-0.1, -0.2] + assert choice.message.tool_calls[0].function.name == "shell" + assert json_loads(choice.message.tool_calls[0].function.arguments) == {"command": "ls"} + assert tokenizer.decode_calls == [ + { + "token_ids": [11, 12], + "skip_special_tokens": False, + "spaces_between_special_tokens": False, + } + ] + + +def json_loads(value: str) -> dict: + import json + + return json.loads(value) + + +def test_followup_prompt_splices_exact_sampled_ids() -> None: + tokenizer = FakeTokenizer([10, 11]) + model = make_model(tokenizer=tokenizer) + request = SimpleNamespace(session={SESSION_ID_KEY: "session-2"}) + first_messages = [{"role": "user", "content": "first"}] + + first_prompt, session_id = model._build_sglang_prompt_ids( + request, + first_messages, + tools=None, + chat_template_kwargs={}, + ) + model._update_sglang_session_seq( + session_id, + first_messages, + first_prompt, + generation_token_ids=[20, 21], + tools=None, + chat_template_kwargs={}, + ) + followup_messages = [ + *first_messages, + {"role": "assistant", "content": "a decode that must not be re-tokenized"}, + {"role": "user", "content": "continue"}, + ] + + followup_prompt, _ = model._build_sglang_prompt_ids( + request, + followup_messages, + tools=None, + chat_template_kwargs={}, + ) + + assert followup_prompt == [10, 11, 20, 21, 90, 91, 30, 31] + + +@pytest.mark.parametrize( + ("generation_token_ids", "expected_sequence"), + [ + ([20], [1, 20, 90, 91]), + ([20, 90], [1, 20, 90, 91]), + ([20, 90, 91], [1, 20, 90, 91]), + ], +) +def test_session_cache_appends_only_missing_eos_suffix( + generation_token_ids: list[int], + expected_sequence: list[int], +) -> None: + model = make_model(tokenizer=FakeTokenizer([1])) + + model._update_sglang_session_seq( + "session-eos", + [{"role": "user", "content": "first"}], + prompt_token_ids=[1], + generation_token_ids=generation_token_ids, + tools=None, + chat_template_kwargs={}, + ) + + assert model._sglang_session_seq["session-eos"]["seq"] == expected_sequence + + +@pytest.mark.parametrize("max_tokens_field", ["max_completion_tokens", "max_tokens"]) +async def test_explicit_max_tokens_is_clamped_to_remaining_context( + max_tokens_field: str, +) -> None: + tokenizer = FakeTokenizer([1, 2, 3]) + client = FakeSGLangClient( + { + "meta_info": { + "output_ids": [11], + "output_token_logprobs": [ + {"token_id": 11, "logprob": -0.1}, + ], + } + } + ) + model = make_model( + context_length=10, + tokenizer=tokenizer, + client=client, + ) + body = NeMoGymChatCompletionCreateParamsNonStreaming( + messages=[{"role": "user", "content": "inspect"}], + **{max_tokens_field: 10}, + ) + + await model.chat_completions( + SimpleNamespace(session={SESSION_ID_KEY: "session-clamp"}), + body, + ) + + assert client.calls[0]["sampling_params"]["max_new_tokens"] == 7 + + +async def test_over_context_terminates_without_calling_generate() -> None: + tokenizer = FakeTokenizer([1, 2, 3, 4]) + client = FakeSGLangClient({"must": "not be used"}) + model = make_model(context_length=4, tokenizer=tokenizer, client=client) + body = NeMoGymChatCompletionCreateParamsNonStreaming( + messages=[{"role": "user", "content": "too long"}], + ) + request = SimpleNamespace(session={SESSION_ID_KEY: "session-3"}) + + response = await model.chat_completions(request, body) + + assert client.calls == [] + assert response.choices[0].finish_reason == "length" + assert response.choices[0].message.prompt_token_ids == [1, 2, 3, 4] + assert response.choices[0].message.generation_token_ids == [] + + +def test_followup_prompt_rejects_changed_rendering_inputs() -> None: + tokenizer = FakeTokenizer([10, 11]) + model = make_model(tokenizer=tokenizer) + request = SimpleNamespace(session={SESSION_ID_KEY: "session-rendering"}) + first_messages = [{"role": "user", "content": "first"}] + first_prompt, session_id = model._build_sglang_prompt_ids( + request, + first_messages, + tools=[{"type": "function", "function": {"name": "old"}}], + chat_template_kwargs={"enable_thinking": True}, + ) + model._update_sglang_session_seq( + session_id, + first_messages, + first_prompt, + generation_token_ids=[20, 21], + tools=[{"type": "function", "function": {"name": "old"}}], + chat_template_kwargs={"enable_thinking": True}, + ) + followup_messages = [ + *first_messages, + {"role": "assistant", "content": "cached"}, + {"role": "user", "content": "continue"}, + ] + + with pytest.raises(RuntimeError, match="session tools or chat-template"): + model._build_sglang_prompt_ids( + request, + followup_messages, + tools=[{"type": "function", "function": {"name": "new"}}], + chat_template_kwargs={"enable_thinking": True}, + ) + with pytest.raises(RuntimeError, match="session tools or chat-template"): + model._build_sglang_prompt_ids( + request, + followup_messages, + tools=[{"type": "function", "function": {"name": "old"}}], + chat_template_kwargs={"enable_thinking": False}, + ) + + +def test_sglang_config_owns_context_and_tool_format() -> None: + model = make_model(context_length=128) + + assert model.config.context_length == 128 + assert model.config.sglang_tool_format == "hermes" + assert not hasattr(model.config, "engine") + + +def test_inline_chat_template_is_used_directly() -> None: + model = make_model() + model.config.sglang_chat_template = "inline-template" + model.config.sglang_chat_template_path = "/must/not/be/read" + + assert model._get_sglang_chat_template() == "inline-template" diff --git a/responses_api_models/sglang_model/tests/test_logic.py b/responses_api_models/sglang_model/tests/test_logic.py new file mode 100644 index 0000000000..c57c2ca1fe --- /dev/null +++ b/responses_api_models/sglang_model/tests/test_logic.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from responses_api_models.sglang_model._logic import extract_generated_tokens_and_logprobs + + +def test_extracts_tuple_entries() -> None: + result = { + "meta_info": { + "output_ids": [11, 12], + "output_token_logprobs": [ + [-0.1, 11, "a"], + (-0.2, 12, "b"), + ], + } + } + + assert extract_generated_tokens_and_logprobs(result) == ([11, 12], [-0.1, -0.2]) + + +def test_extracts_mapping_entries() -> None: + result = { + "meta_info": { + "output_ids": [21, 22], + "output_token_logprobs": [ + {"token_id": 21, "logprob": -0.3}, + {"id": 22, "logprob": -0.4}, + ], + } + } + + assert extract_generated_tokens_and_logprobs(result) == ([21, 22], [-0.3, -0.4]) + + +@pytest.mark.parametrize("location", ["meta", "result"]) +def test_extracts_split_value_index_fallback(location: str) -> None: + fields = { + "output_ids": [31, 32], + "output_token_logprobs_val": [-0.5, -0.6], + "output_token_logprobs_idx": [31, 32], + } + result = {"meta_info": fields} if location == "meta" else fields + + assert extract_generated_tokens_and_logprobs(result) == ([31, 32], [-0.5, -0.6]) + + +def test_extracts_output_ids_fallback() -> None: + result = { + "output_ids": [41, 42], + "meta_info": {"output_token_logprobs_val": [-0.7, -0.8]}, + } + + assert extract_generated_tokens_and_logprobs(result) == ([41, 42], [-0.7, -0.8]) + + +@pytest.mark.parametrize( + "result", + [ + {"meta_info": {"output_ids": [], "output_token_logprobs": []}}, + { + "meta_info": { + "output_ids": [], + "output_token_logprobs_val": [], + "output_token_logprobs_idx": [], + } + }, + ], +) +def test_present_empty_arrays_are_valid_zero_token_completion(result: dict) -> None: + assert extract_generated_tokens_and_logprobs(result) == ([], []) + + +@pytest.mark.parametrize( + "result", + [ + {}, + {"meta_info": {"output_token_logprobs": [[-0.1]]}}, + {"meta_info": {"output_token_logprobs": [{"token_id": 1}]}}, + { + "meta_info": { + "output_token_logprobs_val": [-0.1, -0.2], + "output_token_logprobs_idx": [1], + } + }, + { + "meta_info": { + "output_ids": [1], + "output_token_logprobs_val": [], + } + }, + ], +) +def test_rejects_missing_or_malformed_training_data(result: dict) -> None: + with pytest.raises(RuntimeError): + extract_generated_tokens_and_logprobs(result) + + +@pytest.mark.parametrize( + "result", + [ + { + "meta_info": { + "output_ids": [1, 2], + "output_token_logprobs": [{"token_id": 1, "logprob": -0.1}], + } + }, + { + "meta_info": { + "output_ids": [1], + "output_token_logprobs": [{"token_id": 2, "logprob": -0.1}], + } + }, + { + "meta_info": { + "output_ids": [1], + "output_token_logprobs_val": [-0.1], + "output_token_logprobs_idx": [2], + } + }, + ], +) +def test_rejects_logprob_ids_that_do_not_match_output_ids(result: dict) -> None: + with pytest.raises(RuntimeError, match="do not match output_ids"): + extract_generated_tokens_and_logprobs(result) diff --git a/responses_api_models/sglang_model/tests/test_tool_parsers.py b/responses_api_models/sglang_model/tests/test_tool_parsers.py new file mode 100644 index 0000000000..338fa6470a --- /dev/null +++ b/responses_api_models/sglang_model/tests/test_tool_parsers.py @@ -0,0 +1,166 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +import json + +from responses_api_models.sglang_model.tool_parsers import ( + normalize_tool_call_arguments, + parse_qwen3_coder_tool_calls, +) + + +TOOL_CALL = ( + "\n" + "\n" + "\n42\n\n" + "\n" + "text spanning\nmultiple lines\n" + "\n" + "\n" + "" +) + + +def test_parses_multiline_tool_call_and_keeps_content() -> None: + tool_calls, content = parse_qwen3_coder_tool_calls(f"reasoning\n{TOOL_CALL}") + + assert content == "reasoning" + assert tool_calls[0]["function"]["name"] == "editor" + assert json.loads(tool_calls[0]["function"]["arguments"]) == { + "line_number": "42", + "command": "text spanning\nmultiple lines", + } + + +def test_uses_tool_schema_for_argument_types() -> None: + tools = [ + { + "type": "function", + "function": { + "name": "editor", + "parameters": { + "properties": { + "line_number": {"type": "integer"}, + "command": {"type": "string"}, + } + }, + }, + } + ] + + tool_calls, _ = parse_qwen3_coder_tool_calls(TOOL_CALL, tools) + + assert json.loads(tool_calls[0]["function"]["arguments"])["line_number"] == 42 + + +def test_no_tool_call_is_passthrough() -> None: + assert parse_qwen3_coder_tool_calls("plain text") == ([], "plain text") + + +def test_parses_multiple_tool_calls() -> None: + second = TOOL_CALL.replace("editor", "shell") + + tool_calls, content = parse_qwen3_coder_tool_calls(f"{TOOL_CALL}\n{second}") + + assert [call["function"]["name"] for call in tool_calls] == ["editor", "shell"] + assert content == "" + + +def test_invalid_typed_value_falls_back_to_string() -> None: + tools = [ + { + "type": "function", + "function": { + "name": "editor", + "parameters": {"properties": {"line_number": {"type": "integer"}}}, + }, + } + ] + text = TOOL_CALL.replace("42", "not-a-number") + + tool_calls, _ = parse_qwen3_coder_tool_calls(text, tools) + + assert json.loads(tool_calls[0]["function"]["arguments"])["line_number"] == "not-a-number" + + +def test_unknown_tool_parameters_remain_strings() -> None: + tool_calls, _ = parse_qwen3_coder_tool_calls(TOOL_CALL, tools=[]) + + assert json.loads(tool_calls[0]["function"]["arguments"])["line_number"] == "42" + + +def test_shell_text_with_angle_brackets_is_preserved() -> None: + text = ( + "\n\n" + "\n" + "grep -rn 'x < y && y > z' src/ | head -5\n" + "\n" + "\n" + ) + + tool_calls, content = parse_qwen3_coder_tool_calls(text) + + assert json.loads(tool_calls[0]["function"]["arguments"]) == { + "command": "grep -rn 'x < y && y > z' src/ | head -5" + } + assert content == "" + + +def test_normalizes_string_arguments_without_mutating_input() -> None: + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [ + { + "type": "function", + "function": {"name": "editor", "arguments": '{"command": "view"}'}, + } + ], + }, + ] + + normalized = normalize_tool_call_arguments(messages) + + assert normalized[1]["tool_calls"][0]["function"]["arguments"] == {"command": "view"} + assert messages[1]["tool_calls"][0]["function"]["arguments"] == '{"command": "view"}' + + +def test_normalizes_parser_output() -> None: + tool_calls, _ = parse_qwen3_coder_tool_calls(TOOL_CALL) + + normalized = normalize_tool_call_arguments([{"role": "assistant", "tool_calls": tool_calls}]) + + assert normalized[0]["tool_calls"][0]["function"]["arguments"] == { + "line_number": "42", + "command": "text spanning\nmultiple lines", + } + + +def test_normalizer_leaves_non_object_arguments_unchanged() -> None: + def message(arguments): + return { + "role": "assistant", + "tool_calls": [ + { + "type": "function", + "function": {"name": "editor", "arguments": arguments}, + } + ], + } + + mapping = message({"command": "view"}) + assert normalize_tool_call_arguments([mapping])[0]["tool_calls"][0]["function"]["arguments"] == {"command": "view"} + for raw in ("not json", "[1, 2]"): + assert normalize_tool_call_arguments([message(raw)])[0]["tool_calls"][0]["function"]["arguments"] == raw diff --git a/responses_api_models/sglang_model/tool_parsers.py b/responses_api_models/sglang_model/tool_parsers.py new file mode 100644 index 0000000000..d8b15674ab --- /dev/null +++ b/responses_api_models/sglang_model/tool_parsers.py @@ -0,0 +1,112 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. +"""Client-side tool-call parsers for the SGLang ``/generate`` path.""" + +import json +import re +from typing import Any, Dict, List, Optional, Tuple +from uuid import uuid4 + + +QWEN3_CODER_TOOL_CALL_PATTERN = re.compile( + r"\s*\n]+)>(.*?)\s*", + re.DOTALL, +) +QWEN3_CODER_PARAM_PATTERN = re.compile( + r"\n]+)>\n?(.*?)\n?", + re.DOTALL, +) + + +def _tool_param_types(tools: Optional[List[Dict[str, Any]]], function_name: str) -> Dict[str, str]: + """Map parameter names to JSON-schema types for one requested tool.""" + for tool in tools or []: + function = tool.get("function", tool) + if function.get("name") != function_name: + continue + properties = (function.get("parameters") or {}).get("properties") or {} + return {key: prop.get("type", "string") for key, prop in properties.items() if isinstance(prop, dict)} + return {} + + +def _coerce_param_value(raw: str, type_str: str) -> Any: + """Best-effort conversion of a raw parameter using its declared type.""" + if type_str in ("integer", "number"): + try: + return int(raw) if type_str == "integer" else float(raw) + except ValueError: + return raw + if type_str == "boolean": + lowered = raw.strip().lower() + if lowered in ("true", "false"): + return lowered == "true" + return raw + if type_str in ("array", "object"): + try: + return json.loads(raw) + except json.JSONDecodeError: + return raw + return raw + + +def parse_qwen3_coder_tool_calls( + text: str, tools: Optional[List[Dict[str, Any]]] = None +) -> Tuple[List[Dict[str, Any]], str]: + """Parse XML-like tool calls into OpenAI-compatible tool-call mappings.""" + tool_calls: List[Dict[str, Any]] = [] + for match in QWEN3_CODER_TOOL_CALL_PATTERN.finditer(text): + function_name = match.group(1).strip() + param_types = _tool_param_types(tools, function_name) + arguments: Dict[str, Any] = {} + for param_match in QWEN3_CODER_PARAM_PATTERN.finditer(match.group(2)): + key = param_match.group(1).strip() + arguments[key] = _coerce_param_value( + param_match.group(2), + param_types.get(key, "string"), + ) + tool_calls.append( + { + "id": f"call_{uuid4().hex}", + "type": "function", + "function": { + "name": function_name, + "arguments": json.dumps(arguments, ensure_ascii=False), + }, + } + ) + return tool_calls, QWEN3_CODER_TOOL_CALL_PATTERN.sub("", text).strip() + + +def normalize_tool_call_arguments(messages: List[Any]) -> List[Any]: + """Decode JSON-object argument strings for chat-template rendering.""" + normalized: List[Any] = [] + for message in messages: + tool_calls = message.get("tool_calls") if isinstance(message, dict) else None + if not tool_calls: + normalized.append(message) + continue + new_tool_calls = [] + for tool_call in tool_calls: + function = tool_call.get("function") if isinstance(tool_call, dict) else None + arguments = function.get("arguments") if isinstance(function, dict) else None + if isinstance(arguments, str): + try: + decoded = json.loads(arguments) + except ValueError: + decoded = None + if isinstance(decoded, dict): + tool_call = dict(tool_call, function=dict(function, arguments=decoded)) + new_tool_calls.append(tool_call) + normalized.append(dict(message, tool_calls=new_tool_calls)) + return normalized