From c024fcdb1ce4c132b649cc38a08921c7e85d77cd Mon Sep 17 00:00:00 2001 From: zhiweiz Date: Thu, 5 Nov 2026 18:21:14 -0800 Subject: [PATCH 1/6] release: 1.2.33 --- CHANGELOG.md | 11 + pyproject.toml | 6 +- .../lib/public_interfaces/sampling_client.py | 19 +- .../training/sdk/_vendored/__init__.py | 1 + .../training/sdk/_vendored/kimi_k3/LICENSE | 52 ++ .../sdk/_vendored/kimi_k3/PROVENANCE.md | 42 + .../sdk/_vendored/kimi_k3/__init__.py | 1 + .../sdk/_vendored/kimi_k3/encoding_k3.py | 647 +++++++++++++++ .../training/sdk/_vendored/kimi_k31/LICENSE | 52 ++ .../sdk/_vendored/kimi_k31/PROVENANCE.md | 35 + .../sdk/_vendored/kimi_k31/__init__.py | 1 + .../sdk/_vendored/kimi_k31/encoding_k3.py | 735 ++++++++++++++++++ src/fireworks/training/sdk/client.py | 17 +- .../training/sdk/kimi_native_tokenizer.py | 408 ++++++++++ .../training/sdk/tokenizer_identity.py | 192 ++++- tests_training/test_kimi_native_tokenizer.py | 214 +++++ tests_training/test_tokenizer_identity.py | 153 ++++ 17 files changed, 2576 insertions(+), 10 deletions(-) create mode 100644 src/fireworks/training/sdk/_vendored/__init__.py create mode 100644 src/fireworks/training/sdk/_vendored/kimi_k3/LICENSE create mode 100644 src/fireworks/training/sdk/_vendored/kimi_k3/PROVENANCE.md create mode 100644 src/fireworks/training/sdk/_vendored/kimi_k3/__init__.py create mode 100644 src/fireworks/training/sdk/_vendored/kimi_k3/encoding_k3.py create mode 100644 src/fireworks/training/sdk/_vendored/kimi_k31/LICENSE create mode 100644 src/fireworks/training/sdk/_vendored/kimi_k31/PROVENANCE.md create mode 100644 src/fireworks/training/sdk/_vendored/kimi_k31/__init__.py create mode 100644 src/fireworks/training/sdk/_vendored/kimi_k31/encoding_k3.py create mode 100644 src/fireworks/training/sdk/kimi_native_tokenizer.py create mode 100644 tests_training/test_kimi_native_tokenizer.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 482969d..461c985 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,16 @@ # Changelog +## 1.2.33 (2026-11-06) + +Full Changelog: [v1.2.32...v1.2.33](https://github.com/fw-ai-external/python-sdk/compare/v1.2.32...v1.2.33) + +### Features +* **training-sdk:** load reviewed Kimi K3 and K3.1 tokenizers without executing checkpoint Python, with shared vocabulary and distinct packaged chat encoders. + +### Bug Fixes +* **training-sdk:** preserve message sampling, saved fast-tokenizer overrides and chat-profile save/reload behavior on code-free Kimi loading paths. +* **training-sdk:** follow the November 5 Kimi K3.1 release encoder, which rejects the invalid `medium` thinking effort without changing vocabulary or image-token IDs. + ## 1.2.32 (2026-11-05) Full Changelog: [v1.2.31...v1.2.32](https://github.com/fw-ai-external/python-sdk/compare/v1.2.31...v1.2.32) diff --git a/pyproject.toml b/pyproject.toml index 60d1f23..fd84655 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "fireworks-ai" -version = "1.2.32" +version = "1.2.33" description = "The official Python library for the fireworks API" dynamic = ["readme"] license = "Apache-2.0" @@ -65,9 +65,9 @@ training-sdk = [ training = [ "fireworks-ai[training-sdk]", "torch>=2.0.0", - "transformers", + "transformers>=4.57.6,<6", "datasets", - "tiktoken", + "tiktoken>=0.12.0", "numpy", "wandb", ] diff --git a/src/fireworks/training/sdk/_core/lib/public_interfaces/sampling_client.py b/src/fireworks/training/sdk/_core/lib/public_interfaces/sampling_client.py index 73e5f82..0dc512b 100644 --- a/src/fireworks/training/sdk/_core/lib/public_interfaces/sampling_client.py +++ b/src/fireworks/training/sdk/_core/lib/public_interfaces/sampling_client.py @@ -15,7 +15,12 @@ from fireworks.training.sdk import types from fireworks.training.sdk._core._utils import lru_cache from fireworks.training.sdk._core.lib.sidecar import SidecarRPC, SidecarHandle, create_sidecar_handle -from fireworks.training.sdk.tokenizer_identity import is_deepseek_v4_model_name, requires_code_free_tokenizer_loading +from fireworks.training.sdk.tokenizer_identity import ( + is_deepseek_v4_model_name, + is_kimi_k3_k31_model_name, + is_kimi_k3_k31_tokenizer_source, + requires_code_free_tokenizer_loading, +) from fireworks.training.sdk._core.lib.telemetry import Telemetry, capture_exceptions from fireworks.training.sdk._core.lib.retry_handler import RetryConfig, RetryHandler from fireworks.training.sdk._core.lib.api_future_impl import QueueState, QueueStateObserver, _APIFuture @@ -538,6 +543,18 @@ def _load_tokenizer_from_model_info( "revision": "b5aabbfb20227ed42becbf5541dbffd213942c58", } + if ( + is_kimi_k3_k31_model_name(model_name) + or model_arch in ("KimiK3ForConditionalGeneration", "KimiK31ForConditionalGeneration") + or is_kimi_k3_k31_tokenizer_source(tokenizer_id, revision=kwargs.get("revision")) + ): + # Recognized reviewed Kimi K3/K3.1 data: build the id-identical fast + # tokenizer from reviewed installed code. Checkpoint Python (the + # auto_map TikTokenTokenizer) is never imported. + from fireworks.training.sdk.kimi_native_tokenizer import load_native_tokenizer + + return load_native_tokenizer(tokenizer_id, revision=kwargs.get("revision")) + if ( is_deepseek_v4_model_name(model_name) or model_arch == "DeepseekV4ForCausalLM" diff --git a/src/fireworks/training/sdk/_vendored/__init__.py b/src/fireworks/training/sdk/_vendored/__init__.py new file mode 100644 index 0000000..4c52c4f --- /dev/null +++ b/src/fireworks/training/sdk/_vendored/__init__.py @@ -0,0 +1 @@ +"""Reviewed runtime assets; no checkpoint module is imported.""" diff --git a/src/fireworks/training/sdk/_vendored/kimi_k3/LICENSE b/src/fireworks/training/sdk/_vendored/kimi_k3/LICENSE new file mode 100644 index 0000000..0de9f06 --- /dev/null +++ b/src/fireworks/training/sdk/_vendored/kimi_k3/LICENSE @@ -0,0 +1,52 @@ +Kimi K3 License + +Copyright (c) 2026 Moonshot AI + +Permission is hereby granted, free of charge, to any person (the "Licensee") +obtaining a copy of this software — including the model weights, parameters, +configuration files, inference and training code, and associated documentation +(collectively, the "Software") — to deal in the Software without restriction. +This includes, without limitation, the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of the Software; to run, +deploy, fine-tune, or otherwise modify the Software and create derivative works +from it; and to permit persons to whom the Software is furnished to do so, in +each case subject to the following conditions: + +1. The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. Licensee's use of the +Software must comply with applicable laws and regulations. + +2. "Model as a Service" means giving a third party access to language model +inference or fine-tuning (e.g., via API) in a manner that allows such third +party to exercise meaningful control over the inputs, parameters, or training +data. This does not include (a) end-user products with model capabilities solely +embedded within specific features or harnesses, or (b) mere relaying of requests +to models hosted by others. + +If the Licensee or any of its affiliates operates a Model as a Service business, +and the aggregate revenue of the Licensee and its affiliates exceeds 20 million +US dollars (or the equivalent in other currencies) in total over any consecutive +12 months, the Licensee must enter into a separate agreement with Moonshot AI +before using the Software or its derivative works for any commercial purpose. + +3. If the Software (or any derivative works thereof) is used for any of the +Licensee's commercial products or services that have more than 100 million +monthly active users, or more than 20 million US dollars (or equivalent in other +currencies) in monthly revenue, "Kimi K3" must be prominently displayed on the +user interface of such product or service. + +4. The requirements set forth in Sections 2 and 3 do not apply to: (a) internal +use of the Software, defined as any use that does not make the Software, its +outputs, or its underlying capabilities available to third parties; or (b) any +use of the Software accessed through Moonshot AI's official products or +certified inference partners. + +5. THE SOFTWARE AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” +BASIS, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT +LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE +AND NONINFRINGEMENT. IN NO EVENT SHALL MOONSHOT AI OR ITS AFFILIATES OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +For any questions regarding this license, please contact . diff --git a/src/fireworks/training/sdk/_vendored/kimi_k3/PROVENANCE.md b/src/fireworks/training/sdk/_vendored/kimi_k3/PROVENANCE.md new file mode 100644 index 0000000..508b634 --- /dev/null +++ b/src/fireworks/training/sdk/_vendored/kimi_k3/PROVENANCE.md @@ -0,0 +1,42 @@ +# SDK packaged-encoder adaptation + +The manifest below describes the reviewed upstream source. The SDK copy differs only by the existing Ruff import order and, for K3, renaming the unused loop index to `_message_index`. No template/token semantics change and no lint rule is disabled. Regenerate using the tokenizer snapshot generator. + +SDK packaged SHA-256: `b04b6ecd223d0b8519c5926f095f9531e97396c915595674effcc39fd4658323`. + +## Original source manifest + +# Vendored Kimi K3 encoder + +`encoding_k3.py` is a byte-for-byte copy of the upstream Kimi K3 XTML chat +encoder, incorporated as reviewed package source so that renderer tokenization +never fetches or executes checkpoint Python at runtime. + +| Field | Value | +|---|---| +| Upstream file | `encoding_k3.py` | +| Source repository | `moonshotai/Kimi-K3` | +| Source revision | `301be1b88c89c0d3a763da6301352cb8fe399e90` | +| sha256 | `b9cb7ae100fed34b9337f80dacee5abbf7e261fe9b74bc0e76366701d46f5333` | +| License | Kimi K3 License (see `LICENSE`, copyright 2026 Moonshot AI) | +| Retrieved | 2026-11-05 | + +The sha256 above is identical to the constant `KIMI_K3_ENCODER_SHA256` in +`training/renderer/kimi_k3.py`; a unit test asserts the vendored file still +hashes to it. + +Static inspection at incorporation (26311 bytes): imports limited to +`__future__`, `json`, `dataclasses`, `typing`; no I/O, no `exec`/`eval`, no +subprocess, no network, no third-party dependencies. + +Do not edit this file. To move to a newer upstream revision: download the +exact new revision, verify its sha256, statically inspect it, replace this +file, and update `KIMI_K3_ENCODER_SHA256` and this document in the same +change. + +Note: the Kimi K3.1 encoder is vendored separately under +`training/renderer/tito/_vendored/kimi_k31/` (pin +`367f0a7ab8800f78f7e09dc31e58544ee7fc46eea992ea05b32d53506a921cc5`). The two +encoders are distinct rendering contracts -- K3.1 renders non-string tool +arguments from their JSON literal bytes -- +and must not be aliased. diff --git a/src/fireworks/training/sdk/_vendored/kimi_k3/__init__.py b/src/fireworks/training/sdk/_vendored/kimi_k3/__init__.py new file mode 100644 index 0000000..28ffd66 --- /dev/null +++ b/src/fireworks/training/sdk/_vendored/kimi_k3/__init__.py @@ -0,0 +1 @@ +"""Kimi K3 XTML chat encoder, vendored verbatim (see PROVENANCE.md).""" diff --git a/src/fireworks/training/sdk/_vendored/kimi_k3/encoding_k3.py b/src/fireworks/training/sdk/_vendored/kimi_k3/encoding_k3.py new file mode 100644 index 0000000..1a63e70 --- /dev/null +++ b/src/fireworks/training/sdk/_vendored/kimi_k3/encoding_k3.py @@ -0,0 +1,647 @@ +"""Kimi K3 XTML encoding helpers. + +This module keeps chat rendering in Python. +Callers that need token IDs should consume ``EncodeSegment`` objects directly: +structural markers may be encoded as tiktoken special tokens, while user/tool +text and attribute values are encoded as ordinary text. +""" + +from __future__ import annotations + +import json +from typing import Any, Iterable, Optional +from dataclasses import dataclass + +OPEN_TOKEN = "<|open|>" +CLOSE_TOKEN = "<|close|>" +SEP_TOKEN = "<|sep|>" +END_OF_MSG_TOKEN = "<|end_of_msg|>" +IMAGE_PLACEHOLDER = "<|kimi_image_placeholder|>" + +_VALID_THINKING_EFFORTS = {"low", "high", "max"} + + +@dataclass(frozen=True) +class EncodeSegment: + text: str + allow_special: bool = False + + +class _ImagePromptState: + def __init__(self, image_prompts: Optional[list[str]] = None): + self.image_prompts = image_prompts + self.index = 0 + + def next_prompt(self) -> str: + if self.image_prompts is None: + return IMAGE_PLACEHOLDER + if self.index >= len(self.image_prompts): + raise ValueError("More image placeholders than image prompts.") + prompt = self.image_prompts[self.index] + self.index += 1 + return prompt + + def assert_consumed(self) -> None: + if self.image_prompts is None: + return + if self.index != len(self.image_prompts): + raise ValueError( + f"image prompt count {len(self.image_prompts)} != " + f"consumed placeholder count {self.index}" + ) + + +def _segment(text: Any, *, allow_special: bool = False) -> list[EncodeSegment]: + text = str(text) + if not text: + return [] + return [EncodeSegment(text, allow_special=allow_special)] + + +def _control(text: str) -> list[EncodeSegment]: + return _segment(text, allow_special=True) + + +def _text(text: Any) -> list[EncodeSegment]: + return _segment(text, allow_special=False) + + +def _append_text( + segments: list[EncodeSegment], + text: Any, + image_state: _ImagePromptState, +) -> None: + text = str(text) + if text == "": + return + if image_state.image_prompts is None or IMAGE_PLACEHOLDER not in text: + segments.extend(_text(text)) + return + + parts = text.split(IMAGE_PLACEHOLDER) + for i, part in enumerate(parts): + segments.extend(_text(part)) + if i < len(parts) - 1: + segments.extend(_segment(image_state.next_prompt(), + allow_special=True)) + + +def _escape_attr_value(value: Any) -> str: + return str(value).replace("&", "&").replace('"', """) + + +def _attr(key: str, value: Any) -> list[EncodeSegment]: + return ( + _text(f" {key}") + + _text('="') + + _text(_escape_attr_value(value)) + + _text('"') + ) + + +def _open_tag(tag: str, attrs: Iterable[tuple[str, Any]] = ()) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(OPEN_TOKEN)) + segments.extend(_text(tag)) + for key, value in attrs: + segments.extend(_attr(key, value)) + segments.extend(_control(SEP_TOKEN)) + return segments + +def _close_tag(tag: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(CLOSE_TOKEN)) + segments.extend(_text(tag)) + segments.extend(_control(SEP_TOKEN)) + return segments + + +def _end_of_msg() -> list[EncodeSegment]: + return _control(END_OF_MSG_TOKEN) + + +def _json_compact(value: Any) -> str: + return json.dumps(value, ensure_ascii=False, separators=(",", ":")) + + +def _is_mapping(value: Any) -> bool: + return isinstance(value, dict) + + +def _xtml_type(value: Any) -> str: + if isinstance(value, bool): + return "boolean" + if value is None: + return "null" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return "number" + if isinstance(value, str): + return "string" + if _is_mapping(value): + return "object" + return "array" + + +def _xtml_value(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + + +def _get_value(obj: Any, key: str, default: Any = None) -> Any: + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + +def extract_response_schema(response_format: Any) -> Any: + if response_format is None: + return None + + json_schema = _get_value(response_format, "json_schema") + if json_schema is None: + return None + + if isinstance(json_schema, dict): + return json_schema.get( + "schema", + json_schema.get("json_schema", json_schema), + ) + + schema = _get_value(json_schema, "schema") + if schema is not None: + return schema + + schema = _get_value(json_schema, "json_schema") + if schema is not None: + return schema + + return json_schema + + +def deep_sort_dict(obj: Any) -> Any: + if isinstance(obj, dict): + return {k: deep_sort_dict(v) for k, v in sorted(obj.items())} + if isinstance(obj, list): + return [deep_sort_dict(item) for item in obj] + return obj + + +def normalize_tool_arguments(arguments: Any) -> tuple[dict[str, Any], Optional[str]]: + if arguments is None: + return {}, None + if isinstance(arguments, dict): + return arguments, None + if isinstance(arguments, str): + if not arguments.strip(): + return {}, None + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + return {}, arguments + if not isinstance(parsed, dict): + raise ValueError("Kimi K3 tool call arguments must be a JSON object.") + return parsed, None + raise TypeError( + "Kimi K3 tool call arguments must be a dict or a JSON object string." + ) + + +def normalize_message(message: Any) -> Any: + if not isinstance(message, dict): + return message + + normalized = dict(message) + + tools = normalized.get("tools") + if tools is not None: + normalized["tools"] = deep_sort_dict(tools) + + tool_calls = normalized.get("tool_calls") + if not tool_calls: + return normalized + + normalized_calls = [] + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + normalized_calls.append(tool_call) + continue + + tc = dict(tool_call) + function = tc.get("function") + if isinstance(function, dict): + fn = dict(function) + arguments, json_block = normalize_tool_arguments(fn.get("arguments")) + fn["arguments"] = arguments + if json_block is None: + fn.pop("_xtml_json_block", None) + else: + fn["_xtml_json_block"] = json_block + tc["function"] = fn + else: + arguments, json_block = normalize_tool_arguments(tc.get("arguments")) + tc["arguments"] = arguments + if json_block is None: + tc.pop("_xtml_json_block", None) + else: + tc["_xtml_json_block"] = json_block + normalized_calls.append(tc) + + normalized["tool_calls"] = normalized_calls + return normalized + + +def normalize_conversation(conversation: Any) -> Any: + if not isinstance(conversation, list): + return conversation + + def normalize_messages(messages: list[Any]) -> list[Any]: + return [normalize_message(message) for message in messages] + + if conversation and isinstance(conversation[0], list): + return [normalize_messages(messages) for messages in conversation] + return normalize_messages(conversation) + + +def _tool_call_id_index(tool_calls: Any) -> dict: + """Map assistant ``tool_calls[].id`` to ``(1-based position, function name)``. + + The position mirrors the chat template's enumeration over ``tool_calls`` + (every entry advances the position, even an id-less one). Duplicate ids keep + their first occurrence. + """ + index: dict = {} + if not isinstance(tool_calls, list): + return index + for position, tool_call in enumerate(tool_calls, start=1): + if not isinstance(tool_call, dict): + continue + call_id = tool_call.get("id") + if call_id is None: + continue + key = str(call_id) + if key in index: + continue + function = tool_call.get("function") + name = ( + function.get("name") if isinstance(function, dict) else tool_call.get("name") + ) + index[key] = (position, name) + return index + + +def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]: + """Re-sort K3 XTML tool results into assistant ``tool_calls`` order. + + Serving frameworks generally deliver tool results already in call order. A + direct Transformers caller, however, may pass OpenAI-style tool messages in any + order, so each run of consecutive tool messages is matched against the most + recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` == + ``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and + sorted by the matched 1-based position. The matched call is authoritative, + so each matched message's ``tool`` is set to that call's function name -- + this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting + out of sync with the reordered position. ``index`` is still derived from the + rendered position by the chat template. A run that cannot be fully matched is + left untouched. Re-running is idempotent. + + This function is side-effect free: matched tool messages are shallow-copied + before their ``tool``/``name`` is rewritten, and every other message is + appended to the output as-is. The input list and its message objects are + never mutated. + """ + if not isinstance(messages, list): + return messages + + output: list[Any] = [] + current_index: dict = {} + i = 0 + n = len(messages) + + while i < n: + message = messages[i] + + if isinstance(message, dict) and message.get("role") == "assistant": + tool_calls = message.get("tool_calls") + current_index = _tool_call_id_index(tool_calls) if tool_calls else {} + output.append(message) + i += 1 + continue + + if not isinstance(message, dict) or message.get("role") != "tool": + output.append(message) + i += 1 + continue + + run: list[tuple] = [] # (position, original_offset, message, name) + unresolved = False + offset = 0 + while ( + i < n and isinstance(messages[i], dict) and messages[i].get("role") == "tool" + ): + tool_message = messages[i] + call_id = tool_message.get("tool_call_id", tool_message.get("id")) + matched = current_index.get(str(call_id)) if call_id is not None else None + if matched is None: + unresolved = True + run.append((None, offset, tool_message, None)) + else: + position, name = matched + run.append((position, offset, tool_message, name)) + offset += 1 + i += 1 + + if unresolved: + output.extend(item[2] for item in run) + else: + run.sort(key=lambda item: (item[0], item[1])) + for _, _, tool_message, name in run: + if name is None: + output.append(tool_message) + continue + # The id-matched call is authoritative: align tool (and any + # explicit name) so the rendered XTML tool attribute cannot + # disagree with the reordered position. Copy first so the + # caller's message object is never mutated. + resolved = dict(tool_message) + resolved["tool"] = name + if "name" in resolved: + resolved["name"] = name + output.append(resolved) + + return output + + +def is_batched_conversation(conversation: Any) -> bool: + return ( + isinstance(conversation, list) + and bool(conversation) + and isinstance(conversation[0], list) + ) + + +def _render_content_segments( + content: Any, + image_state: _ImagePromptState, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + if isinstance(content, str): + _append_text(segments, content, image_state) + elif content is not None: + for part in content: + if part["type"] in ["image", "image_url"]: + segments.extend( + _segment(image_state.next_prompt(), allow_special=True)) + else: + _append_text(segments, part["text"], image_state) + return segments + + +def _internal_system_message(message_type: str, body: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_open_tag("message", [("role", "system"), ("type", message_type)])) + segments.extend(_text(body.strip())) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def _render_assistant_segments( + message: dict[str, Any], + image_state: _ImagePromptState, + thinking: bool = True, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + # The channel is structural: in thinking mode every assistant + # message carries the open/close tags even when there is no reasoning + # content to fill in. In non-thinking mode the channel is dropped + # entirely. + if thinking: + reasoning_content = message.get("reasoning_content") or message.get( + "reasoning" + ) + segments.extend(_open_tag("think")) + if reasoning_content is not None and str(reasoning_content).strip(): + _append_text(segments, reasoning_content, image_state) + segments.extend(_close_tag("think")) + + segments.extend(_open_tag("response")) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("response")) + + tool_calls = message.get("tool_calls") + if tool_calls: + segments.extend(_open_tag("tools")) + for index, tool_call in enumerate(tool_calls, start=1): + fn = tool_call.get("function", tool_call) + segments.extend( + _open_tag("call", [("tool", fn["name"]), ("index", index)]) + ) + args = fn.get("arguments", {}) + json_block = fn.get("_xtml_json_block") + if json_block is not None: + segments.extend(_open_tag("json", [("type", "object")])) + _append_text(segments, json_block, image_state) + segments.extend(_close_tag("json")) + elif _is_mapping(args): + for key, value in args.items(): + segments.extend( + _open_tag( + "argument", + [("key", key), ("type", _xtml_type(value))], + ) + ) + _append_text(segments, _xtml_value(value), image_state) + segments.extend(_close_tag("argument")) + segments.extend(_close_tag("call")) + segments.extend(_close_tag("tools")) + + return segments + + +def _render_tool_declare(tools: Any, *, dynamic: bool = False) -> list[EncodeSegment]: + if dynamic: + body = ( + "## New Tools Available\n" + "The system dynamically extends the toolset via lazy-loading.\n" + "You have access to all existing and extended tools.\n" + "Here are the specs for the extended tools.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```" + ) + else: + body = ( + "# Tools\n" + "Here are the available tools, described in JSONSchema.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```" + ) + segments: list[EncodeSegment] = [] + segments.extend(_open_tag("message", [("role", "system"), ("type", "tool-declare")])) + segments.extend(_text(body)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def build_chat_segments( + messages: list[Any], + tools: Optional[list[dict]] = None, + *, + add_generation_prompt: bool = True, + thinking: bool = True, + image_prompts: Optional[list[str]] = None, + **kwargs: Any, +) -> list[EncodeSegment]: + # Re-sort tool results by tool_call_id at the lowest layer so every caller + # (processor or direct tokenizer) gets correctly ordered XTML. The helper is + # side-effect free, so the caller's message objects are left untouched. + messages = normalize_xtml_tool_result_messages(messages) + messages = normalize_conversation(messages) + tools = deep_sort_dict(tools) + + kwargs = dict(kwargs) + response_format = kwargs.get("response_format") + if "response_schema" not in kwargs: + response_schema = extract_response_schema(response_format) + if response_schema is not None: + kwargs["response_schema"] = response_schema + if kwargs.get("response_schema") is not None: + kwargs["response_schema"] = deep_sort_dict(kwargs["response_schema"]) + + image_state = _ImagePromptState(image_prompts) + segments: list[EncodeSegment] = [] + + tool_calls = None + tool_index = 0 + + if tools: + segments.extend(_render_tool_declare(tools)) + + thinking_effort = kwargs.get("thinking_effort") + if thinking and thinking_effort is not None: + assert thinking_effort in _VALID_THINKING_EFFORTS, ( + f"Unsupported thinking_effort={thinking_effort!r}; " + f"supported values are {sorted(_VALID_THINKING_EFFORTS)}." + ) + if thinking and thinking_effort in _VALID_THINKING_EFFORTS: + segments.extend( + _internal_system_message( + "thinking-effort", + "`thinking_effort` guides on how much to think in your " + "thinking channel (not including the response channel), " + "supported values include `low`, `medium`, `high`, and `max`.\n" + f"Now the system is invoked with `thinking_effort={thinking_effort}`.", + ) + ) + + for _message_index, message in enumerate(messages): + if not isinstance(message, dict): + continue + + role = message["role"] + if role == "user": + attrs = [("role", "user")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "system" and message.get("tools"): + segments.extend(_render_tool_declare(message["tools"], dynamic=True)) + elif role == "system": + attrs = [("role", "system")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "tool": + tool_index += 1 + tool_name = message.get("tool", message.get("name")) + if ( + tool_name is None + and tool_calls is not None + and tool_index <= len(tool_calls) + ): + tc = tool_calls[tool_index - 1] + fn = tc.get("function", tc) + tool_name = fn["name"] + if tool_name is None: + raise ValueError( + "Kimi K3 tool messages need a resolvable tool name: " + "carry `tool`/`name`, or match a preceding assistant " + "tool_call by order." + ) + segments.extend( + _open_tag( + "message", + [("role", "tool"), ("tool", tool_name), ("index", tool_index)], + ) + ) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "assistant": + tool_calls = message.get("tool_calls") + tool_index = 0 + attrs = [("role", "assistant")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_assistant_segments(message, image_state, thinking)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + + tool_choice = kwargs.get("tool_choice") + if tool_choice == "required": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=required`.\n" + "You MUST call tools in the next message.", + ) + ) + elif tool_choice == "none": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=none`.\n" + "You MUST NOT call any tools in the next message.", + ) + ) + + rf = kwargs.get("response_format") + rf_type = _get_value(rf, "type", rf) if isinstance(rf, dict) else rf + if rf_type == "json_object": + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_object`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.", + ) + ) + elif rf_type == "json_schema": + schema = _json_compact(kwargs.get("response_schema")) + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_schema`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.\n" + "The JSON data must match the following schema:\n" + f"```json\n{schema}\n```", + ) + ) + + if add_generation_prompt: + segments.extend(_open_tag("message", [("role", "assistant")])) + segments.extend(_open_tag("think" if thinking else "response")) + + image_state.assert_consumed() + return segments diff --git a/src/fireworks/training/sdk/_vendored/kimi_k31/LICENSE b/src/fireworks/training/sdk/_vendored/kimi_k31/LICENSE new file mode 100644 index 0000000..c33f1d2 --- /dev/null +++ b/src/fireworks/training/sdk/_vendored/kimi_k31/LICENSE @@ -0,0 +1,52 @@ +Kimi K3.1-Pro License + +Copyright (c) 2026 Moonshot AI + +Permission is hereby granted, free of charge, to any person (the "Licensee") +obtaining a copy of this software — including the model weights, parameters, +configuration files, inference and training code, and associated documentation +(collectively, the "Software") — to deal in the Software without restriction. +This includes, without limitation, the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of the Software; to run, +deploy, fine-tune, or otherwise modify the Software and create derivative works +from it; and to permit persons to whom the Software is furnished to do so, in +each case subject to the following conditions: + +1. The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. Licensee's use of the +Software must comply with applicable laws and regulations. + +2. "Model as a Service" means giving a third party access to language model +inference or fine-tuning (e.g., via API) in a manner that allows such third +party to exercise meaningful control over the inputs, parameters, or training +data. This does not include (a) end-user products with model capabilities solely +embedded within specific features or harnesses, or (b) mere relaying of requests +to models hosted by others. + +If the Licensee or any of its affiliates operates a Model as a Service business, +and the aggregate revenue of the Licensee and its affiliates exceeds 20 million +US dollars (or the equivalent in other currencies) in total over any consecutive +12 months, the Licensee must enter into a separate agreement with Moonshot AI +before using the Software or its derivative works for any commercial purpose. + +3. If the Software (or any derivative works thereof) is used for any of the +Licensee's commercial products or services that have more than 100 million +monthly active users, or more than 20 million US dollars (or equivalent in other +currencies) in monthly revenue, "Kimi K3.1-Pro" must be prominently displayed on the +user interface of such product or service. + +4. The requirements set forth in Sections 2 and 3 do not apply to: (a) internal +use of the Software, defined as any use that does not make the Software, its +outputs, or its underlying capabilities available to third parties; or (b) any +use of the Software accessed through Moonshot AI's official products or +certified inference partners. + +5. THE SOFTWARE AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” +BASIS, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT +LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE +AND NONINFRINGEMENT. IN NO EVENT SHALL MOONSHOT AI OR ITS AFFILIATES OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +For any questions regarding this license, please contact . \ No newline at end of file diff --git a/src/fireworks/training/sdk/_vendored/kimi_k31/PROVENANCE.md b/src/fireworks/training/sdk/_vendored/kimi_k31/PROVENANCE.md new file mode 100644 index 0000000..42c02af --- /dev/null +++ b/src/fireworks/training/sdk/_vendored/kimi_k31/PROVENANCE.md @@ -0,0 +1,35 @@ +# SDK packaged-encoder adaptation + +The manifest below describes the reviewed upstream source. The SDK copy differs only by the existing Ruff import order and, for K3, renaming the unused loop index to `_message_index`. No template/token semantics change and no lint rule is disabled. Regenerate using the tokenizer snapshot generator. + +SDK packaged SHA-256: `5ca26282a0f72a6786addd95958a8117605d02ecbec9951ae9b67e54a121ddd5`. + +## Original source manifest + +# Vendored Kimi K3.1 encoder + +`encoding_k3.py` is a byte-for-byte copy of the upstream Kimi K3.1 XTML chat +encoder, incorporated as reviewed package source so that TITO rendering never +fetches or executes checkpoint Python at runtime. + +| Field | Value | +|---|---| +| Upstream file | `encoding_k3.py` | +| Source repository | `moonshotai/Kimi-K3.1-Pro` | +| Source revision | `d9e99b0792a5e93b78ea9bea62a68167e36ddd9e` | +| sha256 | `367f0a7ab8800f78f7e09dc31e58544ee7fc46eea992ea05b32d53506a921cc5` | +| License | Kimi K3.1-Pro License (see `LICENSE`, copyright 2026 Moonshot AI) | +| Retrieved | 2026-11-05 | + +The sha256 above is identical to the TITO constant `KIMI_K31_ENCODER_SHA256` +in `training/renderer/tito/kimi_k31.py`; a unit test asserts the vendored file +still hashes to it. + +Static inspection at incorporation (735 lines): imports limited to +`__future__`, `json`, `dataclasses`, `typing`; no I/O, no `exec`/`eval`, no +subprocess, no network, no third-party dependencies. + +Do not edit this file. To move to a newer upstream revision: download the +exact new revision, verify its sha256, statically inspect it, replace this +file, and update `KIMI_K31_ENCODER_SHA256` and this document in the same +change. diff --git a/src/fireworks/training/sdk/_vendored/kimi_k31/__init__.py b/src/fireworks/training/sdk/_vendored/kimi_k31/__init__.py new file mode 100644 index 0000000..c87a380 --- /dev/null +++ b/src/fireworks/training/sdk/_vendored/kimi_k31/__init__.py @@ -0,0 +1 @@ +"""Kimi K3.1 XTML chat encoder, vendored verbatim (see PROVENANCE.md).""" diff --git a/src/fireworks/training/sdk/_vendored/kimi_k31/encoding_k3.py b/src/fireworks/training/sdk/_vendored/kimi_k31/encoding_k3.py new file mode 100644 index 0000000..c21b520 --- /dev/null +++ b/src/fireworks/training/sdk/_vendored/kimi_k31/encoding_k3.py @@ -0,0 +1,735 @@ +"""Kimi K3 XTML encoding helpers. + +This module keeps chat rendering in Python. +Callers that need token IDs should consume ``EncodeSegment`` objects directly: +structural markers may be encoded as tiktoken special tokens, while user/tool +text and attribute values are encoded as ordinary text. +""" + +from __future__ import annotations + +import json +from typing import Any, Iterable, Optional +from dataclasses import dataclass + +OPEN_TOKEN = "<|open|>" +CLOSE_TOKEN = "<|close|>" +SEP_TOKEN = "<|sep|>" +END_OF_MSG_TOKEN = "<|end_of_msg|>" +IMAGE_PLACEHOLDER = "<|kimi_image_placeholder|>" + +_VALID_THINKING_EFFORTS = {"low", "high", "max"} + + +@dataclass(frozen=True) +class EncodeSegment: + text: str + allow_special: bool = False + + +class _ImagePromptState: + + def __init__(self, image_prompts: Optional[list[str]] = None): + self.image_prompts = image_prompts + self.index = 0 + + def next_prompt(self) -> str: + if self.image_prompts is None: + return IMAGE_PLACEHOLDER + if self.index >= len(self.image_prompts): + raise ValueError("More image placeholders than image prompts.") + prompt = self.image_prompts[self.index] + self.index += 1 + return prompt + + def assert_consumed(self) -> None: + if self.image_prompts is None: + return + if self.index != len(self.image_prompts): + raise ValueError( + f"image prompt count {len(self.image_prompts)} != " + f"consumed placeholder count {self.index}") + + +def _segment(text: Any, *, allow_special: bool = False) -> list[EncodeSegment]: + text = str(text) + if not text: + return [] + return [EncodeSegment(text, allow_special=allow_special)] + + +def _control(text: str) -> list[EncodeSegment]: + return _segment(text, allow_special=True) + + +def _text(text: Any) -> list[EncodeSegment]: + return _segment(text, allow_special=False) + + +def _append_text( + segments: list[EncodeSegment], + text: Any, + image_state: _ImagePromptState, +) -> None: + text = str(text) + if text == "": + return + if image_state.image_prompts is None or IMAGE_PLACEHOLDER not in text: + segments.extend(_text(text)) + return + + parts = text.split(IMAGE_PLACEHOLDER) + for i, part in enumerate(parts): + segments.extend(_text(part)) + if i < len(parts) - 1: + segments.extend( + _segment(image_state.next_prompt(), allow_special=True)) + + +def _escape_attr_value(value: Any) -> str: + return str(value).replace("&", "&").replace('"', """) + + +def _attr(key: str, value: Any) -> list[EncodeSegment]: + return (_text(f" {key}") + _text('="') + _text(_escape_attr_value(value)) + + _text('"')) + + +def _open_tag( + tag: str, attrs: Iterable[tuple[str, Any]] = ()) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(OPEN_TOKEN)) + segments.extend(_text(tag)) + for key, value in attrs: + segments.extend(_attr(key, value)) + segments.extend(_control(SEP_TOKEN)) + return segments + + +def _close_tag(tag: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(CLOSE_TOKEN)) + segments.extend(_text(tag)) + segments.extend(_control(SEP_TOKEN)) + return segments + + +def _end_of_msg() -> list[EncodeSegment]: + return _control(END_OF_MSG_TOKEN) + + +def _json_compact(value: Any) -> str: + return json.dumps(value, ensure_ascii=False, separators=(",", ":")) + + +def _is_mapping(value: Any) -> bool: + return isinstance(value, dict) + + +def _xtml_type(value: Any) -> str: + if isinstance(value, bool): + return "boolean" + if value is None: + return "null" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return "number" + if isinstance(value, str): + return "string" + if _is_mapping(value): + return "object" + return "array" + + +def _xtml_value(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + + +def _get_value(obj: Any, key: str, default: Any = None) -> Any: + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + +def extract_response_schema(response_format: Any) -> Any: + if response_format is None: + return None + + json_schema = _get_value(response_format, "json_schema") + if json_schema is None: + return None + + if isinstance(json_schema, dict): + return json_schema.get( + "schema", + json_schema.get("json_schema", json_schema), + ) + + schema = _get_value(json_schema, "schema") + if schema is not None: + return schema + + schema = _get_value(json_schema, "json_schema") + if schema is not None: + return schema + + return json_schema + + +def deep_sort_dict(obj: Any) -> Any: + if isinstance(obj, dict): + return {k: deep_sort_dict(v) for k, v in sorted(obj.items())} + if isinstance(obj, list): + return [deep_sort_dict(item) for item in obj] + return obj + + +# One normalized tool-call argument: (key, XTML type, rendered text). The +# text keeps the original JSON literal for non-string values and is the +# decoded string for string values. +XtmlArgument = tuple[str, str, str] + + +def _parse_arguments_object(s: str) -> list[XtmlArgument]: + """Parse a JSON object string one level deep. + + Each top-level key-value pair yields one ``(key, type, text)`` triple. The + text keeps the original JSON literal for non-string values (``1e2`` stays + ``1e2``, ``[1,2]`` keeps its exact bytes) and is the decoded (unescaped) + string for string values. Nested values are never re-serialized. Anything + after the closing ``}`` is ignored. Raises ``ValueError`` on malformed + input (including a valid non-object JSON document). + """ + idx = 0 + + def _skip_whitespaces() -> None: + nonlocal idx + while idx < len(s) and s[idx] in (" ", "\t", "\n", "\r"): + idx += 1 + + def _next_char() -> Optional[str]: + nonlocal idx + if idx < len(s): + c = s[idx] + idx += 1 + return c + return None + + _skip_whitespaces() + if _next_char() != "{": + raise ValueError("JSON arguments must be an object") + _skip_whitespaces() + + parsed: list[XtmlArgument] = [] + + if idx >= len(s): + raise ValueError("Unexpected end of JSON object") + if s[idx] == "}": + return parsed + + json_decoder = json.JSONDecoder(strict=False) + + def _raw_decode(start_idx: int) -> tuple[Any, int]: + try: + return json_decoder.raw_decode(s, idx=start_idx) + except json.JSONDecodeError as e: + raise ValueError(str(e)) from e + + while True: + decoded_key, idx = _raw_decode(idx) + if not isinstance(decoded_key, str): + raise ValueError( + f"JSON object key must be a string, got {decoded_key!r}") + _skip_whitespaces() + if _next_char() != ":": + raise ValueError(f"Expects ':' after {decoded_key}") + _skip_whitespaces() + + value_start_idx = idx + decoded_value, idx = _raw_decode(idx) + value_end_idx = idx + + text = (decoded_value if isinstance(decoded_value, str) else + s[value_start_idx:value_end_idx]) + parsed.append((decoded_key, _xtml_type(decoded_value), text)) + + _skip_whitespaces() + c = _next_char() + _skip_whitespaces() + + if c == "}": + break + elif c == ",": + continue + else: + raise ValueError(f"Expect '}}' or ',', got {c!r}") + + return parsed + + +def normalize_tool_arguments( + arguments: Any, ) -> tuple[list[XtmlArgument], Optional[str]]: + """Normalize tool call arguments for XTML rendering. + + Returns ``(argument_triples, raw_json_block)``. String arguments are parsed + one level deep so non-string values keep their original JSON literal text, + and anything after the closing ``}`` is discarded; dict arguments (already + Python objects) are serialized on the spot. Any string that is not a + well-formed JSON object -- unparsable, whitespace only, or valid + non-object JSON -- falls back to the raw ```` block. + """ + if arguments is None: + return [], None + if isinstance(arguments, dict): + return [(str(key), _xtml_type(value), _xtml_value(value)) + for key, value in arguments.items()], None + if isinstance(arguments, str): + if arguments == "": + return [], None + try: + return _parse_arguments_object(arguments), None + except ValueError: + return [], arguments + raise TypeError( + "Kimi K3 tool call arguments must be a dict or a JSON object string.") + + +def normalize_message(message: Any) -> Any: + if not isinstance(message, dict): + return message + + normalized = dict(message) + + tools = normalized.get("tools") + if tools is not None: + normalized["tools"] = deep_sort_dict(tools) + + tool_calls = normalized.get("tool_calls") + if not tool_calls: + return normalized + + normalized_calls = [] + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + normalized_calls.append(tool_call) + continue + + tc = dict(tool_call) + function = tc.get("function") + if isinstance(function, dict): + fn = dict(function) + arguments, json_block = normalize_tool_arguments( + fn.get("arguments")) + fn["arguments"] = arguments + if json_block is None: + fn.pop("_xtml_json_block", None) + else: + fn["_xtml_json_block"] = json_block + tc["function"] = fn + else: + arguments, json_block = normalize_tool_arguments( + tc.get("arguments")) + tc["arguments"] = arguments + if json_block is None: + tc.pop("_xtml_json_block", None) + else: + tc["_xtml_json_block"] = json_block + normalized_calls.append(tc) + + normalized["tool_calls"] = normalized_calls + return normalized + + +def normalize_conversation(conversation: Any) -> Any: + if not isinstance(conversation, list): + return conversation + + def normalize_messages(messages: list[Any]) -> list[Any]: + return [normalize_message(message) for message in messages] + + if conversation and isinstance(conversation[0], list): + return [normalize_messages(messages) for messages in conversation] + return normalize_messages(conversation) + + +def _tool_call_id_index(tool_calls: Any) -> dict: + """Map assistant ``tool_calls[].id`` to ``(1-based position, function name)``. + + The position mirrors the chat template's enumeration over ``tool_calls`` + (every entry advances the position, even an id-less one). Duplicate ids keep + their first occurrence. + """ + index: dict = {} + if not isinstance(tool_calls, list): + return index + for position, tool_call in enumerate(tool_calls, start=1): + if not isinstance(tool_call, dict): + continue + call_id = tool_call.get("id") + if call_id is None: + continue + key = str(call_id) + if key in index: + continue + function = tool_call.get("function") + name = (function.get("name") + if isinstance(function, dict) else tool_call.get("name")) + index[key] = (position, name) + return index + + +def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]: + """Re-sort K3 XTML tool results into assistant ``tool_calls`` order. + + Serving frameworks generally deliver tool results already in call order. A + direct Transformers caller, however, may pass OpenAI-style tool messages in any + order, so each run of consecutive tool messages is matched against the most + recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` == + ``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and + sorted by the matched 1-based position. The matched call is authoritative, + so each matched message's ``tool`` is set to that call's function name -- + this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting + out of sync with the reordered position. ``index`` is still derived from the + rendered position by the chat template. A run that cannot be fully matched is + left untouched. Re-running is idempotent. + + This function is side-effect free: matched tool messages are shallow-copied + before their ``tool``/``name`` is rewritten, and every other message is + appended to the output as-is. The input list and its message objects are + never mutated. + """ + if not isinstance(messages, list): + return messages + + output: list[Any] = [] + current_index: dict = {} + i = 0 + n = len(messages) + + while i < n: + message = messages[i] + + if isinstance(message, dict) and message.get("role") == "assistant": + tool_calls = message.get("tool_calls") + current_index = _tool_call_id_index( + tool_calls) if tool_calls else {} + output.append(message) + i += 1 + continue + + if not isinstance(message, dict) or message.get("role") != "tool": + output.append(message) + i += 1 + continue + + run: list[tuple] = [] # (position, original_offset, message, name) + unresolved = False + offset = 0 + while (i < n and isinstance(messages[i], dict) + and messages[i].get("role") == "tool"): + tool_message = messages[i] + call_id = tool_message.get("tool_call_id", tool_message.get("id")) + matched = current_index.get( + str(call_id)) if call_id is not None else None + if matched is None: + unresolved = True + run.append((None, offset, tool_message, None)) + else: + position, name = matched + run.append((position, offset, tool_message, name)) + offset += 1 + i += 1 + + if unresolved: + output.extend(item[2] for item in run) + else: + run.sort(key=lambda item: (item[0], item[1])) + for _, _, tool_message, name in run: + if name is None: + output.append(tool_message) + continue + # The id-matched call is authoritative: align tool (and any + # explicit name) so the rendered XTML tool attribute cannot + # disagree with the reordered position. Copy first so the + # caller's message object is never mutated. + resolved = dict(tool_message) + resolved["tool"] = name + if "name" in resolved: + resolved["name"] = name + output.append(resolved) + + return output + + +def is_batched_conversation(conversation: Any) -> bool: + return (isinstance(conversation, list) and bool(conversation) + and isinstance(conversation[0], list)) + + +def _render_content_segments( + content: Any, + image_state: _ImagePromptState, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + if isinstance(content, str): + _append_text(segments, content, image_state) + elif content is not None: + for part in content: + if part["type"] in ["image", "image_url"]: + segments.extend( + _segment(image_state.next_prompt(), allow_special=True)) + else: + _append_text(segments, part["text"], image_state) + return segments + + +def _internal_system_message(message_type: str, + body: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend( + _open_tag("message", [("role", "system"), ("type", message_type)])) + segments.extend(_text(body.strip())) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def _render_assistant_segments( + message: dict[str, Any], + image_state: _ImagePromptState, + thinking: bool = True, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + # The channel is structural: in thinking mode every assistant + # message carries the open/close tags even when there is no reasoning + # content to fill in. In non-thinking mode the channel is dropped + # entirely. + if thinking: + reasoning_content = message.get("reasoning_content") or message.get( + "reasoning") + segments.extend(_open_tag("think")) + # Only an empty string counts as no reasoning, so whitespace-only + # reasoning is still rendered into the think channel. + if reasoning_content is not None and str(reasoning_content) != "": + _append_text(segments, reasoning_content, image_state) + segments.extend(_close_tag("think")) + + segments.extend(_open_tag("response")) + segments.extend( + _render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("response")) + + tool_calls = message.get("tool_calls") + if tool_calls: + segments.extend(_open_tag("tools")) + for index, tool_call in enumerate(tool_calls, start=1): + fn = tool_call.get("function", tool_call) + segments.extend( + _open_tag("call", [("tool", fn["name"]), ("index", index)])) + args = fn.get("arguments", []) + json_block = fn.get("_xtml_json_block") + if json_block is not None: + segments.extend(_open_tag("json", [("type", "object")])) + _append_text(segments, json_block, image_state) + segments.extend(_close_tag("json")) + else: + for key, arg_type, arg_text in args: + segments.extend( + _open_tag("argument", [("key", key), + ("type", arg_type)])) + _append_text(segments, arg_text, image_state) + segments.extend(_close_tag("argument")) + segments.extend(_close_tag("call")) + segments.extend(_close_tag("tools")) + + return segments + + +def _render_tool_declare(tools: Any, + *, + dynamic: bool = False) -> list[EncodeSegment]: + if dynamic: + body = ( + "## New Tools Available\n" + "The system dynamically extends the toolset via lazy-loading.\n" + "You have access to all existing and extended tools.\n" + "Here are the specs for the extended tools.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```") + else: + body = ("# Tools\n" + "Here are the available tools, described in JSONSchema.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```") + segments: list[EncodeSegment] = [] + segments.extend( + _open_tag("message", [("role", "system"), ("type", "tool-declare")])) + segments.extend(_text(body)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def build_chat_segments( + messages: list[Any], + tools: Optional[list[dict]] = None, + *, + add_generation_prompt: bool = True, + thinking: bool = True, + image_prompts: Optional[list[str]] = None, + **kwargs: Any, +) -> list[EncodeSegment]: + # Re-sort tool results by tool_call_id at the lowest layer so every caller + # (processor or direct tokenizer) gets correctly ordered XTML. The helper is + # side-effect free, so the caller's message objects are left untouched. + messages = normalize_xtml_tool_result_messages(messages) + messages = normalize_conversation(messages) + tools = deep_sort_dict(tools) + + kwargs = dict(kwargs) + response_format = kwargs.get("response_format") + if "response_schema" not in kwargs: + response_schema = extract_response_schema(response_format) + if response_schema is not None: + kwargs["response_schema"] = response_schema + if kwargs.get("response_schema") is not None: + kwargs["response_schema"] = deep_sort_dict(kwargs["response_schema"]) + + image_state = _ImagePromptState(image_prompts) + segments: list[EncodeSegment] = [] + + tool_calls = None + tool_index = 0 + + if tools: + segments.extend(_render_tool_declare(tools)) + + thinking_effort = kwargs.get("thinking_effort") + if thinking and thinking_effort is not None: + assert thinking_effort in _VALID_THINKING_EFFORTS, ( + f"Unsupported thinking_effort={thinking_effort!r}; " + f"supported values are {sorted(_VALID_THINKING_EFFORTS)}.") + if thinking and thinking_effort in _VALID_THINKING_EFFORTS: + segments.extend( + _internal_system_message( + "thinking-effort", + "`thinking_effort` guides on how much to think in your " + "thinking channel (not including the response channel), " + "supported values include `low`, `medium`, `high`, and `max`.\n" + f"Now the system is invoked with `thinking_effort={thinking_effort}`.", + )) + + for message_index, message in enumerate(messages): + # Malformed messages are rejected instead of being silently skipped. + if not isinstance(message, dict): + raise ValueError( + f"Kimi K3 messages must be dicts, got {type(message).__name__} " + f"at index {message_index}.") + + role = message.get("role") + if role == "user": + attrs = [("role", "user")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend( + _render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "system" and message.get("tools"): + segments.extend( + _render_tool_declare(message["tools"], dynamic=True)) + elif role == "system": + attrs = [("role", "system")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend( + _render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "tool": + tool_index += 1 + tool_name = message.get("tool", message.get("name")) + if (tool_name is None and tool_calls is not None + and tool_index <= len(tool_calls)): + tc = tool_calls[tool_index - 1] + fn = tc.get("function", tc) + tool_name = fn["name"] + if tool_name is None: + raise ValueError( + "Kimi K3 tool messages need a resolvable tool name: " + "carry `tool`/`name`, or match a preceding assistant " + "tool_call by order.") + segments.extend( + _open_tag( + "message", + [("role", "tool"), ("tool", tool_name), + ("index", tool_index)], + )) + segments.extend( + _render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "assistant": + tool_calls = message.get("tool_calls") + tool_index = 0 + attrs = [("role", "assistant")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend( + _render_assistant_segments(message, image_state, thinking)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + else: + raise ValueError( + f"Unknown message role {role!r} at index {message_index}.") + + tool_choice = kwargs.get("tool_choice") + if tool_choice == "required": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=required`.\n" + "You MUST call tools in the next message.", + )) + elif tool_choice == "none": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=none`.\n" + "You MUST NOT call any tools in the next message.", + )) + + rf = kwargs.get("response_format") + rf_type = _get_value(rf, "type", rf) if isinstance(rf, dict) else rf + if rf_type == "json_object": + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_object`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.", + )) + elif rf_type == "json_schema": + schema = _json_compact(kwargs.get("response_schema")) + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_schema`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.\n" + "The JSON data must match the following schema:\n" + f"```json\n{schema}\n```", + )) + + if add_generation_prompt: + segments.extend(_open_tag("message", [("role", "assistant")])) + segments.extend(_open_tag("think" if thinking else "response")) + + image_state.assert_consumed() + return segments diff --git a/src/fireworks/training/sdk/client.py b/src/fireworks/training/sdk/client.py index f1ebca9..b9cd51a 100644 --- a/src/fireworks/training/sdk/client.py +++ b/src/fireworks/training/sdk/client.py @@ -41,7 +41,12 @@ from fireworks.training.sdk import types from fireworks.training.sdk._core import APIStatusError, SamplingClient from fireworks.training.sdk._core._compat import model_dump as _training_model_dump -from fireworks.training.sdk.tokenizer_identity import is_deepseek_v4_model_name, requires_code_free_tokenizer_loading +from fireworks.training.sdk.tokenizer_identity import ( + is_deepseek_v4_model_name, + is_kimi_k3_k31_model_name, + is_kimi_k3_k31_tokenizer_source, + requires_code_free_tokenizer_loading, +) if TYPE_CHECKING: from transformers import PreTrainedTokenizerBase @@ -3107,6 +3112,16 @@ def get_tokenizer(self): # lazy: tokenizer downloads are optional; transport needs no transformers. from transformers import AutoTokenizer + if is_kimi_k3_k31_model_name( + getattr(self, "_tokenizer_base_model", None) + ) or is_kimi_k3_k31_tokenizer_source(self._tokenizer_model): + # Recognized reviewed Kimi K3/K3.1 data: build the id-identical + # fast tokenizer from reviewed installed code. Checkpoint Python + # (the auto_map TikTokenTokenizer) is never imported, and an + # explicit trust_remote_code elsewhere cannot reopen it. + from fireworks.training.sdk.kimi_native_tokenizer import load_native_tokenizer + + return load_native_tokenizer(self._tokenizer_model) if is_deepseek_v4_model_name( getattr(self, "_tokenizer_base_model", None) ) or requires_code_free_tokenizer_loading(self._tokenizer_model): diff --git a/src/fireworks/training/sdk/kimi_native_tokenizer.py b/src/fireworks/training/sdk/kimi_native_tokenizer.py new file mode 100644 index 0000000..63dbad6 --- /dev/null +++ b/src/fireworks/training/sdk/kimi_native_tokenizer.py @@ -0,0 +1,408 @@ +"""Code-free native tokenizer loading for reviewed Kimi K3/K3.1 sources. + +Kimi K3 and K3.1 ship a tiktoken vocabulary (``tiktoken.model``) behind a +remote-code ``TikTokenTokenizer`` and no ``tokenizer.json``. This module +converts the pinned data files into an id-identical fast tokenizer entirely +from reviewed installed code, so host loading never imports +checkpoint-staged Python and ``auto_map`` metadata stays inert. + +Only data files are read. ``tokenization_kimi.py`` is hash-checked as +provenance evidence (the token pattern and reserved-token count below are +copied from it) but never imported. The two models share one byte-identical +vocabulary; their XTML chat encoders are distinct reviewed contracts, +vendored separately by the cookbook renderers and never aliased here. + +The cookbook wrapper ``training.utils.kimi_native_tokenizer`` delegates to +this module when ``KIMI_NATIVE_TOKENIZER_CAPABILITY`` matches and otherwise +falls back to its generated snapshot, so older installed SDKs keep the same +behavior. +""" + +from __future__ import annotations + +import json +import hashlib +from typing import Any +from pathlib import Path +from functools import cached_property +from importlib import import_module + +# Capability marker for cookbook delegation, same contract as +# TOKENIZER_IDENTITY_POLICY_CAPABILITY: the wrapper defers to this module +# only when the constant matches exactly, so an intermediate SDK falls back +# to the cookbook's pinned snapshot instead of inheriting older behavior. +KIMI_NATIVE_TOKENIZER_CAPABILITY = "kimi-k3-k31-native-tokenizer/chat-and-saved-bundles-v3" + +KIMI_K3_PROFILE = "kimi_k3" +KIMI_K31_PROFILE = "kimi_k31" + +KIMI_K3_REPOSITORY = "moonshotai/Kimi-K3" +KIMI_K3_REVISION = "301be1b88c89c0d3a763da6301352cb8fe399e90" +KIMI_K31_REPOSITORY = "moonshotai/Kimi-K3.1-Pro" +KIMI_K31_REVISION = "d9e99b0792a5e93b78ea9bea62a68167e36ddd9e" + +# Pinned source files per profile. ``tiktoken.model`` is byte-identical +# across both; ``tokenizer_config.json`` differs only in formatting, and +# ``tokenization_kimi.py`` only in formatting (the token pattern is the +# same), but each profile pins its own reviewed bytes. +KIMI_K3_SOURCE_SHA256 = { + "tiktoken.model": "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103", + "tokenizer_config.json": "5d0803c94db9cd78763499e0956c95fd5a225c14a727e5a6cf5db3f96f010a6e", + "tokenization_kimi.py": "f28ea66e2d862a2a5814970b2ce40c2f7d8296ff09aed90a7e7def689b906944", +} +KIMI_K31_SOURCE_SHA256 = { + "tiktoken.model": "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103", + "tokenizer_config.json": "9fc322181b8ec44236e148916aea4b6c2467b66ded7bfaff0efe623ad750b852", + "tokenization_kimi.py": "204ee39fbf13a6d6608b0d3ca1b4b04a14ac5f0a062a7fdee02e61f5d936f9c6", +} +KIMI_SOURCE_SHA256_BY_PROFILE = { + KIMI_K3_PROFILE: KIMI_K3_SOURCE_SHA256, + KIMI_K31_PROFILE: KIMI_K31_SOURCE_SHA256, +} +KIMI_REPOSITORY_BY_PROFILE = { + KIMI_K3_PROFILE: KIMI_K3_REPOSITORY, + KIMI_K31_PROFILE: KIMI_K31_REPOSITORY, +} +KIMI_REVISION_BY_PROFILE = { + KIMI_K3_PROFILE: KIMI_K3_REVISION, + KIMI_K31_PROFILE: KIMI_K31_REVISION, +} +_PROFILE_BY_REPOSITORY = {repository: profile for profile, repository in KIMI_REPOSITORY_BY_PROFILE.items()} +_ENCODER_MODULES = { + KIMI_K3_PROFILE: "fireworks.training.sdk._vendored.kimi_k3.encoding_k3", + KIMI_K31_PROFILE: "fireworks.training.sdk._vendored.kimi_k31.encoding_k3", +} +_PROFILE_KEY = "kimi_native_profile" + +# ``TikTokenTokenizer.pat_str`` and ``num_reserved_special_tokens`` from the +# pinned ``tokenization_kimi.py`` (identical in both profiles). +KIMI_K3_K31_PAT_STR = "|".join( + [ + r"""[\p{Han}]+""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""\p{N}{1,3}""", + r""" ?[^\s\p{L}\p{N}]+[\r\n]*""", + r"""\s*[\r\n]+""", + r"""\s+(?!\S)""", + r"""\s+""", + ] +) +KIMI_NUM_RESERVED_SPECIAL_TOKENS = 256 +KIMI_BASE_VOCAB_SIZE = 163_584 + +_BOS, _EOS, _UNK, _PAD = "[BOS]", "[EOS]", "[UNK]", "[PAD]" + + +class _NativeChatTemplate: + """The reviewed tokenizer's chat API, without importing checkpoint Python. + + Keep the standard fast tokenizer class/backend serializable. The callable + is installed on load; only its profile (data, not code) is saved in config. + The two private backends avoid mutating a tokenizer shared by other threads. + """ + + def __init__(self, tokenizer: Any, profile: str) -> None: + self.tokenizer = tokenizer + self.profile = profile + + def __getstate__(self) -> dict[str, Any]: + # tokenizers' pickle does not preserve encode_special_tokens. Rebuild + # our private backends so literal control-looking text stays ordinary. + state = self.__dict__.copy() + state.pop("backends", None) + return state + + @property + def encoder(self) -> Any: + # Do not store a module on the tokenizer: ordinary tokenizer pickling + # must continue to work. import_module already caches installed code. + return import_module(_ENCODER_MODULES[self.profile]) + + @cached_property + def backends(self) -> tuple[Any, Any]: + # lazy: tokenize=False renders without copying a large vocabulary. + from tokenizers import Tokenizer + + serialized = self.tokenizer.backend_tokenizer.to_str() + special, text = Tokenizer.from_str(serialized), Tokenizer.from_str(serialized) + for backend in (special, text): + backend.no_padding() + backend.no_truncation() + special.encode_special_tokens = False + text.encode_special_tokens = True + return special, text + + def encode_segments(self, segments: list[Any]) -> list[int]: + tokens: list[int] = [] + for segment in segments: + backend = self.backends[0 if segment.allow_special else 1] + # Same chunk boundaries as the release TikTokenTokenizer. + text = segment.text + for offset in range(0, len(text), 400_000): + chunk = text[offset : offset + 400_000] + start, length = 0, 0 + is_space = chunk[0].isspace() if chunk else False + for index, char in enumerate(chunk): + current = char.isspace() + length = length + 1 if current == is_space else 1 + is_space = current + if length > 25_000: + tokens.extend(backend.encode(chunk[start:index], add_special_tokens=False).ids) + start, length = index, 1 + if chunk[start:]: + tokens.extend(backend.encode(chunk[start:], add_special_tokens=False).ids) + return tokens + + def __call__( + self, + conversation: Any, + tools: list[dict[str, Any]] | None = None, + tokenize: bool = False, + add_generation_prompt: bool = True, + thinking: bool = True, + padding: Any = False, + truncation: bool = False, + max_length: int | None = None, + return_tensors: Any = None, + return_dict: bool = False, + **kwargs: Any, + ) -> Any: + is_batched = self.encoder.is_batched_conversation(conversation) + conversations = conversation if is_batched else [conversation] + image_prompts = kwargs.pop("image_prompts", None) + if is_batched and image_prompts is not None: + raise ValueError("image_prompts is only supported for one chat.") + kwargs.setdefault("thinking_effort", "max") + batches = [ + self.encoder.build_chat_segments( + messages, + tools=tools, + add_generation_prompt=add_generation_prompt, + thinking=thinking, + image_prompts=image_prompts, + **kwargs, + ) + for messages in conversations + ] + if not tokenize: + rendered = ["".join(segment.text for segment in segments) for segments in batches] + return rendered if is_batched else rendered[0] + encoded = [self.encode_segments(segments) for segments in batches] + if truncation and max_length is not None: + encoded = [ids[:max_length] for ids in encoded] + if not (is_batched or padding or return_tensors is not None or return_dict): + return encoded[0] + batch = self.tokenizer.pad( + [{"input_ids": ids, "attention_mask": [1] * len(ids)} for ids in encoded], + padding=padding, + max_length=max_length if padding else None, + return_attention_mask=True, + return_tensors=return_tensors, + ) + if return_dict: + return batch + if is_batched or return_tensors is not None: + return batch["input_ids"] + return batch["input_ids"][0] + + +def _attach_chat_template(tokenizer: Any, profile: str | None) -> Any: + if profile is not None: + if not isinstance(profile, str) or profile not in _ENCODER_MODULES: + raise ValueError(f"unknown Kimi tokenizer profile {profile!r}") + tokenizer.init_kwargs[_PROFILE_KEY] = profile + tokenizer.apply_chat_template = _NativeChatTemplate(tokenizer, profile) + return tokenizer + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def verify_source_dir(source_dir: str | Path, *, profile: str | None = None) -> str: + """Fail closed unless every pinned source file is present and unchanged. + + With ``profile`` unset, either reviewed profile is accepted; the matched + profile name is returned. The tiktoken vocabulary digest is shared, so a + directory mixing one profile's vocabulary with the other's + ``tokenizer_config.json``/``tokenization_kimi.py`` is rejected. + """ + root = Path(source_dir) + profiles = [profile] if profile is not None else sorted(KIMI_SOURCE_SHA256_BY_PROFILE) + mismatched: dict[str, list[str]] = {} + for candidate in profiles: + pins = KIMI_SOURCE_SHA256_BY_PROFILE.get(candidate) + if pins is None: + raise ValueError(f"unknown Kimi tokenizer profile {candidate!r}") + problems = [] + for filename, expected in sorted(pins.items()): + path = root / filename + if not path.is_file(): + problems.append(f"{filename}: missing") + elif _sha256(path) != expected: + problems.append(f"{filename}: sha256 {_sha256(path)} != {expected}") + if not problems: + return candidate + mismatched[candidate] = problems + details = "; ".join(f"{name}: " + ", ".join(problems) for name, problems in mismatched.items()) + raise ValueError(f"Kimi K3/K3.1 tokenizer sources do not match the reviewed pins: {details}") + + +def special_tokens(source_dir: str | Path) -> list[tuple[str, int]]: + """Every special token in id order, named exactly as ``TikTokenTokenizer`` names them.""" + root = Path(source_dir) + config = json.loads((root / "tokenizer_config.json").read_text(encoding="utf-8")) + named = {int(token_id): entry["content"] for token_id, entry in config.get("added_tokens_decoder", {}).items()} + first = KIMI_BASE_VOCAB_SIZE + return [ + (named.get(token_id, f"<|reserved_token_{token_id}|>"), token_id) + for token_id in range(first, first + KIMI_NUM_RESERVED_SPECIAL_TOKENS) + ] + + +def build_fast_tokenizer(source_dir: str | Path, *, profile: str | None = None) -> Any: + """Convert the pinned tiktoken vocabulary into an id-identical fast tokenizer. + + ``TikTokenConverter`` appends special tokens after the base vocabulary in + the order given, so every one of the 256 reserved ids is added explicitly in + id order; passing only the named tokens would shift their ids. + """ + # lazy: conversion needs the optional tokenizers/transformers stack. + from tokenizers import AddedToken + from transformers import PreTrainedTokenizerFast as TokenizersBackend + from transformers.convert_slow_tokenizer import TikTokenConverter + + root = Path(source_dir) + profile = verify_source_dir(root, profile=profile) + backend = TikTokenConverter( + vocab_file=str(root / "tiktoken.model"), + pattern=KIMI_K3_K31_PAT_STR, + additional_special_tokens=[], + ).converted() + if backend.get_vocab_size() != KIMI_BASE_VOCAB_SIZE: + raise ValueError( + f"Kimi K3/K3.1 base vocabulary has {backend.get_vocab_size()} tokens, expected {KIMI_BASE_VOCAB_SIZE}" + ) + ordered = special_tokens(root) + backend.add_special_tokens([AddedToken(name, special=True, normalized=False) for name, _ in ordered]) + for name, token_id in ordered: + if backend.token_to_id(name) != token_id: + raise ValueError(f"special token {name!r} landed at {backend.token_to_id(name)}, expected {token_id}") + tokenizer = TokenizersBackend( + tokenizer_object=backend, + bos_token=_BOS, + eos_token=_EOS, + unk_token=_UNK, + pad_token=_PAD, + clean_up_tokenization_spaces=False, + ) + # Renderer pickling reconstructs the tokenizer from ``name_or_path``; the + # verified source dir reloads through the same code-free path. + tokenizer.name_or_path = str(root) + return _attach_chat_template(tokenizer, profile) + + +def _download_source( + tokenizer_model: str, + *, + revision: str | None, + local_files_only: bool, +) -> Path: + """Fetch only tokenizer data/provenance files, never model weights.""" + # lazy: Hugging Face access is needed only for hub sources. + from huggingface_hub import snapshot_download + + return Path( + snapshot_download( + repo_id=tokenizer_model, + revision=revision, + allow_patterns=["tiktoken.model", "tokenizer.json", "tokenizer_config.json", "tokenization_kimi.py"], + local_files_only=local_files_only, + ) + ) + + +def load_native_tokenizer( + tokenizer_model: str, + *, + revision: str | None = None, + local_files_only: bool = False, + profile: str | None = None, +) -> Any: + """Load raw reviewed Kimi data or an installed-code tokenizer override. + + This is the code-free replacement for + ``AutoTokenizer.from_pretrained(..., trust_remote_code=True)`` on + recognized Kimi sources. It reads only data files, verifies every pinned + byte digest, and builds the tokenizer from reviewed installed code. + Raw tiktoken inputs always retain all pin checks, even if tokenizer.json + is also present. Saved fast bundles are already converted: load their JSON + directly with the installed backend, ignoring auto_map. Selected Kimi + context constrains an unrelated override to installed code, but does not + force it through raw conversion or assign it a Kimi chat profile. + + Official repositories default to their pinned reviewed revisions so the + byte digests below certify immutable data, not a mutable ``main``. + """ + source = Path(tokenizer_model) + if not source.is_dir(): + if profile is None: + profile = _PROFILE_BY_REPOSITORY.get(tokenizer_model) + if revision is None and profile is not None: + revision = KIMI_REVISION_BY_PROFILE[profile] + source = _download_source(tokenizer_model, revision=revision, local_files_only=local_files_only) + if not any((source / name).is_file() for name in ("tiktoken.model", "tokenizer.json")): + # A built-in slow override may need vocab files outside our raw + # conversion allowlist. Keep its Hub identity/revision, but never + # allow checkpoint code to resolve its class. + # lazy: only built-in slow overrides need AutoTokenizer. + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained( + tokenizer_model, revision=revision, local_files_only=local_files_only, trust_remote_code=False + ) + tokenizer = load_native_tokenizer(str(source), profile=profile) + # Restore through the repository identity (re-pinned and re-verified) + # rather than a local cache path that may not exist on another host. + tokenizer.name_or_path = tokenizer_model + return tokenizer + if (source / "tiktoken.model").is_file(): + return build_fast_tokenizer(source, profile=profile) + if (source / "tokenizer.json").is_file(): + # lazy: direct backend loading cannot resolve checkpoint auto_map. + from transformers import PreTrainedTokenizerFast as TokenizersBackend + + config_path = source / "tokenizer_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) if config_path.is_file() else {} + tokenizer = TokenizersBackend.from_pretrained( + str(source), tokenizer_file=str(source / "tokenizer.json"), trust_remote_code=False + ) + # A saved profile is explicit data. An unrelated override gets no + # inferred Kimi chat encoder, even when the base model is Kimi. + return _attach_chat_template(tokenizer, config.get(_PROFILE_KEY)) + # lazy: built-in slow overrides remain valid, with checkpoint Python denied. + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(str(source), trust_remote_code=False, local_files_only=local_files_only) + + +__all__ = [ + "KIMI_BASE_VOCAB_SIZE", + "KIMI_K3_K31_PAT_STR", + "KIMI_K3_PROFILE", + "KIMI_K3_REPOSITORY", + "KIMI_K3_REVISION", + "KIMI_K3_SOURCE_SHA256", + "KIMI_K31_PROFILE", + "KIMI_K31_REPOSITORY", + "KIMI_K31_REVISION", + "KIMI_K31_SOURCE_SHA256", + "KIMI_NATIVE_TOKENIZER_CAPABILITY", + "KIMI_NUM_RESERVED_SPECIAL_TOKENS", + "KIMI_SOURCE_SHA256_BY_PROFILE", + "build_fast_tokenizer", + "load_native_tokenizer", + "special_tokens", + "verify_source_dir", +] diff --git a/src/fireworks/training/sdk/tokenizer_identity.py b/src/fireworks/training/sdk/tokenizer_identity.py index b991113..48363fe 100644 --- a/src/fireworks/training/sdk/tokenizer_identity.py +++ b/src/fireworks/training/sdk/tokenizer_identity.py @@ -1,9 +1,18 @@ -"""Restrict-only signals for code-free DeepSeek V4 tokenizer loading. +"""Restrict-only signals for code-free tokenizer loading of reviewed models. -An exact reviewed canonical V4 identifier or recognized config/tokenizer +An exact reviewed canonical identifier or recognized config/tokenizer data forces loaders through installed code. These signals never establish artifact provenance or authorize Python execution. Without one, failed data reads leave the caller's existing policy unchanged. + +Currently covered families: + +* DeepSeek V4 (``deepseek_v4``), recognized by ``tokenizer.json`` digests; +* Kimi K3 / K3.1 (``kimi_k3`` / ``kimi_k31``), recognized by the shared + tiktoken vocabulary digest plus a pinned ``tokenizer_config.json`` digest. + The two models ship byte-identical ``tiktoken.model`` vocabularies and + equivalent tokenization; their ``encoding_k3.py`` chat encoders are + distinct reviewed contracts and are never aliased. """ from __future__ import annotations @@ -24,7 +33,9 @@ # module only when the constant is present with this exact value, so an # intermediate SDK carrying the module with an older policy revision falls # back to the pinned snapshot instead of inheriting the older behavior. -TOKENIZER_IDENTITY_POLICY_CAPABILITY = "deepseek-v4-tokenizer-identity/separate-family-and-code-policy" +# Capability history: "deepseek-v4-tokenizer-identity/separate-family-and-code-policy" +# covered DeepSeek V4 only; the current value adds the Kimi K3/K3.1 family. +TOKENIZER_IDENTITY_POLICY_CAPABILITY = "deepseek-v4+kimi-k3-k31-tokenizer-identity/saved-native-profiles-v2" # Official DeepSeek V4 ``tokenizer.json`` content digests: # 1. The hub snapshot form (identical for DeepSeek-V4-Flash-0731 and @@ -85,6 +96,55 @@ } ) +# --------------------------------------------------------------------------- +# Kimi K3 / K3.1 (Moonshot AI) +# --------------------------------------------------------------------------- +# +# Both checkpoints ship a tiktoken vocabulary (``tiktoken.model``) behind a +# remote-code ``TikTokenTokenizer`` (``auto_map`` on ``tokenizer_config.json``) +# and no ``tokenizer.json``, so the DeepSeek V4 digest file does not apply. +# The recognition data signal is the shared vocabulary digest plus a pinned +# ``tokenizer_config.json`` digest. The checkpoint Python +# (``tokenization_kimi.py``) is never read for identity and never imported on +# any path. +KIMI_K3_MODEL_TYPE = "kimi_k3" +KIMI_K31_MODEL_TYPE = "kimi_k31" +KIMI_K3_K31_MODEL_TYPES = frozenset({KIMI_K3_MODEL_TYPE, KIMI_K31_MODEL_TYPE}) + +# Digest-reviewed Hugging Face repositories, recognized without a read. +KIMI_K3_REPOSITORY = "moonshotai/Kimi-K3" +KIMI_K31_REPOSITORY = "moonshotai/Kimi-K3.1-Pro" +KIMI_K3_K31_REPOSITORIES = frozenset({KIMI_K3_REPOSITORY, KIMI_K31_REPOSITORY}) + +# Byte-identical in ``moonshotai/Kimi-K3@301be1b8`` and +# ``moonshotai/Kimi-K3.1-Pro@6d0a0202``: one shared BPE vocabulary. +KIMI_K3_K31_TIKTOKEN_MODEL_SHA256 = "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103" + +# The reviewed ``tokenizer_config.json`` of each model. Both declare +# ``TikTokenTokenizer`` with an ``auto_map`` entry, which is exactly the +# remote-code surface code-free loading bypasses. +KIMI_K3_TOKENIZER_CONFIG_SHA256 = "5d0803c94db9cd78763499e0956c95fd5a225c14a727e5a6cf5db3f96f010a6e" +KIMI_K31_TOKENIZER_CONFIG_SHA256 = "9fc322181b8ec44236e148916aea4b6c2467b66ded7bfaff0efe623ad750b852" +KIMI_K3_K31_TOKENIZER_CONFIG_SHA256 = frozenset({KIMI_K3_TOKENIZER_CONFIG_SHA256, KIMI_K31_TOKENIZER_CONFIG_SHA256}) + +# Reviewed canonical Kimi K3/K3.1 identifiers (normalized lowercase form): +# the digest-pinned Hugging Face repositories and the Fireworks resource IDs +# onboarded in the renderer capability registry. Exact membership only, like +# the DeepSeek V4 names. +KIMI_K3_MODEL_NAMES = frozenset( + { + "moonshotai/kimi-k3", + "accounts/fireworks/models/kimi-k3", + } +) +KIMI_K31_MODEL_NAMES = frozenset( + { + "moonshotai/kimi-k3.1-pro", + "accounts/fireworks/models/kimi-k3p1-pro", + } +) +KIMI_K3_K31_MODEL_NAMES = KIMI_K3_MODEL_NAMES | KIMI_K31_MODEL_NAMES + def is_deepseek_v4_model_name(model_name: str | None) -> bool: """Whether ``model_name`` is reviewed DeepSeek V4 context. @@ -111,6 +171,22 @@ def is_deepseek_v4_model_name(model_name: str | None) -> bool: return normalized == DEEPSEEK_V4_MODEL_TYPE or normalized in DEEPSEEK_V4_MODEL_NAMES +def is_kimi_k3_k31_model_name(model_name: str | None) -> bool: + """Whether ``model_name`` is reviewed Kimi K3/K3.1 context. + + Same contract as :func:`is_deepseek_v4_model_name`: the exact + ``model_type`` architecture strings (``kimi_k3`` / ``kimi_k31``) as + independently selected control-plane context, or exact membership in + the reviewed canonical IDs (``KIMI_K3_K31_MODEL_NAMES``). Basenames, + ancestors and substrings never contribute identity. This signal can + only restrict code loading; it never authorizes execution. + """ + if not model_name: + return False + normalized = model_name.strip().rstrip("/").lower() + return normalized in KIMI_K3_K31_MODEL_TYPES or normalized in KIMI_K3_K31_MODEL_NAMES + + def _read_json(path: Path) -> dict[str, Any] | None: try: parsed = json.loads(path.read_text(encoding="utf-8")) @@ -331,15 +407,45 @@ def is_deepseek_v4_tokenizer_source( ) +def is_kimi_k3_k31_tokenizer_source( + tokenizer_model: str | None, + *, + revision: str | None = None, + local_files_only: bool = False, +) -> bool: + """Whether a Kimi K3/K3.1 signal restricts ``tokenizer_model`` to installed code. + + Same data-first contract as :func:`is_deepseek_v4_tokenizer_source`, + with the Kimi data signal: a local artifact directory or hub source is + recognized through ``config.json`` ``model_type`` in + ``KIMI_K3_K31_MODEL_TYPES`` (read as JSON; never an ``AutoConfig`` load, + never code execution), or through the shared ``tiktoken.model`` + vocabulary digest combined with a pinned ``tokenizer_config.json`` + digest. An explicit different model family takes precedence over shared + tokenizer bytes for family recognition; missing/empty model types and + the generic ``"model"`` saved-tokenizer marker permit the digest + fallback. Directory basenames and ancestors never contribute identity. + + This is semantic family recognition, not permission to load code, and + never establishes provenance. Any read or transport failure returns + ``False`` (not recognized), leaving the caller's current family + decision unchanged. + """ + return _kimi_k3_k31_source_matches( + tokenizer_model, revision=revision, local_files_only=local_files_only, explicit_family_wins=True + ) + + def requires_code_free_tokenizer_loading( tokenizer_model: str | None, *, revision: str | None = None, local_files_only: bool = False, ) -> bool: - """Restrict recognized V4 data to installed tokenizer/processor code. + """Restrict recognized reviewed data to installed tokenizer/processor code. - A shared reviewed tokenizer digest is restrictive even when mutable + Covers DeepSeek V4 and Kimi K3/K3.1. A shared reviewed tokenizer digest + is restrictive even when mutable ``model_type`` names another family. Unlike semantic family recognition, this must not select a renderer or encoder: it only keeps checkpoint Python inert. Unknown data retains the caller's existing policy. @@ -347,7 +453,83 @@ def requires_code_free_tokenizer_loading( """ return _tokenizer_source_matches( tokenizer_model, revision=revision, local_files_only=local_files_only, explicit_family_wins=False + ) or _kimi_k3_k31_source_matches( + tokenizer_model, revision=revision, local_files_only=local_files_only, explicit_family_wins=False + ) + + +def _local_kimi_k3_k31_data_matches(root: Path) -> bool: + """Raw pinned data or a saved native profile; neither authorizes code.""" + config = _read_json(root / "tokenizer_config.json") or {} + profile = config.get("kimi_native_profile") + if isinstance(profile, str) and profile in KIMI_K3_K31_MODEL_TYPES and (root / "tokenizer.json").is_file(): + return True + return ( + _sha256(root / "tiktoken.model") == KIMI_K3_K31_TIKTOKEN_MODEL_SHA256 + and _sha256(root / "tokenizer_config.json") in KIMI_K3_K31_TOKENIZER_CONFIG_SHA256 + ) + + +def _hub_kimi_k3_k31_data_matches(model: str, *, revision: str | None, local_files_only: bool) -> bool: + """Byte-digest recognition of the reviewed Kimi data on the Hub. + + No metadata prefilter: ``tiktoken.model`` is not LFS-pinned by digest on + every mirror, and the file is small enough (2.8 MB) that a cached byte + read is the reliable signal. Both reads are data-only downloads. + """ + config_path = _hub_download(model, "tokenizer_config.json", revision=revision, local_files_only=local_files_only) + config = _read_json(config_path) if config_path is not None else None + profile = config.get("kimi_native_profile") if config else None + if isinstance(profile, str) and profile in KIMI_K3_K31_MODEL_TYPES: + return _hub_download(model, "tokenizer.json", revision=revision, local_files_only=local_files_only) is not None + vocab_path = _hub_download(model, "tiktoken.model", revision=revision, local_files_only=local_files_only) + if vocab_path is None or _sha256(vocab_path) != KIMI_K3_K31_TIKTOKEN_MODEL_SHA256: + return False + config_path = _hub_download(model, "tokenizer_config.json", revision=revision, local_files_only=local_files_only) + return config_path is not None and _sha256(config_path) in KIMI_K3_K31_TOKENIZER_CONFIG_SHA256 + + +def _kimi_k3_k31_source_matches( + tokenizer_model: str | None, + *, + revision: str | None, + local_files_only: bool, + explicit_family_wins: bool, +) -> bool: + """Kimi K3/K3.1 leg of source matching; see the public wrappers.""" + if not tokenizer_model: + return False + model = tokenizer_model + local_dir = Path(model) + if local_dir.is_dir(): + config = _read_json(local_dir / "config.json") + model_type = config.get("model_type") if config is not None else None + if isinstance(model_type, str) and model_type not in {"", "model"}: + if model_type in KIMI_K3_K31_MODEL_TYPES: + return True + if explicit_family_wins: + return False + return _local_kimi_k3_k31_data_matches(local_dir) + # Only the digest-reviewed pair skips the advisory read; reviewed + # Fireworks aliases are selected-context identities + # (is_kimi_k3_k31_model_name), not no-read source shortcuts. + if model in KIMI_K3_K31_REPOSITORIES: + return True + config_path = _hub_download( + model, + "config.json", + revision=revision, + local_files_only=local_files_only, ) + if config_path is not None: + config = _read_json(config_path) + model_type = config.get("model_type") if config is not None else None + if isinstance(model_type, str) and model_type not in {"", "model"}: + if model_type in KIMI_K3_K31_MODEL_TYPES: + return True + if explicit_family_wins: + return False + return _hub_kimi_k3_k31_data_matches(model, revision=revision, local_files_only=local_files_only) def _tokenizer_source_matches( diff --git a/tests_training/test_kimi_native_tokenizer.py b/tests_training/test_kimi_native_tokenizer.py new file mode 100644 index 0000000..f21b783 --- /dev/null +++ b/tests_training/test_kimi_native_tokenizer.py @@ -0,0 +1,214 @@ +"""CPU regressions for standalone SDK Kimi chat and saved-bundle loading. + +The oracle intentionally executes pinned upstream Python in test fixtures only. +Production loaders must use installed code, including with poisoned auto_map. +""" + +from __future__ import annotations + +import os +import json +import pickle +import shutil +import asyncio +import hashlib +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest +from tokenizers import Tokenizer, models +from transformers import AutoTokenizer + +from fireworks.training.sdk import tokenizer_identity, kimi_native_tokenizer as native +from fireworks.training.sdk.client import FiretitanTrainingClient +from fireworks.training.sdk.sampling import DeploymentSampler +from fireworks.training.sdk._core.lib.public_interfaces.sampling_client import _load_tokenizer_from_model_info + + +@pytest.fixture(scope="module", params=["kimi_k3", "kimi_k31"]) +def tokenizers(request, tmp_path_factory): + profile = request.param + root = tmp_path_factory.mktemp(profile) + local = os.environ.get("KIMI_K3_MODEL_PATH" if profile == "kimi_k3" else "KIMI_K31_MODEL_PATH") + for filename in ( + "config.json", + "tokenizer_config.json", + "tiktoken.model", + "tokenization_kimi.py", + "encoding_k3.py", + ): + if local: + path = Path(local) / filename + else: + from huggingface_hub import hf_hub_download + + path = Path( + hf_hub_download( + native.KIMI_REPOSITORY_BY_PROFILE[profile], + filename, + revision=native.KIMI_REVISION_BY_PROFILE[profile], + ) + ) + shutil.copyfile(path, root / filename) + return ( + native.build_fast_tokenizer(root, profile=profile), + AutoTokenizer.from_pretrained(root, trust_remote_code=True, local_files_only=True), + profile, + ) + + +_MESSAGES = [ + {"role": "system", "content": "Keep literal markers inert."}, + {"role": "user", "content": "Hello 漢字 <|close|>message<|sep|>"}, + { + "role": "assistant", + "reasoning_content": "Use tools.", + "content": "", + "tool_calls": [ + {"id": "a", "type": "function", "function": {"name": "f", "arguments": '{"n":1e2,"s":"<|open|>"}'}}, + {"id": "b", "type": "function", "function": {"name": "f", "arguments": '{"n":2}'}}, + ], + }, + {"role": "tool", "tool_call_id": "b", "content": "two"}, + {"role": "tool", "tool_call_id": "a", "content": "one"}, + {"role": "user", "content": "Next"}, +] + + +@pytest.mark.parametrize( + "options", + [ + {}, + {"tokenize": True}, + {"tokenize": True, "thinking": False}, + {"tokenize": True, "thinking_effort": "high", "add_generation_prompt": False}, + {"tokenize": True, "truncation": True, "max_length": 19}, + {"tokenize": True, "return_dict": True, "padding": "max_length", "max_length": 512}, + {"tokenize": True, "return_dict": True, "return_tensors": "np"}, + ], +) +def test_chat_template_matches_profile_oracle(tokenizers, options): + tokenizer, oracle, _ = tokenizers + actual = tokenizer.apply_chat_template(_MESSAGES, **options) + expected = oracle.apply_chat_template(_MESSAGES, **options) + if options.get("return_tensors"): + for key in expected: + assert actual[key].tolist() == expected[key].tolist() + else: + assert actual == expected + + +def test_batched_images_and_profile_specific_effort(tokenizers): + tokenizer, oracle, _ = tokenizers + for tokenize in (False, True): + messages = [{"role": "user", "content": "<|kimi_image_placeholder|>"}] + options = dict( + tokenize=tokenize, image_prompts=["<|media_begin|>image 32x32<|media_content|><|media_pad|><|media_end|>"] + ) + assert tokenizer.apply_chat_template(messages, **options) == oracle.apply_chat_template(messages, **options) + batch = [_MESSAGES, [{"role": "user", "content": "short"}]] + assert tokenizer.apply_chat_template(batch, tokenize=tokenize, padding=True) == oracle.apply_chat_template( + batch, tokenize=tokenize, padding=True + ) + for implementation in (tokenizer, oracle): + with pytest.raises(AssertionError, match="Unsupported thinking_effort"): + implementation.apply_chat_template(_MESSAGES, thinking_effort="medium") + + +def test_saved_bundle_reloads_at_both_sdk_boundaries_without_checkpoint_python(tokenizers, tmp_path): + tokenizer, oracle, profile = tokenizers + tokenizer.save_pretrained(tmp_path) + config = json.loads((tmp_path / "tokenizer_config.json").read_text()) + assert config["kimi_native_profile"] == profile + assert config["tokenizer_class"] in {"TokenizersBackend", "PreTrainedTokenizerFast"} + marker = tmp_path / "executed" + config["auto_map"] = {"AutoTokenizer": ["canary.Canary", "canary.Canary"]} + (tmp_path / "tokenizer_config.json").write_text(json.dumps(config)) + (tmp_path / "config.json").write_text( + json.dumps({"model_type": profile, "auto_map": {"AutoConfig": "canary.CanaryConfig"}}) + ) + (tmp_path / "canary.py").write_text( + f"from pathlib import Path\nPath({str(marker)!r}).touch()\nraise RuntimeError('ran')\n" + ) + assert tokenizer_identity.is_kimi_k3_k31_tokenizer_source(str(tmp_path)) + client = object.__new__(FiretitanTrainingClient) + client._tokenizer_model = str(tmp_path) + client._tokenizer_base_model = native.KIMI_REPOSITORY_BY_PROFILE[profile] + for loaded in ( + native.load_native_tokenizer(str(tmp_path)), + client.get_tokenizer(), + _load_tokenizer_from_model_info(client._tokenizer_base_model, str(tmp_path)), + ): + assert loaded.apply_chat_template(_MESSAGES, tokenize=True) == oracle.apply_chat_template( + _MESSAGES, tokenize=True + ) + assert not marker.exists() + restored = pickle.loads(pickle.dumps(tokenizer)) + assert restored.apply_chat_template(_MESSAGES, tokenize=True) == oracle.apply_chat_template( + _MESSAGES, tokenize=True + ) + + +def test_selected_kimi_does_not_force_unrelated_fast_override_through_raw_conversion(tmp_path): + backend = Tokenizer(models.WordLevel({"[UNK]": 0, "hello": 1}, unk_token="[UNK]")) + backend.save(str(tmp_path / "tokenizer.json")) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "PreTrainedTokenizerFast", + "chat_template": "{% for message in messages %}{{ message.content }}{% endfor %}", + "auto_map": {"AutoTokenizer": ["missing.Unsafe", "missing.Unsafe"]}, + } + ) + ) + loaded = _load_tokenizer_from_model_info("moonshotai/Kimi-K3", str(tmp_path)) + assert loaded.encode("hello", add_special_tokens=False) == [1] + assert loaded.apply_chat_template([{"role": "user", "content": "hello"}], tokenize=True, return_dict=False) == [1] + + +def test_builtin_slow_hub_override_retains_source_and_denies_checkpoint_code(tmp_path, monkeypatch): + from unittest.mock import Mock + + download = Mock(return_value=tmp_path) + load = Mock(return_value=object()) + monkeypatch.setattr(native, "_download_source", download) + monkeypatch.setattr(AutoTokenizer, "from_pretrained", load) + assert native.load_native_tokenizer("acme/builtin-override", revision="pinned", local_files_only=True) is load.return_value + load.assert_called_once_with( + "acme/builtin-override", revision="pinned", local_files_only=True, trust_remote_code=False + ) + + +@pytest.mark.parametrize("profile", [["kimi_k3"], {}, "unknown"]) +def test_invalid_saved_profile_cannot_select_an_encoder(tmp_path, profile): + backend = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + backend.save(str(tmp_path / "tokenizer.json")) + (tmp_path / "tokenizer_config.json").write_text(json.dumps({"kimi_native_profile": profile})) + assert not tokenizer_identity.is_kimi_k3_k31_tokenizer_source(str(tmp_path)) + with pytest.raises(ValueError, match="unknown Kimi tokenizer profile"): + native.load_native_tokenizer(str(tmp_path)) + + +def test_message_sampler_uses_restored_template_and_token_sampler_stays_primitive(tokenizers): + tokenizer, oracle, _ = tokenizers + sampler = object.__new__(DeploymentSampler) + sampler.tokenizer = tokenizer + sampler._do_one_completion = AsyncMock(return_value=[]) + assert asyncio.run(sampler.sample_with_tokens(_MESSAGES)) == [] + assert sampler._do_one_completion.await_args.args[0] == oracle.apply_chat_template( + _MESSAGES, tokenize=True, add_generation_prompt=True, return_dict=False + ) + sampler._do_one_completion.reset_mock() + assert asyncio.run(sampler.sample_with_prompt_tokens([123, 456])) == [] + assert sampler._do_one_completion.await_args.args[0] == [123, 456] + + +def test_sdk_vendors_match_reviewed_lint_adapted_profile_sources(): + for profile, digest in ( + ("kimi_k3", "b04b6ecd223d0b8519c5926f095f9531e97396c915595674effcc39fd4658323"), + ("kimi_k31", "5ca26282a0f72a6786addd95958a8117605d02ecbec9951ae9b67e54a121ddd5"), + ): + directory = Path(native.__file__).parent / "_vendored" / profile + assert hashlib.sha256((directory / "encoding_k3.py").read_bytes()).hexdigest() == digest + assert (directory / "LICENSE").is_file() + assert (directory / "PROVENANCE.md").is_file() diff --git a/tests_training/test_tokenizer_identity.py b/tests_training/test_tokenizer_identity.py index 6197ed0..a1ba83b 100644 --- a/tests_training/test_tokenizer_identity.py +++ b/tests_training/test_tokenizer_identity.py @@ -505,3 +505,156 @@ def test_other_model_kimi_policies_and_revisions_remain_unchanged(monkeypatch, t sampling_client._load_tokenizer_from_model_info.cache_clear() sampling_client._load_tokenizer_from_model_info("unmigrated-model", tokenizer) assert calls == [{"fast": True, "trust_remote_code": True, "revision": revision}] + + +# --------------------------------------------------------------------------- +# Kimi K3/K3.1 identity (local only) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "model_name,expected", + [ + (None, False), + ("", False), + # Reviewed canonical IDs: the digest-pinned HF repositories and the + # registry-onboarded Fireworks resource IDs. + ("moonshotai/Kimi-K3", True), + ("moonshotai/Kimi-K3.1-Pro", True), + ("accounts/fireworks/models/kimi-k3", True), + ("accounts/fireworks/models/kimi-k3p1-pro", True), + # Registry-style normalization: strip, trailing slashes, case. + ("accounts/fireworks/models/kimi-k3/", True), + ("MOONSHOTAI/KIMI-K3.1-PRO", True), + (" moonshotai/Kimi-K3 ", True), + # The exact architecture model_type strings are independently + # selected control-plane context. + ("kimi_k3", True), + ("kimi_k31", True), + ("Kimi_K31/", True), + # Substrings, basenames, ancestors, namespaces and forks never + # contribute identity. + ("accounts/team/models/kimi-k3-finetune", False), + ("accounts/fireworks/models/kimi-k3p1-pro/finetune-x", False), + ("acme/Kimi-K3", False), + ("moonshotai/Kimi-K30", False), + ("accounts/fireworks/models/kimi-k3p10", False), + ("kimi", False), + ("kimi_k3_k31", False), + # Kimi K2 stays outside the K3/K3.1 set. + ("moonshotai/Kimi-K2-Thinking", False), + ("moonshotai/Kimi-K2.6", False), + ("kimi_k2", False), + ("Qwen/Qwen3-8B", False), + ], +) +def test_kimi_model_name_policy_exact_reviewed_ids(model_name, expected): + assert identity.is_kimi_k3_k31_model_name(model_name) is expected + + +def test_kimi_reviewed_hub_sources_cannot_fail_open_on_metadata_fetch(monkeypatch): + def unavailable(*_args, **_kwargs): + raise AssertionError("reviewed Kimi Hub sources must not need advisory reads") + + monkeypatch.setattr(identity, "_hub_download", unavailable) + for source in identity.KIMI_K3_K31_REPOSITORIES: + assert identity.is_kimi_k3_k31_tokenizer_source(source) + + +def test_kimi_empty_and_missing_inputs_are_not_recognized() -> None: + assert identity.is_kimi_k3_k31_tokenizer_source(None) is False + assert identity.is_kimi_k3_k31_tokenizer_source("") is False + assert identity.is_kimi_k3_k31_tokenizer_source("/nonexistent/path/xyz") is False + + +@pytest.mark.parametrize("model_type", ["kimi_k3", "kimi_k31"]) +def test_kimi_config_model_type_branch(tmp_path: Path, model_type: str) -> None: + (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) + assert identity.is_kimi_k3_k31_tokenizer_source(str(tmp_path)) is True + + +def test_kimi_other_model_type_is_not_recognized(tmp_path: Path) -> None: + (tmp_path / "config.json").write_text(json.dumps({"model_type": "qwen3"})) + assert identity.is_kimi_k3_k31_tokenizer_source(str(tmp_path)) is False + + +@pytest.mark.parametrize("model_type", ["qwen3", "llama", "new_model_family"]) +def test_kimi_explicit_other_model_type_precedes_shared_vocab_digest( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, model_type: str +) -> None: + """Shared vocabulary bytes must not change an explicitly selected family.""" + (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) + vocab = tmp_path / "tiktoken.model" + vocab.write_bytes(b"shared vocab bytes") + config = tmp_path / "tokenizer_config.json" + config.write_bytes(b"shared tokenizer config") + monkeypatch.setattr(identity, "KIMI_K3_K31_TIKTOKEN_MODEL_SHA256", hashlib.sha256(vocab.read_bytes()).hexdigest()) + monkeypatch.setattr( + identity, + "KIMI_K3_K31_TOKENIZER_CONFIG_SHA256", + frozenset({hashlib.sha256(config.read_bytes()).hexdigest()}), + ) + assert identity.is_kimi_k3_k31_tokenizer_source(str(tmp_path)) is False + assert identity.requires_code_free_tokenizer_loading(str(tmp_path)) is True + + +def test_kimi_configless_digest_branch(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + vocab = tmp_path / "tiktoken.model" + vocab.write_bytes(b"shared vocab bytes") + config = tmp_path / "tokenizer_config.json" + config.write_bytes(b"shared tokenizer config") + monkeypatch.setattr(identity, "KIMI_K3_K31_TIKTOKEN_MODEL_SHA256", hashlib.sha256(vocab.read_bytes()).hexdigest()) + monkeypatch.setattr( + identity, + "KIMI_K3_K31_TOKENIZER_CONFIG_SHA256", + frozenset({hashlib.sha256(config.read_bytes()).hexdigest()}), + ) + assert identity.is_kimi_k3_k31_tokenizer_source(str(tmp_path)) is True + vocab.write_bytes(b"different bytes") + assert identity.is_kimi_k3_k31_tokenizer_source(str(tmp_path)) is False + + +def test_kimi_vocab_digest_alone_is_not_enough(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Both the shared vocabulary and a pinned tokenizer_config must match.""" + vocab = tmp_path / "tiktoken.model" + vocab.write_bytes(b"shared vocab bytes") + (tmp_path / "tokenizer_config.json").write_bytes(b"unreviewed config bytes") + monkeypatch.setattr(identity, "KIMI_K3_K31_TIKTOKEN_MODEL_SHA256", hashlib.sha256(vocab.read_bytes()).hexdigest()) + assert identity.is_kimi_k3_k31_tokenizer_source(str(tmp_path)) is False + + +def test_kimi_same_artifact_renamed_has_same_decision(tmp_path: Path) -> None: + """Directory names never contribute identity: the same artifact keeps the + decision its data earns under every name.""" + for name in ("totally-unrelated-name", "Kimi-K3.1-clone", "kimi-k3p1-pro"): + directory = tmp_path / name + directory.mkdir() + (directory / "config.json").write_text(json.dumps({"model_type": "kimi_k31"})) + assert identity.is_kimi_k3_k31_tokenizer_source(str(directory)) is True, name + (directory / "config.json").write_text(json.dumps({"model_type": "qwen3"})) + assert identity.is_kimi_k3_k31_tokenizer_source(str(directory)) is False, name + + +def test_kimi_absent_metadata_renamed_does_not_gain_identity(tmp_path: Path) -> None: + for name in ("Kimi-K3.1-Pro", "kimi-k3", "kimi_k31"): + directory = tmp_path / name + directory.mkdir() + assert identity.is_kimi_k3_k31_tokenizer_source(str(directory)) is False, name + + +def test_kimi_model_type_sentinel_is_context_but_never_source_identity(monkeypatch) -> None: + """``kimi_k31`` (the exact model_type) is valid selected context, yet as a + tokenizer SOURCE it is recognized only through data: with advisory reads + unavailable it must not be recognized by the string alone.""" + assert identity.is_kimi_k3_k31_model_name(identity.KIMI_K31_MODEL_TYPE) is True + + reads = [] + + def unavailable(repo_id, filename, **_kwargs): + reads.append((repo_id, filename)) + return None + + monkeypatch.setattr(identity, "_hub_download", unavailable) + monkeypatch.setattr(identity, "_hub_kimi_k3_k31_data_matches", lambda *_args, **_kwargs: False) + assert identity.is_kimi_k3_k31_tokenizer_source("kimi_k31") is False + assert reads == [("kimi_k31", "config.json")] From 92291224a0d5f739b4a60b121b26c561f685d88c Mon Sep 17 00:00:00 2001 From: Cheng Lyu Date: Fri, 6 Nov 2026 09:34:25 -0800 Subject: [PATCH 2/6] release: 1.2.33 Co-authored-by: yoyoyocmu-fw <349039719+yoyoyocmu-fw@users.noreply.github.com> Co-authored-by: Cheng Lyu --- CHANGELOG.md | 1 + src/fireworks/training/sdk/client.py | 4 +++- src/fireworks/training/sdk/managed.py | 6 ++++- tests_training/test_comms_negotiation.py | 29 ++++++++++++++++-------- tests_training/test_managed_reference.py | 6 +++-- 5 files changed, 32 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 461c985..470a613 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ Full Changelog: [v1.2.32...v1.2.33](https://github.com/fw-ai-external/python-sdk/compare/v1.2.32...v1.2.33) ### Features +* **training-sdk:** support dedicated LoRA full-weight synchronization over RDMA, with capability negotiation and managed provisioning checks * **training-sdk:** load reviewed Kimi K3 and K3.1 tokenizers without executing checkpoint Python, with shared vocabulary and distinct packaged chat encoders. ### Bug Fixes diff --git a/src/fireworks/training/sdk/client.py b/src/fireworks/training/sdk/client.py index b9cd51a..a04d762 100644 --- a/src/fireworks/training/sdk/client.py +++ b/src/fireworks/training/sdk/client.py @@ -1683,7 +1683,7 @@ def _refresh_weight_sync_backend(self) -> Any: if required: raise RuntimeError("RDMA weight sync is enabled but the sampler backend does not support RDMA") return sampler - eligible = sampler.allow_rdma and self._lora_rank == 0 and self.run_name is None and not sampler.cmek_resource + eligible = sampler.allow_rdma and self.run_name is None and not sampler.cmek_resource if not eligible and not required: # Legacy LoRA/session deployments need not expose base-model hotload # status. The rollout enforces its flag at FILE publication as well. @@ -1691,6 +1691,7 @@ def _refresh_weight_sync_backend(self) -> Any: state = sampler.deploy_mgr.hotload_check_status(sampler.deployment_id, sampler.base_model) replicas = state.get("replicas", []) receiver_support = [replica.get("supports_rdma_weight_sync") is True for replica in replicas] + receiver_lora_support = [replica.get("supports_rdma_lora_weight_sync") is True for replica in replicas] receiver_enabled = [replica.get("rdma_weight_sync_enabled") for replica in replicas] required = required or any(receiver_support) or any(flag is True for flag in receiver_enabled) if not required: @@ -1701,6 +1702,7 @@ def _refresh_weight_sync_backend(self) -> Any: or not trainer_support or not receiver_support or not all(receiver_support) + or (self._lora_rank != 0 and not all(receiver_lora_support)) or any(flag is False for flag in receiver_enabled) ): raise RuntimeError( diff --git a/src/fireworks/training/sdk/managed.py b/src/fireworks/training/sdk/managed.py index 2f8bed9..4d56115 100644 --- a/src/fireworks/training/sdk/managed.py +++ b/src/fireworks/training/sdk/managed.py @@ -829,7 +829,11 @@ def _uses_manual_training_infra(config: _ManagedTinkerConfig) -> bool: def _rdma_candidate(config: _ManagedTinkerConfig) -> bool: """Allow capability negotiation without changing platform-owned resources.""" return ( - config.create_deployment and config.lora_rank == 0 and config.max_lora_rank is None and not config.forward_only + config.create_deployment + # LoRA capacity does not imply a multi-session trainer: cookbook also + # uses max_lora_rank for dedicated jobs. Each created client must still + # negotiate explicit trainer and rollout support before using RDMA. + and not config.forward_only ) diff --git a/tests_training/test_comms_negotiation.py b/tests_training/test_comms_negotiation.py index 0d16f34..ed390f7 100644 --- a/tests_training/test_comms_negotiation.py +++ b/tests_training/test_comms_negotiation.py @@ -46,9 +46,7 @@ def receive(request): model_id = ( f"base-{index}" if body.get("base_only") - else f"run-{index}:train:0" - if route == "serverless" - else f"model-{index}" + else f"run-{index}:train:0" if route == "serverless" else f"model-{index}" ) pending[request_id] = {"type": "create_model", "model_id": model_id, **next(models)} return httpx.Response(200, json={"request_id": request_id}) @@ -549,11 +547,13 @@ def test_old_response_model_accepts_router_replay_capability(): assert response.model_id == "test" +@pytest.mark.parametrize("capacity", [None, 16], ids=["legacy-config", "cookbook-capacity"]) @pytest.mark.parametrize("lora_rank", [0, 8], ids=["full", "lora"]) @pytest.mark.parametrize("trainer_support", [False, True], ids=["old-trainer", "new-trainer"]) @pytest.mark.parametrize("inference_support", [False, True], ids=["old-inference", "new-inference"]) -def test_weight_sync_transport_matrix(connect, lora_rank, trainer_support, inference_support): - from fireworks.training.sdk.managed import _TinkerSamplerBackend +@pytest.mark.parametrize("lora_support", [False, True]) +def test_weight_sync_transport_matrix(connect, lora_rank, trainer_support, inference_support, lora_support, capacity): + from fireworks.training.sdk.managed import _rdma_candidate, _ManagedTinkerConfig, _TinkerSamplerBackend # Missing fields reproduce the response of old images, rather than requiring # those images to learn how to advertise a negative capability. @@ -568,24 +568,33 @@ def test_weight_sync_transport_matrix(connect, lora_rank, trainer_support, infer ) manager._hotload_headers.return_value = {} replica = {"supports_rdma_weight_sync": True} if inference_support else {} + replica["supports_rdma_lora_weight_sync"] = lora_support manager.hotload_check_status.return_value = {"replicas": [replica]} manager.hotload_and_wait.return_value = True - if (lora_rank == 0 and trainer_support != inference_support) or (lora_rank != 0 and trainer_support): + config = _ManagedTinkerConfig(base_model="test/model", max_lora_rank=capacity) + backend = _TinkerSamplerBackend( + manager, + "rollout", + "test/model", + lora_rank=lora_rank, + allow_rdma=_rdma_candidate(config), + ) + expected_rdma = trainer_support and inference_support and (lora_rank == 0 or lora_support) + if (trainer_support or inference_support) and not expected_rdma: with pytest.raises(RuntimeError, match="FILE fallback is not allowed"): - client._attach_sampler_backend(_TinkerSamplerBackend(manager, "rollout", "test/model", lora_rank=lora_rank)) + client._attach_sampler_backend(backend) assert not any(op in ("weight_sync", "save_weights_for_sampler") for op, _ in calls) manager.hotload_and_wait.assert_not_called() return - client._attach_sampler_backend(_TinkerSamplerBackend(manager, "rollout", "test/model", lora_rank=lora_rank)) + client._attach_sampler_backend(backend) - expected_rdma = lora_rank == 0 and trainer_support and inference_support assert client.supports_rdma_weight_sync is expected_rdma result = client.weight_sync().result(timeout=5) submitted = [op for op, _ in calls if op in ("weight_sync", "save_weights_for_sampler")] assert submitted == (["weight_sync"] if expected_rdma else ["save_weights_for_sampler"]) assert manager.hotload_and_wait.call_count == (0 if expected_rdma else 1) assert (result.optimizer_version is not None) is expected_rdma - assert manager.hotload_check_status.call_count == (2 if lora_rank == 0 else 0) + assert manager.hotload_check_status.call_count == 2 def _weight_sync_client(connect, *, weight_sync_status=200, rdma=True): diff --git a/tests_training/test_managed_reference.py b/tests_training/test_managed_reference.py index e0718ee..d5ba887 100644 --- a/tests_training/test_managed_reference.py +++ b/tests_training/test_managed_reference.py @@ -55,8 +55,10 @@ def test_managed_config_defaults_use_reservation_true_and_accepts_opt_out(): "overrides,enabled", [ ({}, True), - ({"lora_rank": 8}, False), - ({"max_lora_rank": 8}, False), + ({"lora_rank": 8}, True), + ({"lora_rank": 8, "forward_only": True}, False), + ({"lora_rank": 8, "extra_values": {"testValue": "test-value"}}, True), + ({"max_lora_rank": 8}, True), ({"forward_only": True}, False), ({"create_deployment": False}, False), ({"extra_values": {"testValue": "test-value"}}, True), From c20d17c79214c25eed18299e0fc8e44c4552fd0c Mon Sep 17 00:00:00 2001 From: greg-fw Date: Fri, 6 Nov 2026 18:07:48 -0800 Subject: [PATCH 3/6] release: 1.2.33 Co-authored-by: Claude --- CHANGELOG.md | 3 ++- src/fireworks/types/deployment.py | 6 ++++++ src/fireworks/types/deployment_shape.py | 6 ++++++ src/fireworks/types/model.py | 6 ++++++ src/fireworks/types/shared/deployed_model.py | 6 ++++++ 5 files changed, 26 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 470a613..b24e74b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,10 +1,11 @@ # Changelog -## 1.2.33 (2026-11-06) +## 1.2.33 (2026-11-07) Full Changelog: [v1.2.32...v1.2.33](https://github.com/fw-ai-external/python-sdk/compare/v1.2.32...v1.2.33) ### Features +* **deployments, models, deployed models, deployment shapes:** return `created_by` (the creator's email) on `Deployment`, `Model`, `DeployedModel` and `DeploymentShape`. Only members of the resource's own account see it; for other callers it is empty. * **training-sdk:** support dedicated LoRA full-weight synchronization over RDMA, with capability negotiation and managed provisioning checks * **training-sdk:** load reviewed Kimi K3 and K3.1 tokenizers without executing checkpoint Python, with shared vocabulary and distinct packaged chat encoders. diff --git a/src/fireworks/types/deployment.py b/src/fireworks/types/deployment.py index 8cd7942..b544b66 100644 --- a/src/fireworks/types/deployment.py +++ b/src/fireworks/types/deployment.py @@ -78,6 +78,12 @@ class Deployment(BaseModel): create_time: Optional[datetime] = FieldInfo(alias="createTime", default=None) """The creation time of the deployment.""" + created_by: Optional[str] = FieldInfo(alias="createdBy", default=None) + """The email address of the user who created this deployment. + + Only returned to members of the deployment's account. + """ + delete_time: Optional[datetime] = FieldInfo(alias="deleteTime", default=None) """The time at which the resource will be soft deleted.""" diff --git a/src/fireworks/types/deployment_shape.py b/src/fireworks/types/deployment_shape.py index e1d04a5..87e1785 100644 --- a/src/fireworks/types/deployment_shape.py +++ b/src/fireworks/types/deployment_shape.py @@ -43,6 +43,12 @@ class DeploymentShape(BaseModel): create_time: Optional[datetime] = FieldInfo(alias="createTime", default=None) """The creation time of the deployment shape.""" + created_by: Optional[str] = FieldInfo(alias="createdBy", default=None) + """The email address of the user who created this deployment shape. + + Only returned to members of the deployment shape's account. + """ + description: Optional[str] = None """The description of the deployment shape. diff --git a/src/fireworks/types/model.py b/src/fireworks/types/model.py index c6bea3f..c9855f8 100644 --- a/src/fireworks/types/model.py +++ b/src/fireworks/types/model.py @@ -43,6 +43,12 @@ class Model(BaseModel): create_time: Optional[datetime] = FieldInfo(alias="createTime", default=None) """The creation time of the model.""" + created_by: Optional[str] = FieldInfo(alias="createdBy", default=None) + """The email address of the user who created this model. + + Only returned to members of the model's account. + """ + default_draft_model: Optional[str] = FieldInfo(alias="defaultDraftModel", default=None) """The default draft model to use when creating a deployment. diff --git a/src/fireworks/types/shared/deployed_model.py b/src/fireworks/types/shared/deployed_model.py index 7db95ee..110bdbb 100644 --- a/src/fireworks/types/shared/deployed_model.py +++ b/src/fireworks/types/shared/deployed_model.py @@ -16,6 +16,12 @@ class DeployedModel(BaseModel): create_time: Optional[datetime] = FieldInfo(alias="createTime", default=None) """The creation time of the resource.""" + created_by: Optional[str] = FieldInfo(alias="createdBy", default=None) + """The email address of the user who created this deployed model. + + Only returned to members of the deployed model's account. + """ + default: Optional[bool] = None """ If true, this is the default target when querying this model without the From e8e430337961a93ca97b3b48cdd628271e5b9110 Mon Sep 17 00:00:00 2001 From: zhiweiz Date: Fri, 6 Nov 2026 22:09:15 -0800 Subject: [PATCH 4/6] release: 1.2.33 --- CHANGELOG.md | 1 + .../lib/public_interfaces/sampling_client.py | 8 +- src/fireworks/training/sdk/client.py | 10 +- src/fireworks/training/sdk/managed.py | 3 +- .../training/sdk/tokenizer_identity.py | 101 ++++++++++++++---- tests_training/test_tokenizer_identity.py | 30 ++++-- 6 files changed, 120 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b24e74b..a8f03e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ Full Changelog: [v1.2.32...v1.2.33](https://github.com/fw-ai-external/python-sdk ### Bug Fixes * **training-sdk:** preserve message sampling, saved fast-tokenizer overrides and chat-profile save/reload behavior on code-free Kimi loading paths. * **training-sdk:** follow the November 5 Kimi K3.1 release encoder, which rejects the invalid `medium` thinking effort without changing vocabulary or image-token IDs. +* **training-sdk:** keep DeepSeek V4.1 tokenizer and renderer loading in installed code, including explicit base-model context and opaque managed tokenizer paths. ## 1.2.32 (2026-11-05) diff --git a/src/fireworks/training/sdk/_core/lib/public_interfaces/sampling_client.py b/src/fireworks/training/sdk/_core/lib/public_interfaces/sampling_client.py index 0dc512b..7857adb 100644 --- a/src/fireworks/training/sdk/_core/lib/public_interfaces/sampling_client.py +++ b/src/fireworks/training/sdk/_core/lib/public_interfaces/sampling_client.py @@ -16,7 +16,7 @@ from fireworks.training.sdk._core._utils import lru_cache from fireworks.training.sdk._core.lib.sidecar import SidecarRPC, SidecarHandle, create_sidecar_handle from fireworks.training.sdk.tokenizer_identity import ( - is_deepseek_v4_model_name, + is_code_free_model_name, is_kimi_k3_k31_model_name, is_kimi_k3_k31_tokenizer_source, requires_code_free_tokenizer_loading, @@ -556,11 +556,11 @@ def _load_tokenizer_from_model_info( return load_native_tokenizer(tokenizer_id, revision=kwargs.get("revision")) if ( - is_deepseek_v4_model_name(model_name) - or model_arch == "DeepseekV4ForCausalLM" + is_code_free_model_name(model_name) + or model_arch in {"DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM", "DeepseekV41ForConditionalGeneration"} or requires_code_free_tokenizer_loading(tokenizer_id, revision=kwargs.get("revision")) ): - # Recognized official DeepSeek V4 data ships no remote code; an + # Reviewed DeepSeek V4/V4.1 data ships no remote code; an # explicit False resolves the identical tokenizer class and keeps any # checkpoint-staged Python inert. kwargs["trust_remote_code"] = False diff --git a/src/fireworks/training/sdk/client.py b/src/fireworks/training/sdk/client.py index a04d762..b243254 100644 --- a/src/fireworks/training/sdk/client.py +++ b/src/fireworks/training/sdk/client.py @@ -42,7 +42,7 @@ from fireworks.training.sdk._core import APIStatusError, SamplingClient from fireworks.training.sdk._core._compat import model_dump as _training_model_dump from fireworks.training.sdk.tokenizer_identity import ( - is_deepseek_v4_model_name, + is_code_free_model_name, is_kimi_k3_k31_model_name, is_kimi_k3_k31_tokenizer_source, requires_code_free_tokenizer_loading, @@ -3124,10 +3124,10 @@ def get_tokenizer(self): from fireworks.training.sdk.kimi_native_tokenizer import load_native_tokenizer return load_native_tokenizer(self._tokenizer_model) - if is_deepseek_v4_model_name( + if is_code_free_model_name( getattr(self, "_tokenizer_base_model", None) ) or requires_code_free_tokenizer_loading(self._tokenizer_model): - # Recognized official DeepSeek V4 data ships no remote code; an + # Reviewed DeepSeek V4/V4.1 data ships no remote code; an # explicit False resolves the identical tokenizer class and keeps # any checkpoint-staged Python inert. return AutoTokenizer.from_pretrained(self._tokenizer_model, trust_remote_code=False) @@ -4160,7 +4160,9 @@ def create_training_client( training_client = managed_handle.service_client.create_training_client(**create_model_kwargs) training_client._tokenizer_model = managed_config.tokenizer_model training_client._tokenizer_base_model = ( - "deepseek_v4" if managed_config.base_model_type == "deepseek_v4" else managed_config.base_model + managed_config.base_model_type + if is_code_free_model_name(managed_config.base_model_type) + else managed_config.base_model ) if managed_handle.sampler_backend is not None: training_client._attach_sampler_backend(managed_handle.sampler_backend) diff --git a/src/fireworks/training/sdk/managed.py b/src/fireworks/training/sdk/managed.py index 4d56115..1ff25ab 100644 --- a/src/fireworks/training/sdk/managed.py +++ b/src/fireworks/training/sdk/managed.py @@ -49,6 +49,7 @@ from fireworks.training.sdk._snapshot_chain import ( build_incremental_metadata, ) +from fireworks.training.sdk.tokenizer_identity import is_code_free_model_name if TYPE_CHECKING: from fireworks.training.sdk._rdma import _RdmaSamplerBackend @@ -632,7 +633,7 @@ def _create_managed_tinker_client( # Let get_tokenizer() load from the HF tokenizer name without a get_info RPC. training_client._tokenizer_model = config.tokenizer_model training_client._tokenizer_base_model = ( - "deepseek_v4" if config.base_model_type == "deepseek_v4" else config.base_model + config.base_model_type if is_code_free_model_name(config.base_model_type) else config.base_model ) handle.training_client = training_client handle.training_clients.append(training_client) diff --git a/src/fireworks/training/sdk/tokenizer_identity.py b/src/fireworks/training/sdk/tokenizer_identity.py index 48363fe..743a0d8 100644 --- a/src/fireworks/training/sdk/tokenizer_identity.py +++ b/src/fireworks/training/sdk/tokenizer_identity.py @@ -8,6 +8,7 @@ Currently covered families: * DeepSeek V4 (``deepseek_v4``), recognized by ``tokenizer.json`` digests; +* DeepSeek V4.1, restricted separately without selecting a V4 renderer; * Kimi K3 / K3.1 (``kimi_k3`` / ``kimi_k31``), recognized by the shared tiktoken vocabulary digest plus a pinned ``tokenizer_config.json`` digest. The two models ship byte-identical ``tiktoken.model`` vocabularies and @@ -34,8 +35,8 @@ # intermediate SDK carrying the module with an older policy revision falls # back to the pinned snapshot instead of inheriting the older behavior. # Capability history: "deepseek-v4-tokenizer-identity/separate-family-and-code-policy" -# covered DeepSeek V4 only; the current value adds the Kimi K3/K3.1 family. -TOKENIZER_IDENTITY_POLICY_CAPABILITY = "deepseek-v4+kimi-k3-k31-tokenizer-identity/saved-native-profiles-v2" +# covered DeepSeek V4 only; the current value also covers V4.1 and saved Kimi profiles. +TOKENIZER_IDENTITY_POLICY_CAPABILITY = "deepseek-v4-v41+kimi-k3-k31-tokenizer-identity/saved-native-profiles-v3" # Official DeepSeek V4 ``tokenizer.json`` content digests: # 1. The hub snapshot form (identical for DeepSeek-V4-Flash-0731 and @@ -55,6 +56,26 @@ DEEPSEEK_V41_MODEL_RE = re.compile(r"(?:^|[/_.-])deepseek[-_]?v?4(?:[._-]?1|p1)(?=$|[/_.-])") +DEEPSEEK_V41_MODEL_TYPE = "deepseek_v41" +DEEPSEEK_V41_MODEL_TYPES = frozenset({"deepseek_v41", "deepseek_v41_text", "deepseek_v4.1", "deepseek_v4.1_text"}) +DEEPSEEK_V41_REPOSITORIES = frozenset({"deepseek-ai/DeepSeek-V4.1-Flash"}) +# Real upstream, AutoTokenizer-saved, and declared-postprocessor cookbook forms. +# These restrict loading only; they must never select the V4 renderer. +DEEPSEEK_V41_TOKENIZER_JSON_SHA256 = frozenset( + { + "c90dfa01249db1be4245780a052ede752e1361c612ac6d08e2bdada7d599476b", + "23136b659802e7e16fc849c551ac4a05a85e708732adeed010610b11dce85290", + "f8c19e2bc45ba524379b42c29ea07e83c9f856c15c2f7ebb69e36ddcfa1cf650", + } +) +_DEEPSEEK_V41_TOKENIZER_JSON_SIZES = frozenset({6_367_257, 10_071_774, 10_072_050}) +_DEEPSEEK_V41_TOKENIZER_JSON_GIT_BLOB_SHA1 = frozenset( + { + "6a15814dd25c934028034531da744c689f87ff21", + "2345dbdaeeaaf36cb6b01e0a2b282befd18d08ec", + "48ae10170e965f3508778b75842eb08d1244ffdf", + } +) # Byte profiles of the three official ``tokenizer.json`` forms, used only by # the Hub metadata prefilter: exact byte sizes, and git blob sha1 digests as @@ -187,6 +208,21 @@ def is_kimi_k3_k31_model_name(model_name: str | None) -> bool: return normalized in KIMI_K3_K31_MODEL_TYPES or normalized in KIMI_K3_K31_MODEL_NAMES +def is_code_free_model_name(model_name: str | None) -> bool: + """Restrictive selected context, never semantic renderer identity.""" + if is_deepseek_v4_model_name(model_name): + return True + if not model_name: + return False + return model_name.strip().rstrip("/").lower() in DEEPSEEK_V41_MODEL_TYPES | { + "deepseek_v41_flash", + "deepseek-ai/deepseek-v4.1-flash", + "accounts/fireworks/models/deepseek-v41-flash", + "accounts/fireworks/models/deepseek-v4.1-flash", + "accounts/fireworks/models/deepseek-v4p1-flash", + } + + def _read_json(path: Path) -> dict[str, Any] | None: try: parsed = json.loads(path.read_text(encoding="utf-8")) @@ -267,7 +303,7 @@ def _head_tokenizer_json_metadata(model: str, revision: str | None) -> Any: return None -def _metadata_decision(size: Any, etag: Any) -> bool | None: +def _metadata_decision(size: Any, etag: Any, *, include_v41: bool = False) -> bool | None: """Prefilter decision from HEAD metadata: True/False, or None for bytes. ``False`` is returned only when byte inequality is proven by a supported @@ -286,14 +322,21 @@ def _metadata_decision(size: Any, etag: Any) -> bool | None: # A missing etag header is a missing capability: full byte reads. return None etag = etag.lower() - size_matches = (size in _DEEPSEEK_V4_TOKENIZER_JSON_SIZES) if size is not None else None + sizes = _DEEPSEEK_V4_TOKENIZER_JSON_SIZES + digests = DEEPSEEK_V4_TOKENIZER_JSON_SHA256 + blobs = _DEEPSEEK_V4_TOKENIZER_JSON_GIT_BLOB_SHA1 + if include_v41: + sizes = sizes | _DEEPSEEK_V41_TOKENIZER_JSON_SIZES + digests = digests | DEEPSEEK_V41_TOKENIZER_JSON_SHA256 + blobs = blobs | _DEEPSEEK_V41_TOKENIZER_JSON_GIT_BLOB_SHA1 + size_matches = (size in sizes) if size is not None else None if len(etag) == 64 and _HEX_RE.match(etag): - if etag in DEEPSEEK_V4_TOKENIZER_JSON_SHA256: + if etag in digests: # Two proven signals contradicting each other: let bytes decide. return None if size_matches is False else True return False if len(etag) == 40 and _HEX_RE.match(etag): - if etag in _DEEPSEEK_V4_TOKENIZER_JSON_GIT_BLOB_SHA1: + if etag in blobs: # A candidate even when the size contradicts it: bytes decide. return None # The blob sha1 alone proves inequality; size agreement is irrelevant. @@ -325,14 +368,21 @@ def _hub_offline() -> bool: return bool(getattr(constants, "HF_HUB_OFFLINE", False)) -def _download_tokenizer_json_matches(model: str, *, revision: str | None, local_files_only: bool) -> bool: +def _download_tokenizer_json_matches( + model: str, *, revision: str | None, local_files_only: bool, include_v41: bool = False +) -> bool: path = _hub_download(model, "tokenizer.json", revision=revision, local_files_only=local_files_only) if path is None: return False - return _sha256(path) in DEEPSEEK_V4_TOKENIZER_JSON_SHA256 + digests = DEEPSEEK_V4_TOKENIZER_JSON_SHA256 + if include_v41: + digests = digests | DEEPSEEK_V41_TOKENIZER_JSON_SHA256 + return _sha256(path) in digests -def _hub_tokenizer_json_matches(model: str, *, revision: str | None, local_files_only: bool) -> bool: +def _hub_tokenizer_json_matches( + model: str, *, revision: str | None, local_files_only: bool, include_v41: bool = False +) -> bool: """Whether the Hub ``tokenizer.json`` is a pinned official form. Online against the official endpoint, one HEAD request answers from @@ -344,7 +394,9 @@ def _hub_tokenizer_json_matches(model: str, *, revision: str | None, local_files the HEAD observed so the two reads are coherent. No result is cached. """ if local_files_only or _hub_offline() or not _official_hub_endpoint(): - return _download_tokenizer_json_matches(model, revision=revision, local_files_only=local_files_only) + return _download_tokenizer_json_matches( + model, revision=revision, local_files_only=local_files_only, include_v41=include_v41 + ) metadata = _head_tokenizer_json_metadata(model, revision) if metadata is _HEAD_ABSENT: return False @@ -352,13 +404,17 @@ def _hub_tokenizer_json_matches(model: str, *, revision: str | None, local_files if metadata is not None: pinned = _pinned_commit(metadata) if pinned is not None: - decision = _metadata_decision(getattr(metadata, "size", None), getattr(metadata, "etag", None)) + decision = _metadata_decision( + getattr(metadata, "size", None), getattr(metadata, "etag", None), include_v41=include_v41 + ) if decision is not None: return decision revision_for_download = pinned # Without a valid commit pin the metadata answer is unusable: no # metadata negative, full byte reads at the caller's revision. - return _download_tokenizer_json_matches(model, revision=revision_for_download, local_files_only=False) + return _download_tokenizer_json_matches( + model, revision=revision_for_download, local_files_only=False, include_v41=include_v41 + ) def is_deepseek_v4_tokenizer_source( @@ -444,7 +500,7 @@ def requires_code_free_tokenizer_loading( ) -> bool: """Restrict recognized reviewed data to installed tokenizer/processor code. - Covers DeepSeek V4 and Kimi K3/K3.1. A shared reviewed tokenizer digest + Covers DeepSeek V4/V4.1 and Kimi K3/K3.1. A shared reviewed tokenizer digest is restrictive even when mutable ``model_type`` names another family. Unlike semantic family recognition, this must not select a renderer or encoder: it only keeps checkpoint @@ -543,20 +599,27 @@ def _tokenizer_source_matches( if not tokenizer_model: return False model = tokenizer_model + model_types = {DEEPSEEK_V4_MODEL_TYPE} + digests = DEEPSEEK_V4_TOKENIZER_JSON_SHA256 + repositories = DEEPSEEK_V4_REPOSITORIES + if not explicit_family_wins: + model_types.update(DEEPSEEK_V41_MODEL_TYPES) + digests = digests | DEEPSEEK_V41_TOKENIZER_JSON_SHA256 + repositories = repositories | DEEPSEEK_V41_REPOSITORIES local_dir = Path(model) if local_dir.is_dir(): config = _read_json(local_dir / "config.json") model_type = config.get("model_type") if config is not None else None if isinstance(model_type, str) and model_type not in {"", "model"}: - if model_type == DEEPSEEK_V4_MODEL_TYPE: + if model_type in model_types: return True if explicit_family_wins: return False - return _sha256(local_dir / "tokenizer.json") in DEEPSEEK_V4_TOKENIZER_JSON_SHA256 + return _sha256(local_dir / "tokenizer.json") in digests # Only the digest-reviewed pair skips the advisory read; reviewed # mirror/Fireworks aliases are selected-context identities # (is_deepseek_v4_model_name), not no-read source shortcuts. - if model in DEEPSEEK_V4_REPOSITORIES: + if model in repositories: return True config_path = _hub_download( model, @@ -568,8 +631,10 @@ def _tokenizer_source_matches( config = _read_json(config_path) model_type = config.get("model_type") if config is not None else None if isinstance(model_type, str) and model_type not in {"", "model"}: - if model_type == DEEPSEEK_V4_MODEL_TYPE: + if model_type in model_types: return True if explicit_family_wins: return False - return _hub_tokenizer_json_matches(model, revision=revision, local_files_only=local_files_only) + return _hub_tokenizer_json_matches( + model, revision=revision, local_files_only=local_files_only, include_v41=not explicit_family_wins + ) diff --git a/tests_training/test_tokenizer_identity.py b/tests_training/test_tokenizer_identity.py index a1ba83b..aa19cea 100644 --- a/tests_training/test_tokenizer_identity.py +++ b/tests_training/test_tokenizer_identity.py @@ -282,8 +282,18 @@ def load(model: Any, **kwargs: Any) -> Any: @pytest.mark.parametrize("override", ["local", "kimi"]) +@pytest.mark.parametrize( + "selected_model", + [ + "accounts/fireworks/models/deepseek-v4-flash-0731", + "accounts/fireworks/models/deepseek-v41-flash", + "accounts/fireworks/models/deepseek-v4p1-flash", + "deepseek_v41_text", + "deepseek_v4.1_text", + ], +) def test_sdk_v4_model_context_cannot_lose_policy_on_tokenizer_override( - unidentified_v4_tokenizer, monkeypatch, override + unidentified_v4_tokenizer, monkeypatch, override, selected_model ): """Synthetic Hub staging keeps the actual HF loader/code execution intact.""" import transformers @@ -306,7 +316,7 @@ def load(model, **kwargs): tokenizer = None try: tokenizer = sampling_client._load_tokenizer_from_model_info( - "accounts/fireworks/models/deepseek-v4-flash-0731", + selected_model, str(directory) if override == "local" else "moonshotai/Kimi-K2.6", ) except ValueError as exc: @@ -322,6 +332,11 @@ def load(model, **kwargs): [ ("accounts/fireworks/models/deepseek-v4-flash-0731", ""), ("accounts/customer/models/production-model", "deepseek_v4"), + ("accounts/fireworks/models/deepseek-v41-flash", ""), + ("accounts/customer/models/production-model", "deepseek_v41"), + ("accounts/customer/models/production-model", "deepseek_v41_text"), + ("accounts/customer/models/production-model", "deepseek_v4.1"), + ("accounts/customer/models/production-model", "deepseek_v4.1_text"), ], ) def test_managed_training_client_retains_base_identity_for_tokenizer( @@ -439,7 +454,12 @@ def unavailable(repo_id, filename, **_kwargs): assert reads == [("deepseek_v4", "config.json")] -def test_core_training_tokenizer_uses_existing_architecture_metadata(unidentified_v4_tokenizer, monkeypatch): +@pytest.mark.parametrize( + "model_arch", ["DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM", "DeepseekV41ForConditionalGeneration"] +) +def test_core_training_tokenizer_uses_existing_architecture_metadata( + unidentified_v4_tokenizer, monkeypatch, model_arch +): """An opaque model name can still constrain policy from get_info metadata.""" import asyncio from types import SimpleNamespace @@ -454,9 +474,7 @@ def test_core_training_tokenizer_uses_existing_architecture_metadata(unidentifie directory, original, marker = unidentified_v4_tokenizer info = types.GetInfoResponse( model_id="test-model", - model_data=types.ModelData( - arch="DeepseekV4ForCausalLM", model_name="renamed-model", tokenizer_id=str(directory) - ), + model_data=types.ModelData(arch=model_arch, model_name="renamed-model", tokenizer_id=str(directory)), ) class Holder: From 918e6a259fa32dfceb78a5955471149a135d198f Mon Sep 17 00:00:00 2001 From: Hershal Bhatia Date: Sat, 7 Nov 2026 10:57:21 -0800 Subject: [PATCH 5/6] release: 1.2.33 --- src/fireworks/training/sdk/client.py | 29 ++++- tests_training/test_release_channel.py | 151 +++++++++++++++++++++++++ 2 files changed, 179 insertions(+), 1 deletion(-) create mode 100644 tests_training/test_release_channel.py diff --git a/src/fireworks/training/sdk/client.py b/src/fireworks/training/sdk/client.py index b243254..c95087f 100644 --- a/src/fireworks/training/sdk/client.py +++ b/src/fireworks/training/sdk/client.py @@ -112,6 +112,8 @@ class LoadAdapterResponse(BaseModel): SAMPLER_DRAIN_TIMEOUT_S: float = 3600.0 DEFAULT_PARALLEL_CHUNK_SEND_CONCURRENCY: int = 128 PARALLEL_CHUNK_SEND_CONCURRENCY_ENV = "FIREWORKS_TRAINING_PARALLEL_CHUNK_SEND_CONCURRENCY" +TRAINING_RELEASE_CHANNEL_ENV = "FIREWORKS_TRAINING_RELEASE_CHANNEL" +TRAINING_RELEASE_CHANNEL_HEADER = "Fireworks-Release-Channel" FIRETITAN_TINKER_CLIENT_CONFIG: dict[str, bool] = { "parallel_fwdbwd_chunks": True, "proto_write_fwdbwd": False, @@ -3290,9 +3292,18 @@ class FiretitanServiceClient(ServiceClient): Each training client automatically selects optimized requests when its model-creation response advertises ``comms="v2"``. Older trainers retain the legacy request path. Inspect ``training_client.comms`` for the selected version. + + ``release_channel`` selects an authorized serverless training session release. When omitted, + an explicit default header or ``FIREWORKS_TRAINING_RELEASE_CHANNEL`` is used. + Pass an empty string for the default release. Dedicated clients ignore the + environment setting and reject an explicit nonempty selection. """ - def __init__(self, *args, api_key: str | None = None, **kwargs): + def __init__( + self, *args: Any, api_key: str | None = None, release_channel: str | None = None, **kwargs: Any + ) -> None: + if release_channel is not None and not isinstance(release_channel, str): + raise ValueError("release_channel must be a string or None") project_id = kwargs.get("project_id", args[1] if len(args) > 1 else None) if project_id is not None: raise NotImplementedError( @@ -3329,9 +3340,25 @@ def __init__(self, *args, api_key: str | None = None, **kwargs): self._default_user_metadata: dict[str, str] | None = kwargs.get("user_metadata") if self._managed_config is not None: + if release_channel: + raise ValueError("release_channel is only supported for serverless training") return self._training_base_url = str(resolve_training_base_url(constructor_base_url)) + serverless = self._training_base_url.rstrip("/").endswith("/training/v1/serverless") + if release_channel and not serverless: + raise ValueError("release_channel is only supported for serverless training") + if serverless: + headers = dict(kwargs.get("default_headers") or {}) + channel_header_names = [name for name in headers if name.lower() == TRAINING_RELEASE_CHANNEL_HEADER.lower()] + if release_channel is None and not channel_header_names: + release_channel = os.environ.get(TRAINING_RELEASE_CHANNEL_ENV) + if release_channel is not None: + for name in channel_header_names: + del headers[name] + if release_channel: + headers[TRAINING_RELEASE_CHANNEL_HEADER] = release_channel + kwargs["default_headers"] = headers self._serverless_caller_headers = _serverless_caller_headers(kwargs.get("default_headers")) if self._fireworks_api_key is not None: headers = dict(kwargs.pop("default_headers", None) or {}) diff --git a/tests_training/test_release_channel.py b/tests_training/test_release_channel.py new file mode 100644 index 0000000..ef62508 --- /dev/null +++ b/tests_training/test_release_channel.py @@ -0,0 +1,151 @@ +"""Release selection is resolved at service construction without changing default routing.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch +from collections.abc import Iterator + +import pytest + +from fireworks.training.sdk.client import ( + TRAINING_RELEASE_CHANNEL_ENV, + TRAINING_RELEASE_CHANNEL_HEADER, + FiretitanServiceClient, + FiretitanSamplingClient, +) +from fireworks.training.sdk._core.lib.public_interfaces.service_client import ServiceClient + + +@pytest.fixture +def base_client_init(monkeypatch: pytest.MonkeyPatch) -> Iterator[MagicMock]: + monkeypatch.delenv(TRAINING_RELEASE_CHANNEL_ENV, raising=False) + monkeypatch.delenv("TINKER_BASE_URL", raising=False) + with patch.object(ServiceClient, "__init__", return_value=None) as constructor: + yield constructor + + +def test_explicit_channel_is_forwarded_without_mutating_headers(base_client_init: MagicMock) -> None: + headers = {"X-Test": "preserved", "fireworks-release-channel": "canary-old"} + service = FiretitanServiceClient(api_key="fw-test", release_channel="canary-selected", default_headers=headers) + + sent_headers = base_client_init.call_args.kwargs["default_headers"] + assert sent_headers[TRAINING_RELEASE_CHANNEL_HEADER] == "canary-selected" + assert sent_headers["X-Test"] == "preserved" + assert "fireworks-release-channel" not in sent_headers + assert headers["fireworks-release-channel"] == "canary-old" + assert service._serverless_sampling_headers()[TRAINING_RELEASE_CHANNEL_HEADER] == "canary-selected" + + +def test_environment_channel_is_resolved_once(base_client_init: MagicMock, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, "canary-original") + service = FiretitanServiceClient(api_key="fw-test") + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, "canary-later") + + assert base_client_init.call_args.kwargs["default_headers"][TRAINING_RELEASE_CHANNEL_HEADER] == "canary-original" + assert service._serverless_sampling_headers()[TRAINING_RELEASE_CHANNEL_HEADER] == "canary-original" + + +def test_sampler_handle_preserves_channel_after_environment_changes( + base_client_init: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, "canary-original") + service = FiretitanServiceClient(api_key="fw-test") + sampler = FiretitanSamplingClient.create( + inference_url="https://training.example.test/training/v1/serverless", + model="accounts/test-account/trainingSessions/ts-test", + api_key="fw-test", + additional_headers=service._serverless_sampling_headers(), + ) + try: + handle = sampler.get_sampler_handle() + finally: + sampler.close() + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, "canary-later") + restored = FiretitanSamplingClient.from_sampler_handle(handle, api_key="fw-restored") + try: + assert restored.deployment_sampler.additional_headers[TRAINING_RELEASE_CHANNEL_HEADER] == "canary-original" + finally: + restored.close() + + +def test_explicit_channel_overrides_environment(base_client_init: MagicMock, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, "canary-environment") + FiretitanServiceClient(api_key="fw-test", release_channel="canary-explicit") + + assert base_client_init.call_args.kwargs["default_headers"][TRAINING_RELEASE_CHANNEL_HEADER] == "canary-explicit" + + +def test_explicit_empty_channel_selects_default(base_client_init: MagicMock, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, "canary-environment") + service = FiretitanServiceClient( + api_key="fw-test", release_channel="", default_headers={"fireworks-release-channel": "canary-header"} + ) + + assert not any("release-channel" in name.lower() for name in base_client_init.call_args.kwargs["default_headers"]) + assert not any("release-channel" in name.lower() for name in service._serverless_sampling_headers()) + + +def test_omitted_channel_preserves_default_behavior(base_client_init: MagicMock) -> None: + service = FiretitanServiceClient(api_key="fw-test") + + assert TRAINING_RELEASE_CHANNEL_HEADER not in base_client_init.call_args.kwargs["default_headers"] + assert TRAINING_RELEASE_CHANNEL_HEADER not in service._serverless_sampling_headers() + assert "release_channel" not in base_client_init.call_args.kwargs + + +def test_non_string_channel_cannot_silently_select_default(base_client_init: MagicMock) -> None: + with pytest.raises(ValueError, match="must be a string or None"): + FiretitanServiceClient(api_key="fw-test", release_channel=False) # type: ignore[arg-type] + base_client_init.assert_not_called() + + +@pytest.mark.parametrize("env_channel", [None, "canary-environment"]) +def test_manual_header_is_preserved( + base_client_init: MagicMock, monkeypatch: pytest.MonkeyPatch, env_channel: str | None +) -> None: + if env_channel is not None: + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, env_channel) + FiretitanServiceClient(api_key="fw-test", default_headers={"fireworks-release-channel": "canary-header"}) + + sent_headers = base_client_init.call_args.kwargs["default_headers"] + assert sent_headers["fireworks-release-channel"] == "canary-header" + assert TRAINING_RELEASE_CHANNEL_HEADER not in sent_headers + + +def test_dedicated_client_ignores_environment(base_client_init: MagicMock, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, "canary-environment") + FiretitanServiceClient(api_key="fw-test", base_url="https://trainer.example.test") + + assert TRAINING_RELEASE_CHANNEL_HEADER not in base_client_init.call_args.kwargs["default_headers"] + + +def test_dedicated_client_rejects_explicit_channel(base_client_init: MagicMock) -> None: + with pytest.raises(ValueError, match="only supported for serverless training"): + FiretitanServiceClient( + api_key="fw-test", base_url="https://trainer.example.test", release_channel="canary-selected" + ) + base_client_init.assert_not_called() + + +def test_managed_client_ignores_environment(base_client_init: MagicMock, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, "canary-environment") + service = FiretitanServiceClient(api_key="fw-test", managed_config=object()) + + base_client_init.assert_not_called() + assert TRAINING_RELEASE_CHANNEL_HEADER not in service._serverless_sampling_headers() + + +def test_managed_client_rejects_explicit_channel(base_client_init: MagicMock) -> None: + with pytest.raises(ValueError, match="only supported for serverless training"): + FiretitanServiceClient(api_key="fw-test", managed_config=object(), release_channel="canary-selected") + base_client_init.assert_not_called() + + +def test_environment_endpoint_is_resolved_before_channel( + base_client_init: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("TINKER_BASE_URL", "https://training.example.test/training/v1/serverless/") + monkeypatch.setenv(TRAINING_RELEASE_CHANNEL_ENV, "canary-environment") + FiretitanServiceClient(api_key="fw-test") + + assert base_client_init.call_args.kwargs["default_headers"][TRAINING_RELEASE_CHANNEL_HEADER] == "canary-environment" From 7045fadd60aa09f9c0777ac17453e59103d9d417 Mon Sep 17 00:00:00 2001 From: chengxilifw Date: Sat, 7 Nov 2026 21:54:53 -0800 Subject: [PATCH 6/6] release: 1.2.33 --- CHANGELOG.md | 2 +- .../training/sdk/_core/_streaming.py | 16 +- src/fireworks/training/sdk/_sse.py | 8 +- src/fireworks/training/sdk/sampling.py | 97 +++++++-- tests_training/test_comms_negotiation.py | 23 +- tests_training/test_deployment.py | 102 ++++----- tests_training/test_routing_references.py | 6 +- tests_training/test_sampling.py | 17 +- tests_training/test_sampling_client_close.py | 6 +- tests_training/test_sampling_stream_errors.py | 31 ++- tests_training/test_sampling_token_ids.py | 203 ++++++++++++++++++ tests_training/test_sse_chunks.py | 36 ++++ tests_training/test_tito.py | 2 +- 13 files changed, 434 insertions(+), 115 deletions(-) create mode 100644 tests_training/test_sampling_token_ids.py create mode 100644 tests_training/test_sse_chunks.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a8f03e9..20ea652 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Changelog -## 1.2.33 (2026-11-07) +## 1.2.33 (2026-11-08) Full Changelog: [v1.2.32...v1.2.33](https://github.com/fw-ai-external/python-sdk/compare/v1.2.32...v1.2.33) diff --git a/src/fireworks/training/sdk/_core/_streaming.py b/src/fireworks/training/sdk/_core/_streaming.py index 79ed9b0..6528533 100644 --- a/src/fireworks/training/sdk/_core/_streaming.py +++ b/src/fireworks/training/sdk/_core/_streaming.py @@ -148,15 +148,15 @@ def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[ServerSentEvent]: def _iter_chunks(self, iterator: Iterator[bytes]) -> Iterator[bytes]: """Given an iterator that yields raw binary data, iterate over it and yield individual SSE chunks""" - data = b"" + data = bytearray() for chunk in iterator: for line in chunk.splitlines(keepends=True): data += line if data.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")): - yield data - data = b"" + yield bytes(data) + data = bytearray() if data: - yield data + yield bytes(data) async def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]: """Given an iterator that yields raw binary data, iterate over it & yield every event encountered""" @@ -170,15 +170,15 @@ async def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[Ser async def _aiter_chunks(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[bytes]: """Given an iterator that yields raw binary data, iterate over it and yield individual SSE chunks""" - data = b"" + data = bytearray() async for chunk in iterator: for line in chunk.splitlines(keepends=True): data += line if data.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")): - yield data - data = b"" + yield bytes(data) + data = bytearray() if data: - yield data + yield bytes(data) def decode(self, line: str) -> ServerSentEvent | None: # See: https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation # noqa: E501 diff --git a/src/fireworks/training/sdk/_sse.py b/src/fireworks/training/sdk/_sse.py index 103e2d2..f9751b8 100644 --- a/src/fireworks/training/sdk/_sse.py +++ b/src/fireworks/training/sdk/_sse.py @@ -84,15 +84,15 @@ def __init__(self) -> None: @staticmethod async def _aiter_chunks(stream: Any) -> Any: """Reassemble raw bytes into SSE chunks (delimited by blank lines).""" - buf = b"" + buf = bytearray() async for raw in stream.aiter_bytes(): for line in raw.splitlines(keepends=True): buf += line if buf.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")): - yield buf - buf = b"" + yield bytes(buf) + buf = bytearray() if buf: - yield buf + yield bytes(buf) async def aiter_events(self, response: Any) -> Any: """Yield :class:`_SSEEvent` objects from an ``httpx.Response``.""" diff --git a/src/fireworks/training/sdk/sampling.py b/src/fireworks/training/sdk/sampling.py index 9654337..f1929af 100644 --- a/src/fireworks/training/sdk/sampling.py +++ b/src/fireworks/training/sdk/sampling.py @@ -226,7 +226,7 @@ class SampledCompletion: Contains the full token sequence (prompt + completion) needed for training, with prompt tokens from client-side tokenization and completion tokens from - the deployment's ``raw_output`` response. + the deployment's exact token-ID response. """ text: str @@ -288,7 +288,7 @@ class DeploymentSampler(_RestClient): Uses a local HuggingFace tokenizer to apply chat templates and tokenize prompts, then sends token IDs to the ``/inference/v1/completions`` endpoint - (token-in, token-out). Completion token IDs come back via ``raw_output``. + (token-in, token-out). Completion token IDs come back via SSE token deltas. BOS and special tokens are handled by the tokenizer's chat template -- no manual prepend is needed. @@ -554,8 +554,12 @@ async def async_completions_stream( if "routing_matrix_format" in kwargs: raise ValueError("routing_matrix_format is selected by the trainer-bound sampler") http_timeout = kwargs.pop("http_timeout", 600) - if kwargs.get("images"): - kwargs.setdefault("return_token_ids", True) + if kwargs.get("raw_output"): + raise ValueError( + "DeploymentSampler requires raw_output=False; use the inference client for raw diagnostics" + ) + kwargs["raw_output"] = False + kwargs["return_token_ids"] = True payload: dict[str, Any] = { "model": self.model, "prompt": prompt, @@ -647,9 +651,11 @@ async def async_completions_stream( with _sampling_response_evidence(resp, prompt): accumulated_text = "" accumulated_logprobs: list[dict] = [] + accumulated_token_ids: list[int] | None = None + server_prompt_ids: list[int] | None = None + missing_token_ids = False finish_reason = None usage_info = None - raw_output = None routing_metadata = {} perf_metrics_dict: dict[str, str] | None = None first_token_time: float | None = None @@ -698,6 +704,8 @@ async def async_completions_stream( raise _SSEProtocolError("completion choices must be a list of objects") for choice in chunk.get("choices", []): + if choice.get("index", 0) != 0: + raise _SSEProtocolError("streaming completions require n=1") routing_metadata.update( {key: choice[key] for key in _STORAGE_REFERENCE_FIELDS if key in choice} ) @@ -708,11 +716,38 @@ async def async_completions_stream( accumulated_text += text_delta lp = choice.get("logprobs") + content = lp.get("content") if isinstance(lp, dict) else None if lp and isinstance(lp, dict): - content = lp.get("content") if isinstance(content, list): accumulated_logprobs.extend(content) + token_ids = choice.get("token_ids") + if token_ids is not None: + if not isinstance(token_ids, list) or any(type(token) is not int for token in token_ids): + raise _SSEProtocolError("completion token_ids must be a list of integers") + if content is not None: + if not isinstance(content, list) or len(content) != len(token_ids): + raise _SSEProtocolError("completion token IDs and logprobs are not aligned") + if any( + not isinstance(row, dict) + or (row.get("token_id") is not None and row["token_id"] != token) + for token, row in zip(token_ids, content) + ): + raise _SSEProtocolError("completion token IDs disagree with logprob token IDs") + if accumulated_token_ids is None: + accumulated_token_ids = [] + accumulated_token_ids.extend(token_ids) + elif text_delta or content: + missing_token_ids = True + + prompt_ids = choice.get("prompt_token_ids") + if prompt_ids is not None: + if not isinstance(prompt_ids, list) or any(type(token) is not int for token in prompt_ids): + raise _SSEProtocolError("prompt_token_ids must be a list of integers") + if server_prompt_ids is not None and prompt_ids != server_prompt_ids: + raise _SSEProtocolError("prompt token IDs changed within the completion stream") + server_prompt_ids = prompt_ids + fr = choice.get("finish_reason") if fr: if fr not in ("stop", "length", "tool_calls"): @@ -721,10 +756,6 @@ async def async_completions_stream( raise error finish_reason = fr - ro = choice.get("raw_output") - if ro: - raw_output = ro - if "usage" in chunk: usage_info = chunk["usage"] @@ -736,6 +767,21 @@ async def async_completions_stream( if finish_reason is None: raise _SSETruncationError("inference stream ended without a successful finish reason") + if missing_token_ids or accumulated_token_ids is None or server_prompt_ids is None: + raise _SSEProtocolError( + "deployment did not return complete token IDs; serving must support return_token_ids and prompt_token_ids" + ) + if payload.get("logprobs") and len(accumulated_logprobs) != len(accumulated_token_ids): + raise _SSEProtocolError("completion token IDs and logprobs are not aligned") + generated_count = usage_info.get("completion_tokens") if usage_info else None + echo_count = payload.get("echo_last") + if echo_count is None: + echo_count = len(server_prompt_ids) if payload.get("echo") else 0 + if type(generated_count) is int and type(echo_count) is int: + expected_count = generated_count + min(echo_count, len(server_prompt_ids)) + if len(accumulated_token_ids) != expected_count: + raise _SSEProtocolError("completion token count disagrees with usage") + client_ttft = (first_token_time - t0) if first_token_time else None elapsed = time.monotonic() - t0 completion_tokens = usage_info.get("completion_tokens") if usage_info else None @@ -761,12 +807,12 @@ async def async_completions_stream( assembled_choice: dict[str, Any] = { "text": accumulated_text, "finish_reason": finish_reason, + "token_ids": accumulated_token_ids, + "prompt_token_ids": server_prompt_ids, } assembled_choice.update(routing_metadata) if accumulated_logprobs: assembled_choice["logprobs"] = {"content": accumulated_logprobs} - if raw_output: - assembled_choice["raw_output"] = raw_output result: dict[str, Any] = {"choices": [assembled_choice]} if upstream_response_id is not None: result["id"] = upstream_response_id @@ -938,10 +984,15 @@ async def sample_with_prompt_tokens( sampler entry point for renderer-backed RL rollouts where the prompt has already been built from a token-native renderer. - Concurrency control, ``raw_output`` echo semantics, ``max_seq_len`` filtering, + Concurrency control, echo semantics, ``max_seq_len`` filtering, and logprob extraction are inherited from :meth:`_do_one_completion` and :meth:`_parse_completions_result`. + By default, collect exact token IDs from SSE deltas and + omit the duplicate raw snapshot. This requires a deployment that returns + incremental token IDs. ``raw_output=True`` is not supported by this sampler; + use the inference client directly for raw diagnostics. + Pass ``images`` (base64 data URLs) for vision models with an unexpanded token prompt (one ``<|image_pad|>`` token per image). The server expands image pads; ``return_token_ids`` is set automatically so @@ -1092,7 +1143,7 @@ def drain_metrics(self) -> list[ServerMetrics]: # Per-completion retry covers two transient server-side classes: # 1. SSE truncation: the deployment closes the stream mid-generation - # without [DONE]/finish_reason/raw_output. Surfaces as a RuntimeError + # without [DONE]/finish_reason/token_ids. Surfaces as a RuntimeError # raised from async_completions_stream. # 2. Transient HTTP / connection errors right after a fresh deployment # or hotload: 408/429/502/503/504 plus httpx connection-level errors. @@ -1415,6 +1466,8 @@ async def _do_one_completion_result( diagnostic_context = kwargs.pop("timeout_diagnostic_context", None) sampling_context = kwargs.pop("sampling_context", None) raw_logprobs_match_sampling = self._raw_logprobs_match_sampling_params(temperature, kwargs) + kwargs.setdefault("raw_output", False) + kwargs["return_token_ids"] = True # One stable id for this logical request, constant across every retry and # sent as X-Request-Id so a failure is searchable in server logs. logical_request_id = logical_request_id or uuid.uuid4().hex @@ -1433,7 +1486,6 @@ async def _do_one_completion_result( prompt=prompt_ids, max_tokens=max_tokens, temperature=temperature, - raw_output=True, logical_request_id=logical_request_id, **kwargs, ) @@ -1702,16 +1754,15 @@ def _parse_completions_result( for choice in result.get("choices", []): text = choice.get("text", "") finish_reason = choice.get("finish_reason", "unknown") - raw = choice.get("raw_output") or {} - completion_ids = raw.get("completion_token_ids") + completion_ids = choice.get("token_ids") if completion_ids is None: raise RuntimeError( format_sdk_error( - "Deployment did not return raw_output token IDs", - f"The API response is missing completion_token_ids. Got choice keys: {list(choice.keys())}", - "The sampler requested raw_output=True, which is required for token-in RL rollouts. " - "Use a deployment path that returns raw_output.completion_token_ids for completions.", + "Deployment did not return completion token IDs", + f"The API response is missing choice.token_ids. Got choice keys: {list(choice.keys())}", + "Token-in RL rollouts require exact completion token IDs. " + "Serving must return incremental token IDs and server prompt IDs.", docs_url=DOCS_SDK, show_support=True, ) @@ -1768,14 +1819,14 @@ def _parse_completions_result( raise ValueError("Sampling references belong to a different shared storage domain") sampling_references[column] = references - expanded_prompt_ids = choice.get("prompt_token_ids") or raw.get("prompt_token_ids") + expanded_prompt_ids = choice.get("prompt_token_ids") if expanded_prompt_ids is not None: prompt_for_full = [int(x) for x in expanded_prompt_ids] else: prompt_for_full = list(prompt_ids) # With echo=True the API returns P+C tokens in - # completion_token_ids and logprobs cover all P+C positions. + # choice.token_ids and logprobs cover all P+C positions. # Strip the prompt prefix (verified by actual content match) # and drop the unconditional first-token logprob to get # P+C-1 training-aligned entries. diff --git a/tests_training/test_comms_negotiation.py b/tests_training/test_comms_negotiation.py index ed390f7..83d5e21 100644 --- a/tests_training/test_comms_negotiation.py +++ b/tests_training/test_comms_negotiation.py @@ -46,7 +46,9 @@ def receive(request): model_id = ( f"base-{index}" if body.get("base_only") - else f"run-{index}:train:0" if route == "serverless" else f"model-{index}" + else f"run-{index}:train:0" + if route == "serverless" + else f"model-{index}" ) pending[request_id] = {"type": "create_model", "model_id": model_id, **next(models)} return httpx.Response(200, json={"request_id": request_id}) @@ -279,7 +281,10 @@ async def receive(request): return httpx.Response(probe_status, json=_r3_probe_result() if probe_status == 200 else error) assert body.get("routing_matrix_format") == ("parquet_v1" if probe_status == 200 else None) assert bool(body.get("include_routing_matrix")) == (probe_status != 200) - return httpx.Response(200, text='data: {"choices":[{"text":"ok","finish_reason":"stop"}]}\n\ndata: [DONE]\n\n') + return httpx.Response( + 200, + text='data: {"choices":[{"text":"ok","token_ids":[3],"prompt_token_ids":[1,2],"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n', + ) def unexpected_probe(request): raise AssertionError(f"Binding must not call {request.url}") @@ -430,7 +435,10 @@ async def receive(request): ) assert body["include_routing_matrix"] is True assert "x-fireworks-r3-store-id" not in request.headers - return httpx.Response(200, text='data: {"choices":[{"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n') + return httpx.Response( + 200, + text='data: {"choices":[{"token_ids":[3],"prompt_token_ids":[1,2],"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n', + ) sampler = DeploymentSampler("https://inference.invalid", "test/model", "fw-test") sampler.negotiate_routing_matrix_format("parquet_v1", "nfs-test") @@ -480,7 +488,10 @@ async def receive(request): }, ) ) - return httpx.Response(200, text='data: {"choices":[{"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n') + return httpx.Response( + 200, + text='data: {"choices":[{"token_ids":[3],"prompt_token_ids":[1,2],"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n', + ) monkeypatch.setattr( "fireworks.training.sdk._rest_client._make_async_client", @@ -1018,6 +1029,8 @@ def test_mask_only_sse_preserves_references_and_storage_headers(): async def receive(request): payload = json.loads(request.content) assert payload["sampling_mask"] == "parquet_v1" + assert payload["raw_output"] is False + assert payload["return_token_ids"] is True assert "top_sampling_logprobs" not in payload assert request.headers["x-fireworks-r3-store-id"] == "test" assert request.headers["x-fireworks-r3-ttl-seconds"] == "120" @@ -1026,6 +1039,8 @@ async def receive(request): { "text": "ok", "finish_reason": "stop", + "token_ids": [2], + "prompt_token_ids": [1], "r3_store_id": "test", "sampling_mask_format": "parquet_v1", "sampling_mask_references": references, diff --git a/tests_training/test_deployment.py b/tests_training/test_deployment.py index 64ff11b..e02d9ba 100644 --- a/tests_training/test_deployment.py +++ b/tests_training/test_deployment.py @@ -1378,7 +1378,7 @@ def test_basic_sample(self): { "text": "hello world", "finish_reason": "stop", - "raw_output": {"completion_token_ids": completion_ids}, + "token_ids": completion_ids, } ] }, @@ -1409,7 +1409,7 @@ def test_tokenizer_called_with_messages(self): { "text": "ok", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], } ] }, @@ -1427,11 +1427,9 @@ def test_tokenizer_called_with_messages(self): def test_missing_completion_token_ids_raises(self): sampler = _make_sampler(tokenizer=_make_mock_tokenizer([1, 2, 3])) - _mock_async_completions_stream( - sampler, {"choices": [{"text": "ok", "finish_reason": "stop", "raw_output": {}}]} - ) + _mock_async_completions_stream(sampler, {"choices": [{"text": "ok", "finish_reason": "stop"}]}) - with pytest.raises(RuntimeError, match="missing completion_token_ids"): + with pytest.raises(RuntimeError, match="missing choice.token_ids"): asyncio.run( sampler.sample_with_tokens( messages=[{"role": "user", "content": "hi"}], @@ -1451,7 +1449,7 @@ def test_echo_strips_verified_prefix(self): { "text": "gen", "finish_reason": "stop", - "raw_output": {"completion_token_ids": echoed_completion_ids}, + "token_ids": echoed_completion_ids, } ] }, @@ -1485,7 +1483,7 @@ def test_echo_logprobs_aligned(self): { "text": "gen", "finish_reason": "stop", - "raw_output": {"completion_token_ids": echoed_ids}, + "token_ids": echoed_ids, "logprobs": { "content": [ {"logprob": 0.0, "sampling_logprob": None}, @@ -1539,7 +1537,7 @@ def test_max_seq_len_post_filter(self): { "text": "long", "finish_reason": "length", - "raw_output": {"completion_token_ids": long_completion}, + "token_ids": long_completion, } ] }, @@ -1563,7 +1561,7 @@ def test_logprobs_extracted(self): { "text": "hi", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -1.5, "sampling_logprob": -1.6}]}, } ] @@ -1589,7 +1587,7 @@ def test_legacy_logprobs_fall_back_to_raw_for_full_distribution_sampling(self, c { "text": "hi", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -1.5}]}, } ] @@ -1620,7 +1618,7 @@ def test_legacy_logprobs_do_not_fall_back_when_sampling_distribution_changes(sel { "text": "hi", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -1.5}]}, } ] @@ -1647,7 +1645,7 @@ def test_routing_matrices_extracted(self): { "text": "hi", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -0.5, "routing_matrix": "AQID"}]}, } ] @@ -1699,7 +1697,7 @@ def test_basic_sample(self, native_sdk): { "text": "out", "finish_reason": "stop", - "raw_output": {"completion_token_ids": completion_ids}, + "token_ids": completion_ids, } ] }, @@ -1723,7 +1721,7 @@ def test_basic_sample(self, native_sdk): { "text": "out", "finish_reason": "stop", - "raw_output": {"completion_token_ids": completion_ids}, + "token_ids": completion_ids, "logprobs": { "content": [ {"logprob": -0.4, "sampling_logprob": -0.41}, @@ -1755,7 +1753,7 @@ def test_does_not_call_apply_chat_template(self, native_sdk): { "text": "out", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], } ] }, @@ -1775,7 +1773,7 @@ def test_does_not_call_apply_chat_template(self, native_sdk): { "text": "out", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -0.7, "sampling_logprob": -0.8}]}, } ] @@ -1800,7 +1798,7 @@ async def _fake_stream(*args, **kwargs): { "text": "out", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -0.7, "sampling_logprob": -0.8}]}, } ] @@ -1839,7 +1837,7 @@ async def _fake_stream(*args, **kwargs): { "text": "out", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -0.7, "sampling_logprob": -0.8}]}, } ] @@ -1881,7 +1879,7 @@ def test_logprobs_extracted(self, native_sdk): { "text": "x", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -0.7, "sampling_logprob": -0.8}]}, } ] @@ -1901,7 +1899,7 @@ def test_logprobs_extracted(self, native_sdk): { "text": "x", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -0.7, "sampling_logprob": -0.8}]}, } ] @@ -1924,7 +1922,7 @@ async def _fake_stream(*args, **kwargs): { "text": "x", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -0.7}]}, } ] @@ -1951,7 +1949,7 @@ async def _fake_stream(*args, **kwargs): { "text": "x", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -0.7}]}, } ] @@ -1980,7 +1978,7 @@ async def _fake_stream(*args, **kwargs): { "text": "x", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [99]}, + "token_ids": [99], "logprobs": {"content": [{"logprob": -0.7, "sampling_logprob": -0.8}]}, } ] @@ -2032,7 +2030,7 @@ async def _fake_stream(*args, **kwargs): { "text": "out", "finish_reason": "stop", - "raw_output": {"completion_token_ids": completion_ids}, + "token_ids": completion_ids, "logprobs": { "content": [ {"logprob": -0.3, "sampling_logprob": -0.31}, @@ -2118,10 +2116,8 @@ async def _fake_stream(*args, **kwargs): { "text": "out", "finish_reason": "stop", - "raw_output": { - "prompt_token_ids": expanded, - "completion_token_ids": expanded + completion_ids, - }, + "token_ids": expanded + completion_ids, + "prompt_token_ids": expanded, "logprobs": {"content": prompt_rows + completion_rows}, } ] @@ -2160,7 +2156,7 @@ async def _fake_stream(*args, **kwargs): { "text": "out", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [40, 50]}, + "token_ids": [40, 50], "logprobs": { "content": [ { @@ -2218,7 +2214,7 @@ async def _fake_stream(*args, **kwargs): { "text": "out", "finish_reason": "length", - "raw_output": {"completion_token_ids": prompt_ids + completion_ids}, + "token_ids": prompt_ids + completion_ids, "logprobs": { "content": [ {"logprob": 0.0, "sampling_logprob": None, "routing_matrix": "first"}, @@ -2278,9 +2274,7 @@ async def fake_stream(*_args, **kwargs): { "text": "out", "finish_reason": "length", - "raw_output": { - "completion_token_ids": (prompt_ids if echo else []) + completion_ids, - }, + "token_ids": (prompt_ids if echo else []) + completion_ids, "logprobs": {"content": content}, "top_sampling_format": "parquet_v1", "r3_store_id": "test", @@ -2373,7 +2367,7 @@ async def _fake_stream(*args, **kwargs): { "text": "x", "finish_reason": "length", - "raw_output": {"completion_token_ids": prompt_ids + [40]}, + "token_ids": prompt_ids + [40], "logprobs": { "content": [ {"logprob": 0.0, "sampling_logprob": None}, @@ -2491,7 +2485,7 @@ def test_legacy_path_still_renders_via_chat_template(self): { "text": "ok", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [11]}, + "token_ids": [11], } ] }, @@ -2993,7 +2987,8 @@ def test_streaming_with_adaptive_controller(self): "choices": [ { "text": "hello", - "raw_output": {"completion_token_ids": completion_ids}, + "token_ids": completion_ids, + "prompt_token_ids": prompt_ids, "finish_reason": "stop", } ], @@ -3030,7 +3025,8 @@ def test_drained_metrics_include_client_latency_and_usage(self): "choices": [ { "text": "hello", - "raw_output": {"completion_token_ids": [400, 500]}, + "token_ids": [400, 500], + "prompt_token_ids": prompt_ids, "finish_reason": "stop", } ], @@ -3092,9 +3088,9 @@ def _http_error(status: int) -> httpx.HTTPStatusError: @staticmethod def _success(): - return { - "choices": [{"text": "ok", "finish_reason": "stop", "raw_output": {"completion_token_ids": [7]}}] - }, ServerMetrics(server_processing_time=0.1) + return {"choices": [{"text": "ok", "finish_reason": "stop", "token_ids": [7]}]}, ServerMetrics( + server_processing_time=0.1 + ) def _run(self, monkeypatch, failures, *, controller=None): self._no_backoff(monkeypatch) @@ -3203,9 +3199,19 @@ async def _acquire(): controller.release = lambda metrics=None: released.append(metrics) sampler = _make_sampler(tokenizer=None) sampler._concurrency_controller = controller - # HTTP 200 whose choice lacks raw_output.completion_token_ids: the real - # stream succeeds and the real parser rejects it. - chunk = {"choices": [{"text": "ok", "finish_reason": "stop"}]} + # Complete token deltas pass stream validation, but a reordered echo + # prefix must fail in the completion parser and mark the same attempt. + chunk = { + "choices": [ + { + "text": "ok", + "finish_reason": "stop", + "token_ids": [3, 2, 1, 7], + "prompt_token_ids": [1, 2, 3], + } + ], + "usage": {"completion_tokens": 1}, + } raw = f"data: {_json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode() async def _post(payload, headers, timeout): @@ -3221,8 +3227,8 @@ async def aiter_bytes(): return resp sampler._post_completion = _post - with pytest.raises(RuntimeError): - asyncio.run(sampler.sample_with_prompt_tokens([1, 2, 3])) + with pytest.raises(RuntimeError, match="echoed suffix"): + asyncio.run(sampler.sample_with_prompt_tokens([1, 2, 3], echo=True)) records = sampler.drain_metrics() sampler.close() assert len(records) == 1 @@ -3255,7 +3261,7 @@ async def _probe(headers, timeout): sampler.routing_matrix_format = "base64_inline" chunk = { - "choices": [{"text": "ok", "finish_reason": "stop", "raw_output": {"completion_token_ids": [7]}}], + "choices": [{"text": "ok", "finish_reason": "stop", "token_ids": [7], "prompt_token_ids": [1, 2, 3]}], "usage": {"completion_tokens": 1}, } raw = f"data: {_json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode() @@ -3312,7 +3318,7 @@ async def _fake(*args, **kwargs): { "text": "ok", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [7]}, + "token_ids": [7], } ] }, ServerMetrics() diff --git a/tests_training/test_routing_references.py b/tests_training/test_routing_references.py index 998aaed..edc4d97 100644 --- a/tests_training/test_routing_references.py +++ b/tests_training/test_routing_references.py @@ -80,7 +80,7 @@ def test_parquet_full_and_partial_echo_alignment(echo_last): routes = refs(len(ids), 100) choice = { "text": "ok", - "raw_output": {"completion_token_ids": ids}, + "token_ids": ids, "logprobs": {"content": [{"logprob": -1.0, "sampling_logprob": -1.0} for _ in ids]}, "routing_references": routes.to_dict(), "r3_store_id": "test", @@ -110,7 +110,7 @@ def test_parquet_top_sampling_references_are_completion_only(echo_last, echoed_r rows = len(ids) if echoed_rows else 2 choice = { "text": "ok", - "raw_output": {"completion_token_ids": ids}, + "token_ids": ids, "logprobs": {"content": [{"logprob": -1.0, "sampling_logprob": -0.5} for _ in ids]}, f"{column}_references": refs(rows, 100).to_dict(), f"{column}_format": "parquet_v1", @@ -216,7 +216,7 @@ def test_sampling_parser_accepts_independent_mask_head_files(): sampler.r3_store_id = "test" choice = { "text": "x", - "raw_output": {"completion_token_ids": [1, 2]}, + "token_ids": [1, 2], "logprobs": {"content": [{"logprob": -0.5, "sampling_logprob": -0.5}] * 2}, "top_sampling_format": "parquet_v1", "sampling_mask_format": "parquet_v1", diff --git a/tests_training/test_sampling.py b/tests_training/test_sampling.py index 298a9f4..b424dec 100644 --- a/tests_training/test_sampling.py +++ b/tests_training/test_sampling.py @@ -57,7 +57,7 @@ def _http_error( _SUCCESS = ( { "id": "cmpl-upstream-test", - "choices": [{"text": "hi", "finish_reason": "stop", "raw_output": {"completion_token_ids": [40, 50]}}], + "choices": [{"text": "hi", "finish_reason": "stop", "token_ids": [40, 50]}], }, ServerMetrics(), ) @@ -110,7 +110,7 @@ def test_completion_parser_retains_top_logprob_token_ids() -> None: { "text": "hi", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [40, 50]}, + "token_ids": [40, 50], "logprobs": { "content": [ { @@ -548,7 +548,7 @@ def test_error_record_has_no_secrets_or_prompt(self, no_sleep, no_jitter, monkey def _sse_success_bytes() -> bytes: - chunk = '{"id":"cmpl-upstream-test","choices":[{"text":"hi","finish_reason":"stop","raw_output":{"completion_token_ids":[40,50]}}]}' + chunk = '{"id":"cmpl-upstream-test","choices":[{"text":"hi","finish_reason":"stop","token_ids":[40,50],"prompt_token_ids":[1,2,3]}]}' return f"data: {chunk}\n\ndata: [DONE]\n\n".encode("utf-8") @@ -571,10 +571,8 @@ def _handler(request: httpx.Request) -> httpx.Response: sampler = _make_sampler() sampler._async_client = httpx.AsyncClient(transport=httpx.MockTransport(_handler)) - result, metrics = asyncio.run( - sampler.async_completions_stream(prompt=[1, 2, 3], raw_output=True, logical_request_id="lr-xyz") - ) - assert result["choices"][0]["raw_output"]["completion_token_ids"] == [40, 50] + result, metrics = asyncio.run(sampler.async_completions_stream(prompt=[1, 2, 3], logical_request_id="lr-xyz")) + assert result["choices"][0]["token_ids"] == [40, 50] assert result["id"] == "cmpl-upstream-test" assert seen["x-request-id"] == "lr-xyz" assert metrics.response_request_id == "lr-xyz-cmpl-server" @@ -593,7 +591,7 @@ def _handler(request: httpx.Request) -> httpx.Response: sampler = _make_sampler() sampler._async_client = httpx.AsyncClient(transport=httpx.MockTransport(_handler)) with pytest.raises(httpx.HTTPStatusError): - asyncio.run(sampler.async_completions_stream(prompt=[1, 2, 3], raw_output=True, logical_request_id="lr-1")) + asyncio.run(sampler.async_completions_stream(prompt=[1, 2, 3], logical_request_id="lr-1")) # errors.py opt-out means the transport is hit exactly once per stream call. assert posts["n"] == 1 sampler.close() @@ -613,7 +611,6 @@ def _handler(request: httpx.Request) -> httpx.Response: result, _ = asyncio.run( sampler.async_completions_stream( prompt=[1, 2, 3], - raw_output=True, prompt_cache_key="affinity-1", additional_headers_snapshot=snapshot, ) @@ -640,7 +637,6 @@ def _handler(request: httpx.Request) -> httpx.Response: asyncio.run( sampler.async_completions_stream( prompt=[1, 2, 3], - raw_output=True, additional_headers_snapshot=snapshot, ) ) @@ -656,7 +652,6 @@ def _handler(request: httpx.Request) -> httpx.Response: asyncio.run( sampler.async_completions_stream( prompt=[1, 2, 3], - raw_output=True, additional_headers_snapshot=snapshot, ) ) diff --git a/tests_training/test_sampling_client_close.py b/tests_training/test_sampling_client_close.py index 56e0cd1..0639b08 100644 --- a/tests_training/test_sampling_client_close.py +++ b/tests_training/test_sampling_client_close.py @@ -23,10 +23,8 @@ _CHOICE: dict[str, Any] = { "text": "out", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [40, 50]}, - "logprobs": { - "content": [{"logprob": -0.3, "sampling_logprob": -0.3}, {"logprob": -0.4, "sampling_logprob": -0.4}] - }, + "token_ids": [40, 50], + "logprobs": {"content": [{"logprob": -0.3, "sampling_logprob": -0.3}, {"logprob": -0.4, "sampling_logprob": -0.4}]}, } diff --git a/tests_training/test_sampling_stream_errors.py b/tests_training/test_sampling_stream_errors.py index b372216..e2abe3d 100644 --- a/tests_training/test_sampling_stream_errors.py +++ b/tests_training/test_sampling_stream_errors.py @@ -16,6 +16,7 @@ from fireworks.training.sdk import TITOError, TITOSidecar, TITOFailureTraceConfig from tests_training.test_tito import FakeSampler, FakeRenderer, _engine, _first_request from fireworks.training.sdk._sse import _SSEDecoder, _SSEProtocolError, _SSETruncationError +from fireworks.training.sdk.tito import TITOChatRequest from fireworks.training.sdk.sampling import ServerMetrics, DeploymentSampler, SamplingRequestError from fireworks.training.sdk.concurrency import AdaptiveConcurrencyController from fireworks.training.sdk.tito._engine import _LinearTrajectoryCore @@ -27,7 +28,9 @@ def _event(value: Any) -> bytes: return ("data: " + json.dumps(value, ensure_ascii=False) + "\r\n\r\n").encode() -def _completion(finish: str | None = "stop", tokens: tuple[int, ...] = (197, 3)) -> bytes: +def _completion( + finish: str | None = "stop", tokens: tuple[int, ...] = (197, 3), *, prompt_ids: tuple[int, ...] = (1,) +) -> bytes: return _event( { "id": "synthetic-completion-id", @@ -35,7 +38,8 @@ def _completion(finish: str | None = "stop", tokens: tuple[int, ...] = (197, 3)) { "text": "a", "finish_reason": finish, - "raw_output": {"completion_token_ids": list(tokens)}, + "token_ids": list(tokens), + "prompt_token_ids": list(prompt_ids), "logprobs": {"content": [{"logprob": -0.2, "sampling_logprob": -0.2} for _ in tokens]}, } ], @@ -120,7 +124,7 @@ async def test_unsuccessful_finish_cannot_become_success(finish: str) -> None: sampler.close() -async def test_done_and_raw_tokens_do_not_invent_a_finish() -> None: +async def test_done_and_token_ids_do_not_invent_a_finish() -> None: async with _upstream([_completion(None) + _DONE]) as (url, calls): sampler = DeploymentSampler(inference_url=url, model="synthetic", api_key="unused") sampler._RETRY_MAX_ATTEMPTS = 1 @@ -157,11 +161,19 @@ async def test_malformed_json_is_not_hidden_by_a_later_complete_chunk() -> None: @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("fault", [_ERROR, b"data: malformed-json\n\n"]) async def test_error_stream_never_commits_a_turn(tmp_path: Any, stream: bool, fault: bytes) -> None: - async with _upstream([_completion() + fault + _DONE, _completion("stop", (198, 3)) + _DONE]) as (url, calls): + request = {"messages": [{"role": "user", "content": "synthetic"}], "stream": stream} + renderer = FakeRenderer() + prompt_ids = tuple(renderer.render_conversation_tokens(TITOChatRequest.from_openai(request))) + async with _upstream( + [ + _completion(prompt_ids=prompt_ids) + fault + _DONE, + _completion("stop", (198, 3), prompt_ids=prompt_ids) + _DONE, + ] + ) as (url, calls): sampler = DeploymentSampler(inference_url=url, model="synthetic", api_key="unused") sidecar = TITOSidecar.from_deployment_sampler( sampler, - renderer=FakeRenderer(), + renderer=renderer, max_context_tokens=4096, max_output_tokens=128, failure_trace_config=TITOFailureTraceConfig(root_dir=tmp_path, min_free_bytes=0), @@ -170,7 +182,6 @@ async def test_error_stream_never_commits_a_turn(tmp_path: Any, stream: bool, fa try: endpoint = await sidecar.create_trajectory_async() headers = {"authorization": "Bearer " + endpoint.api_key} - request = {"messages": [{"role": "user", "content": "synthetic"}], "stream": stream} async with httpx.AsyncClient(trust_env=False) as client: response = await client.post( endpoint.openai_base_url + "/chat/completions", headers=headers, json=request @@ -186,7 +197,9 @@ async def test_error_stream_never_commits_a_turn(tmp_path: Any, stream: bool, fa ) local_decode = fault != _ERROR assert response.status_code == (200 if stream else (422 if local_decode else 502)) - assert error["code"] == ("tito_sdk_sampling_parse_error" if local_decode else "tito_server_stream_error") + assert error["code"] == ( + "tito_sdk_sampling_parse_error" if local_decode else "tito_server_stream_error" + ) assert error["origin"] == ("sdk_parser" if local_decode else "server") assert error["wire_status"] == 200 assert error["should_retry"] is False @@ -396,7 +409,9 @@ async def test_semantic_stream_status_controls_admission_without_rewriting_wire_ controller = AdaptiveConcurrencyController(initial_window=8, min_window=1, adjustment_interval=1) error = _event({"error": {"status_code": None, "status": status}}) async with _upstream([_completion() + error + _DONE]) as (url, calls): - sampler = DeploymentSampler(inference_url=url, model="synthetic", api_key="unused", concurrency_controller=controller) + sampler = DeploymentSampler( + inference_url=url, model="synthetic", api_key="unused", concurrency_controller=controller + ) sampler._RETRY_MAX_ATTEMPTS = 1 try: with pytest.raises(SamplingRequestError) as failure: diff --git a/tests_training/test_sampling_token_ids.py b/tests_training/test_sampling_token_ids.py new file mode 100644 index 0000000..30d5e5a --- /dev/null +++ b/tests_training/test_sampling_token_ids.py @@ -0,0 +1,203 @@ +"""Exact-token SSE contracts; constructed fixtures, not incident captures.""" + +from __future__ import annotations + +import json +from typing import Any + +import httpx +import pytest + +from fireworks.training.sdk._sse import _SSEProtocolError +from fireworks.training.sdk.sampling import DeploymentSampler + + +def _wire(choices: list[dict[str, Any]], generated: int = 2) -> bytes: + events = [{"choices": [choice]} for choice in choices] + events.append({"choices": [{"finish_reason": "stop"}], "usage": {"completion_tokens": generated}}) + return b"".join(b"data: " + json.dumps(event).encode() + b"\n\n" for event in events) + b"data: [DONE]\n\n" + + +def _delta(tokens: list[int], text: str = "") -> dict[str, Any]: + return { + "text": text, + "token_ids": tokens, + "logprobs": {"content": [{"token_id": token, "logprob": -0.5, "sampling_logprob": -0.2} for token in tokens]}, + } + + +@pytest.mark.parametrize("echo_last", [None, 0, 2, 3]) +@pytest.mark.parametrize("raw_output", [None, False]) +async def test_token_deltas_preserve_echo_and_server_prompt(echo_last: int | None, raw_output: bool | None) -> None: + # Server-expanded prompt IDs must win over the submitted IDs. + prompt = [10, 11, 12] + echoed = prompt[-echo_last:] if echo_last else [] + tokens = echoed + [20, 21] + first = _delta(tokens[:-1], text="你好") + first["prompt_token_ids"] = prompt + last = _delta(tokens[-1:]) # A token with no display text must survive. + requests = [] + + async def handle(request: httpx.Request) -> httpx.Response: + requests.append(json.loads(request.content)) + return httpx.Response(200, content=_wire([first, last])) + + sampler = DeploymentSampler("https://example.invalid", "model", "unused") + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + sampler._get_async_client = lambda: client + kwargs: dict[str, Any] = {"logprobs": True} + if raw_output is not None: + kwargs["raw_output"] = raw_output + if echo_last == 3: + kwargs["echo"] = True + elif echo_last is not None: + kwargs.update(echo=True, echo_last=echo_last) + try: + result = await sampler.sample_with_prompt_tokens_result([1], **kwargs) + completion = result.completions[0] + assert completion.full_tokens == prompt + [20, 21] + assert completion.prompt_len == 3 + assert completion.completion_len == 2 + expected_echo = 2 if echo_last == 3 else (echo_last or 0) + assert completion.echoed_prompt_logprob_count == expected_echo + assert completion.inference_logprobs == [-0.5] * (2 + expected_echo) + assert completion.sampling_logprobs == [-0.2] * (2 + expected_echo) + assert len(requests) == 1 + assert requests[0]["return_token_ids"] is True + assert requests[0]["raw_output"] is False + finally: + sampler.close() + + +@pytest.mark.parametrize("width", [2, 128]) +async def test_raw_off_stream_preserves_top_logprobs_and_sampling_controls(width: int) -> None: + deltas = [_delta([20]), _delta([21])] + deltas[0]["prompt_token_ids"] = [1] + heads = [[{"token_id": token + i, "logprob": -0.5 - i} for i in range(width)] for token in (20, 21)] + for delta, head in zip(deltas, heads): + delta["logprobs"]["content"][0]["top_logprobs"] = head + + async def handle(request: httpx.Request) -> httpx.Response: + payload = json.loads(request.content) + assert payload["raw_output"] is False + assert payload["return_token_ids"] is True + assert payload["top_k"] == 10 + assert payload["top_p"] == 0.9 + assert payload["temperature"] == 0.7 + assert payload["top_logprobs"] == width + return httpx.Response(200, content=_wire(deltas)) + + sampler = DeploymentSampler("https://example.invalid", "model", "unused") + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + sampler._get_async_client = lambda: client + try: + result = await sampler.sample_with_prompt_tokens_result( + [1], logprobs=True, top_logprobs=width, top_k=10, top_p=0.9, temperature=0.7 + ) + completion = result.completions[0] + assert completion.full_tokens == [1, 20, 21] + assert completion.inference_logprobs == [-0.5, -0.5] + assert completion.sampling_logprobs == [-0.2, -0.2] + assert completion.inference_topk_token_ids == [[row["token_id"] for row in head] for head in heads] + assert completion.inference_topk_logprobs == [[row["logprob"] for row in head] for head in heads] + finally: + sampler.close() + + +@pytest.mark.parametrize( + "fault", ["missing_ids", "missing_prompt", "row_count", "row_id", "usage", "prompt_change", "index"] +) +async def test_incomplete_token_stream_fails_without_retry(fault: str) -> None: + first = _delta([20]) + first["prompt_token_ids"] = [1] + last = _delta([21], text="x") + generated = 2 + if fault == "missing_ids": + del first["token_ids"] + elif fault == "missing_prompt": + del first["prompt_token_ids"] + elif fault == "row_count": + first["logprobs"]["content"] = [] + elif fault == "row_id": + first["logprobs"]["content"][0]["token_id"] = 99 + elif fault == "usage": + generated = 3 + elif fault == "prompt_change": + last["prompt_token_ids"] = [2] + else: + last["index"] = 1 + calls = 0 + + async def handle(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response(200, content=_wire([first, last], generated=generated)) + + sampler = DeploymentSampler("https://example.invalid", "model", "unused") + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + sampler._get_async_client = lambda: client + try: + with pytest.raises(_SSEProtocolError): + await sampler.sample_with_prompt_tokens_result([1], logprobs=True, raw_output=False) + assert calls == 1 + finally: + sampler.close() + + +async def test_explicit_raw_mode_is_rejected_before_request() -> None: + calls = 0 + + async def handle(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response(500) + + sampler = DeploymentSampler("https://example.invalid", "model", "unused") + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + sampler._get_async_client = lambda: client + try: + with pytest.raises(ValueError, match="requires raw_output=False"): + await sampler.sample_with_prompt_tokens_result([1], raw_output=True) + assert calls == 0 + finally: + sampler.close() + + +async def test_raw_snapshot_does_not_replace_missing_incremental_ids() -> None: + choice = {"text": "hi", "raw_output": {"completion_token_ids": [20, 21]}} + calls = 0 + + async def handle(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + assert json.loads(request.content)["raw_output"] is False + return httpx.Response(200, content=_wire([choice])) + + sampler = DeploymentSampler("https://example.invalid", "model", "unused") + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + sampler._get_async_client = lambda: client + try: + with pytest.raises(_SSEProtocolError, match="complete token IDs"): + await sampler.sample_with_prompt_tokens_result([1]) + assert calls == 1 + finally: + sampler.close() + + +@pytest.mark.parametrize( + "options", [{}, {"images": ["synthetic"]}, {"include_routing_matrix": True}, {"sampling_mask": "parquet_v1"}] +) +async def test_default_paths_use_incremental_ids(options: dict[str, Any]) -> None: + sampler = DeploymentSampler("https://example.invalid", "model", "unused") + + async def capture(**kwargs: Any) -> Any: + assert kwargs["raw_output"] is False + assert kwargs["return_token_ids"] is True + raise ValueError("captured request") + + sampler.async_completions_stream = capture + try: + with pytest.raises(ValueError, match="captured request"): + await sampler.sample_with_prompt_tokens_result([1], **options) + finally: + sampler.close() diff --git a/tests_training/test_sse_chunks.py b/tests_training/test_sse_chunks.py new file mode 100644 index 0000000..66a231b --- /dev/null +++ b/tests_training/test_sse_chunks.py @@ -0,0 +1,36 @@ +"""Chunk framing and immutable output for all training SSE splitters.""" + +from __future__ import annotations + +from typing import AsyncIterator + +import pytest + +from fireworks.training.sdk._sse import _SSEDecoder +from fireworks.training.sdk._core._streaming import SSEDecoder + + +class _Fragments: + def __init__(self, wire: bytes, size: int) -> None: + self.wire = wire + self.size = size + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + for offset in range(0, len(self.wire), self.size): + yield self.wire[offset : offset + self.size] + + +@pytest.mark.parametrize("delimiter", [b"\n\n", b"\r\r", b"\r\n\r\n"]) +@pytest.mark.parametrize("size", [1, 7, 16384]) +async def test_fragmented_large_events_preserve_bytes(delimiter: bytes, size: int) -> None: + events = [b"data: " + "你好🙂".encode() + delimiter, b"data: " + b"x" * 65536 + delimiter, b"data: tail"] + wire = b"".join(events) + fragments = _Fragments(wire=wire, size=size) + outputs = [ + [chunk async for chunk in _SSEDecoder._aiter_chunks(fragments)], + list(SSEDecoder()._iter_chunks(iter(wire[i : i + size] for i in range(0, len(wire), size)))), + [chunk async for chunk in SSEDecoder()._aiter_chunks(fragments.aiter_bytes())], + ] + for chunks in outputs: + assert chunks == events + assert all(type(chunk) is bytes for chunk in chunks) diff --git a/tests_training/test_tito.py b/tests_training/test_tito.py index c57a726..c783bea 100644 --- a/tests_training/test_tito.py +++ b/tests_training/test_tito.py @@ -1034,7 +1034,7 @@ async def controlled_transport(*args: Any, **kwargs: Any): { "text": "a", "finish_reason": "stop", - "raw_output": {"completion_token_ids": [197, 3]}, + "token_ids": [197, 3], } ], },