diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index fa84471d36..d1bea819eb 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -268,6 +268,7 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_fix_sentencepiece_tokenizer_guard.py \ tests/saving/test_compressed_export_schemes.py \ tests/saving/test_export_api_surface.py \ tests/saving/test_export_dispatch.py \ @@ -358,6 +359,7 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_fix_sentencepiece_tokenizer_guard.py \ tests/saving/test_compressed_export_schemes.py \ tests/saving/test_export_api_surface.py \ tests/saving/test_export_dispatch.py \ diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index cf8c021e38..58ef2558f3 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -729,6 +729,7 @@ jobs: content, events = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": prompt}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": enabled, "session_id": f"{session}-att{attempt_i}", "temperature": TOOL_PROBE_TEMP, @@ -818,6 +819,7 @@ jobs: content, events = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["web_search"], "session_id": "ci-tool-calling-web", "temperature": 0.0, diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index d3d765aa84..946681706a 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -612,6 +612,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["python"], "session_id": "ci-tool-calling-py", "temperature": TEMP, @@ -647,6 +648,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["web_search"], "session_id": "ci-tool-calling-web", "temperature": TEMP, diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 233292f7a3..63a7e9dc8f 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -791,6 +791,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["python"], "session_id": "ci-tool-calling-py", "temperature": TEMP, @@ -816,6 +817,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Use the terminal tool to run `echo hello-bash-tool` and tell me the exact output."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["terminal"], "session_id": "ci-tool-calling-bash", "temperature": TEMP, @@ -840,6 +842,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["web_search"], "session_id": "ci-tool-calling-web", "temperature": TEMP, diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 929cc37bda..1f7bc8dcc0 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -95,8 +95,8 @@ "file": "fastapi/routing.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", - "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" + "evidence": "L587: while True: sha256:06c2c7f15d73bf192e5e3272c5ff5fcaeff7f6774fef5f4eca6ef473ae50e2b3", + "evidence_hash": "57acd497f404c203e4450d0580ad85aa8a33406e8d64ad06fbac6cf47d97b24d" }, { "package": "fastmcp-slim", diff --git a/studio/MCP.md b/studio/MCP.md new file mode 100644 index 0000000000..91b39fcc69 --- /dev/null +++ b/studio/MCP.md @@ -0,0 +1,34 @@ +# Unsloth Studio MCP server + +Studio can expose a local MCP server so an MCP client can inspect models and +GPU state, validate recipes, start or stop training, inspect recipe output, and +export a loaded model. + +The server is disabled by default. Enable it for a local Studio process with: + +```bash +UNSLOTH_STUDIO_ENABLE_MCP=1 \ +UNSLOTH_STUDIO_MCP_TOKEN='use-a-local-secret' \ +unsloth studio +``` + +The endpoint is `http://127.0.0.1:8888/mcp/` when Studio uses its default port +(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual Studio +port when it is configured differently. + +The high-impact tools are: + +- `studio_status` and `list_local_models` for discovery +- `get_training_status`, `start_training`, `stop_training`, and `list_training_runs` +- `validate_recipe`, `get_recipe_job_status`, and `get_recipe_job_dataset` +- `load_checkpoint` and `export_gguf` + +`start_training` accepts the same fields as the Studio `TrainingStartRequest`. +The request is validated by the existing Pydantic model before a subprocess is +started. Export paths use the existing Studio validation as well. + +The endpoint always requires `UNSLOTH_STUDIO_MCP_TOKEN` and checks an exact +Bearer token for both HTTP and WebSocket connections. Keep it on localhost +unless the deployment has an authenticated reverse proxy. The MCP endpoint is +intentionally opt-in because tools can consume GPU memory, write model +artifacts, and stop active work. \ No newline at end of file diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index 897db8262d..5113eebb36 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -10,10 +10,242 @@ native-chat-template fallback used by the transformers and MLX backends. import copy import json import logging +from dataclasses import dataclass from typing import Optional _THINK_OPEN = "" _THINK_CLOSE = "" +_GEMMA_CHANNEL_START = "<|channel>" +_GEMMA_THOUGHT_OPEN = "<|channel>thought" +_GEMMA_THOUGHT_CLOSE = "" +_GEMMA_TEMPLATE_OPENERS = ( + _GEMMA_THOUGHT_OPEN + "\n", + _GEMMA_THOUGHT_OPEN + "\\n", + _GEMMA_THOUGHT_OPEN + _GEMMA_THOUGHT_CLOSE, +) + + +def _tokenizer_objects(tokenizer) -> tuple: + """Return a processor/tokenizer and its distinct nested tokenizer.""" + if tokenizer is None: + return () + nested = getattr(tokenizer, "tokenizer", None) + return (tokenizer,) if nested is None or nested is tokenizer else (tokenizer, nested) + + +def _selected_template_strings_from_value( + template, + tools = None, + *, + prefer_tool_use: bool = True, +) -> tuple[str, ...]: + """Return the named chat template matching HF's default selection rules.""" + tools = tools or None + if isinstance(template, str): + return (template,) + if not isinstance(template, dict): + return () + if prefer_tool_use and tools and isinstance(template.get("tool_use"), str): + return (template["tool_use"],) + if isinstance(template.get("default"), str): + return (template["default"],) + values = tuple(value for value in template.values() if isinstance(value, str)) + return values if len(values) == 1 else () + + +def _selected_chat_template_strings(tokenizer, tools = None) -> tuple[str, ...]: + """Return the active chat template selected for this request.""" + tools = tools or None + getter = getattr(tokenizer, "get_chat_template", None) + if callable(getter): + for kwargs in ({"chat_template": None, "tools": tools}, {"tools": tools}, {}): + try: + selected = getter(**kwargs) + except Exception: + continue + if isinstance(selected, str): + return (selected,) + # ProcessorMixin.apply_chat_template does not switch to "tool_use" implicitly; + # it uses "default" unless chat_template= names another template. + is_processor = getattr(tokenizer, "tokenizer", None) is not None and callable( + getattr(tokenizer, "apply_chat_template", None) + ) + return _selected_template_strings_from_value( + getattr(tokenizer, "chat_template", None), + tools, + prefer_tool_use = not is_processor, + ) + + +def _detect_reasoning_channel_markers_from_templates( + templates: tuple[str, ...], +) -> Optional[tuple[str, str]]: + """Return Gemma native reasoning markers only when a template emits them.""" + if any(opener in template for template in templates for opener in _GEMMA_TEMPLATE_OPENERS): + return _GEMMA_THOUGHT_OPEN, _GEMMA_THOUGHT_CLOSE + return None + + +def detect_reasoning_channel_markers(tokenizer, tools = None) -> Optional[tuple[str, str]]: + """Return native Gemma thought-channel markers supported by a tokenizer. + + Detection uses the active chat template rather than model names or vocabulary + membership. Some models expose Gemma control tokens without using the native + thought-channel response protocol, and those must keep normal + ``skip_special_tokens`` streaming. + """ + for obj in _tokenizer_objects(tokenizer): + templates = _selected_chat_template_strings(obj, tools) + if templates: + return _detect_reasoning_channel_markers_from_templates(templates) + return None + + +def detect_reasoning_channel_markers_from_template( + template, tools = None +) -> Optional[tuple[str, str]]: + """Return native Gemma thought-channel markers from a raw template value.""" + return _detect_reasoning_channel_markers_from_templates( + _selected_template_strings_from_value(template, tools) + ) + + +def detect_reasoning_channel_markers_from_model_info( + tokenizer, + model_info: Optional[dict] = None, + tools = None, +) -> Optional[tuple[str, str]]: + """Return reasoning markers from the active or cached native template.""" + markers = detect_reasoning_channel_markers(tokenizer, tools = tools) + if markers is not None or not isinstance(model_info, dict): + return markers + + native_templates = ( + model_info.get("native_chat_template"), + (model_info.get("chat_template_info") or {}).get("template"), + ) + for template in native_templates: + markers = detect_reasoning_channel_markers_from_template(template, tools) + if markers is not None: + return markers + return None + + +@dataclass(frozen = True) +class ChatTemplateRenderResult: + """Prompt plus response-protocol metadata selected by the renderer.""" + + prompt: str + reasoning_channel_markers: Optional[tuple[str, str]] = None + + +def _split_partial_marker(text: str, marker: str) -> tuple[str, str]: + """Hold the longest suffix that may become ``marker`` in the next chunk.""" + for length in range(min(len(text), len(marker) - 1), 0, -1): + if text.endswith(marker[:length]): + return text[:-length], text[-length:] + return text, "" + + +class ReasoningChannelNormalizer: + """Incrementally convert one native reasoning channel to ````. + + The parser follows mlx-vlm's streaming boundary behavior but emits Studio's + established canonical text contract. Only the configured opening and + closing markers are consumed; tool-call and other control markers remain + available to downstream parsers. + """ + + def __init__(self, opening_marker: str, closing_marker: str): + self._opening_marker = opening_marker + self._closing_marker = closing_marker + self._buffer = "" + self._in_reasoning = False + self._reasoning_done = False + self._skip_opening_newline = False + + def feed(self, text: str) -> str: + """Consume a raw text delta and return the stable canonical delta.""" + self._buffer += text or "" + output: list[str] = [] + while self._buffer: + if self._reasoning_done: + output.append(self._buffer) + self._buffer = "" + break + + if self._in_reasoning and self._skip_opening_newline: + if self._buffer.startswith("\n"): + self._buffer = self._buffer[1:] + self._skip_opening_newline = False + if not self._buffer: + break + + marker = self._closing_marker if self._in_reasoning else self._opening_marker + index = self._buffer.find(marker) + if index < 0: + stable, self._buffer = _split_partial_marker(self._buffer, marker) + output.append(stable) + break + + output.append(self._buffer[:index]) + self._buffer = self._buffer[index + len(marker) :] + if self._in_reasoning: + output.append(_THINK_CLOSE) + self._in_reasoning = False + self._reasoning_done = True + else: + output.append(_THINK_OPEN) + self._in_reasoning = True + self._skip_opening_newline = True + return "".join(output) + + def finish(self) -> str: + """Flush a naturally completed stream and close an open think block.""" + output = self.drain() + if self._in_reasoning: + output += _THINK_CLOSE + self._in_reasoning = False + self._reasoning_done = True + return output + + def drain(self) -> str: + """Flush buffered literal text without synthesizing a closing tag.""" + output = self._buffer + self._buffer = "" + return output + + +def normalize_reasoning_snapshots( + stream, + tokenizer = None, + cancel_event = None, + markers: Optional[tuple[str, str]] = None, + tools = None, +): + """Normalize a prefix-monotonic cumulative text stream when supported.""" + markers = markers or detect_reasoning_channel_markers(tokenizer, tools = tools) + if markers is None: + yield from stream + return + + normalizer = ReasoningChannelNormalizer(*markers) + raw_output = "" + normalized_output = "" + for snapshot in stream: + if not snapshot.startswith(raw_output): + raise RuntimeError("Reasoning normalization requires cumulative text snapshots") + delta = normalizer.feed(snapshot[len(raw_output) :]) + raw_output = snapshot + if delta: + normalized_output += delta + yield normalized_output + + cancelled = cancel_event is not None and cancel_event.is_set() + tail = normalizer.drain() if cancelled else normalizer.finish() + if tail: + normalized_output += tail + yield normalized_output def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str: @@ -166,7 +398,8 @@ def render_native_template( preserve_thinking: Optional[bool] = None, apply_fn = None, hf_token: Optional[str] = None, -) -> Optional[str]: + return_metadata: bool = False, +): """Render ``messages`` + ``tools`` with the model's NATIVE chat template. Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit @@ -175,7 +408,9 @@ def render_native_template( tool-calling syntax. It is loaded straight from the repo (bypassing any override on the live tokenizer) and cached on ``model_info``. Returns the rendered prompt only if the native template actually emits the tools (render - differs with vs without tools); otherwise ``None``. + differs with vs without tools); otherwise ``None``. With ``return_metadata``, + returns ``ChatTemplateRenderResult`` so callers can stream with the response + protocol selected by this request's template. ``hf_token`` is the token the model was loaded with -- passed to the repo load so a gated/private model's native template can still be fetched (otherwise the @@ -261,7 +496,16 @@ def render_native_template( exc, ) return None - return with_tools if with_tools != no_tools else None + if with_tools == no_tools: + return None + if return_metadata: + return ChatTemplateRenderResult( + with_tools, + _detect_reasoning_channel_markers_from_templates( + _selected_template_strings_from_value(native_tpl, tools) + ), + ) + return with_tools def render_with_native_template_fallback( @@ -277,7 +521,8 @@ def render_with_native_template_fallback( preserve_thinking: Optional[bool] = None, apply_fn = None, hf_token: Optional[str] = None, -) -> str: + return_metadata: bool = False, +): """Return ``formatted_prompt``, swapping in a native-template render when an override template dropped the ``tools`` schema. @@ -285,9 +530,27 @@ def render_with_native_template_fallback( them (detected by comparison, robust against tool names in the system prompt), re-render with the model's native template. Shared by the transformers and MLX backends so both advertise tools consistently. ``hf_token`` is forwarded so a - gated/private model's native template can still be fetched.""" + gated/private model's native template can still be fetched. With + ``return_metadata``, returns the selected prompt plus reasoning-channel markers + for the exact template used by this request.""" + live_markers = detect_reasoning_channel_markers(tokenizer, tools = tools) + + def _result(prompt: str, markers = live_markers): + if return_metadata: + return ChatTemplateRenderResult(prompt, markers) + return prompt + if not tools: - return formatted_prompt + # Gemma 4 can emit its native reasoning protocol even when a generation-time + # Unsloth override rendered a marker-free prompt. Preserve the live-verified + # no-tools thinking behavior without letting cached native metadata describe + # unrelated tool prompts that kept the active override. + markers = live_markers + if markers is None: + markers = detect_reasoning_channel_markers_from_model_info( + tokenizer, model_info, tools = None + ) + return _result(formatted_prompt, markers) if apply_fn is None: apply_fn = apply_chat_template_for_generation # Probe whether the live template dropped the schema. A tools-requiring template @@ -307,9 +570,9 @@ def render_with_native_template_fallback( active_model_name, exc, ) - return formatted_prompt + return _result(formatted_prompt) if formatted_prompt != probe_no_tools: - return formatted_prompt # template already emits the tools schema + return _result(formatted_prompt) # template already emits the tools schema native_prompt = render_native_template( model_info = model_info, active_model_name = active_model_name, @@ -320,6 +583,7 @@ def render_with_native_template_fallback( preserve_thinking = preserve_thinking, apply_fn = apply_fn, hf_token = hf_token, + return_metadata = return_metadata, ) if native_prompt: logger.info( @@ -328,4 +592,4 @@ def render_with_native_template_fallback( active_model_name, ) return native_prompt - return formatted_prompt + return _result(formatted_prompt) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 172e9e5546..8d262bbb0f 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -5,7 +5,7 @@ from unsloth import FastLanguageModel, FastVisionModel from unsloth.chat_templates import get_chat_template -from transformers import TextStreamer +from transformers import TextIteratorStreamer, TextStreamer from peft import PeftModel, PeftModelForCausalLM import json @@ -32,6 +32,11 @@ from core.inference.chat_eos import ( chat_eos_repair, resolve_chat_turn_end_eos_ids_using, ) +from core.inference.chat_template_helpers import ( + ReasoningChannelNormalizer, + detect_reasoning_channel_markers, + detect_think_prefill, +) from core.inference.presence_penalty import _make_presence_penalty_processor from io import StringIO import structlog @@ -187,6 +192,53 @@ class HarmonyTextStreamer: self._queue.put(new_content) +class ReasoningTextIteratorStreamer(TextIteratorStreamer): + """TextIteratorStreamer that preserves native channel tokens until parsed.""" + + def __init__( + self, + tokenizer, + *, + markers: tuple[str, str], + skip_prompt: bool = True, + timeout: float = 0.2, + cancel_event = None, + **decode_kwargs, + ): + decode_kwargs["skip_special_tokens"] = False + super().__init__(tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs) + self._normalizer = ReasoningChannelNormalizer(*markers) + self._cancel_event = cancel_event + self._aborted = False + + def abort(self): + """Mark generation as failed so ``end`` drains without closing.""" + self._aborted = True + + def on_finalized_text( + self, + text: str, + stream_end: bool = False, + ): + """Queue canonical deltas, closing only on natural stream completion.""" + delta = self._normalizer.feed(text) + if delta: + self.text_queue.put(delta, timeout = self.timeout) + + if stream_end: + cancelled = self._aborted or ( + self._cancel_event is not None and self._cancel_event.is_set() + ) + tail = self._normalizer.drain() if cancelled else self._normalizer.finish() + if tail: + self.text_queue.put(tail, timeout = self.timeout) + self.text_queue.put(self.stop_signal, timeout = self.timeout) + + +class _GenerationThreadError(RuntimeError): + """Generation worker failures that should propagate through stream routes.""" + + class InferenceBackend: """Unified inference backend supporting text, vision, and LoRA models""" @@ -836,6 +888,7 @@ class InferenceBackend: thread_id: Optional[str] = None, rag_scope: Optional[dict] = None, presence_penalty: float = 0.0, + reasoning_prefilled: bool = False, ): """Run an agentic tool loop on top of ``generate_chat_response``. @@ -889,6 +942,7 @@ class InferenceBackend: session_id = session_id, thread_id = thread_id, rag_scope = rag_scope, + reasoning_prefilled = reasoning_prefilled, ) def generate_chat_response( @@ -960,8 +1014,7 @@ class InferenceBackend: thread can toggle adapters under the generation lock. """ if not self.active_model_name: - yield "Error: No active model" - return + raise RuntimeError("No active model") model_info = self.models[self.active_model_name] is_vision = model_info.get("is_vision", False) @@ -1049,6 +1102,7 @@ class InferenceBackend: template_messages = [{"role": "system", "content": system_prompt}] + messages else: template_messages = messages + reasoning_channel_markers_resolved = False try: if not (hasattr(tokenizer, "chat_template") and tokenizer.chat_template): raise ValueError( @@ -1058,6 +1112,7 @@ class InferenceBackend: f"Please use a model that includes a chat template, or manually set " f"one via tokenizer.chat_template before inference." ) + reasoning_channel_markers = None formatted_prompt = self._apply_chat_template_for_generation( tokenizer, template_messages, @@ -1073,7 +1128,7 @@ class InferenceBackend: render_with_native_template_fallback, ) - formatted_prompt = render_with_native_template_fallback( + render_result = render_with_native_template_fallback( formatted_prompt = formatted_prompt, tokenizer = tokenizer, model_info = model_info, @@ -1085,13 +1140,19 @@ class InferenceBackend: preserve_thinking = preserve_thinking, apply_fn = self._apply_chat_template_for_generation, hf_token = model_info.get("hf_token"), + return_metadata = True, ) + formatted_prompt = render_result.prompt + reasoning_channel_markers = render_result.reasoning_channel_markers + reasoning_channel_markers_resolved = True logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") except Exception as e: logger.error(f"Error applying chat template: {e}") # Fall back to manual formatting formatted_prompt = self.format_chat_prompt(messages, system_prompt) + reasoning_channel_markers = None + reasoning_channel_markers_resolved = True # Step 3: generate yield from self.generate_stream( @@ -1105,6 +1166,8 @@ class InferenceBackend: cancel_event = cancel_event, _adapter_state = _adapter_state, presence_penalty = presence_penalty, + reasoning_channel_markers = reasoning_channel_markers, + reasoning_channel_markers_resolved = reasoning_channel_markers_resolved, ) def _generate_vision_response( @@ -1190,21 +1253,27 @@ class InferenceBackend: # Stream with TextIteratorStreamer + background thread try: - from core.inference.chat_template_helpers import detect_think_prefill - # Re-emit an open prefill swallowed by skip_prompt (see # generate_stream). think_prefix = detect_think_prefill( prompt_text, getattr(raw_tokenizer, "all_special_tokens", None) ) - from transformers import TextIteratorStreamer import threading - streamer = TextIteratorStreamer( + streamer = self._make_text_streamer( raw_tokenizer, + protocol_source = processor, + # The text-only VLM fallback above did not render with the + # processor template, so its native markers do not describe + # this request's response protocol. + reasoning_channel_markers = detect_reasoning_channel_markers(processor) + if image + else None, + reasoning_channel_markers_resolved = True, skip_prompt = True, - skip_special_tokens = True, timeout = 0.2, + cancel_event = cancel_event, + use_harmony = self._is_gpt_oss_model(), ) generation_kwargs = dict( @@ -1226,6 +1295,10 @@ class InferenceBackend: ) if _pp is not None: generation_kwargs["logits_processor"] = _pp + stopping_criteria = self._cancel_stopping_criteria(cancel_event) + if stopping_criteria is not None: + generation_kwargs["stopping_criteria"] = stopping_criteria + active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs) err: dict[str, str] = {} @@ -1235,6 +1308,8 @@ class InferenceBackend: model.generate(**generation_kwargs) except Exception as e: err["msg"] = str(e) + if hasattr(streamer, "abort"): + streamer.abort() logger.error(f"Vision generation error in thread: {e}") finally: try: @@ -1251,12 +1326,17 @@ class InferenceBackend: if think_prefix: yield think_prefix from queue import Empty + import time generation_complete = False + cancel_deadline = None try: while True: if cancel_event is not None and cancel_event.is_set(): - break + if cancel_deadline is None: + cancel_deadline = time.monotonic() + 10 + elif time.monotonic() >= cancel_deadline: + break try: new_token = next(streamer) except StopIteration: @@ -1265,27 +1345,48 @@ class InferenceBackend: except Empty: if not thread.is_alive(): generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) + break + if cancel_deadline is not None: + remaining = cancel_deadline - time.monotonic() + if remaining <= 0: + break + thread.join(timeout = remaining) + if thread.is_alive(): + break + generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) break continue if new_token: - output += new_token - cleaned = self._clean_generated_text(output) + output, cleaned = self._append_stream_delta( + output, new_token, active_stop_token_ids + ) yield cleaned finally: if cancel_event is not None and not generation_complete: cancel_event.set() - thread.join(timeout = 10) + join_timeout = 10 + if cancel_deadline is not None: + join_timeout = max(0, cancel_deadline - time.monotonic()) + thread.join(timeout = join_timeout) if thread.is_alive(): logger.warning( "Vision generation thread did not exit after cancel/join timeout" ) if err.get("msg"): - yield f"Error: {err['msg']}" + raise _GenerationThreadError(err["msg"]) + except _GenerationThreadError: + raise except Exception as e: logger.error(f"Vision generation error: {e}") - yield f"Error: {str(e)}" + raise def generate_audio_input_response( self, @@ -1410,11 +1511,13 @@ class InferenceBackend: ) if err.get("msg"): - yield f"Error: {err['msg']}" + raise _GenerationThreadError(err["msg"]) + except _GenerationThreadError: + raise except Exception as e: logger.error(f"Audio input generation error: {e}") - yield f"Error: {str(e)}" + raise def generate_whisper_response( self, @@ -1447,6 +1550,86 @@ class InferenceBackend: from utils.datasets import is_gpt_oss_model_name return is_gpt_oss_model_name(model_name or self.active_model_name or "") + def _make_text_streamer( + self, + tokenizer, + *, + protocol_source = None, + reasoning_channel_markers = None, + reasoning_channel_markers_resolved: bool = False, + skip_prompt: bool = True, + timeout: float = 0.2, + cancel_event = None, + use_harmony: bool = False, + ): + """Create the streamer matching this model's native response protocol.""" + if use_harmony: + try: + return HarmonyTextStreamer( + tokenizer, + skip_prompt = skip_prompt, + timeout = timeout, + ) + except Exception as e: + logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}") + return TextIteratorStreamer( + tokenizer, + skip_prompt = skip_prompt, + skip_special_tokens = True, + timeout = timeout, + ) + + markers = ( + reasoning_channel_markers + if reasoning_channel_markers_resolved + else reasoning_channel_markers + or detect_reasoning_channel_markers(protocol_source or tokenizer) + ) + if markers is not None: + return ReasoningTextIteratorStreamer( + tokenizer, + markers = markers, + skip_prompt = skip_prompt, + timeout = timeout, + cancel_event = cancel_event, + ) + return TextIteratorStreamer( + tokenizer, + skip_prompt = skip_prompt, + skip_special_tokens = True, + timeout = timeout, + ) + + def _append_stream_delta( + self, + output: str, + new_token: str, + stop_token_ids = None, + ): + """Append a streamer delta and apply response-boundary cleanup.""" + output += new_token + return output, self._clean_generated_text(output, stop_token_ids = stop_token_ids) + + def _drain_streamer_tail( + self, + streamer, + output: str, + stop_token_ids = None, + ): + """Drain queued streamer text after the producer exits.""" + while True: + try: + new_token = next(streamer) + except StopIteration: + return output + except Exception: + return output + if new_token: + output, cleaned = self._append_stream_delta( + output, new_token, stop_token_ids = stop_token_ids + ) + yield cleaned + def generate_stream( self, prompt: str, @@ -1459,6 +1642,8 @@ class InferenceBackend: cancel_event = None, _adapter_state = None, presence_penalty: float = 0.0, + reasoning_channel_markers = None, + reasoning_channel_markers_resolved: bool = False, ) -> Generator[str, None, None]: """Generate a streaming text response (text models only). @@ -1467,8 +1652,7 @@ class InferenceBackend: ``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it). """ if not self.active_model_name: - yield "Error: No active model" - return + raise RuntimeError("No active model") model_info = self.models[self.active_model_name] model = model_info["model"] @@ -1481,9 +1665,7 @@ class InferenceBackend: try: inputs = tokenizer(prompt, return_tensors = "pt").to(model.device) - from transformers import TextIteratorStreamer import threading - from core.inference.chat_template_helpers import detect_think_prefill # skip_prompt swallows an open prefilled by the template; # re-emit it so the frontend can render the thinking block. @@ -1494,30 +1676,16 @@ class InferenceBackend: else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None)) ) - # gpt-oss models: HarmonyTextStreamer parses the multi-channel - # harmony protocol into tags - if self._is_gpt_oss_model(): - try: - streamer = HarmonyTextStreamer( - tokenizer, - skip_prompt = True, - timeout = 0.2, - ) - except Exception as e: - logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}") - streamer = TextIteratorStreamer( - tokenizer, - skip_prompt = True, - skip_special_tokens = True, - timeout = 0.2, - ) - else: - streamer = TextIteratorStreamer( - tokenizer, - skip_prompt = True, - skip_special_tokens = True, - timeout = 0.2, - ) + streamer = self._make_text_streamer( + tokenizer, + protocol_source = model_info.get("tokenizer"), + reasoning_channel_markers = reasoning_channel_markers, + reasoning_channel_markers_resolved = reasoning_channel_markers_resolved, + skip_prompt = True, + timeout = 0.2, + cancel_event = cancel_event, + use_harmony = self._is_gpt_oss_model(), + ) generation_kwargs = dict( **inputs, @@ -1535,27 +1703,16 @@ class InferenceBackend: if tokenizer.pad_token_id is None else tokenizer.pad_token_id, ) + active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs) # Presence penalty (GGUF parity); prompt_len excludes prompt tokens. _pp = _make_presence_penalty_processor( presence_penalty, int(inputs["input_ids"].shape[1]) ) if _pp is not None: generation_kwargs["logits_processor"] = _pp - if cancel_event is not None: - from transformers.generation.stopping_criteria import ( - StoppingCriteria, - StoppingCriteriaList, - ) - class _CancelCriteria(StoppingCriteria): - def __init__(self, ev): - self.ev = ev - - def __call__(self, input_ids, scores, **kwargs): - return self.ev.is_set() - - generation_kwargs["stopping_criteria"] = StoppingCriteriaList( - [_CancelCriteria(cancel_event)] - ) + stopping_criteria = self._cancel_stopping_criteria(cancel_event) + if stopping_criteria is not None: + generation_kwargs["stopping_criteria"] = stopping_criteria def generate_fn(): with self._generation_lock: @@ -1565,6 +1722,8 @@ class InferenceBackend: model.generate(**generation_kwargs) except Exception as e: err["msg"] = str(e) + if hasattr(streamer, "abort"): + streamer.abort() logger.error(f"Generation error: {e}") finally: try: @@ -1582,12 +1741,17 @@ class InferenceBackend: if think_prefix: yield think_prefix from queue import Empty + import time generation_complete = False + cancel_deadline = None try: while True: if cancel_event is not None and cancel_event.is_set(): - break + if cancel_deadline is None: + cancel_deadline = time.monotonic() + 10 + elif time.monotonic() >= cancel_deadline: + break try: new_token = next(streamer) except StopIteration: @@ -1596,11 +1760,27 @@ class InferenceBackend: except Empty: if not thread.is_alive(): generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) + break + if cancel_deadline is not None: + remaining = cancel_deadline - time.monotonic() + if remaining <= 0: + break + thread.join(timeout = remaining) + if thread.is_alive(): + break + generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) break continue if new_token: - output += new_token - cleaned = self._clean_generated_text(output) + output, cleaned = self._append_stream_delta( + output, new_token, active_stop_token_ids + ) yield cleaned finally: # Set cancel_event only on early exit (user cancel), NOT on @@ -1609,16 +1789,21 @@ class InferenceBackend: # disrupt the next serialized request (e.g. compare mode). if cancel_event is not None and not generation_complete: cancel_event.set() - thread.join(timeout = 10) + join_timeout = 10 + if cancel_deadline is not None: + join_timeout = max(0, cancel_deadline - time.monotonic()) + thread.join(timeout = join_timeout) if thread.is_alive(): logger.warning("Generation thread did not exit after cancel/join timeout") if err.get("msg"): - yield f"Error: {err['msg']}" + raise _GenerationThreadError(err["msg"]) + except _GenerationThreadError: + raise except Exception as e: logger.error(f"Error during generation: {e}") - yield f"Error: {str(e)}" + raise # ── Audio (TTS) Generation ──────────────────────────────────── @@ -2107,8 +2292,42 @@ class InferenceBackend: return img.resize(new_size, Image.Resampling.LANCZOS) return img - def _clean_generated_text(self, text: str) -> str: - """Strip leaked special tokens using the tokenizer's own token list.""" + def _generation_stop_token_ids(self, model, generation_kwargs: dict): + """Return the stop-token ids active for a ``generate`` call.""" + if "eos_token_id" in generation_kwargs: + return generation_kwargs.get("eos_token_id") + generation_config = getattr(model, "generation_config", None) + eos_token_id = getattr(generation_config, "eos_token_id", None) + if eos_token_id is not None: + return eos_token_id + config = getattr(model, "config", None) + return getattr(config, "eos_token_id", None) + + def _cancel_stopping_criteria(self, cancel_event): + """Build a Transformers stopping criteria list for user cancellation.""" + if cancel_event is None: + return None + from transformers.generation.stopping_criteria import ( + StoppingCriteria, + StoppingCriteriaList, + ) + + class _CancelCriteria(StoppingCriteria): + def __init__(self, ev): + self.ev = ev + + def __call__(self, input_ids, scores, **kwargs): + return self.ev.is_set() + + return StoppingCriteriaList([_CancelCriteria(cancel_event)]) + + def _clean_generated_text( + self, + text: str, + *, + stop_token_ids = None, + ) -> str: + """Strip leaked response-boundary tokens after streaming.""" if self._is_gpt_oss_model(): # HarmonyTextStreamer emits clean .... Strip any # harmony protocol tokens and other gpt-oss tokens (e.g. @@ -2118,10 +2337,28 @@ class InferenceBackend: return text.strip() tokenizer = self.models.get(self.active_model_name, {}).get("tokenizer") + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) if tokenizer: - for token in getattr(tokenizer, "all_special_tokens", []): - if token in text: - text = text.replace(token, "") + if stop_token_ids is None: + stop_token_ids = self.models.get(self.active_model_name, {}).get( + "chat_turn_end_eos_ids" + ) + if isinstance(stop_token_ids, int): + stop_token_ids = (stop_token_ids,) + for token_id in stop_token_ids or (): + try: + token = tokenizer.convert_ids_to_tokens(int(token_id)) + except Exception: + token = None + if isinstance(token, str) and token and text.endswith(token): + text = text[: -len(token)] + elif ( + isinstance(token, str) + and token + and text.endswith("") + and text[: -len("")].endswith(token) + ): + text = text[: -len("") - len(token)] + "" return text.strip() def _load_chat_template_info(self, model_name: str): diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e797cfe22a..fec5214cf9 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -9,6 +9,7 @@ OpenAI-compatible /v1/chat/completions endpoint. import atexit import contextlib +import functools import json import os import re @@ -82,6 +83,7 @@ from core.inference.tool_call_parser import ( ) from core.inference.tool_loop_controller import ( ToolLoopController, + append_deferred_nudges, tool_event_provenance, ) from state.tool_approvals import ( @@ -998,6 +1000,283 @@ def _cached_colocated_split_main( return None +def _cached_variant_resolution(repo_id: str, hf_variant: str) -> tuple[Optional[str], list[str]]: + """Find a cached main GGUF and its shards for a variant.""" + candidate = next(_cached_variant_candidates(repo_id, hf_variant), None) + if candidate is None: + return None, [] + _, main, shards, _ = candidate + return main, shards + + +def _cached_variant_candidates( + repo_id: str, + hf_variant: str, + *, + require_mmproj: bool = False, +) -> Generator[tuple[str, str, list[str], Path], None, None]: + """Yield complete cached variant copies in snapshot preference order.""" + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + cached_files = _gguf_snapshot_files(snap) + matches = _gguf_files_for_variant(cached_files, hf_variant) + if not matches: + continue + main = matches[0] + shards = _gguf_extra_shards(matches, main) + split = _SHARD_FULL_RE.match(main) + if split: + numbers = { + int(match.group(2)) + for path in [main, *shards] + if (match := _SHARD_FULL_RE.match(path)) + } + if numbers != set(range(1, int(split.group(3)) + 1)): + continue + main_path = snap.joinpath(*main.replace("\\", "/").split("/")) + if not main_path.is_file() or not _snapshot_has_all_shards( + str(main_path), main, shards, {} + ): + continue + if require_mmproj and not _pick_mmproj(cached_files): + continue + yield str(main_path), main, shards, snap + except Exception as e: + logger.debug(f"Cache lookup for variant failed: {e}") + + +def _cached_candidate_matches_revision_size( + repo_id: str, candidate: tuple[str, str, list[str], Path], hf_token: Optional[str] +) -> bool: + """Check cached byte sizes against the snapshot's own Hub revision. + + A snapshot pointer is normally published only after its blob is complete. + When the old revision is still queryable, also compare every weight file's + size so a manually truncated cache entry is not treated as reusable. If + metadata cannot be reached, retain the cache's normal offline semantics. + """ + main_path, main, shards, snap = candidate + paths = [main, *shards] + try: + from huggingface_hub import get_paths_info + infos = list( + get_paths_info( + repo_id, + paths, + revision = snap.name, + token = hf_token, + ) + ) + except Exception as e: + logger.debug( + "Could not size-check cached GGUF %s at revision %s: %s", + repo_id, + snap.name, + e, + ) + return True + + if not infos: + # The Hub answers an unknown (e.g. force-pushed away) revision with an + # empty result, not an error; treat it like unreachable metadata. + return True + expected_sizes = {info.path: info.size for info in infos if info.size is not None} + if any(path not in expected_sizes for path in paths): + return False + try: + if os.path.getsize(main_path) < expected_sizes[main]: + return False + except OSError: + return False + return _snapshot_has_all_shards(main_path, main, shards, expected_sizes) + + +def _cached_complete_candidate( + repo_id: str, gguf_filename: Optional[str], shards: list[str] +) -> Optional[tuple[str, str, list[str], Path]]: + """Return one complete exact-filename cache candidate with snapshot context.""" + if not gguf_filename: + return None + if shards: + main_path = _cached_colocated_split_main(repo_id, gguf_filename, shards, {}) + else: + m = _SHARD_FULL_RE.match(gguf_filename) + if m and int(m.group(3)) > 1: + return None + main_path = _cached_hf_snapshot_file(repo_id, gguf_filename) + if main_path is None: + return None + snap = _snapshot_dir_of(main_path) + if snap is None: + return None + return main_path, gguf_filename, shards, snap + + +def cached_gguf_for_load( + hf_repo: str, + hf_variant: Optional[str], + *, + require_mmproj: bool = False, + verify_sizes: bool = False, + hf_token: Optional[str] = None, +) -> Optional[str]: + """Return a cached GGUF that can be loaded without downloading.""" + if not hf_variant: + return None + hf_repo = _resolve_repo_id_casing(hf_repo) + for candidate in _cached_variant_candidates( + hf_repo, + hf_variant, + require_mmproj = require_mmproj, + ): + if verify_sizes and not _cached_candidate_matches_revision_size( + hf_repo, candidate, hf_token + ): + continue + return candidate[0] + return None + + +def _snapshot_dir_of(path: str) -> Optional[Path]: + """Return the HF cache snapshot containing path, if any.""" + try: + p = Path(os.path.abspath(path)) + except OSError: + return None + for ancestor in p.parents: + if ancestor.parent.name == "snapshots": + return ancestor + return None + + +def _companion_snapshot_sibling( + near_path: str, pick: Callable[[list[str]], Optional[str]] +) -> Optional[str]: + """Find a companion in the same snapshot as near_path.""" + snap = _snapshot_dir_of(near_path) + if snap is None: + return None + try: + sibling = pick(_gguf_snapshot_files(snap)) + except Exception: + return None + if not sibling: + return None + candidate = snap / sibling + return str(candidate) if candidate.is_file() else None + + +def _pick_mmproj(candidates: list[str]) -> Optional[str]: + mmproj_files = sorted( + f for f in candidates if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() + ) + if not mmproj_files: + return None + return next((f for f in mmproj_files if f.lower().endswith("-f16.gguf")), mmproj_files[0]) + + +def _hub_download_in_flight(hf_repo: str) -> bool: + try: + from hub.utils.download_registry import get_models_registry + return bool(get_models_registry().active_job_refs(hf_repo)) + except Exception: + return False + + +def _hub_download_blocks_gguf_load( + hf_repo: str, + hf_variant: Optional[str], + *, + require_mmproj: bool = False, + hf_token: Optional[str] = None, +) -> bool: + """Whether an active Hub job makes this GGUF load unsafe. + + Same-variant jobs can reclaim the stale snapshot a load would reuse, so + they always block. Other jobs block only when this load lacks a complete + cached copy and would write to the shared cache itself. + """ + try: + from hub.utils.download_registry import get_models_registry + + registry = get_models_registry() + if not registry.active_job_refs(hf_repo): + return False + if registry.has_active_variant(hf_repo, hf_variant): + return True + except Exception: + return False + return ( + cached_gguf_for_load( + hf_repo, + hf_variant, + require_mmproj = require_mmproj, + verify_sizes = True, + hf_token = hf_token, + ) + is None + ) + + +# Active GGUF loads by normalized repo ID. +_LOADS_IN_FLIGHT: dict[str, int] = {} +_LOADS_IN_FLIGHT_LOCK = threading.Lock() + + +@contextlib.contextmanager +def gguf_load_in_flight(hf_repo: Optional[str]): + """Track an HF GGUF load until the context exits.""" + key = (hf_repo or "").strip().lower() + if not key: + yield + return + with _LOADS_IN_FLIGHT_LOCK: + _LOADS_IN_FLIGHT[key] = _LOADS_IN_FLIGHT.get(key, 0) + 1 + try: + yield + finally: + with _LOADS_IN_FLIGHT_LOCK: + remaining = _LOADS_IN_FLIGHT.get(key, 1) - 1 + if remaining <= 0: + _LOADS_IN_FLIGHT.pop(key, None) + else: + _LOADS_IN_FLIGHT[key] = remaining + + +def hf_gguf_load_in_flight(hf_repo: str) -> bool: + """Return whether a GGUF load is active for hf_repo.""" + key = (hf_repo or "").strip().lower() + if not key: + return False + with _LOADS_IN_FLIGHT_LOCK: + return _LOADS_IN_FLIGHT.get(key, 0) > 0 + + +def _with_gguf_load_marker(load: Callable): + """Keep an HF repo marked for the full synchronous load call.""" + + @functools.wraps(load) + def wrapped(self, *args, **kwargs): + hf_repo = kwargs.get("hf_repo") + with gguf_load_in_flight(hf_repo): + if hf_repo and _hub_download_blocks_gguf_load( + hf_repo, + kwargs.get("hf_variant"), + require_mmproj = bool( + kwargs.get("is_vision") + and not extra_args_disable_mmproj(kwargs.get("extra_args")) + ), + hf_token = kwargs.get("hf_token"), + ): + raise RuntimeError( + f"'{hf_repo}' is currently being downloaded by the download manager" + ) + return load(self, *args, **kwargs) + + return wrapped + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -1541,6 +1820,31 @@ def _is_external_link(path: Path) -> bool: return False +# Inkling's template takes a numeric thinking-effort dial (0..0.99) and its +# float() coercion turns unrecognized named levels into 0, i.e. no thinking. +# Map OpenAI-style names to the values the model was trained on. Module-level +# so duck-typed engine stand-ins in tests do not need the attribute. +_INKLING_REASONING_EFFORT = { + "none": 0.0, + "minimal": 0.1, + "low": 0.2, + "medium": 0.7, + "high": 0.9, + "xhigh": 0.99, + "max": 0.99, +} + + +def _coerce_reasoning_effort(architecture, kwargs: dict) -> dict: + if architecture == "inkling": + effort = kwargs.get("reasoning_effort") + if isinstance(effort, str): + mapped = _INKLING_REASONING_EFFORT.get(effort.strip().lower()) + if mapped is not None: + kwargs["reasoning_effort"] = mapped + return kwargs + + class LlamaCppBackend: """Manages a llama-server subprocess for GGUF model inference. @@ -1941,36 +2245,15 @@ class LlamaCppBackend: def reasoning_default(self) -> bool: return self._reasoning_default - # Inkling's template takes a numeric thinking-effort dial (0..0.99) and its - # float() coercion turns unrecognized named levels into 0, i.e. no thinking. - # Map OpenAI-style names to the values the model was trained on. - _INKLING_REASONING_EFFORT = { - "none": 0.0, - "minimal": 0.2, - "low": 0.2, - "medium": 0.7, - "high": 0.9, - "xhigh": 0.99, - "max": 0.99, - } - - def _coerce_reasoning_effort(self, kwargs: dict) -> dict: - if getattr(self, "_architecture", None) == "inkling": - effort = kwargs.get("reasoning_effort") - if isinstance(effort, str): - mapped = self._INKLING_REASONING_EFFORT.get(effort.strip().lower()) - if mapped is not None: - kwargs["reasoning_effort"] = mapped - return kwargs - def _reasoning_kwargs(self, enable_thinking: bool) -> dict: if self._reasoning_style == "enable_thinking_effort": # GLM-5.2-style: enable_thinking is the on/off gate; when on, leave # the template's default effort (max) in place. return {"enable_thinking": enable_thinking} if self._reasoning_style == "reasoning_effort": - return self._coerce_reasoning_effort( - {"reasoning_effort": "high" if enable_thinking else "low"} + return _coerce_reasoning_effort( + getattr(self, "_architecture", None), + {"reasoning_effort": "high" if enable_thinking else "low"}, ) return {"enable_thinking": enable_thinking} @@ -2019,7 +2302,7 @@ class LlamaCppBackend: kwargs["enable_thinking"] = enable_thinking if self._supports_preserve_thinking and preserve_thinking is not None: kwargs["preserve_thinking"] = preserve_thinking - self._coerce_reasoning_effort(kwargs) + _coerce_reasoning_effort(getattr(self, "_architecture", None), kwargs) return kwargs or None @property @@ -3801,13 +4084,13 @@ class LlamaCppBackend: hf_repo: str, free_bytes: int, hf_token: Optional[str] = None, - ) -> Optional[tuple[str, int]]: + ) -> Optional[tuple[str, int, list[str]]]: """Find the smallest GGUF variant (including all shards) that fits. Groups split shards by variant prefix and sums their sizes (e.g. UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total). - Returns (first_shard_filename, total_size_bytes) or None. + Returns (first_shard_filename, total_size_bytes, extra_shards) or None. """ try: from huggingface_hub import get_paths_info, list_repo_files @@ -3843,9 +4126,13 @@ class LlamaCppBackend: # Smallest that fits variant_sizes.sort(key = lambda x: x[1]) - for first_file, total_size, _ in variant_sizes: + for first_file, total_size, shard_files in variant_sizes: if total_size > 0 and total_size <= free_bytes: - return first_file, total_size + return ( + first_file, + total_size, + [path for path in sorted(shard_files) if path != first_file], + ) return None except Exception: @@ -4441,42 +4728,52 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"Could not list repo files: {e}") - # Offline: resolve variant -> filename from the local HF cache. - # The heuristic below assumes filenames echo the repo name, which - # breaks for e.g. Qwen3.6-27B-MTP-GGUF (no "MTP" in file). Match - # against the rel path (not just basename) so subdir layouts like - # ``BF16/foo.gguf`` are findable. + # Fall back to the local cache when the repo listing is unavailable. if not gguf_filename: - try: - from utils.models.model_config import _iter_hf_cache_snapshots - for snap in _iter_hf_cache_snapshots(hf_repo): - cached_files = _gguf_snapshot_files(snap) - matches = _gguf_files_for_variant(cached_files, hf_variant) - if not matches: - continue - gguf_filename = matches[0] - gguf_extra_shards = _gguf_extra_shards(matches, gguf_filename) - logger.info( - "Resolved variant %s -> %s from local HF cache", - hf_variant, - gguf_filename, - ) - break - except Exception as e: - logger.debug(f"Offline cache lookup for variant failed: {e}") + cached_name, cached_shards = _cached_variant_resolution(hf_repo, hf_variant) + if cached_name: + gguf_filename = cached_name + gguf_extra_shards = cached_shards + logger.info( + "Resolved variant %s -> %s from local HF cache", + hf_variant, + gguf_filename, + ) if not gguf_filename: repo_name = hf_repo.split("/")[-1].replace("-GGUF", "") gguf_filename = f"{repo_name}-{hf_variant}.gguf" + # Prefer the existing model. Updates use force=True to fetch a new revision. + if not force: + if hf_variant: + # Resolve by variant so a newer revision's filename does not hide + # the complete older copy. Size-check against that older snapshot's + # own revision when its metadata remains available. + cached_main = cached_gguf_for_load( + hf_repo, + hf_variant, + verify_sizes = True, + hf_token = hf_token, + ) + else: + candidate = _cached_complete_candidate(hf_repo, gguf_filename, gguf_extra_shards) + cached_main = ( + candidate[0] + if candidate is not None + and _cached_candidate_matches_revision_size(hf_repo, candidate, hf_token) + else None + ) + if cached_main is not None: + logger.info(f"Reusing cached GGUF: {cached_main}") + return cached_main + # Check disk space; fall back to a smaller variant if needed all_gguf_files = [gguf_filename] + gguf_extra_shards - expected_sizes: dict[str, int] = {} try: from huggingface_hub import get_paths_info, try_to_load_from_cache path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token)) - expected_sizes = {p.path: p.size for p in path_infos if p.size} total_bytes = sum((p.size or 0) for p in path_infos) # Subtract bytes already in the HF cache so we only preflight @@ -4485,25 +4782,10 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - # Cross-snapshot / case-variant cache reuse is offline-only (see the download - # path below); online, hf_hub_download fetches the current revision and - # resumes partials, so an old snapshot must not be counted as cached here or - # the preflight would under-count the download and skip the disk fallback. + # Count only files that can resume this download. offline = _hf_env_offline() - # A split GGUF whose shards are not co-located in a single snapshot is - # refetched as a whole set later, so it must not be counted as cached here. - split_needs_refetch = False - if offline and not force and gguf_extra_shards: - # Scan all snapshots for one that holds the whole set co-located, so a - # newer snapshot with only the first shard does not mask an older - # complete one and needlessly trip the disk fallback. - if ( - _cached_colocated_split_main( - hf_repo, gguf_filename, gguf_extra_shards, expected_sizes - ) - is None - ): - split_needs_refetch = True + # Offline split sets are reusable only when every shard shares a snapshot. + split_needs_refetch = bool(offline and not force and gguf_extra_shards) if not force and not split_needs_refetch: for p in path_infos: if not p.size: @@ -4564,32 +4846,26 @@ class LlamaCppBackend: hf_token, ) if smaller: - fallback_file, fallback_size = smaller + fallback_file, fallback_size, fallback_shards = smaller logger.info( f"Selected variant too large ({total_gb:.1f} GB), " f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)" ) gguf_filename = fallback_file - _m = _SHARD_RE.match(gguf_filename) - _prefix = _m.group(1) if _m else None - if _prefix: - prefix_lower = _prefix.lower() - gguf_extra_shards = sorted( - f - for f in all_gguf_files - if f.lower().startswith(prefix_lower) - and f != gguf_filename - and not _is_companion_gguf_path(f) + gguf_extra_shards = fallback_shards + + # The selected fallback is a new load target. Apply the + # same any-revision reuse policy before starting a fetch. + fallback_candidate = _cached_complete_candidate( + hf_repo, gguf_filename, gguf_extra_shards + ) + if fallback_candidate is not None and ( + _cached_candidate_matches_revision_size( + hf_repo, fallback_candidate, hf_token ) - else: - gguf_extra_shards = [] - # Record the fallback's size so the later cache-reuse probe can - # size-verify it; only for a single-file fallback, since - # _find_smallest_fitting_variant returns the whole-variant size - # and using that as the first shard's expected size would reject - # a valid cached first shard of a split fallback. - if not gguf_extra_shards: - expected_sizes[fallback_file] = fallback_size + ): + logger.info(f"Reusing cached fallback GGUF: {fallback_candidate[0]}") + return fallback_candidate[0] else: raise RuntimeError( f"Not enough disk space to download any variant. " @@ -4609,45 +4885,25 @@ class LlamaCppBackend: raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. - local_path = None - # Reuse a cached copy from another snapshot / case-variant repo dir only when - # offline. Online, fall through to hf_hub_download so its revision/etag check - # fetches the current file (and resumes a partial) instead of serving a stale - # same-name blob from an older revision. - if not force and _hf_env_offline(): - if gguf_extra_shards: - # A split GGUF must load every shard from one snapshot; reuse only a - # snapshot that holds the whole set co-located, scanning past a newer - # snapshot that has just the first shard while an older one is complete. - local_path = _cached_colocated_split_main( - hf_repo, gguf_filename, gguf_extra_shards, expected_sizes - ) - else: - local_path = _cached_hf_snapshot_file( - hf_repo, - gguf_filename, - expected_size = expected_sizes.get(gguf_filename), - ) - if local_path is None: - local_path = hf_hub_download_with_xet_fallback( + local_path = hf_hub_download_with_xet_fallback( + hf_repo, + gguf_filename, + hf_token, + cancel_event = cancel_event, + on_status = lambda m: logger.info(m), + force_download = force, + ) + for shard in gguf_extra_shards: + if cancel_event.is_set(): + raise RuntimeError("Cancelled") + logger.info(f"Resolving GGUF shard: {shard}") + hf_hub_download_with_xet_fallback( hf_repo, - gguf_filename, + shard, hf_token, cancel_event = cancel_event, - on_status = lambda m: logger.info(m), force_download = force, ) - for shard in gguf_extra_shards: - if cancel_event.is_set(): - raise RuntimeError("Cancelled") - logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download_with_xet_fallback( - hf_repo, - shard, - hf_token, - cancel_event = cancel_event, - force_download = force, - ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): raise @@ -4670,10 +4926,12 @@ class LlamaCppBackend: pick: Callable[[list[str]], Optional[str]], label: str, cancel_event: Optional[threading.Event] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name. - Tries the live repo file list, then the local HF cache snapshots + Prefers a companion co-located with ``near_path``'s cache snapshot, + then tries the live repo file list, then the local HF cache snapshots (offline, same fallback as _download_gguf), then hf_hub_download. Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so an /unload between the main download and here skips the fetch. @@ -4683,6 +4941,17 @@ class LlamaCppBackend: if cancel_event.is_set(): return None + # Keep companion files in the main GGUF's snapshot. + if near_path: + cached = _companion_snapshot_sibling(near_path, pick) + if cached: + logger.info("Reusing cached %s: %s", label, cached) + return cached + + if _hub_download_in_flight(hf_repo): + logger.info("Skipping %s download while a hub download is active", label) + return None + target: Optional[str] = None from huggingface_hub import list_repo_files @@ -4755,33 +5024,23 @@ class LlamaCppBackend: hf_repo: str, hf_token: Optional[str] = None, cancel_event: Optional[threading.Event] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Download the mmproj (vision projection) file from a GGUF repo. Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local path, or None if none exists. ``cancel_event`` overrides - ``self._cancel_event`` (defaults to it). + ``self._cancel_event`` (defaults to it). ``near_path`` prefers a + copy co-located with the main GGUF's cache snapshot. """ - def _pick_mmproj(candidates: list[str]) -> Optional[str]: - mmproj_files = sorted( - f - for f in candidates - if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() - ) - if not mmproj_files: - return None - for f in mmproj_files: - if f.lower().endswith("-f16.gguf"): - return f - return mmproj_files[0] - return self._download_companion_gguf( hf_repo = hf_repo, hf_token = hf_token, pick = _pick_mmproj, label = "mmproj", cancel_event = cancel_event, + near_path = near_path, ) def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]: @@ -4812,6 +5071,7 @@ class LlamaCppBackend: *, hf_repo: str, hf_token: Optional[str] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Download the separate MTP drafter (speculative head) from a GGUF repo. @@ -4823,16 +5083,6 @@ class LlamaCppBackend: are intentionally skipped. Returns the local path, or None. """ - # Offline, reuse any drafter already on disk (a fresh copy can't be - # fetched). Online, _download_companion_gguf/hf_hub_download reuse the - # current cached file and refetch a changed one, so skip the probe here - # rather than pair new weights with a stale draft. - if _hf_env_offline(): - cached = self._cached_repo_mtp_drafter(hf_repo) - if cached: - logger.info(f"Reusing cached MTP drafter (offline): {cached}") - return cached - def _pick_mtp(candidates: list[str]) -> Optional[str]: # Root-level only: MTP/ subdir copies now share the mtp- prefix but # are explicit-selection, not auto-fetch (they'd sort ahead of root). @@ -4845,11 +5095,28 @@ class LlamaCppBackend: ) return mtp_files[0] if mtp_files else None + if near_path: + cached = _companion_snapshot_sibling(near_path, _pick_mtp) + if cached: + logger.info("Reusing cached MTP drafter: %s", cached) + return cached + + # Offline, reuse any drafter already on disk (a fresh copy can't be + # fetched). Online, _download_companion_gguf/hf_hub_download reuse the + # current cached file and refetch a changed one, so skip the probe here + # rather than pair new weights with a stale draft. + if _hf_env_offline(): + cached = self._cached_repo_mtp_drafter(hf_repo) + if cached: + logger.info(f"Reusing cached MTP drafter (offline): {cached}") + return cached + return self._download_companion_gguf( hf_repo = hf_repo, hf_token = hf_token, pick = _pick_mtp, label = "MTP drafter", + near_path = near_path, ) def _resolve_launch_mmproj_path( @@ -5431,6 +5698,7 @@ class LlamaCppBackend: ) self._stdout_thread.start() + @_with_gguf_load_marker def load_model( self, *, @@ -5576,6 +5844,7 @@ class LlamaCppBackend: mmproj_path = self._download_mmproj( hf_repo = hf_repo, hf_token = hf_token, + near_path = model_path, ) # Auto-download the separate MTP drafter (e.g. Gemma) when # the requested spec mode can use it. Repos with the head @@ -5593,6 +5862,7 @@ class LlamaCppBackend: mtp_draft_path = self._download_mtp( hf_repo = hf_repo, hf_token = hf_token, + near_path = model_path, ) elif gguf_path: if not Path(gguf_path).is_file(): @@ -8227,6 +8497,11 @@ class LlamaCppBackend: if not is_ours: continue + # A live parent means a running Studio (or the user's + # shell) still owns it -- not an orphan. + if LlamaCppBackend._pid_parent_is_alive(proc.info["pid"]): + continue + proc.kill() killed += 1 logger.info( @@ -8279,6 +8554,9 @@ class LlamaCppBackend: if not owned: continue + if LlamaCppBackend._pid_parent_is_alive(pid): + continue + try: os.kill(pid, signal.SIGKILL) killed += 1 @@ -10107,6 +10385,9 @@ class LlamaCppBackend: assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False + # Collect no-op nudges and flush them after the batch, so a no-op + # doesn't abort it and drop the parallel calls that follow. + deferred_noop_msgs: list = [] # The text-path provisional card uses the parser's default id ("call_0"); # a Mistral-style call carries its own id and would open a duplicate. Reuse @@ -10149,14 +10430,14 @@ class LlamaCppBackend: "provenance": decision.provenance, } completion = tool_controller.record_noop(decision) - conversation.append(completion.model_message()) + deferred_noop_msgs.append(completion.model_message()) if _forced_tool_call_pending: _forced_tool_call_pending = False logger.info( "Suppressed local GGUF tool call as internal no-op: " f"action={decision.action} tool={decision.tool_name}" ) - break + continue if not assistant_appended: assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()] @@ -10275,6 +10556,8 @@ class LlamaCppBackend: if _forced_tool_call_pending: _forced_tool_call_pending = False + append_deferred_nudges(conversation, deferred_noop_msgs) + # Close provisional cards not resolved by execution/no-op handling. for _pid, _pname in provisional_started_tool_calls.items(): if _pid not in resolved_provisional_tool_call_ids: diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 163ade10c4..e78c93b6f3 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -8,14 +8,76 @@ instead of torch/transformers for model loading and generation. import json import os import threading +from contextlib import contextmanager from typing import Optional, Generator from core.inference.message_content import content_to_text from core.inference.runtime_context import runtime_context_length +from core.inference.chat_template_helpers import ( + ReasoningChannelNormalizer, + normalize_reasoning_snapshots, +) from loggers import get_logger logger = get_logger(__name__) +def _mlx_adapter_modules(model): + """Return bypassable adapter entries and unsupported wrapper paths.""" + adapters = [] + unsupported = [] + for path, module in model.named_modules(): + if not path or not (hasattr(module, "lora_a") and hasattr(module, "lora_b")): + continue + base = getattr(module, "linear", None) + if base is None: + base = getattr(module, "embedding", None) + if base is None: + unsupported.append(path) + else: + adapters.append((path, module, base)) + return adapters, unsupported + + +@contextmanager +def _temporary_mlx_adapter_state(model, use_adapter): + """Select base or adapter modules for one request, then restore the tree.""" + if use_adapter is None: + yield + return + if isinstance(use_adapter, str): + raise NotImplementedError( + "Unsloth MLX: named adapter selection is not supported; use True for " + "the loaded adapter or False for the base model." + ) + if use_adapter is not True and use_adapter is not False: + raise TypeError("Unsloth MLX: use_adapter must be None, True, False, or a string.") + + adapters, unsupported = _mlx_adapter_modules(model) + if use_adapter is True: + if not adapters and not unsupported: + logger.warning("MLX adapter requested, but the active model has no adapter layers") + yield + return + if unsupported: + raise RuntimeError( + "Unsloth MLX: cannot disable adapter layers without their base modules: " + + ", ".join(unsupported[:5]) + ) + if not adapters: + yield + return + + from mlx.utils import tree_unflatten + + base_modules = tree_unflatten([(path, base) for path, _, base in adapters]) + adapter_modules = tree_unflatten([(path, wrapper) for path, wrapper, _ in adapters]) + try: + model.update_modules(base_modules) + yield + finally: + model.update_modules(adapter_modules) + + def _mlx_vlm_model_config(model): """Return the loaded MLX model config and its type, preferring whichever of config / _config actually carries a model_type.""" @@ -504,6 +566,7 @@ class MLXInferenceBackend: reasoning_effort = None, preserve_thinking = None, presence_penalty = 0.0, + _adapter_state = None, ) -> Generator[str, None, None]: if self._model is None: raise RuntimeError("No model loaded") @@ -533,7 +596,7 @@ class MLXInferenceBackend: break if self._is_vlm: - yield from self._generate_vlm( + stream = self._generate_vlm( full_messages, image, temperature, @@ -548,9 +611,10 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, presence_penalty = presence_penalty, + _adapter_state = _adapter_state, ) else: - yield from self._generate_text( + stream = self._generate_text( full_messages, temperature, top_p, @@ -564,7 +628,9 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, presence_penalty = presence_penalty, + _adapter_state = _adapter_state, ) + yield from stream def _generate_text( self, @@ -582,6 +648,7 @@ class MLXInferenceBackend: reasoning_effort = None, preserve_thinking = None, presence_penalty = 0.0, + _adapter_state = None, ): from mlx_lm import stream_generate from mlx_lm.sample_utils import make_sampler, make_logits_processors @@ -609,7 +676,7 @@ class MLXInferenceBackend: # probe and native render share a renderer. (VLM renders via the # processor for image tokens and is not wired here.) model_info = self.models.get(self.active_model_name, {}) - prompt = render_with_native_template_fallback( + render_result = render_with_native_template_fallback( formatted_prompt = prompt, tokenizer = self._tokenizer, model_info = model_info, @@ -620,17 +687,16 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, hf_token = model_info.get("hf_token"), + return_metadata = True, ) + prompt = render_result.prompt + reasoning_channel_markers = render_result.reasoning_channel_markers # An open prefilled by the template lives in the prompt, not # the generated tokens; re-emit it so the frontend renders the block. think_prefix = detect_think_prefill( prompt, getattr(self._tokenizer, "all_special_tokens", None) ) - # Emit it before the first token so the block renders during prefill. - if think_prefix: - yield think_prefix - sampler = make_sampler( temp = temperature, top_p = top_p, @@ -654,7 +720,17 @@ class MLXInferenceBackend: if not logits_processors: logits_processors = None + preserve_native_channels = reasoning_channel_markers is not None token_ids = [] + normalizer = ( + ReasoningChannelNormalizer(*reasoning_channel_markers) + if reasoning_channel_markers is not None + else None + ) + # MLX consumers diff cumulative snapshots. Keep a prompt-prefilled + # prefix on every native-protocol snapshot just as the normal + # decoding path does below. + normalized_output = think_prefix logger.info( "Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s", len(prompt), @@ -662,9 +738,12 @@ class MLXInferenceBackend: type(self._model).__name__, type(self._tokenizer).__name__, ) - with self._generation_lock: + with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): final_response = None try: + # Enter request-scoped model state before yielding any response. + if think_prefix: + yield think_prefix gen_kwargs = dict( prompt = prompt, max_tokens = max_new_tokens, @@ -678,12 +757,19 @@ class MLXInferenceBackend: **gen_kwargs, ): final_response = response - token_ids.append(response.token) - cumulative = self._tokenizer.decode( - token_ids, - skip_special_tokens = True, - ) - yield think_prefix + cumulative + if preserve_native_channels: + piece = getattr(response, "text", None) or "" + delta = normalizer.feed(piece) + if delta: + normalized_output += delta + yield normalized_output + else: + token_ids.append(response.token) + cumulative = self._tokenizer.decode( + token_ids, + skip_special_tokens = True, + ) + yield think_prefix + cumulative if cancel_event and cancel_event.is_set(): break @@ -700,6 +786,12 @@ class MLXInferenceBackend: getattr(final_response, "generation_tokens", 0), getattr(final_response, "generation_tps", 0.0), ) + if normalizer is not None: + cancelled = cancel_event is not None and cancel_event.is_set() + tail = normalizer.drain() if cancelled else normalizer.finish() + if tail: + normalized_output += tail + yield normalized_output def _generate_vlm( self, @@ -718,6 +810,7 @@ class MLXInferenceBackend: reasoning_effort = None, preserve_thinking = None, presence_penalty = 0.0, + _adapter_state = None, ): from mlx_vlm import stream_generate as vlm_stream @@ -821,9 +914,6 @@ class MLXInferenceBackend: # Re-emit an open prefill from the prompt (see _generate_text). cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None)) - # Emit it before the first token so the block renders during prefill. - if cumulative: - yield cumulative logger.info( "VLM generating: prompt_len=%d, has_image=%s", len(prompt), @@ -858,31 +948,46 @@ class MLXInferenceBackend: elif _rep_active: vlm_kwargs["repetition_penalty"] = float(repetition_penalty) - with self._generation_lock: - final_response = None - try: - for response in vlm_stream( - self._model, - self._processor, - prompt, - images, - **vlm_kwargs, - ): - final_response = response - token_text = response.text if hasattr(response, "text") else str(response) - cumulative += token_text - yield cumulative - if cancel_event and cancel_event.is_set(): - break - finally: - # mlx_vlm exposes the same stats fields as mlx_lm. - if final_response is not None: - self.last_generation_stats = _build_generation_stats( - getattr(final_response, "prompt_tokens", 0), - getattr(final_response, "prompt_tps", 0.0), - getattr(final_response, "generation_tokens", 0), - getattr(final_response, "generation_tps", 0.0), - ) + def _stream_vlm_snapshots(): + nonlocal cumulative + # Hold the generation lock AND the request-scoped adapter state for the + # whole stream so Base-vs-LoRA compare mode honors use_adapter and the + # wrapper tree is restored on completion, cancellation, or close. + with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): + final_response = None + try: + # Emit any prefilled block before the first token so the + # UI renders it during prefill, matching _generate_text. Done + # inside the adapter context so an unsupported request raises + # before any output escapes. + if cumulative: + yield cumulative + for response in vlm_stream( + self._model, + self._processor, + prompt, + images, + **vlm_kwargs, + ): + final_response = response + token_text = response.text if hasattr(response, "text") else str(response) + cumulative += token_text + yield cumulative + if cancel_event and cancel_event.is_set(): + break + finally: + # mlx_vlm exposes the same stats fields as mlx_lm. + if final_response is not None: + self.last_generation_stats = _build_generation_stats( + getattr(final_response, "prompt_tokens", 0), + getattr(final_response, "prompt_tps", 0.0), + getattr(final_response, "generation_tokens", 0), + getattr(final_response, "generation_tps", 0.0), + ) + + yield from normalize_reasoning_snapshots( + _stream_vlm_snapshots(), chat_target, cancel_event, tools = tools + ) def generate_with_adapter_control( self, @@ -890,8 +995,11 @@ class MLXInferenceBackend: cancel_event = None, **gen_kwargs, ) -> Generator[str, None, None]: - # MLX LoRA adapter toggling not yet supported; generate normally - yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) + yield from self.generate_chat_response( + cancel_event = cancel_event, + _adapter_state = use_adapter, + **gen_kwargs, + ) def reset_generation_state(self): import mlx.core as mx diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index c2082bc198..eaa474d9b8 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -59,7 +59,32 @@ class GenStreamError(str): "Error:" by checking isinstance(chunk, GenStreamError). """ - __slots__ = () + __slots__ = ("public",) + + def __new__( + cls, + value, + *, + public: bool = False, + ): + obj = str.__new__(cls, value) + obj.public = bool(public) + return obj + + +class GenStreamErrorRaised(RuntimeError): + """Internal exception form of ``GenStreamError`` for generator boundaries.""" + + __slots__ = ("public",) + + def __init__( + self, + value, + *, + public: bool = False, + ): + super().__init__(value) + self.public = bool(public) class InferenceOrchestrator: @@ -531,13 +556,19 @@ class InferenceOrchestrator: initial_resp_queue = self._resp_queue while True: if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue: - yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") + yield GenStreamError( + f"Error: {self._subprocess_crash_message(crash_context)}", + public = True, + ) return resp = read_one(read_timeout) if resp is None: # Check subprocess health if not self._ensure_subprocess_alive(): - yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") + yield GenStreamError( + f"Error: {self._subprocess_crash_message(crash_context)}", + public = True, + ) return continue @@ -689,11 +720,11 @@ class InferenceOrchestrator: GPU work stays serialized; this only avoids orchestrator lock contention. """ if not self._ensure_subprocess_alive(): - yield GenStreamError("Error: Inference subprocess is not running") + yield GenStreamError("Error: Inference subprocess is not running", public = True) return if not self.active_model_name: - yield GenStreamError("Error: No active model") + yield GenStreamError("Error: No active model", public = True) return # Latch the target model so the recheck below can detect a switch that completed # between _start_dispatcher and mailbox registration (mirrors the locked path's @@ -704,7 +735,7 @@ class InferenceOrchestrator: # so without this early-out a compare request would enqueue a generate on the # outgoing model and delay the switch. if self._unload_pending: - yield GenStreamError("Error: model is being unloaded") + yield GenStreamError("Error: model is being unloaded", public = True) return # Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under @@ -776,7 +807,7 @@ class InferenceOrchestrator: # _stop_dispatcher joins the dispatcher, which itself takes that lock. if orphaned_dispatcher: self._stop_dispatcher() - yield GenStreamError("Error: model is being unloaded") + yield GenStreamError("Error: model is being unloaded", public = True) return try: @@ -1376,6 +1407,7 @@ class InferenceOrchestrator: use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, presence_penalty: float = 0.0, + reasoning_prefilled: bool = False, **_unused, ): """Run the safetensors agentic tool loop in the parent process, @@ -1414,12 +1446,27 @@ class InferenceOrchestrator: presence_penalty = presence_penalty, ) if use_adapter is not None: - yield from self.generate_with_adapter_control( + stream = self.generate_with_adapter_control( use_adapter = use_adapter, **common_kwargs, ) else: - yield from self.generate_chat_response(**common_kwargs) + stream = self.generate_chat_response(**common_kwargs) + close_stream = False + try: + for chunk in stream: + if isinstance(chunk, GenStreamError): + close_stream = True + raise GenStreamErrorRaised(str(chunk), public = chunk.public) + yield chunk + finally: + if close_stream: + close = getattr(stream, "close", None) + if callable(close): + try: + close() + except Exception: + logger.debug("failed to close errored generation stream", exc_info = True) initial = list(messages) if system_prompt: @@ -1441,6 +1488,7 @@ class InferenceOrchestrator: confirm_tool_calls = confirm_tool_calls, bypass_permissions = bypass_permissions, permission_mode = permission_mode, + reasoning_prefilled = reasoning_prefilled, ) def generate_with_adapter_control( @@ -1454,14 +1502,27 @@ class InferenceOrchestrator: Uses the dispatcher path (no _gen_lock) so compare-mode requests don't block each other; the subprocess serializes them via its - sequential command loop. + sequential command loop. Backend failures raise instead of becoming + assistant text. """ - yield from self._generate_dispatched( + stream = self._generate_dispatched( use_adapter = use_adapter, cancel_event = cancel_event, stats_holder = stats_holder, **gen_kwargs, ) + try: + for chunk in stream: + if isinstance(chunk, GenStreamError): + # Preserve the public/operational flag so the route can surface + # the real message (e.g. "model is being unloaded") instead of a + # generic error. Mirrors the safetensors tool loop's _single_turn. + raise GenStreamErrorRaised(str(chunk), public = chunk.public) + yield chunk + finally: + close = getattr(stream, "close", None) + if callable(close): + close() def _generate_inner( self, @@ -1489,11 +1550,11 @@ class InferenceOrchestrator: readers don't consume each other's tokens off the shared resp_queue. """ if not self._ensure_subprocess_alive(): - yield GenStreamError("Error: Inference subprocess is not running") + yield GenStreamError("Error: Inference subprocess is not running", public = True) return if not self.active_model_name: - yield GenStreamError("Error: No active model") + yield GenStreamError("Error: No active model", public = True) return expected_model = self.active_model_name @@ -1510,7 +1571,7 @@ class InferenceOrchestrator: # so we never generate on the wrong one. if self._unload_pending or self.active_model_name != expected_model: # Won the lock handoff during a switch; don't start on the outgoing model. - yield GenStreamError("Error: model is being unloaded") + yield GenStreamError("Error: model is being unloaded", public = True) return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None @@ -1695,10 +1756,10 @@ class InferenceOrchestrator: ) -> Generator[str, None, None]: """Shared inner logic for audio input generation (Whisper + ASR).""" if not self._ensure_subprocess_alive(): - yield GenStreamError("Error: Inference subprocess is not running") + yield GenStreamError("Error: Inference subprocess is not running", public = True) return if not self.active_model_name: - yield GenStreamError("Error: No active model") + yield GenStreamError("Error: No active model", public = True) return expected_model = self.active_model_name @@ -1707,7 +1768,7 @@ class InferenceOrchestrator: # cleared or swapped the model while we waited. if self._unload_pending or self.active_model_name != expected_model: # Won the lock handoff during a switch; don't start on the outgoing model. - yield GenStreamError("Error: model is being unloaded") + yield GenStreamError("Error: model is being unloaded", public = True) return request_id = str(uuid.uuid4()) diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index 5b72373c03..d3bffc2f3d 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -276,8 +276,9 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { "auth_header": "Authorization", "auth_prefix": "Bearer ", "notes": ( - "Local Ollama server. OpenAI-compatible /v1/chat/completions; " - "no API key. Surfaced via CUSTOM_PROVIDER_PRESETS in the frontend." + "Ollama server (local or cloud). OpenAI-compatible " + "/v1/chat/completions; API key optional (required by Ollama " + "cloud). Surfaced via CUSTOM_PROVIDER_PRESETS in the frontend." ), "hidden": True, }, diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index f4c243d1bf..9110315815 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -50,6 +50,7 @@ from core.inference.tool_call_parser import ( # pattern lists, so the safetensors streaming strip stays aligned with the parser. from core.tool_healing import ( _REHEARSAL_TAIL_STRIP_RE, + _THINK_CLOSE_RE, _strip_bracket_tag_calls, _think_spans_outside_tool_markup, apply_tool_strip_patterns, @@ -57,6 +58,7 @@ from core.tool_healing import ( ) from core.inference.tool_loop_controller import ( ToolLoopController, + append_deferred_nudges, coerce_tool_arguments, status_for_tool, tool_event_provenance, @@ -303,6 +305,45 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str: return status_for_tool(tool_name, arguments) +def _reprompt_intent_text(text: str, *, reasoning_prefilled: bool = False) -> str: + """Return visible answer text for the plan-without-action classifier. + + Safetensors reasoning shares the cumulative text channel with the answer. + Forward-looking phrases inside ```` / ``[THINK]`` are private + planning, not a user-visible promise to call a tool. Match GGUF's behavior: + classify visible content when present and fall back to reasoning only for a + reasoning-only stall. + """ + prefilled_reasoning = "" + if reasoning_prefilled: + close = _THINK_CLOSE_RE.search(text) + if close is None: + return text.strip() + prefilled_reasoning = text[: close.end()].strip() + text = text[close.end() :].strip() + if not text: + return prefilled_reasoning + + spans = _think_spans_outside_tool_markup(text) + if not spans: + return text.strip() + + visible: list[str] = [] + reasoning: list[str] = [] + cursor = 0 + for start, end in spans: + visible.append(text[cursor:start]) + reasoning.append(text[start:end]) + cursor = end + visible.append(text[cursor:]) + + visible_text = "".join(visible).strip() + reasoning_text = "".join(reasoning).strip() + if visible_text: + return visible_text + return "\n".join(part for part in (prefilled_reasoning, reasoning_text) if part).strip() + + def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool: """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" probe = strip_llama3_leading_sentinels(text.lstrip()) @@ -447,6 +488,7 @@ def run_safetensors_tool_loop( confirm_tool_calls: bool = False, bypass_permissions: bool = False, permission_mode: Optional[str] = None, + reasoning_prefilled: bool = False, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -955,7 +997,10 @@ def run_safetensors_tool_loop( # (GGUF loop parity). The retry is gated on nudge_tool_calls so # Studio callers (which send True) always nudge, while API callers # who omit the flag keep today's no-reprompt behavior (opt-in). - stripped_answer = content_accum.strip() + intent_text = _reprompt_intent_text( + content_accum, + reasoning_prefilled = reasoning_prefilled, + ) if ( auto_heal_tool_calls and nudge_tool_calls @@ -964,7 +1009,7 @@ def run_safetensors_tool_loop( and not rag_autoinjected and not tool_denied and not any(record.executed for record in tool_controller.history) - and is_short_intent_without_action(stripped_answer) + and is_short_intent_without_action(intent_text) ): reprompt_count += 1 logger.info( @@ -972,9 +1017,9 @@ def run_safetensors_tool_loop( "calling tools (%d chars)", reprompt_count, MAX_ACT_REPROMPTS, - len(stripped_answer), + len(intent_text), ) - conversation.append({"role": "assistant", "content": stripped_answer}) + conversation.append({"role": "assistant", "content": intent_text}) tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool" conversation.append( { @@ -1099,6 +1144,9 @@ def run_safetensors_tool_loop( assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False + # Collect no-op nudges and flush them after the batch, so a no-op doesn't + # abort it and drop the parallel calls that follow. + deferred_noop_msgs: list = [] for tc in tool_calls or []: func = tc.get("function", {}) or {} @@ -1127,12 +1175,12 @@ def run_safetensors_tool_loop( "provenance": decision.provenance, } completion = tool_controller.record_noop(decision) - conversation.append(completion.model_message()) + deferred_noop_msgs.append(completion.model_message()) logger.info( "Suppressed local safetensors tool call as internal no-op: " f"action={decision.action} tool={decision.tool_name}" ) - break + continue if not assistant_appended: assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()] @@ -1243,6 +1291,8 @@ def run_safetensors_tool_loop( yield completion.tool_end_event() conversation.append(completion.tool_message()) + append_deferred_nudges(conversation, deferred_noop_msgs) + # Clear the status badge before the next turn. yield {"type": "status", "text": ""} diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py index f595531b90..f7ed450d11 100644 --- a/studio/backend/core/inference/tool_loop_controller.py +++ b/studio/backend/core/inference/tool_loop_controller.py @@ -266,6 +266,17 @@ def strip_result_for_model(result: str) -> str: return result +def append_deferred_nudges(conversation: list, msgs: Sequence[dict]) -> None: + """Append a batch's no-op nudges as one deduped ``role=user`` message. + + Deferred to after the batch's tool results so a no-op never splits an + assistant's ``tool_calls`` from their ``role=tool`` results. + """ + contents = list(dict.fromkeys(msg["content"] for msg in msgs)) + if contents: + conversation.append({"role": "user", "content": "\n\n".join(contents)}) + + def _tool_name_from_schema(tool: Mapping[str, Any]) -> str: function = tool.get("function") if not isinstance(function, Mapping): @@ -277,8 +288,9 @@ def _tool_name_from_schema(tool: Mapping[str, Any]) -> str: def _noop_result(reason: NoopReason, tool_name: str) -> str: if reason == "duplicate": return ( - "The previous tool request was not executed because this exact " - "tool call already completed successfully. Do not repeat the same " + f"One earlier request to call tool '{tool_name}' in this batch was " + "not executed because an identical call had already completed " + "successfully. Do not repeat the same " "tool call. Continue with a different enabled tool if that would " "materially help, or provide the final answer if you have enough " "information." @@ -291,8 +303,8 @@ def _noop_result(reason: NoopReason, tool_name: str) -> str: "the requested final note or answer." ) return ( - f"The previous tool request was not executed because tool " - f"'{tool_name}' is not enabled for this request. Provide the " + f"One earlier request to call tool '{tool_name}' in this batch was " + "not executed because that tool is not enabled for this request. Provide the " "final answer now without calling more tools." ) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index cfe46d4190..0f1ed464cc 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -4194,14 +4194,18 @@ _MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion) # Raw download cap > _MAX_PAGE_CHARS since SSR pages embed large sections # stripped during conversion; 512 KB still reaches article content. _MAX_FETCH_BYTES = 512 * 1024 +# PDF cross-reference data lives at EOF, so extraction needs the whole body. +_MAX_PDF_FETCH_BYTES = 10 * 1024 * 1024 +_MAX_WEB_PDF_PAGES = 50 # Control/undecodable chars, excluding text whitespace and ESC (for ANSI logs). # Binary when they exceed 12.5%, after allowing 16 minor encoding glitches. _BINARY_CHAR_RE = re.compile("[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]") _MIN_BINARY_CHARS = 16 _BINARY_CHAR_DIVISOR = 8 # Common binary signatures that can otherwise look text-heavy when mislabeled. +_PDF_MAGIC = b"%PDF-" _BINARY_MAGIC = ( - b"%PDF-", # PDF + _PDF_MAGIC, b"PK\x03\x04", # zip / docx / xlsx / pptx / epub / jar b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", # OLE / legacy Office b"\x89PNG\r\n\x1a\n", # PNG @@ -4235,14 +4239,22 @@ def _looks_binary(text: str) -> bool: ) -def _has_binary_magic(data: bytes) -> bool: - """Whether a common binary signature follows optional BOM or whitespace.""" +def _magic_head(data: bytes) -> bytes: head = data[:1024].lstrip() for bom, _codec in _UNICODE_BOM_CODECS: if head.startswith(bom): head = head.removeprefix(bom).lstrip() break - return head.startswith(_BINARY_MAGIC) + return head + + +def _has_pdf_magic(data: bytes) -> bool: + return _magic_head(data).startswith(_PDF_MAGIC) + + +def _has_binary_magic(data: bytes) -> bool: + """Whether a common binary signature follows optional BOM or whitespace.""" + return _magic_head(data).startswith(_BINARY_MAGIC) def _has_single_byte_text_evidence(data: bytes) -> bool: @@ -4253,6 +4265,45 @@ def _has_single_byte_text_evidence(data: bytes) -> bool: return ascii_text_bytes / len(data) >= _MIN_SINGLE_BYTE_ASCII_RATIO +def _extract_pdf_text(data: bytes) -> str: + """Extract page-delimited text with the same parser used by RAG ingestion.""" + from ..rag.parsers import parse_pdf_bytes + + pages, total_pages = parse_pdf_bytes(data, max_pages = _MAX_WEB_PDF_PAGES) + page_limit_reached = total_pages > _MAX_WEB_PDF_PAGES + parts: list[str] = [] + length = 0 + text_limited = False + for page in pages: + page_text = page.text.strip() + if not page_text: + continue + section = f"## Page {page.page_number}\n\n{page_text}" + piece = ("\n\n" if parts else "") + section + remaining = _MAX_PAGE_CHARS - length + if len(piece) > remaining: + parts.append(piece[:remaining]) + text_limited = True + break + parts.append(piece) + length += len(piece) + + text = "".join(parts).rstrip() + if not text: + if page_limit_reached: + return f"(PDF contains no extractable text in the first {_MAX_WEB_PDF_PAGES} pages)" + return "" + limits = [] + if text_limited: + limits.append(f"text limited to {_MAX_PAGE_CHARS:,} characters") + if page_limit_reached: + limits.append(f"page processing capped at {_MAX_WEB_PDF_PAGES} pages") + if limits: + marker = f"\n\n... (PDF extraction {'; '.join(limits)})" + text = text[: _MAX_PAGE_CHARS - len(marker)].rstrip() + marker + return text + + _USER_AGENTS = ( "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36", @@ -4648,29 +4699,71 @@ def _fetch_url_raw( return reason2, "", "" current_host = rp.hostname continue + + # get_content_type() defaults to "text/plain" when the header is + # absent (RFC 2045); report "" instead so callers can tell a missing + # header apart from a server that really declared text/plain. + if resp.headers.get("Content-Type") is None: + content_type = "" + else: + content_type = (resp.headers.get_content_type() or "").lower() + # Success: read the capped body enforcing the budget between chunks # (see _read_capped_body), so a slow-drip server can't stretch a # single resp.read past the deadline. + declared_pdf = content_type == "application/pdf" + read_limit = _MAX_PDF_FETCH_BYTES + 1 if declared_pdf else max_bytes body_error, raw_bytes = _read_capped_body( resp, - max_bytes, + read_limit, timeout, deadline, cancel_event, ) if body_error is not None: return body_error, "", "" + + # A missing or wrong PDF MIME type is common: once the initial text-sized + # read identifies PDF magic, finish the bounded download to reach the EOF xref. + if not declared_pdf and len(raw_bytes) == max_bytes and _has_pdf_magic(raw_bytes): + tail_error, tail = _read_capped_body( + resp, + _MAX_PDF_FETCH_BYTES - max_bytes + 1, + timeout, + deadline, + cancel_event, + ) + if tail_error is not None: + return tail_error, "", "" + raw_bytes += tail break else: return "Failed to fetch URL: too many redirects.", "", "" - # get_content_type() defaults to "text/plain" when the header is - # absent (RFC 2045); report "" instead so callers can tell a missing - # header apart from a server that really declared text/plain. - if resp.headers.get("Content-Type") is None: - content_type = "" - else: - content_type = (resp.headers.get_content_type() or "").lower() + is_pdf = declared_pdf or _has_pdf_magic(raw_bytes) + if is_pdf: + if len(raw_bytes) > _MAX_PDF_FETCH_BYTES: + return ( + "(PDF content exceeds the download limit; not readable as text)", + "", + content_type, + ) + budget_error = _fetch_budget_exceeded(deadline, cancel_event) + if budget_error is not None: + return budget_error, "", content_type + try: + pdf_text = _extract_pdf_text(raw_bytes) + except Exception as exc: + logger.debug("web PDF text extraction failed (%s)", type(exc).__name__) + return "(PDF content could not be read as text)", "", content_type + budget_error = _fetch_budget_exceeded(deadline, cancel_event) + if budget_error is not None: + return budget_error, "", content_type + if not pdf_text: + pdf_text = "(PDF contains no extractable text)" + # Report the true type even for a mislabeled body so the caller's "html" + # check routes the extracted text to the plain-text path, not html_to_markdown. + return None, pdf_text, "application/pdf" # Reject known-binary MIME types before decoding. Binary is returned as the # error string so the caller surfaces the placeholder, not replacement chars. diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index e4628dcea8..9f301ba37e 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -513,20 +513,25 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: logger.info("Starting text generation for request_id=%s", request_id) - for cumulative_text in generator: - # cancel_event is an mp.Event — checked instantly, no queue polling. - if cancel_event.is_set(): - logger.info("Generation cancelled for request %s", request_id) - break + try: + for cumulative_text in generator: + # cancel_event is an mp.Event — checked instantly, no queue polling. + if cancel_event.is_set(): + logger.info("Generation cancelled for request %s", request_id) + break - _send_response( - resp_queue, - { - "type": "token", - "request_id": request_id, - "text": cumulative_text, - }, - ) + _send_response( + resp_queue, + { + "type": "token", + "request_id": request_id, + "text": cumulative_text, + }, + ) + finally: + close = getattr(generator, "close", None) + if callable(close): + close() _send_response( resp_queue, diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index 9afddf1d9e..0b42906b85 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -103,7 +103,7 @@ def _markdown_incomplete(markdown: str, plain: str) -> bool: return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters -def _pdf_markdown(doc) -> list[str] | None: +def _pdf_markdown(doc, pages: range | None = None) -> list[str] | None: """Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index i maps to page i+1. Returns None when the lib is missing, extraction fails, or the page count does not line up, so the caller falls back to plain PyMuPDF text.""" @@ -112,28 +112,44 @@ def _pdf_markdown(doc) -> list[str] | None: except Exception: return None try: - chunks = pymupdf4llm.to_markdown( - doc, - page_chunks = True, - show_progress = False, - ) + kwargs = {"page_chunks": True, "show_progress": False} + if pages is not None: + kwargs["pages"] = list(pages) + chunks = pymupdf4llm.to_markdown(doc, **kwargs) except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True) return None - if not isinstance(chunks, list) or len(chunks) != doc.page_count: + expected_pages = doc.page_count if pages is None else len(pages) + if not isinstance(chunks, list) or len(chunks) != expected_pages: return None return [str(c.get("text") or "") for c in chunks] -def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: +def _pdf( + source: str | bytes, + want_images: bool, + max_pages: int | None = None, +) -> tuple[list[Page], list[ParsedImage], int]: import fitz # PyMuPDF pages: list[Page] = [] images: list[ParsedImage] = [] - doc = fitz.open(path) + doc = ( + fitz.open(stream = source, filetype = "pdf") if isinstance(source, bytes) else fitz.open(source) + ) try: - md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None - for i, page in enumerate(doc): + if doc.needs_pass: + raise ValueError("encrypted PDF requires a password") + total_pages = doc.page_count + page_numbers = range(total_pages if max_pages is None else min(total_pages, max_pages)) + if not config.PDF_MARKDOWN: + md = None + elif max_pages is None: + md = _pdf_markdown(doc) + else: + md = _pdf_markdown(doc, page_numbers) + for i, page_number in enumerate(page_numbers): + page = doc[page_number] plain = page.get_text("text") or "" candidate = md[i] if md else "" # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval), @@ -147,7 +163,7 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: text = candidate else: text = plain - pages.append(_page(text, i + 1)) + pages.append(_page(text, page_number + 1)) if want_images: for img in page.get_images(full = True): xref = img[0] @@ -161,13 +177,22 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: images.append( ParsedImage( image_bytes = image_bytes, - page_number = i + 1, + page_number = page_number + 1, xref = xref, ) ) finally: doc.close() - return pages, images + return pages, images, total_pages + + +def parse_pdf_bytes(data: bytes, *, max_pages: int | None = None) -> tuple[list[Page], int]: + """Extract PDF pages from an in-memory download using the ingestion parser. + + Returns the (capped) pages plus the document's full page count, so a caller + that set ``max_pages`` can tell a fully-read short PDF from a truncated one.""" + pages, _images, total_pages = _pdf(data, want_images = False, max_pages = max_pages) + return pages, total_pages def _merge_rects(boxes: list) -> list: @@ -416,7 +441,7 @@ def parse(path: str, *, want_images: bool = False): ext = os.path.splitext(path)[1].lower() if ext == ".pdf": - pages, images = _pdf(path, want_images) + pages, images, _total = _pdf(path, want_images) return (pages, images) if want_images else pages if ext == ".docx": diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py index 50e20fc13f..e5d48872c1 100644 --- a/studio/backend/hub/services/download_lifecycle.py +++ b/studio/backend/hub/services/download_lifecycle.py @@ -8,6 +8,7 @@ import os import signal import subprocess import sys +import time import threading from pathlib import Path from typing import Callable, Optional @@ -210,7 +211,9 @@ def finalize_worker_exit( repo_type: Optional[RepoType] = None, repo_id: Optional[str] = None, transport: Optional[str] = None, -) -> None: + cancel_marker_transport: Optional[str] = None, + defer_error: bool = False, +) -> str: """Block until *proc* exits, then record the job's terminal state in *registry*. Drains and scrubs stderr first, then classifies the exit code. A no-op when the process was already dropped (e.g. superseded). @@ -222,7 +225,7 @@ def finalize_worker_exit( rc = proc.wait() cancel_requested = registry.cancel_requested(key) if not registry.drop_process(key, proc): - return + return "idle" stderr_text = download_registry.scrub_secrets( (stderr_data or b"").decode("utf-8", "replace").strip(), hf_token = hf_token, @@ -230,6 +233,8 @@ def finalize_worker_exit( state = classify_exit(rc, cancel_requested = cancel_requested) if state == "complete": registry.set_job(key, "complete") + if transport == download_registry.TRANSPORT_HTTP: + registry.update_job_transport(key, download_registry.TRANSPORT_HTTP) if stderr_text: if download_manifest.MANIFEST_DEGRADED_MARKER in stderr_text: logger.warning( @@ -262,18 +267,226 @@ def finalize_worker_exit( metadata.variant if metadata is not None and metadata.variant else download_registry.variant_from_key(key), - transport, + cancel_marker_transport or transport, logger = logger, ) else: - registry.set_job( - key, - "error", - stderr_text or f"worker exited with code {rc}", - ) + if not defer_error: + registry.set_job( + key, + "error", + stderr_text or f"worker exited with code {rc}", + ) logger.error( f"{log_prefix} failed for {label} (rc={rc}): {stderr_text}", ) + return state + + +def _set_retry_failure_state( + registry: download_registry.DownloadRegistry, + key: str, + error: str, + *, + repo_type: RepoType, + repo_id: str, + fallback_variant: Optional[str], + fallback_transport: Optional[str], + logger, +) -> str: + state, metadata = registry.set_error_unless_cancelled(key, error) + if state == "cancelled": + download_registry.persist_cancel_marker( + repo_type, + repo_id, + metadata.variant if metadata is not None and metadata.variant else fallback_variant, + metadata.transport + if metadata is not None and metadata.transport + else fallback_transport, + logger = logger, + ) + return state + + +def _try_http_retry( + registry: download_registry.DownloadRegistry, + key: str, + *, + hf_token: Optional[str], + label: str, + log_prefix: str, + logger, + repo_type: RepoType, + repo_id: str, + watch_name: str, +) -> bool: + """Reclaim *key* with HTTP transport and spawn a recovery worker. + + Returns ``True`` when the HTTP worker was successfully registered. + Caller is responsible for ensuring this is only called when: the job is + in ``"error"`` state, the original transport was XET, and HTTP is available. + + Derives variant and blob-hash metadata from the registry entry written by + the original XET claim so callers do not re-construct worker arguments. + Re-queries peer protection hashes at spawn time to reflect any concurrent + sibling changes between the XET failure and this call. + """ + original_metadata = registry.get_job_metadata(key) + if original_metadata is None: + logger.debug("%s XET retry skipped for %s; metadata unavailable", log_prefix, label) + _set_retry_failure_state( + registry, + key, + "XET retry skipped: metadata unavailable", + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = download_registry.variant_from_key(key), + fallback_transport = download_registry.TRANSPORT_XET, + logger = logger, + ) + return False + if original_metadata.transport != download_registry.TRANSPORT_XET: + logger.debug( + "%s XET retry skipped for %s; original transport was %s", + log_prefix, + label, + original_metadata.transport, + ) + _set_retry_failure_state( + registry, + key, + f"XET retry skipped: original transport was {original_metadata.transport}", + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = original_metadata.variant, + fallback_transport = original_metadata.transport, + logger = logger, + ) + return False + variant = original_metadata.variant + blob_hashes = original_metadata.blob_hashes + progress_blob_hashes = original_metadata.progress_blob_hashes + completed_baseline_bytes = ( + download_registry.completed_blob_bytes( + repo_type, + repo_id, + progress_blob_hashes, + ) + if progress_blob_hashes + else 0 + ) + generation = registry.current_generation(key) + registry.release_active_slot(key) + while True: + if registry.cancel_requested(key): + _set_retry_failure_state( + registry, + key, + "HTTP retry cancelled before reclaiming the download slot", + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = variant, + fallback_transport = original_metadata.transport, + logger = logger, + ) + return False + + claimed, conflict_state = registry.claim( + key, + download_registry.TRANSPORT_HTTP, + repo_type = repo_type, + repo_id = repo_id, + variant = variant, + blob_hashes = blob_hashes, + progress_blob_hashes = progress_blob_hashes, + completed_baseline_bytes = completed_baseline_bytes, + generation = generation, + replace_active = True, + cancel_marker_transport = original_metadata.transport, + ) + if claimed: + break + if conflict_state == "deleting": + logger.debug( + "%s XET retry claim rejected for %s; repo is being deleted", + log_prefix, + label, + ) + _set_retry_failure_state( + registry, + key, + "HTTP retry could not reclaim the download slot", + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = variant, + fallback_transport = original_metadata.transport, + logger = logger, + ) + return False + logger.debug( + "%s XET retry claim blocked for %s by active sibling state %s; waiting", + log_prefix, + label, + conflict_state, + ) + time.sleep(0.05) + + args: list[str] = ["--repo-id", repo_id] + if repo_type == "dataset": + args.append("--dataset") + elif variant: + args.extend(["--variant", variant]) + + # Re-query at spawn time: sibling state may have changed since XET failed. + peer_hashes = registry.peer_blob_hashes(key) if variant else frozenset() + + logger.warning( + "%s XET worker failed for %s; retrying over HTTP", + log_prefix, + label, + ) + try: + proc = spawn_worker( + args, + hf_token, + use_xet = False, + protected_blob_hashes = peer_hashes or None, + ) + except Exception as exc: + scrubbed = download_registry.scrub_secrets(str(exc), hf_token = hf_token) + logger.error( + "%s HTTP retry spawn failed for %s: %s", + log_prefix, + label, + scrubbed, + ) + registry.update_job_transport(key, original_metadata.transport) + _set_retry_failure_state( + registry, + key, + scrubbed, + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = variant, + fallback_transport = original_metadata.transport, + logger = logger, + ) + return False + + return register_worker( + registry, + key, + proc, + hf_token = hf_token, + label = label, + log_prefix = log_prefix, + logger = logger, + repo_type = repo_type, + repo_id = repo_id, + transport = download_registry.TRANSPORT_HTTP, + cancel_marker_transport = original_metadata.transport, + watch_name = watch_name, + ) def kill_and_reap_process( @@ -309,6 +522,7 @@ def register_worker( repo_type: RepoType, repo_id: str, transport: str, + cancel_marker_transport: Optional[str] = None, watch_name: str, ) -> bool: if not registry.register_process(key, proc): @@ -319,7 +533,14 @@ def register_worker( def _watch() -> None: try: - finalize_worker_exit( + can_retry_http = ( + transport == download_registry.TRANSPORT_XET + and download_registry.download_transport_unavailable_reason( + download_registry.TRANSPORT_HTTP + ) + is None + ) + state = finalize_worker_exit( registry, key, proc, @@ -330,7 +551,25 @@ def register_worker( repo_type = repo_type, repo_id = repo_id, transport = transport, + cancel_marker_transport = cancel_marker_transport, + defer_error = can_retry_http, ) + # XET-to-HTTP recovery: when a non-cancelled XET worker fails and + # HTTP is available, attempt one automatic retry over HTTP. The + # transport check is the recursion guard: an HTTP worker that errors + # never satisfies `transport == TRANSPORT_XET`, so it stays terminal. + if can_retry_http and state == "error": + _try_http_retry( + registry, + key, + hf_token = worker_token, + label = label, + log_prefix = log_prefix, + logger = logger, + repo_type = repo_type, + repo_id = repo_id, + watch_name = watch_name, + ) except Exception: # finalize_worker_exit is the only thing that clears running/cancelling; # if it raises, force a terminal state so claim() isn't blocked until restart. @@ -426,8 +665,19 @@ def cancel_worker( return "cancelling" return registry.get_job(key).state # Worker already exited; let its watcher classify the real return code. - # Arming a pending cancel here could mislabel a genuine failure as a cancel. if proc.poll() is not None: + get_metadata = getattr(registry, "get_job_metadata", None) + metadata = get_metadata(key) if get_metadata is not None else None + can_retry_http = ( + metadata is not None + and metadata.transport == download_registry.TRANSPORT_XET + and download_registry.download_transport_unavailable_reason( + download_registry.TRANSPORT_HTTP + ) + is None + ) + if can_retry_http and registry.mark_pending_cancel(key, generation): + return "cancelling" return registry.get_job(key).state if not registry.request_cancel(key, proc, generation): diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py index c2ffe7bffc..862af0141a 100644 --- a/studio/backend/hub/services/models/downloads.py +++ b/studio/backend/hub/services/models/downloads.py @@ -60,6 +60,30 @@ def _job_status( return DownloadJobStatus(state = state, error = error, generation = generation) +def _load_in_flight(repo_id: str) -> bool: + try: + from core.inference.llama_cpp import hf_gguf_load_in_flight + return hf_gguf_load_in_flight(repo_id) + except Exception: + return False + + +def _load_in_flight_error(repo_id: str) -> HTTPException: + return HTTPException( + status_code = 409, + detail = ( + f"A model load for '{repo_id}' is in progress and may be " + "downloading it. Wait for the load to finish (or cancel it), " + "then start the download." + ), + ) + + +def _reject_if_load_in_flight(repo_id: str) -> None: + if _load_in_flight(repo_id): + raise _load_in_flight_error(repo_id) + + def _spawn_download_worker( repo_id: str, variant: Optional[str], @@ -89,6 +113,9 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional # Canonicalize so two different-cased paste-ins share one job + cache dir. repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + # Avoid concurrent writers to the same HF cache files. + _reject_if_load_in_flight(repo_id) + variant = (body.gguf_variant or "").strip() or None if variant is not None and not _is_valid_gguf_variant(variant): raise HTTPException( @@ -147,9 +174,12 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional blob_hashes = variant_blob_hashes, progress_blob_hashes = variant_progress_blob_hashes, completed_baseline_bytes = completed_baseline_bytes, + admission_check = lambda: not _load_in_flight(repo_id), ) generation = _registry.current_generation(key) if not claimed: + if claim_state == "admission_blocked": + raise _load_in_flight_error(repo_id) # claim_state is the blocking job's state. The client can attach only # when the blocker is this key's own in-flight job (adoptable); a # cross-variant conflict or in-progress delete is not accepted. diff --git a/studio/backend/hub/tests/test_download_lifecycle.py b/studio/backend/hub/tests/test_download_lifecycle.py index a4baafa317..87346573b0 100644 --- a/studio/backend/hub/tests/test_download_lifecycle.py +++ b/studio/backend/hub/tests/test_download_lifecycle.py @@ -1,27 +1,147 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import io +import logging + from hub.services import download_lifecycle +from hub.utils import download_registry, state_dir -def _set_xet_reason(monkeypatch, reason): +class _Proc: + pid = 4242 + + def __init__( + self, + rc, + stderr = b"", + ): + self.rc = rc + self.stderr = io.BytesIO(stderr) + self.waited = False + + def poll(self): + return self.rc if self.waited else None + + def wait(self, timeout = None): + self.waited = True + return self.rc + + def kill(self): + pass + + +class _ImmediateThread: + def __init__(self, *, target, **_kwargs): + self.target = target + + def start(self): + self.target() + + +def test_resolve_effective_use_xet(monkeypatch): + for requested, unavailable_reason, expected in ( + (False, "unused", False), + (True, None, True), + (True, "hf_xet is not installed", False), + ): + monkeypatch.setattr( + download_lifecycle.download_registry, + "download_transport_unavailable_reason", + lambda _transport, reason = unavailable_reason: reason, + ) + assert download_lifecycle.resolve_effective_use_xet(requested) is expected + + +def test_xet_failure_retries_over_http_for_model_and_dataset(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr(download_lifecycle.threading, "Thread", _ImmediateThread) + register_worker = download_lifecycle.register_worker + + for repo_type, repo_id, variant, expected_args in ( + ("model", "Org/Model", "Q4_K_M", ["--repo-id", "Org/Model", "--variant", "Q4_K_M"]), + ("dataset", "Org/Data", None, ["--repo-id", "Org/Data", "--dataset"]), + ): + registry = download_registry.DownloadRegistry() + key = download_registry.normalize_job_key(f"{repo_id}::{variant}" if variant else repo_id) + assert registry.claim( + key, + download_registry.TRANSPORT_XET, + repo_type = repo_type, + repo_id = repo_id, + variant = variant, + blob_hashes = frozenset({"blob"}), + )[0] + generation = registry.current_generation(key) + spawned = [] + + def fake_spawn( + args, + _token, + *, + use_xet, + protected_blob_hashes = None, + ): + spawned.append((args, use_xet, protected_blob_hashes)) + return _Proc(0) + + def fake_retry_register(*_args, **kwargs): + assert kwargs["transport"] == download_registry.TRANSPORT_HTTP + return True + + monkeypatch.setattr(download_lifecycle, "spawn_worker", fake_spawn) + monkeypatch.setattr(download_lifecycle, "register_worker", fake_retry_register) + assert register_worker( + registry, + key, + _Proc(1, b"xet failed"), + hf_token = None, + label = repo_id, + log_prefix = "Download", + logger = logging.getLogger("test"), + repo_type = repo_type, + repo_id = repo_id, + transport = download_registry.TRANSPORT_XET, + watch_name = f"{repo_type}-watch", + ) + + metadata = registry.get_job_metadata(key) + assert spawned == [(expected_args, False, None)] + assert metadata.transport == download_registry.TRANSPORT_HTTP + assert metadata.blob_hashes == frozenset({"blob"}) + assert registry.current_generation(key) == generation + + +def test_http_failure_remains_terminal(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr(download_lifecycle.threading, "Thread", _ImmediateThread) + register_worker = download_lifecycle.register_worker + registry = download_registry.DownloadRegistry() + key = download_registry.normalize_repo_key("Org/Data") + assert registry.claim( + key, + download_registry.TRANSPORT_HTTP, + repo_type = "dataset", + repo_id = "Org/Data", + )[0] monkeypatch.setattr( - download_lifecycle.download_registry, - "download_transport_unavailable_reason", - lambda _transport: reason, + download_lifecycle, + "register_worker", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("HTTP failures must not retry") + ), ) - - -def test_resolve_effective_use_xet_keeps_http_when_not_requested(monkeypatch): - _set_xet_reason(monkeypatch, "should not be consulted") - assert download_lifecycle.resolve_effective_use_xet(False) is False - - -def test_resolve_effective_use_xet_keeps_xet_when_available(monkeypatch): - _set_xet_reason(monkeypatch, None) - assert download_lifecycle.resolve_effective_use_xet(True) is True - - -def test_resolve_effective_use_xet_downgrades_when_xet_unavailable(monkeypatch): - _set_xet_reason(monkeypatch, "Xet transport is unavailable because hf_xet is not installed.") - assert download_lifecycle.resolve_effective_use_xet(True) is False + assert register_worker( + registry, + key, + _Proc(1, b"http failed"), + hf_token = None, + label = "Org/Data", + log_prefix = "Download", + logger = logging.getLogger("test"), + repo_type = "dataset", + repo_id = "Org/Data", + transport = download_registry.TRANSPORT_HTTP, + watch_name = "dataset-watch", + ) + assert registry.get_job(key).state == "error" diff --git a/studio/backend/hub/utils/download_registry.py b/studio/backend/hub/utils/download_registry.py index 777d63e1b5..b6bdee3bce 100644 --- a/studio/backend/hub/utils/download_registry.py +++ b/studio/backend/hub/utils/download_registry.py @@ -45,9 +45,9 @@ import sys import threading import time import weakref -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Iterator, Literal, Optional +from typing import Callable, Iterator, Literal, Optional from loggers import get_logger @@ -126,6 +126,9 @@ def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMeta "repo_id": metadata.repo_id if metadata is not None else None, "variant": metadata.variant if metadata is not None else None, "transport": metadata.transport if metadata is not None else None, + "cancel_marker_transport": metadata.cancel_marker_transport + if metadata is not None + else None, } tmp = path.with_name(f".{path.name}.tmp-{pid}") try: @@ -305,7 +308,7 @@ def reap_orphan_workers() -> None: data.get("repo_type"), repo_id, data.get("variant"), - data.get("transport"), + data.get("cancel_marker_transport") or data.get("transport"), ) except Exception as exc: logger.debug("Reaper failed for breadcrumb %s: %s", entry, exc) @@ -699,6 +702,7 @@ class DownloadMetadata: repo_id: str variant: Optional[str] transport: Optional[str] + cancel_marker_transport: Optional[str] = None # GGUF variant main/writable hashes, identifying the variant-specific shards # for concurrency decisions. blob_hashes: frozenset[str] = field(default_factory = frozenset) @@ -801,6 +805,7 @@ class DownloadRegistry: self._processes: dict[str, subprocess.Popen] = {} self._repo_active: dict[str, set[str]] = {} self._metadata: dict[str, DownloadMetadata] = {} + self._cancel_marker_transports: dict[str, str] = {} self._pending_cancel: dict[str, Optional[int]] = {} self._generations: dict[str, int] = {} # Monotonic across keys so an evicted then re-claimed key never reuses a @@ -839,6 +844,7 @@ class DownloadRegistry: if state in TERMINAL_STATES: self._put_terminal_job_locked(key, state, error) self._pending_cancel.pop(key, None) + self._cancel_marker_transports.pop(key, None) repo = _repo_of_key(key) active = self._repo_active.get(repo) if active is not None: @@ -848,6 +854,57 @@ class DownloadRegistry: else: self._jobs[key] = DownloadState(state, error) + def set_error_unless_cancelled( + self, key: str, error: str + ) -> tuple[JobState, Optional[DownloadMetadata]]: + key = normalize_job_key(key) + with self._lock: + current = self._jobs.get(key, DownloadState("idle")).state + has_pending_cancel = key in self._pending_cancel + pending_generation = self._pending_cancel.get(key) + metadata = self._metadata.get(key) + should_cancel = current == "cancelling" or ( + has_pending_cancel and self._generation_matches_locked(key, pending_generation) + ) + terminal_state: JobState = "cancelled" if should_cancel else "error" + marker_transport = self._cancel_marker_transports.pop(key, None) + if marker_transport is None and metadata is not None: + marker_transport = metadata.cancel_marker_transport + self._put_terminal_job_locked( + key, + terminal_state, + None if should_cancel else error, + ) + self._pending_cancel.pop(key, None) + repo = _repo_of_key(key) + active = self._repo_active.get(repo) + if active is not None: + active.discard(key) + if not active: + self._repo_active.pop(repo, None) + if should_cancel and metadata is not None and marker_transport is not None: + metadata = replace(metadata, transport = marker_transport) + return terminal_state, metadata + + def update_job_transport(self, key: str, transport: str) -> None: + key = normalize_job_key(key) + with self._lock: + metadata = self._metadata.get(key) + if metadata is None or metadata.transport == transport: + return + self._metadata[key] = replace(metadata, transport = transport) + + def release_active_slot(self, key: str) -> None: + key = normalize_job_key(key) + repo = _repo_of_key(key) + with self._lock: + active = self._repo_active.get(repo) + if active is None: + return + active.discard(key) + if not active: + self._repo_active.pop(repo, None) + def get_job(self, key: str) -> DownloadState: key = normalize_job_key(key) with self._lock: @@ -884,6 +941,14 @@ class DownloadRegistry: ): self._put_terminal_job_locked(key, "cancelled") metadata_to_persist = self._metadata.pop(key, None) + marker_transport = self._cancel_marker_transports.pop(key, None) + if marker_transport is None and metadata_to_persist is not None: + marker_transport = metadata_to_persist.cancel_marker_transport + if metadata_to_persist is not None and marker_transport is not None: + metadata_to_persist = replace( + metadata_to_persist, + transport = marker_transport, + ) repo = _repo_of_key(key) active = self._repo_active.get(repo) if active is not None: @@ -963,12 +1028,24 @@ class DownloadRegistry: blob_hashes: Optional[frozenset[str]] = None, progress_blob_hashes: Optional[frozenset[str]] = None, completed_baseline_bytes: int = 0, + admission_check: Optional[Callable[[], bool]] = None, + generation: Optional[int] = None, + replace_active: bool = False, + metadata_transport: Optional[str] = None, + cancel_marker_transport: Optional[str] = None, ) -> tuple[bool, str]: key = normalize_job_key(key) repo = _repo_of_key(key) requested_hashes = blob_hashes or frozenset() requested_progress_hashes = progress_blob_hashes or frozenset() with self._lock: + # Run the final external admission check while the registry lock is + # held, immediately before inspecting and publishing active state. + # The GGUF load path establishes its marker before calling + # its active-job probe, so either this claim observes that marker + # or the load's later probe observes this claim. + if admission_check is not None and not admission_check(): + return False, "admission_blocked" deleting_scopes = self._deleting.get(repo) if deleting_scopes is not None and ( None in deleting_scopes or variant_from_key(key) in deleting_scopes @@ -1007,10 +1084,13 @@ class DownloadRegistry: if conflict_state is not None: return False, conflict_state current = self._jobs.get(key, DownloadState("idle")).state - if current in _ACTIVE_STATES: + if current in _ACTIVE_STATES and not replace_active: return False, current - self._generation_seq += 1 - self._generations[key] = self._generation_seq + if generation is None: + self._generation_seq += 1 + self._generations[key] = self._generation_seq + else: + self._generations[key] = generation self._jobs[key] = DownloadState("running") self._repo_active.setdefault(repo, active).add(key) if repo_type and repo_id: @@ -1018,7 +1098,8 @@ class DownloadRegistry: repo_type = repo_type, repo_id = repo_id, variant = variant, - transport = transport, + transport = metadata_transport if metadata_transport is not None else transport, + cancel_marker_transport = cancel_marker_transport, blob_hashes = requested_hashes, progress_blob_hashes = requested_progress_hashes, completed_baseline_bytes = max( @@ -1026,8 +1107,13 @@ class DownloadRegistry: int(completed_baseline_bytes or 0), ), ) + if cancel_marker_transport is not None: + self._cancel_marker_transports[key] = cancel_marker_transport + else: + self._cancel_marker_transports.pop(key, None) else: self._metadata.pop(key, None) + self._cancel_marker_transports.pop(key, None) return True, "running" def adoptable(self, key: str) -> bool: @@ -1053,7 +1139,8 @@ class DownloadRegistry: download. A variant delete conflicts only with that same variant or a whole-repo download writing the shared snapshot; other quantizations download concurrently and never block it.""" - for key in self._repo_active.get(repo_id, set()): + active_keys = self._repo_active.get(repo_id, set()) + for key in active_keys: job = self._jobs.get(key) if job is None or job.state not in _ACTIVE_STATES: continue @@ -1062,6 +1149,16 @@ class DownloadRegistry: other_variant = self._active_job_variant_locked(key) if other_variant is None or other_variant == variant: return True + for key, job in self._jobs.items(): + if key in active_keys or _repo_of_key(key) != repo_id: + continue + if job.state not in _ACTIVE_STATES: + continue + if variant is None: + return True + other_variant = self._active_job_variant_locked(key) + if other_variant is None or other_variant == variant: + return True return False def peer_blob_hashes(self, key: str) -> frozenset[str]: @@ -1108,6 +1205,16 @@ class DownloadRegistry: candidate_keys = list(self._repo_active.get(repo_key, set())) else: candidate_keys = [key for active in self._repo_active.values() for key in active] + # An XET->HTTP retry handoff briefly drops its key from _repo_active + # while its job stays active; include those released-but-active jobs + # so the waiting retry still lists and can be adopted or cancelled. + seen = set(candidate_keys) + for key, job in self._jobs.items(): + if key in seen or job.state not in _ACTIVE_STATES: + continue + if repo_key is not None and _repo_of_key(key) != repo_key: + continue + candidate_keys.append(key) refs: list[ActiveDownloadRef] = [] for key in candidate_keys: job = self._jobs.get(key) @@ -1123,6 +1230,23 @@ class DownloadRegistry: ) return refs + def has_active_variant(self, repo_id: str, variant: Optional[str]) -> bool: + """Whether an active model job targets this exact GGUF variant. + + Scans the job table rather than only ``_repo_active`` so an XET-to-HTTP + retry handoff remains visible while it has temporarily released its + active slot. + """ + repo_key = normalize_repo_key(repo_id) + target = (variant or "").strip().lower() or None + with self._lock: + for key, job in self._jobs.items(): + if _repo_of_key(key) != repo_key or job.state not in _ACTIVE_STATES: + continue + if self._active_job_variant_locked(key) == target: + return True + return False + def begin_delete( self, repo_id: str, @@ -1169,12 +1293,25 @@ class DownloadRegistry: repo_id = normalize_repo_key(repo_id) target = (variant or "").strip().lower() or None with self._lock: - for key in self._repo_active.get(repo_id, set()): + active_keys = self._repo_active.get(repo_id, set()) + for key in active_keys: job = self._jobs.get(key) if job is None or job.state not in _ACTIVE_STATES: continue if self._active_job_variant_locked(key) != target: return True + # An XET->HTTP retry peer between release_active_slot() and its reclaim + # is briefly absent from _repo_active while its job stays active and + # still owns the shared companion; mirror the released-but-active scan + # used by _delete_blocked_by_active_locked so it still blocks companion + # deletion of a different variant. + for key, job in self._jobs.items(): + if key in active_keys or _repo_of_key(key) != repo_id: + continue + if job.state not in _ACTIVE_STATES: + continue + if self._active_job_variant_locked(key) != target: + return True return False def request_cancel( @@ -1198,17 +1335,58 @@ class DownloadRegistry: return True def terminate_all(self, kind: str = "download") -> None: + settled_no_proc: list[Optional[DownloadMetadata]] = [] with self._lock: live = [ (key, proc, self._metadata.get(key)) for key, proc in self._processes.items() if proc.poll() is None ] + live_keys = {key for key, _proc, _metadata in live} # Flag as an intentional stop so the watcher's exit classification # reports them cancelled rather than an OOM/crash once SIGKILL lands. for key, _proc, _metadata in live: if self._jobs.get(key, DownloadState("idle")).state == "running": self._jobs[key] = DownloadState("cancelling") + # Settle active jobs without a live worker too. Two cases: an + # XET->HTTP retry parked in the reclaim wait loop has dropped its + # worker and slot guard, so it is absent from `live`; and a + # registered worker that already exited with an error but whose + # watcher has not yet run would otherwise stay `running` and spawn an + # HTTP retry after this shutdown snapshot. Skip a registered worker + # that exited cleanly (rc == 0): it completed and the watcher will + # mark it done, so marking it cancelling would strand a stale marker. + for key, job in list(self._jobs.items()): + if job.state not in _ACTIVE_STATES or key in live_keys: + continue + proc = self._processes.get(key) + if proc is not None: + if proc.poll() == 0: + continue + # A registered worker that exited nonzero on its own over HTTP + # is a genuine terminal download failure, not a shutdown cancel + # and not retry-capable: leave its error status intact rather + # than persisting a cancel marker that would read as + # cancelled/resumable after restart. Only an exited XET worker + # could still spawn a post-shutdown HTTP retry, so only that + # needs settling here. + metadata = self._metadata.get(key) + if metadata is not None and metadata.transport == TRANSPORT_HTTP: + continue + self._pending_cancel[key] = self._generations.get(key) + self._jobs[key] = DownloadState("cancelling") + settled_no_proc.append(self._metadata.get(key)) + # Persist a cancel marker for each settled no-live-worker job outside the + # lock (mirroring the reaped path) so shutdown records resumable/cancelled + # state even if it returns before the daemon watcher wakes to do so. + for metadata in settled_no_proc: + if metadata is not None: + persist_cancel_marker( + metadata.repo_type, + metadata.repo_id, + metadata.variant, + metadata.cancel_marker_transport or metadata.transport, + ) reaped: list[tuple[str, subprocess.Popen, Optional[DownloadMetadata]]] = [] for key, proc, metadata in live: try: @@ -1222,7 +1400,7 @@ class DownloadRegistry: metadata.repo_type, metadata.repo_id, metadata.variant, - metadata.transport, + metadata.cancel_marker_transport or metadata.transport, ) continue reaped.append((key, proc, metadata)) @@ -1242,7 +1420,7 @@ class DownloadRegistry: metadata.repo_type, metadata.repo_id, metadata.variant, - metadata.transport, + metadata.cancel_marker_transport or metadata.transport, ) diff --git a/studio/backend/main.py b/studio/backend/main.py index e64048dc00..6e16dc00ca 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -612,6 +612,22 @@ app = FastAPI( lifespan = lifespan, ) +# The MCP surface is opt-in because it can start GPU jobs and write model +# artifacts. Mount it only when explicitly enabled by the Studio process. +if os.environ.get("UNSLOTH_STUDIO_ENABLE_MCP") == "1": + from fastmcp.utilities.lifespan import combine_lifespans + + from mcp_server import BearerTokenMiddleware, create_studio_mcp + + _studio_mcp_app = create_studio_mcp().http_app(path = "/") + _studio_mcp_lifespan = _studio_mcp_app.lifespan + _mcp_token = os.environ.get("UNSLOTH_STUDIO_MCP_TOKEN") + if not _mcp_token: + raise RuntimeError("UNSLOTH_STUDIO_MCP_TOKEN is required when MCP is enabled") + _studio_mcp_app = BearerTokenMiddleware(_studio_mcp_app, _mcp_token) + app.router.lifespan_context = combine_lifespans(lifespan, _studio_mcp_lifespan) + app.mount("/mcp", _studio_mcp_app) + from loggers.config import LogConfig from loggers.handlers import LoggingMiddleware @@ -752,6 +768,7 @@ _BODY_PROTECTED_PREFIXES = ( "/api/settings", "/api/train", "/api/export", + "/mcp", ) _DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload" _DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = ( diff --git a/studio/backend/mcp_server.py b/studio/backend/mcp_server.py new file mode 100644 index 0000000000..f837f46425 --- /dev/null +++ b/studio/backend/mcp_server.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Curated MCP tools for driving an Unsloth Studio instance. + +The MCP surface deliberately wraps the existing Studio services instead of +duplicating training or export logic. It is opt-in because several tools can +start GPU work or write model artifacts. +""" + +from __future__ import annotations + +import hmac +from typing import Any + +from fastmcp import FastMCP + + +class BearerTokenMiddleware: + """Require an exact bearer token when Studio MCP is exposed remotely.""" + + def __init__(self, app: Any, token: str) -> None: + if not token or not token.strip(): + raise ValueError("Studio MCP bearer token must be a non-empty value") + if not token.isascii(): + # A non-ASCII token cannot be sent in an HTTP header; reject it here. + raise ValueError("Studio MCP bearer token must contain ASCII characters only") + self.app = app + # Compare on raw header bytes: str hmac.compare_digest raises on non-ASCII + # input, which would surface as a 500 instead of a clean 401. + self.expected = token.encode("utf-8") + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + scope_type = scope.get("type") + if scope_type not in ("http", "websocket"): + await self.app(scope, receive, send) + return + + headers = dict(scope.get("headers", [])) + raw_auth = headers.get(b"authorization", b"") + scheme, _, supplied = raw_auth.partition(b" ") + if scheme.lower() != b"bearer" or not hmac.compare_digest(supplied, self.expected): + await _send_unauthorized(send, scope_type) + return + + await self.app(scope, receive, send) + + +async def _send_unauthorized(send: Any, scope_type: str) -> None: + if scope_type == "websocket": + await send({"type": "websocket.close", "code": 4401}) + return + + await send( + { + "type": "http.response.start", + "status": 401, + "headers": [(b"content-type", b"application/json"), (b"www-authenticate", b"Bearer")], + } + ) + await send( + { + "type": "http.response.body", + "body": b'{"detail":"MCP bearer token required"}', + } + ) + + +def _dump(value: Any) -> Any: + """Convert Pydantic responses to plain JSON values for MCP clients.""" + if hasattr(value, "model_dump"): + return value.model_dump(mode = "json") + return value + + +def _clamp(value: int, low: int, high: int) -> int: + """Clamp an MCP-supplied integer into an inclusive range. + + MCP tools call the Studio route functions directly, which skips FastAPI's + Query(ge=, le=) validation, so we re-apply the same bounds here. + """ + return max(low, min(value, high)) + + +def create_studio_mcp() -> FastMCP: + """Create the Studio MCP server and register the high-value tools.""" + mcp = FastMCP( + "Unsloth Studio", + instructions = ( + "Use read tools to inspect the local Studio state before starting GPU work. " + "Training and export tools can consume substantial VRAM and write files. " + "Never expose tokens or local paths from tool results unless the user asks." + ), + ) + + @mcp.tool + async def studio_status() -> dict[str, Any]: + """Return the current training, export, inference, and GPU state.""" + from routes.export import get_export_status + from routes.inference import get_status as get_inference_status + from routes.training import get_training_status + + from utils.hardware import get_gpu_utilization + + training, export, inference = await _gather_status( + get_training_status(current_subject = "mcp"), + get_export_status(current_subject = "mcp"), + get_inference_status(current_subject = "mcp"), + ) + return { + "training": _dump(training), + "export": _dump(export), + "inference": _dump(inference), + "hardware": get_gpu_utilization(), + } + + @mcp.tool + async def list_local_models(models_dir: str = "./models") -> dict[str, Any]: + """List local and cached models available to Studio.""" + from routes.models import list_local_models as list_models + return _dump(await list_models(models_dir = models_dir, current_subject = "mcp")) + + @mcp.tool + async def get_training_status() -> dict[str, Any]: + """Read the active training job, phase, progress, and recent metrics.""" + from routes.training import get_training_status as get_status + return _dump(await get_status(current_subject = "mcp")) + + @mcp.tool + async def start_training(config: dict[str, Any]) -> dict[str, Any]: + """Start a validated Studio training job from a TrainingStartRequest-shaped object. + + The config is validated by the same Pydantic model used by the Studio UI. + Call get_training_status first and do not start work while another job runs. + """ + from models import TrainingStartRequest + from routes.training import start_training as start + + request = TrainingStartRequest.model_validate(config) + # Pass via_api_key explicitly (a direct call leaves it a Depends object). + # MCP drives Studio like the UI session, so it coexists and frees VRAM. + return _dump(await start(request, current_subject = "mcp", via_api_key = False)) + + @mcp.tool + async def stop_training(save: bool = True) -> dict[str, Any]: + """Ask the active training process to stop at its next safe checkpoint.""" + from routes.training import TrainingStopRequest, stop_training as stop + return _dump(await stop(TrainingStopRequest(save = save), current_subject = "mcp")) + + @mcp.tool + async def list_training_runs(limit: int = 50, offset: int = 0) -> dict[str, Any]: + """List completed and stopped training runs, newest first.""" + from routes.training_history import list_training_runs as list_runs + + # Clamp here (direct call skips Query bounds); a negative LIMIT = no limit. + limit = _clamp(limit, 1, 200) + offset = max(0, offset) + return _dump(await list_runs(limit = limit, offset = offset, current_subject = "mcp")) + + @mcp.tool + def validate_recipe(recipe: dict[str, Any]) -> dict[str, Any]: + """Validate a Data Recipe with the same validator used by Studio.""" + from models.data_recipe import RecipePayload + from routes.data_recipe.validate import validate + + return _dump(validate(RecipePayload(recipe = recipe))) + + @mcp.tool + def get_recipe_job_status(job_id: str) -> dict[str, Any]: + """Read the status of a Data Recipe job.""" + from routes.data_recipe.jobs import job_status + return _dump(job_status(job_id)) + + @mcp.tool + def get_recipe_job_dataset( + job_id: str, + limit: int = 20, + offset: int = 0, + ) -> dict[str, Any]: + """Read a bounded page of generated Data Recipe rows.""" + from routes.data_recipe.jobs import job_dataset + + # Clamp here (direct call skips FastAPI's Query bounds). + limit = _clamp(limit, 1, 500) + offset = max(0, offset) + return _dump(job_dataset(job_id, limit = limit, offset = offset)) + + @mcp.tool + async def load_checkpoint( + checkpoint_path: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + trust_remote_code: bool = False, + approved_remote_code_fingerprint: str | None = None, + hf_token: str | None = None, + ) -> dict[str, Any]: + """Load a checkpoint into the export backend. + + Export runs in its own subprocess and coexists with training and + inference; it does not unload them, so a load can fail with a clear + out-of-memory error if the GPU is already full. Pass hf_token to load a + gated checkpoint, and approved_remote_code_fingerprint to retry a + trust_remote_code load that was blocked pending review. + """ + from models import LoadCheckpointRequest + from routes.export import load_checkpoint as load + + request = LoadCheckpointRequest( + checkpoint_path = checkpoint_path, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + trust_remote_code = trust_remote_code, + approved_remote_code_fingerprint = approved_remote_code_fingerprint, + hf_token = hf_token, + ) + return _dump(await load(request, current_subject = "mcp")) + + @mcp.tool + async def export_gguf( + save_directory: str, + quantization_method: str | list[str] = "Q4_K_M", + push_to_hub: bool = False, + repo_id: str | None = None, + hf_token: str | None = None, + imatrix: bool = False, + imatrix_path: str | None = None, + ) -> dict[str, Any]: + """Export the loaded model to GGUF using Studio's existing path validation. + + quantization_method may be a single method or a list to produce several + GGUFs from one load. Pass hf_token when push_to_hub is set (the backend + rejects a Hub upload without it). Set imatrix (or imatrix_path) for the + IQ low-bit quants that require an importance matrix. + """ + from models import ExportGGUFRequest + from routes.export import export_gguf as export + + request = ExportGGUFRequest( + save_directory = save_directory, + quantization_method = quantization_method, + push_to_hub = push_to_hub, + repo_id = repo_id, + hf_token = hf_token, + imatrix = imatrix, + imatrix_path = imatrix_path, + ) + return _dump(await export(request, current_subject = "mcp")) + + return mcp + + +async def _gather_status(*coroutines: Any) -> tuple[Any, ...]: + """Gather independent status calls without letting one optional backend fail all state.""" + import asyncio + + results = await asyncio.gather(*coroutines, return_exceptions = True) + return tuple( + {"error": str(result)} if isinstance(result, Exception) else result for result in results + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0541a81ff2..9299e26d56 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -20,6 +20,7 @@ from loggers import get_logger import asyncio import threading import weakref +from contextlib import ExitStack import re as _re @@ -28,6 +29,7 @@ import re as _re from utils.models import extract_model_size_b as _extract_model_size_b from utils.api_errors import openai_error_body, anthropic_error_body +from core.inference.orchestrator import GenStreamError, GenStreamErrorRaised from core.inference.llama_admission import ( LlamaAdmissionCancelled, LlamaAdmissionConfig, @@ -212,6 +214,14 @@ def _friendly_error(exc: Exception) -> str: return "An internal error occurred" +def _friendly_gen_stream_error(value) -> str: + """Return a client-safe message for typed local generation errors.""" + text = str(value) + if getattr(value, "public", False): + return text + return safe_error_detail(RuntimeError(text), fallback = "An internal error occurred.") + + def _friendly_upstream_error(text: str) -> str: """Rewrite a raw llama-server error body into an actionable message where we can. @@ -998,6 +1008,7 @@ try: from core.inference.llama_server_args import ( _effective_tensor_parallel, _tensor_parallel_matches_loaded, + extra_args_disable_mmproj, parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, @@ -1035,6 +1046,7 @@ except ImportError: from core.inference.llama_server_args import ( _effective_tensor_parallel, _tensor_parallel_matches_loaded, + extra_args_disable_mmproj, parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, @@ -1915,16 +1927,57 @@ async def artifact_preview_frame(allow_network: bool = False): _BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:') -def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: +def _detect_safetensors_features( + backend, + chat_template: Optional[str], + tools = None, +) -> dict: """Classify reasoning/tool capabilities via the GGUF classifier so flags match across backends. gpt-oss is overridden: Harmony routes reasoning and tools through tokenizer channels, not template markup.""" model_id = getattr(backend, "active_model_name", None) + feature_template = chat_template + try: + from core.inference.chat_template_helpers import _selected_template_strings_from_value + selected_templates = _selected_template_strings_from_value(chat_template, tools) + if selected_templates: + feature_template = selected_templates[0] + except Exception: + logger.debug("safetensors_named_template_selection_failed", exc_info = True) flags = detect_reasoning_flags( - chat_template, + feature_template, model_identifier = model_id, log_source = "safetensors", ) + if not flags.get("supports_reasoning"): + try: + from core.inference.chat_template_helpers import ( + detect_reasoning_channel_markers_from_template, + ) + + templates = [chat_template] + models = getattr(backend, "models", None) + model_info = ( + models.get(model_id, {}) + if isinstance(models, dict) and model_id is not None + else {} + ) + if isinstance(model_info, dict): + templates.extend( + ( + model_info.get("native_chat_template"), + (model_info.get("chat_template_info") or {}).get("template"), + ) + ) + if any( + detect_reasoning_channel_markers_from_template(template, tools = tools) is not None + for template in templates + ): + flags["supports_reasoning"] = True + flags["reasoning_always_on"] = True + logger.info("safetensors: model always reasons (native channel markers)") + except Exception: + logger.debug("safetensors_native_reasoning_marker_check_failed", exc_info = True) # Markers any supported parser recognises (template advertises tools but # uses none -> drop the pill). Reuse the parser's own signal list so this # gate never drifts (a hand-maintained copy lost the DeepSeek variants); @@ -1938,9 +1991,9 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: ) if ( flags.get("supports_tools") - and chat_template - and not any(m in chat_template for m in _PARSER_MARKERS) - and not _BARE_JSON_NAME_MARKER_RE.search(chat_template) + and isinstance(feature_template, str) + and not any(m in feature_template for m in _PARSER_MARKERS) + and not _BARE_JSON_NAME_MARKER_RE.search(feature_template) ): logger.info( "safetensors: template advertises tools but uses an " @@ -3968,6 +4021,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre native_grant_backed = False model_log_label = request.model_path + gguf_load_stack = ExitStack() try: # Validate user pass-through args up front so a managed-flag collision # returns 400 before any model work. @@ -4187,15 +4241,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() - # Unload any active Unsloth model to free VRAM (off the event loop: - # unload takes _gen_lock and can wait on an in-flight stream). - if unsloth_backend.active_model_name: - logger.info( - f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" - ) - await asyncio.to_thread( - unsloth_backend.unload_model, unsloth_backend.active_model_name - ) + if config.gguf_hf_repo: + from core.inference.llama_cpp import gguf_load_in_flight + gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) # Inherit llama_extra_args from the previous load when the request # omits the field (the chat-settings Apply path doesn't round-trip @@ -4275,6 +4323,38 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre extra_llama_args, ) + # Block cache writes that would race the download manager. This runs + # after pass-through argument inheritance so a carried --no-mmproj + # changes the companion requirement exactly as it does for the load. + if config.gguf_hf_repo: + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + if await asyncio.to_thread( + _hub_download_blocks_gguf_load, + config.gguf_hf_repo, + config.gguf_variant, + require_mmproj = bool( + config.is_vision and not extra_args_disable_mmproj(extra_llama_args) + ), + hf_token = request.hf_token, + ): + raise HTTPException( + status_code = 409, + detail = ( + f"'{model_log_label}' is currently being downloaded " + "by the download manager. Wait for the download to " + "finish (or cancel it), then load the model." + ), + ) + + # Unload any active Unsloth model only after every hub conflict check. + if unsloth_backend.active_model_name: + logger.info( + f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" + ) + await asyncio.to_thread( + unsloth_backend.unload_model, unsloth_backend.active_model_name + ) + # Route to HF or local mode based on config. Run in a thread so the # event loop stays free for progress polling and other requests # during the (potentially long) GGUF download + llama-server start. @@ -4639,6 +4719,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre logger.error(f"Error loading model: {e}", exc_info = True) msg = _maybe_unsupported_message(redacted_msg) raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}") + finally: + gguf_load_stack.close() def _requires_trust_remote_code_for_model( @@ -5392,6 +5474,10 @@ async def generate_stream( if chunk is _DONE: completed = True break + if isinstance(chunk, GenStreamError): + yield f"data: {json.dumps({'error': _friendly_gen_stream_error(chunk)})}\n\n" + yield "data: [DONE]\n\n" + return yield f"data: {json.dumps({'content': chunk})}\n\n" if completed: yield "data: [DONE]\n\n" @@ -5405,6 +5491,7 @@ async def generate_stream( backend.reset_generation_state() logger.error(f"Error during generation: {e}", exc_info = True) yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n" + yield "data: [DONE]\n\n" finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if not completed and not cancel_event.is_set(): @@ -7028,6 +7115,13 @@ async def openai_chat_completions( chunk_text = await asyncio.to_thread(next, gen, _DONE) if chunk_text is _DONE: break + if isinstance(chunk_text, GenStreamError): + _msg = _friendly_gen_stream_error(chunk_text) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) + return if chunk_text: api_monitor.append_reply(monitor_id, chunk_text) yield _chat_content_chunk( @@ -7043,8 +7137,11 @@ async def openai_chat_completions( raise except Exception as e: logger.error(f"Error during audio input streaming: {e}", exc_info = True) - api_monitor.fail(monitor_id, _friendly_error(e)) - yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" + _msg = _friendly_error(e) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) @@ -7061,7 +7158,15 @@ async def openai_chat_completions( ) else: try: - full_text = "".join(audio_input_generate()) + full_text = "" + for chunk_text in audio_input_generate(): + if isinstance(chunk_text, GenStreamError): + _msg = _friendly_gen_stream_error(chunk_text) + api_monitor.fail(monitor_id, _msg) + raise HTTPException(status_code = 500, detail = _msg) + full_text += chunk_text + except HTTPException: + raise except Exception as e: api_monitor.fail(monitor_id, _friendly_error(e)) raise @@ -8605,19 +8710,33 @@ async def openai_chat_completions( # Classify capability flags from the loaded template. _sf_model_info = backend.models.get(backend.active_model_name, {}) _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template") - _sf_features = _detect_safetensors_features(backend, _sf_tpl) - - # GGUF parity: enable_thinking templates prefill an unclosed ; split into - # reasoning_content deltas so the UI renders the block for safetensors and MLX. - _sf_parse_think = bool( - _sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on") + # Named templates may expose native reasoning only in their ``tool_use`` + # branch. Use a truthy placeholder for Studio-managed tools, whose concrete + # schemas are selected below, and the request schemas for client passthrough. + _sf_server_tool_intent = bool( + _effective_enable_tools(payload) or _explicit_studio_tool_loop_requested(payload) ) - # Prefilled-open only for prefill styles with thinking on; gpt-oss uses the normal mode. - _sf_reasoning_prefilled = _sf_reasoning_prefill_mode( - _sf_features, - payload.enable_thinking, - _sf_tpl, - reasoning_effort = payload.reasoning_effort, + _sf_template_tools = payload.tools if payload.tool_choice != "none" else None + if not _sf_template_tools and _sf_server_tool_intent: + _sf_template_tools = ({},) + + def _sf_response_protocol(tools = None): + features = _detect_safetensors_features(backend, _sf_tpl, tools = tools) + parse_think = bool( + features.get("supports_reasoning") or features.get("reasoning_always_on") + ) + reasoning_prefilled = _sf_reasoning_prefill_mode( + features, + payload.enable_thinking, + _sf_tpl, + reasoning_effort = payload.reasoning_effort, + ) + return features, parse_think, reasoning_prefilled + + # GGUF parity: split canonical output into reasoning_content. The + # selected template branch must match whether this request renders tools. + _sf_features, _sf_parse_think, _sf_reasoning_prefilled = _sf_response_protocol( + _sf_template_tools ) def _new_sf_reasoning_extractor(): @@ -8767,6 +8886,7 @@ async def openai_chat_completions( permission_mode = payload.permission_mode, use_adapter = payload.use_adapter, stats_holder = _sf_stats_holder, + reasoning_prefilled = _sf_reasoning_prefilled, ) _sf_tool_sentinel = object() @@ -8826,6 +8946,18 @@ async def openai_chat_completions( _sf_next_task = None if event is _sf_tool_sentinel: break + if isinstance(event, GenStreamError): + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(event) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) + return + if not isinstance(event, dict): + raise RuntimeError( + f"Invalid safetensors tool event: {type(event).__name__}" + ) if event["type"] == "heartbeat": # Tool-execution wrapper heartbeat -> SSE keepalive. @@ -8913,6 +9045,11 @@ async def openai_chat_completions( backend.reset_generation_state() api_monitor.finish(monitor_id, "cancelled") raise + except GenStreamErrorRaised as exc: + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse({"error": {"message": _msg, "type": "server_error"}}) except Exception: backend.reset_generation_state() # Generic wire message; full trace stays in the log (CWE-209: @@ -8962,6 +9099,15 @@ async def openai_chat_completions( for event in gen: if cancel_event.is_set(): break + if isinstance(event, GenStreamError): + raise HTTPException( + status_code = 500, + detail = _friendly_gen_stream_error(event), + ) + if not isinstance(event, dict): + raise RuntimeError( + f"Invalid safetensors tool event: {type(event).__name__}" + ) if event.get("type") == "content": full_text = _strip_tool_xml_for_display( event.get("text", ""), @@ -9002,6 +9148,15 @@ async def openai_chat_completions( backend.reset_generation_state() api_monitor.finish(monitor_id, "cancelled") raise + except GenStreamErrorRaised as exc: + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + raise HTTPException(status_code = 500, detail = _msg) + except HTTPException as exc: + backend.reset_generation_state() + api_monitor.fail(monitor_id, str(exc.detail)) + raise except Exception: backend.reset_generation_state() # CWE-209: generic detail; full trace in log. @@ -9088,6 +9243,12 @@ async def openai_chat_completions( else: gen_kwargs["tools"] = payload.tools + # The potential tool context above is needed before server/client routing is + # known. This standard path now has the exact schemas that will be rendered, + # so resolve reasoning parsing again to keep empty registries, forced-tool + # misses, and tool_choice="none" on the marker-free template branch. + _, _sf_parse_think, _sf_reasoning_prefilled = _sf_response_protocol(gen_kwargs.get("tools")) + # Request-scoped usage/timings receptacle (filled at gen_done). stats_holder: dict = {} @@ -9168,6 +9329,14 @@ async def openai_chat_completions( _next_task = None if cumulative is _DONE: break + if isinstance(cumulative, GenStreamError): + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(cumulative) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) + return if await request.is_disconnected(): cancel_event.set() backend.reset_generation_state() @@ -9274,6 +9443,13 @@ async def openai_chat_completions( backend.reset_generation_state() api_monitor.finish(monitor_id, "cancelled") raise + except GenStreamErrorRaised as exc: + # Adapter-controlled (compare-mode) backend failure. Honor the + # public flag so operational errors surface their real message. + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse({"error": {"message": _msg, "type": "server_error"}}) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI streaming: {e}", exc_info = True) @@ -9317,6 +9493,11 @@ async def openai_chat_completions( try: full_text = "" for token in generate(): + if isinstance(token, GenStreamError): + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(token) + api_monitor.fail(monitor_id, _msg) + raise HTTPException(status_code = 500, detail = _msg) full_text = token # Split prefilled reasoning (GGUF parity); also covers MLX via @@ -9415,6 +9596,15 @@ async def openai_chat_completions( api_monitor.finish(monitor_id) return _model_json_response(response) + except HTTPException: + raise + except GenStreamErrorRaised as exc: + # Adapter-controlled (compare-mode) backend failure. Honor the public + # flag so operational errors surface their real message. + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + raise HTTPException(status_code = 500, detail = _msg) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI completion: {e}", exc_info = True) diff --git a/studio/backend/run.py b/studio/backend/run.py index 56b9c78343..4f105b53b1 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1354,11 +1354,23 @@ def run_server( if secure: os.environ["UNSLOTH_SECURE"] = "1" - import nest_asyncio - - nest_asyncio.apply() - import asyncio + + # nest_asyncio is for Colab/IPython, where the main thread already runs a loop + # the blocking waits below would collide with. Apply it only with a loop running + # (a plain CLI start has nothing to nest) and only on Python <= 3.13: on 3.14+ + # its global Task patch leaves asyncio.current_task() None (tracking moved into + # C), which also breaks the background uvicorn loop and 500s every request. It + # is archived upstream, so no 3.14 fix is coming; skip it there. + if sys.version_info < (3, 14): + try: + asyncio.get_running_loop() + except RuntimeError: + pass + else: + import nest_asyncio + nest_asyncio.apply() + from threading import Thread, Event import uvicorn diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py new file mode 100644 index 0000000000..15d91cd324 --- /dev/null +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -0,0 +1,740 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for cached GGUF reuse and load/download exclusion. + +No GPU, network, or subprocesses are required. +""" + +from __future__ import annotations + +import asyncio +import sys +import threading +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub optional dependencies before importing the modules under test. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +sys.modules.setdefault("structlog", _structlog_stub) + +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + "HTTPStatusError", + ): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Request = type("Request", (), {}) + + class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + _httpx_stub.Timeout = _FakeTimeout + _httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, + ) + sys.modules.setdefault("httpx", _httpx_stub) + + +from huggingface_hub import constants as hf_constants + +from core.inference.llama_cpp import ( + LlamaCppBackend, + cached_gguf_for_load, + gguf_load_in_flight, + hf_gguf_load_in_flight, +) + + +REPO = "unsloth/gemma-test-GGUF" +VARIANT = "UD-Q4_K_XL" +MAIN = f"gemma-test-{VARIANT}.gguf" + + +def _build_cache( + root: Path, + repo_id: str, + files: dict[str, int], + *, + snapshot_sha: str = "a" * 40, +) -> Path: + """Create ``$root/models--/snapshots//`` for each entry.""" + repo_dir = root / f"models--{repo_id.replace('/', '--')}" + (repo_dir / "blobs").mkdir(parents = True, exist_ok = True) + snap = repo_dir / "snapshots" / snapshot_sha + snap.mkdir(parents = True, exist_ok = True) + for rel, size in files.items(): + full = snap / rel + full.parent.mkdir(parents = True, exist_ok = True) + full.write_bytes(b"\0" * size) + return snap + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + return tmp_path + + +def _fail_download(*_args, **_kwargs): + raise AssertionError("must reuse the cached GGUF instead of downloading") + + +def _fail_get_paths_info(*_args, **_kwargs): + raise AssertionError("cached reuse must return before the sizing preflight") + + +class TestLoadReusesCachedCopy: + def test_online_reuse_after_revision_bump(self, hf_cache): + """A new repo revision does not replace a complete cached model.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / MAIN) + + def test_reuse_size_check_uses_cached_snapshot_revision(self, hf_cache): + """Current-revision size changes do not invalidate an older complete copy.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + revisions: list[str | None] = [] + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + revisions.append(revision) + size = 4 if revision == snap.name else 8 + return [_types.SimpleNamespace(path = path, size = size) for path in paths] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / MAIN) + assert revisions == [snap.name] + + def test_reuse_when_cached_revision_vanished_from_hub(self, hf_cache): + """The Hub answers an unknown revision with an empty result, not an error.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", lambda *_a, **_k: []), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / MAIN) + + def test_truncated_cached_file_is_not_reused(self, hf_cache): + backend = LlamaCppBackend() + _build_cache(hf_cache, REPO, {MAIN: 4}) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 8) for path in paths] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert downloaded == [MAIN] + assert out == f"/fake/{REPO}/{MAIN}" + + def test_truncated_cached_split_shard_is_not_reused(self, hf_cache): + backend = LlamaCppBackend() + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + _build_cache(hf_cache, REPO, {shard1: 8, shard2: 4}) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 8) for path in paths] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert downloaded == [shard1, shard2] + assert out == f"/fake/{REPO}/{shard1}" + + def test_online_reuse_when_reupload_renamed_the_file(self, hf_cache): + """A renamed variant still reuses its cached file.""" + backend = LlamaCppBackend() + old_name = f"gemma-test-old-{VARIANT}.gguf" + snap = _build_cache(hf_cache, REPO, {old_name: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / old_name) + + def test_downloads_when_nothing_cached(self, hf_cache): + backend = LlamaCppBackend() + downloaded: list[str] = [] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert downloaded == [MAIN] + assert out == f"/fake/{REPO}/{MAIN}" + + def test_force_redownloads_despite_cache(self, hf_cache): + """A forced download ignores a complete cached copy.""" + backend = LlamaCppBackend() + _build_cache(hf_cache, REPO, {MAIN: 4}) + downloaded: list[str] = [] + + def fake_download( + repo_id, + filename, + token = None, + **kwargs, + ): + assert kwargs.get("force_download") is True + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT, force = True) + + assert downloaded == [MAIN] + assert out == f"/fake/{REPO}/{MAIN}" + + def test_split_reused_only_when_colocated(self, hf_cache): + backend = LlamaCppBackend() + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + snap = _build_cache(hf_cache, REPO, {shard1: 4, shard2: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / shard1) + + def test_partial_split_set_downloads(self, hf_cache): + """A partial split set is not reused.""" + backend = LlamaCppBackend() + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + _build_cache(hf_cache, REPO, {shard1: 4}) + downloaded: list[str] = [] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p is not None] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert downloaded == [shard1, shard2] + assert out == f"/fake/{REPO}/{shard1}" + + def test_reuse_prefers_newest_snapshot_after_update(self, hf_cache): + """Loads prefer the newest complete snapshot.""" + import os + + backend = LlamaCppBackend() + old_snap = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "a" * 40) + new_snap = _build_cache(hf_cache, REPO, {MAIN: 6}, snapshot_sha = "b" * 40) + os.utime(old_snap, (1_000_000, 1_000_000)) + os.utime(new_snap, (2_000_000, 2_000_000)) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(new_snap / MAIN) + + def test_low_disk_fallback_reuses_cached_copy(self, hf_cache): + backend = LlamaCppBackend() + fallback = "gemma-test-Q2_K.gguf" + snap = _build_cache(hf_cache, REPO, {fallback: 4}) + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + size = 4 if revision == snap.name else 100 + return [_types.SimpleNamespace(path = path, size = size) for path in paths] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("shutil.disk_usage", lambda *_a, **_k: _types.SimpleNamespace(free = 10)), + patch.object( + backend, + "_find_smallest_fitting_variant", + lambda *_a, **_k: (fallback, 4, []), + ), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / fallback) + + def test_companion_prefers_main_snapshot_sibling(self, hf_cache): + """A cached mmproj is reused from the main model's snapshot.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4, "mmproj-F16.gguf": 2}) + + def _fail_list(*_args, **_kwargs): + raise AssertionError("snapshot sibling must resolve without a repo listing") + + with patch("huggingface_hub.list_repo_files", _fail_list): + out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) + + assert out == str(snap / "mmproj-F16.gguf") + + def test_companion_finds_snapshot_through_hf_symlink(self, hf_cache): + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {}) + blobs = snap.parent.parent / "blobs" + main_blob = blobs / "main" + mmproj_blob = blobs / "mmproj" + main_blob.write_bytes(b"main") + mmproj_blob.write_bytes(b"mmproj") + try: + (snap / MAIN).symlink_to(main_blob) + (snap / "mmproj-F16.gguf").symlink_to(mmproj_blob) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + with patch("huggingface_hub.list_repo_files", _fail_download): + out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) + + assert out == str(snap / "mmproj-F16.gguf") + + def test_companion_does_not_download_during_hub_job(self, hf_cache): + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + registry = _types.SimpleNamespace(active_job_refs = lambda _repo: [object()]) + + with ( + patch("huggingface_hub.list_repo_files", _fail_download), + patch("hub.utils.download_registry.get_models_registry", lambda: registry), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) + + assert out is None + + +class TestCachedGgufForLoadProbe: + def test_complete_copy_found(self, hf_cache): + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN) + + def test_absent_copy_is_none(self, hf_cache): + assert cached_gguf_for_load(REPO, VARIANT) is None + + def test_partial_split_is_none(self, hf_cache): + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + _build_cache(hf_cache, REPO, {shard1: 4}) + assert cached_gguf_for_load(REPO, VARIANT) is None + + def test_partial_new_snapshot_does_not_hide_complete_split(self, hf_cache): + import os + + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + old = _build_cache( + hf_cache, + REPO, + {shard1: 4, shard2: 4}, + snapshot_sha = "a" * 40, + ) + new = _build_cache(hf_cache, REPO, {shard1: 4}, snapshot_sha = "b" * 40) + os.utime(old, (1_000_000, 1_000_000)) + os.utime(new, (2_000_000, 2_000_000)) + + assert cached_gguf_for_load(REPO, VARIANT) == str(old / shard1) + + def test_split_requires_every_declared_shard(self, hf_cache): + shard1 = f"gemma-test-{VARIANT}-00001-of-00003.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00003.gguf" + _build_cache(hf_cache, REPO, {shard1: 4, shard2: 4}) + + assert cached_gguf_for_load(REPO, VARIANT) is None + + def test_required_mmproj_must_share_main_snapshot(self, hf_cache): + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN) + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) is None + + (snap / "mmproj-F16.gguf").write_bytes(b"mmproj") + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(snap / MAIN) + + def test_required_mmproj_scans_past_newer_main_only_snapshot(self, hf_cache): + import os + + old = _build_cache( + hf_cache, + REPO, + {MAIN: 4, "mmproj-F16.gguf": 2}, + snapshot_sha = "a" * 40, + ) + new = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "b" * 40) + os.utime(old, (1_000_000, 1_000_000)) + os.utime(new, (2_000_000, 2_000_000)) + + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(old / MAIN) + + +class TestLoadHubDownloadExclusion: + def test_in_flight_marker_counts_and_normalizes_case(self): + assert not hf_gguf_load_in_flight(REPO) + with gguf_load_in_flight(REPO): + assert hf_gguf_load_in_flight(REPO.upper()) + with gguf_load_in_flight(REPO.lower()): + assert hf_gguf_load_in_flight(REPO) + assert hf_gguf_load_in_flight(REPO) + assert not hf_gguf_load_in_flight(REPO) + + def test_marker_noops_for_local_loads(self): + with gguf_load_in_flight(None): + assert not hf_gguf_load_in_flight("") + + def test_marker_cleared_on_exception(self): + with pytest.raises(RuntimeError): + with gguf_load_in_flight(REPO): + raise RuntimeError("boom") + assert not hf_gguf_load_in_flight(REPO) + + def test_hub_download_refused_while_load_in_flight(self): + from fastapi import HTTPException + + from hub.schemas.downloads import DownloadModelRequest + from hub.services.models import downloads as dl + + body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT) + with ( + patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id), + gguf_load_in_flight(REPO), + ): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(dl.download_model_response(body)) + + assert exc_info.value.status_code == 409 + assert "load" in exc_info.value.detail.lower() + + def test_hub_download_rechecks_marker_before_claim(self): + from fastapi import HTTPException + + from hub.schemas.downloads import DownloadModelRequest + from hub.services.models import downloads as dl + + scope = None + + def mark_load(*_args, **_kwargs): + nonlocal scope + if scope is None: + scope = gguf_load_in_flight(REPO) + scope.__enter__() + return frozenset() + + class _Registry: + def claim(self, *_args, admission_check, **_kwargs): + assert admission_check() is False + return False, "admission_blocked" + + def current_generation(self, _key): + return 0 + + registry = _Registry() + body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT) + try: + with ( + patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id), + patch.object(dl.gguf_variants, "gguf_variant_blob_hashes", mark_load), + patch.object(dl, "_registry", registry), + ): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(dl.download_model_response(body)) + finally: + if scope is not None: + scope.__exit__(None, None, None) + + assert exc_info.value.status_code == 409 + + def test_registry_admission_check_prevents_claim(self): + from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP + + registry = DownloadRegistry() + claimed, state = registry.claim( + f"{REPO}::{VARIANT}", + TRANSPORT_HTTP, + repo_type = "model", + repo_id = REPO, + variant = VARIANT, + admission_check = lambda: False, + ) + + assert claimed is False + assert state == "admission_blocked" + assert registry.active_jobs(REPO) == {} + + def test_same_variant_job_stays_visible_during_retry_handoff(self): + from hub.utils.download_registry import DownloadRegistry, TRANSPORT_XET + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + + registry = DownloadRegistry() + key = f"{REPO}::{VARIANT}" + claimed, _ = registry.claim( + key, + TRANSPORT_XET, + repo_type = "model", + repo_id = REPO, + variant = VARIANT, + ) + assert claimed is True + assert registry.has_active_variant(REPO, VARIANT.lower()) is True + + registry.release_active_slot(key) + + assert registry.active_jobs(REPO) == {} + assert registry.active_job_refs(REPO) + assert registry.has_active_variant(REPO, VARIANT) is True + with ( + patch("hub.utils.download_registry.get_models_registry", lambda: registry), + patch( + "core.inference.llama_cpp.cached_gguf_for_load", + side_effect = AssertionError("same-variant jobs must block before cache reuse"), + ), + ): + assert _hub_download_blocks_gguf_load(REPO, VARIANT) is True + + registry.set_job(key, "complete") + assert registry.has_active_variant(REPO, VARIANT) is False + + def test_other_variant_job_still_allows_complete_cached_load(self): + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP + + registry = DownloadRegistry() + registry.claim( + f"{REPO}::Q8_0", + TRANSPORT_HTTP, + repo_type = "model", + repo_id = REPO, + variant = "Q8_0", + ) + with ( + patch("hub.utils.download_registry.get_models_registry", lambda: registry), + patch( + "core.inference.llama_cpp.cached_gguf_for_load", + return_value = "/cached/model.gguf", + ) as cached_probe, + ): + assert _hub_download_blocks_gguf_load(REPO, VARIANT) is False + + cached_probe.assert_called_once_with( + REPO, + VARIANT, + require_mmproj = False, + verify_sizes = True, + hf_token = None, + ) + + def test_cancelled_request_keeps_marker_until_load_thread_finishes(self): + from core.inference.llama_cpp import _with_gguf_load_marker + + started = threading.Event() + release = threading.Event() + finished = threading.Event() + + class FakeBackend: + @_with_gguf_load_marker + def load_model(self, *, hf_repo): + started.set() + release.wait(timeout = 2) + finished.set() + return True + + async def scenario(): + with patch( + "core.inference.llama_cpp._hub_download_blocks_gguf_load", + return_value = False, + ): + task = asyncio.create_task( + asyncio.to_thread(FakeBackend().load_model, hf_repo = REPO) + ) + assert await asyncio.to_thread(started.wait, 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert hf_gguf_load_in_flight(REPO) + + release.set() + assert await asyncio.to_thread(finished.wait, 1) + for _ in range(100): + if not hf_gguf_load_in_flight(REPO): + break + await asyncio.sleep(0.001) + assert not hf_gguf_load_in_flight(REPO) + + asyncio.run(scenario()) + + def test_load_marker_precedes_hub_guard_and_unload(self): + source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() + gguf_branch = source[source.index("if config.is_gguf:") :] + + assert ( + gguf_branch.index("enter_context(gguf_load_in_flight") + < gguf_branch.index("if request.llama_extra_args is None") + < gguf_branch.index("_hub_download_blocks_gguf_load") + < gguf_branch.index("unsloth_backend.unload_model") + ) + llama_source = ( + Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" + ).read_text() + assert "@_with_gguf_load_marker\n def load_model(" in llama_source diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index c3161c5714..bd2c008589 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -1061,6 +1061,80 @@ def test_same_turn_duplicate_web_search_is_internal_noop(monkeypatch): ] +def test_same_turn_duplicate_does_not_drop_later_parallel_call(monkeypatch): + # One batch: search(a), search(a) [duplicate], search(b). The duplicate is an + # internal no-op, but the distinct search(b) after it must still run, and the + # no-op nudge must land after the tool results rather than splitting them. + batch = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_a1", + "type": "function", + "function": {"name": "web_search", "arguments": json.dumps({"query": "a"})}, + }, + { + "index": 1, + "id": "call_a2", + "type": "function", + "function": {"name": "web_search", "arguments": json.dumps({"query": "a"})}, + }, + { + "index": 2, + "id": "call_b", + "type": "function", + "function": {"name": "web_search", "arguments": json.dumps({"query": "b"})}, + }, + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [batch, final_stream], payloads) + + calls: list[dict] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append(arguments) + return "search-result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 3, + ) + ) + + # Both distinct calls ran; the duplicate did not (old `break` dropped search(b)). + assert calls == [{"query": "a"}, {"query": "b"}] + assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [ + "call_a1", + "call_b", + ] + + # The next generation's conversation must be well-formed: the assistant lists + # only the executed calls (no orphan for the duplicate), the two tool results + # follow contiguously, and the no-op nudge lands after them, never between. + conv = payloads[1]["messages"] + asst = next(m for m in conv if m["role"] == "assistant" and m.get("tool_calls")) + assert [tc.get("id") for tc in asst["tool_calls"]] == ["call_a1", "call_b"] + after = conv[conv.index(asst) + 1 :] + assert [m["role"] for m in after[:2]] == ["tool", "tool"] + assert [m.get("tool_call_id") for m in after[:2]] == ["call_a1", "call_b"] + assert after[2]["role"] == "user" # deferred duplicate nudge, after the results + assert after[2]["content"].startswith( + "One earlier request to call tool 'web_search' in this batch was not executed" + ) + assert "previous tool request" not in after[2]["content"].lower() + + def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(monkeypatch): same_turn_render_calls = [ _sse( diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py index 3c21f41701..d0213f6079 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py @@ -373,6 +373,7 @@ def test_kill_orphaned_servers_returns_count(): with ( patch.dict(sys.modules, {"psutil": fake_psutil}), patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}), + patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), ): n = LlamaCppBackend._kill_orphaned_servers() assert n == 1, "only the Studio-owned orphan should be counted" @@ -384,11 +385,53 @@ def test_kill_orphaned_servers_returns_count(): with ( patch.dict(sys.modules, {"psutil": fake_psutil}), patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}), + patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), ): assert LlamaCppBackend._kill_orphaned_servers() == 0 assert killed == [] +def test_kill_orphaned_servers_spares_live_parent(): + """A Studio-owned llama-server whose parent is still running is not an + orphan (a live Studio or the user's shell owns it) and must never be + killed; only the true orphan (parent gone) is reaped.""" + import os + + mypid = os.getpid() + fake_path = "/tmp/unsloth-test-llama/llama-server" + killed: list[int] = [] + + class _FakeProc: + def __init__(self, pid, name, exe): + self.info = {"pid": pid, "name": name, "exe": exe} + + def kill(self): + killed.append(self.info["pid"]) + + live_parent = _FakeProc(mypid + 1, "llama-server", fake_path) + true_orphan = _FakeProc(mypid + 2, "llama-server", fake_path) + + fake_psutil = _types.ModuleType("psutil") + fake_psutil.NoSuchProcess = type("NoSuchProcess", (Exception,), {}) + fake_psutil.AccessDenied = type("AccessDenied", (Exception,), {}) + fake_psutil.ZombieProcess = type("ZombieProcess", (Exception,), {}) + fake_psutil.process_iter = lambda attrs = None: [live_parent, true_orphan] + + with ( + patch.dict(sys.modules, {"psutil": fake_psutil}), + patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}), + patch.object(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0)), + patch.object( + LlamaCppBackend, + "_pid_parent_is_alive", + staticmethod(lambda pid: pid == mypid + 1), + ), + ): + n = LlamaCppBackend._kill_orphaned_servers() + assert n == 1, "only the true orphan should be reaped" + assert killed == [mypid + 2], "the live-parent server must be spared" + + def test_startup_reaper_arms_settle_timestamp(): """__init__ arms ``_last_kill_monotonic`` when the startup reaper kills an orphan (so the first load_model waits for VRAM to settle), and leaves the diff --git a/studio/backend/tests/test_mcp_server.py b/studio/backend/tests/test_mcp_server.py new file mode 100644 index 0000000000..71792605ae --- /dev/null +++ b/studio/backend/tests/test_mcp_server.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import asyncio +import sys +import types + +import pytest + +from mcp_server import BearerTokenMiddleware, _clamp, _dump, create_studio_mcp + + +def _get_tool(name): + tools = asyncio.run(create_studio_mcp().list_tools()) + return {tool.name: tool for tool in tools}[name] + + +def test_studio_mcp_registers_control_plane_tools(): + tools = asyncio.run(create_studio_mcp().list_tools()) + + assert {tool.name for tool in tools} == { + "studio_status", + "list_local_models", + "get_training_status", + "start_training", + "stop_training", + "list_training_runs", + "validate_recipe", + "get_recipe_job_status", + "get_recipe_job_dataset", + "load_checkpoint", + "export_gguf", + } + + +def test_dump_serializes_pydantic_values(): + class Response: + def model_dump(self, *, mode): + assert mode == "json" + return {"ok": True} + + assert _dump(Response()) == {"ok": True} + assert _dump({"already": "json"}) == {"already": "json"} + + +def test_bearer_token_middleware_rejects_wrong_token(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "http", "headers": [(b"authorization", b"Bearer wrong")]}, + None, + send, + ) + ) + + assert events[0]["status"] == 401 + assert "app" not in events + + +def test_bearer_token_middleware_closes_unauthorized_websocket(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "websocket", "headers": []}, + None, + send, + ) + ) + + assert events == [{"type": "websocket.close", "code": 4401}] + + +def test_bearer_token_middleware_rejects_non_ascii_authorization(): + # A non-ASCII bearer value must produce a clean 401, not a 500. Comparing on + # bytes avoids the str hmac.compare_digest TypeError on non-ASCII input. + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "http", "headers": [(b"authorization", b"Bearer \xff\xff")]}, + None, + send, + ) + ) + + assert events[0]["status"] == 401 + assert "app" not in events + + +def test_bearer_token_middleware_accepts_correct_token(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "http", "headers": [(b"authorization", b"Bearer secret")]}, + None, + send, + ) + ) + + assert events == ["app"] + + +def test_bearer_token_middleware_requires_non_empty_token(): + async def app(scope, receive, send): + pass + + for bad in ("", " "): + with pytest.raises(ValueError): + BearerTokenMiddleware(app, bad) + + +def test_bearer_token_middleware_rejects_non_ascii_token(): + async def app(scope, receive, send): + pass + + # non-ASCII tokens cannot be transmitted in an HTTP header by a standard + # client, so they are rejected at construction instead of locking out. + for bad in ("töken", "\U0001f600"): + with pytest.raises(ValueError): + BearerTokenMiddleware(app, bad) + + +def test_bearer_token_middleware_passes_through_non_http_scopes(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run(middleware({"type": "lifespan"}, None, send)) + + assert events == ["app"] + + +def test_clamp_restricts_to_inclusive_bounds(): + assert _clamp(5, 1, 200) == 5 + assert _clamp(-10, 1, 200) == 1 + assert _clamp(10_000, 1, 200) == 200 + assert _clamp(0, 1, 500) == 1 + assert _clamp(1_000, 1, 500) == 500 + + +def test_export_and_checkpoint_tools_expose_forwarded_fields(): + export_props = set(_get_tool("export_gguf").parameters["properties"]) + assert {"hf_token", "imatrix", "imatrix_path"} <= export_props + + checkpoint_props = set(_get_tool("load_checkpoint").parameters["properties"]) + assert {"hf_token", "approved_remote_code_fingerprint"} <= checkpoint_props + + +def _stub_module(monkeypatch, name, **attrs): + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + if "." in name: + module.__path__ = [] # mark package-like so submodule imports resolve + monkeypatch.setitem(sys.modules, name, module) + return module + + +def test_export_gguf_forwards_hf_token_and_imatrix(monkeypatch): + captured = {} + + class FakeExportGGUFRequest: + def __init__(self, **kwargs): + captured.update(kwargs) + + async def fake_export(request, current_subject): + return {"current_subject": current_subject} + + _stub_module(monkeypatch, "models", ExportGGUFRequest = FakeExportGGUFRequest) + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.export", export_gguf = fake_export) + + tool = _get_tool("export_gguf") + result = asyncio.run( + tool.fn( + save_directory = "/tmp/out", + quantization_method = ["Q4_K_M", "Q8_0"], + push_to_hub = True, + repo_id = "me/model", + hf_token = "hf_secret", + imatrix = True, + imatrix_path = "/tmp/imatrix.dat", + ) + ) + + assert captured["hf_token"] == "hf_secret" + assert captured["imatrix"] is True + assert captured["imatrix_path"] == "/tmp/imatrix.dat" + assert captured["quantization_method"] == ["Q4_K_M", "Q8_0"] + assert result["current_subject"] == "mcp" + + +def test_load_checkpoint_forwards_token_and_fingerprint(monkeypatch): + captured = {} + + class FakeLoadCheckpointRequest: + def __init__(self, **kwargs): + captured.update(kwargs) + + async def fake_load(request, current_subject): + return {"current_subject": current_subject} + + _stub_module(monkeypatch, "models", LoadCheckpointRequest = FakeLoadCheckpointRequest) + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.export", load_checkpoint = fake_load) + + tool = _get_tool("load_checkpoint") + asyncio.run( + tool.fn( + checkpoint_path = "/tmp/ckpt", + approved_remote_code_fingerprint = "sha256:abc", + hf_token = "hf_secret", + ) + ) + + assert captured["hf_token"] == "hf_secret" + assert captured["approved_remote_code_fingerprint"] == "sha256:abc" + + +def test_list_training_runs_clamps_pagination(monkeypatch): + captured = {} + + async def fake_list_runs(limit, offset, current_subject): + captured["limit"] = limit + captured["offset"] = offset + return {"ok": True} + + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.training_history", list_training_runs = fake_list_runs) + + tool = _get_tool("list_training_runs") + asyncio.run(tool.fn(limit = 10_000, offset = -5)) + + assert captured["limit"] == 200 + assert captured["offset"] == 0 + + +def test_get_recipe_job_dataset_clamps_pagination(monkeypatch): + captured = {} + + def fake_job_dataset(job_id, limit, offset): + captured["limit"] = limit + captured["offset"] = offset + return {"ok": True} + + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.data_recipe") + _stub_module(monkeypatch, "routes.data_recipe.jobs", job_dataset = fake_job_dataset) + + tool = _get_tool("get_recipe_job_dataset") # this tool is synchronous + tool.fn(job_id = "job-1", limit = -1, offset = -9) + + assert captured["limit"] == 1 + assert captured["offset"] == 0 diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 29fbb45158..7b8aefb722 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -2,6 +2,7 @@ import sys import types +from contextlib import contextmanager from types import SimpleNamespace import pytest @@ -40,12 +41,16 @@ class _DummyModel: def _install_fake_mlx(monkeypatch): mlx_pkg = types.ModuleType("mlx") mlx_core = types.ModuleType("mlx.core") + mlx_utils = types.ModuleType("mlx.utils") mlx_core.metal = _DummyMetal() mlx_core.set_wired_limit = _DummyMX.set_wired_limit mlx_core.device_info = _DummyMX.device_info + mlx_utils.tree_unflatten = dict mlx_pkg.core = mlx_core + mlx_pkg.utils = mlx_utils monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + monkeypatch.setitem(sys.modules, "mlx.utils", mlx_utils) def _install_fake_fast_mlx(monkeypatch, calls): @@ -68,6 +73,99 @@ def _install_fake_fast_mlx(monkeypatch, calls): monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.loader", mlx_loader) +class _AdapterTree: + def __init__(self, modules): + self.modules = dict(modules) + + def named_modules(self): + return list(self.modules.items()) + + def update_modules(self, modules): + self.modules.update(modules) + + +def test_temporary_mlx_adapter_state_bypasses_and_restores_wrappers(monkeypatch): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import _temporary_mlx_adapter_state + + base = object() + wrapper = SimpleNamespace(lora_a = object(), lora_b = object(), linear = base, m = object()) + model = _AdapterTree({"model.layers.0.proj": wrapper}) + + with pytest.raises(RuntimeError, match = "generation failed"): + with _temporary_mlx_adapter_state(model, False): + assert model.modules["model.layers.0.proj"] is base + raise RuntimeError("generation failed") + assert model.modules["model.layers.0.proj"] is wrapper + + +def test_temporary_mlx_adapter_state_validates_requests(): + from core.inference.mlx_inference import _temporary_mlx_adapter_state + + wrapper = SimpleNamespace(lora_a = object(), lora_b = object(), embedding = object()) + model = _AdapterTree({"embed_tokens": wrapper}) + with _temporary_mlx_adapter_state(model, True): + assert model.modules["embed_tokens"] is wrapper + with pytest.raises(NotImplementedError, match = "named adapter"): + with _temporary_mlx_adapter_state(model, "other"): + pass + + base_model = _AdapterTree({"proj": object()}) + with _temporary_mlx_adapter_state(base_model, None): + pass + with _temporary_mlx_adapter_state(base_model, True): + pass + + unsupported = _AdapterTree({"proj": SimpleNamespace(lora_a = object(), lora_b = object())}) + with _temporary_mlx_adapter_state(unsupported, True): + pass + with pytest.raises(RuntimeError, match = "without their base modules"): + with _temporary_mlx_adapter_state(unsupported, False): + pass + + +def test_temporary_mlx_adapter_state_uses_real_mlx_module_tree(): + nn = pytest.importorskip("mlx.nn") + pytest.importorskip("mlx_lm") + from mlx_lm.models.switch_layers import SwitchLinear + from mlx_lm.tuner.dora import DoRALinear + from mlx_lm.tuner.lora import LoRAEmbedding, LoRALinear, LoRASwitchLinear + + from core.inference.mlx_inference import _temporary_mlx_adapter_state + + class _Layer(nn.Module): + def __init__(self): + super().__init__() + quantized = nn.QuantizedLinear.from_linear(nn.Linear(32, 32), group_size = 32, bits = 4) + self.quantized_proj = LoRALinear.from_base(quantized) + self.dora_proj = DoRALinear.from_base(nn.Linear(4, 4)) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.layers = [_Layer()] + self.embed_tokens = LoRAEmbedding.from_base(nn.Embedding(16, 4)) + self.experts = LoRASwitchLinear.from_base(SwitchLinear(4, 4, 2)) + + model = _Model() + wrappers = { + path: module + for path, module in model.named_modules() + if hasattr(module, "lora_a") and hasattr(module, "lora_b") + } + bases = { + path: getattr(module, "linear", getattr(module, "embedding", None)) + for path, module in wrappers.items() + } + + with _temporary_mlx_adapter_state(model, False): + live = dict(model.named_modules()) + assert all(live[path] is base for path, base in bases.items()) + + restored = dict(model.named_modules()) + assert all(restored[path] is wrapper for path, wrapper in wrappers.items()) + + def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): _install_fake_mlx(monkeypatch) calls = [] @@ -333,10 +431,87 @@ def test_mlx_generate_chat_response_accepts_template_kwargs(): ), f"{name!r} must default to None so existing callers stay valid" +def test_mlx_vlm_reemits_think_prefill_inside_adapter_context(monkeypatch): + """A prefilled block must be re-emitted as the first VLM snapshot, + inside the adapter context (so unsupported requests still raise first), so + the UI renders the thinking block during prefill and a pre-first-token + cancel does not drop it. Mirrors _generate_text.""" + from core.inference import mlx_inference + + MLXInferenceBackend = mlx_inference.MLXInferenceBackend + + order = [] + + @contextmanager + def _adapter_state(_model, state): + assert backend._generation_lock.locked() + order.append("adapter_enter") + try: + yield + finally: + order.append("adapter_exit") + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state) + monkeypatch.setattr( + "core.inference.chat_template_helpers.detect_think_prefill", + lambda *_a, **_k: "\n", + ) + + prompt_utils = SimpleNamespace( + MODEL_CONFIG = {"deepseek_vl_v2": object()}, + apply_chat_template = lambda *_a, **_k: " model-aware", + ) + mlx_vlm = types.ModuleType("mlx_vlm") + mlx_vlm.prompt_utils = prompt_utils + + def _vlm_stream(*_a, **_k): + # The prefill must have been emitted before any generated token. + assert order[-1] == "adapter_enter" + yield SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1) + + mlx_vlm.stream_generate = _vlm_stream + monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm) + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda _t, _m, **_k: " model-aware", + ) + + backend = MLXInferenceBackend() + backend._model = SimpleNamespace(config = {"model_type": "deepseek_vl_v2"}) + backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) + args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) + + gen = backend._generate_vlm(*args, _adapter_state = False) + # First snapshot is the prefill alone, emitted after entering the adapter context. + assert next(gen) == "\n" + assert order == ["adapter_enter"] + # Subsequent snapshots are cumulative (prefill + generated text). + assert next(gen) == "\nok" + gen.close() + assert order == ["adapter_enter", "adapter_exit"] + + def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference import mlx_inference + + MLXInferenceBackend = mlx_inference.MLXInferenceBackend calls = {"generic": [], "model": [], "stream": []} + adapter_events = [] + adapter_active = {"value": False} + + @contextmanager + def _adapter_state(_model, state): + assert backend._generation_lock.locked() + adapter_events.append(("enter", state)) + adapter_active["value"] = True + try: + yield + finally: + adapter_active["value"] = False + adapter_events.append(("exit", state)) + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state) state = {"generic": "serialized", "model": " model-aware"} prompt_utils = SimpleNamespace( MODEL_CONFIG = {"deepseek_vl_v2": object()}, @@ -346,10 +521,13 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): ) mlx_vlm = types.ModuleType("mlx_vlm") mlx_vlm.prompt_utils = prompt_utils - mlx_vlm.stream_generate = lambda *_args, **kwargs: ( - calls["stream"].append((_args, kwargs)) - or iter([SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1)]) - ) + + def _vlm_stream(*args, **kwargs): + assert adapter_active["value"] + calls["stream"].append((args, kwargs)) + yield SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1) + + mlx_vlm.stream_generate = _vlm_stream monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm) def generic(_target, _messages, **kwargs): @@ -369,7 +547,11 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) tools = [{"function": {"name": "search"}}] - assert list(backend._generate_vlm(*args)) == ["ok"] + generator = backend._generate_vlm(*args, _adapter_state = False) + assert next(generator) == "ok" + assert adapter_active["value"] and backend._generation_lock.locked() + generator.close() + assert adapter_events == [("enter", False), ("exit", False)] assert calls["model"][0]["num_images"] == 1 assert calls["stream"][0][0][2] == " model-aware" with pytest.raises(RuntimeError, match = "dropping requested tools"): @@ -449,7 +631,10 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): """Mac text path must route through apply_chat_template_for_generation so reasoning / tool kwargs reach the tokenizer.""" _install_fake_mlx(monkeypatch) - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference import mlx_inference + + MLXInferenceBackend = mlx_inference.MLXInferenceBackend + real_adapter_state = mlx_inference._temporary_mlx_adapter_state # The text path renders once with tools, then the native-template fallback makes a second no- # tools probe call (tools=None) to detect whether the template dropped the schema. @@ -474,11 +659,31 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): mlx_lm_sample.make_sampler = lambda **_kw: object() mlx_lm_sample.make_logits_processors = lambda **_kw: None + adapter_events = [] + adapter_active = {"value": False} + stream_state = {"fail": False} + + @contextmanager + def _adapter_state(_model, state): + assert backend._generation_lock.locked() + adapter_events.append(("enter", state)) + adapter_active["value"] = True + try: + yield + finally: + adapter_active["value"] = False + adapter_events.append(("exit", state)) + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state) + class _Resp: def __init__(self, tok): self.token = tok def _stream_generate(_model, _tokenizer, **_kw): + assert adapter_active["value"] + if stream_state["fail"]: + raise RuntimeError("generation failed") yield _Resp(1) mlx_lm_pkg.stream_generate = _stream_generate @@ -500,17 +705,45 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): backend._tokenizer = _Tok() backend._is_vlm = False - out = list( - backend.generate_chat_response( - messages = [{"role": "user", "content": "ping"}], - tools = [{"function": {"name": "web_search"}}], - enable_thinking = True, - reasoning_effort = "medium", - preserve_thinking = True, - max_new_tokens = 1, - ) + generator = backend.generate_with_adapter_control( + use_adapter = False, + messages = [{"role": "user", "content": "ping"}], + tools = [{"function": {"name": "web_search"}}], + enable_thinking = True, + reasoning_effort = "medium", + preserve_thinking = True, + max_new_tokens = 1, ) - assert out == ["hi"] + assert next(generator) == "hi" + assert adapter_active["value"] and backend._generation_lock.locked() + generator.close() + assert adapter_events == [("enter", False), ("exit", False)] + stream_state["fail"] = True + with pytest.raises(RuntimeError, match = "generation failed"): + list( + backend.generate_with_adapter_control( + use_adapter = False, + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 1, + ) + ) + assert adapter_events[-2:] == [("enter", False), ("exit", False)] + assert not backend._generation_lock.locked() + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", real_adapter_state) + monkeypatch.setattr( + "core.inference.chat_template_helpers.detect_think_prefill", + lambda *_args, **_kwargs: "", + ) + stream_state["fail"] = False + named = backend.generate_with_adapter_control( + use_adapter = "named", + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 1, + ) + with pytest.raises(NotImplementedError, match = "named adapter"): + next(named) + assert not adapter_active["value"] and not backend._generation_lock.locked() # The toggled kwargs must reach the chat-template helper on the real render # (one of the calls carries the tools; the fallback probe passes tools=None). tool_renders = [ @@ -523,3 +756,169 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): assert render["kwargs"]["enable_thinking"] is True assert render["kwargs"]["reasoning_effort"] == "medium" assert render["kwargs"]["preserve_thinking"] is True + + +def test_mlx_text_normalizes_native_reasoning_and_close_releases_lock(monkeypatch): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import MLXInferenceBackend + + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda *_args, **_kwargs: "prompt", + raising = True, + ) + monkeypatch.setattr( + "core.inference.chat_template_helpers.render_with_native_template_fallback", + lambda formatted_prompt, **_kwargs: SimpleNamespace( + prompt = formatted_prompt, + reasoning_channel_markers = ("<|channel>thought\n", ""), + ), + raising = True, + ) + + mlx_lm_pkg = types.ModuleType("mlx_lm") + mlx_lm_sample = types.ModuleType("mlx_lm.sample_utils") + mlx_lm_sample.make_sampler = lambda **_kw: object() + mlx_lm_sample.make_logits_processors = lambda **_kw: None + + class _Resp: + def __init__(self, text, tok): + self.text = text + self.token = tok + + def _stream_generate(_model, _tokenizer, **_kw): + yield _Resp("<|channel>thought\n", 10) + yield _Resp("r", 11) + yield _Resp("", 12) + yield _Resp("a", 13) + + mlx_lm_pkg.stream_generate = _stream_generate + monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg) + monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample) + + backend = MLXInferenceBackend() + backend._model = object() + backend._tokenizer = SimpleNamespace(all_special_tokens = []) + backend._is_vlm = False + + assert list( + backend.generate_chat_response( + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 4, + ) + ) == ["", "r", "r", "ra"] + + gen = backend.generate_chat_response( + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 4, + ) + assert next(gen) == "" + assert backend._generation_lock.locked() + gen.close() + assert not backend._generation_lock.locked() + + +def test_mlx_text_native_metadata_preserves_prefilled_think_snapshots(monkeypatch): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import MLXInferenceBackend + + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda *_args, **_kwargs: "prompt\n", + raising = True, + ) + monkeypatch.setattr( + "core.inference.chat_template_helpers.render_with_native_template_fallback", + lambda formatted_prompt, **_kwargs: SimpleNamespace( + prompt = formatted_prompt, + reasoning_channel_markers = ("<|channel>thought", ""), + ), + raising = True, + ) + + mlx_lm_pkg = types.ModuleType("mlx_lm") + mlx_lm_sample = types.ModuleType("mlx_lm.sample_utils") + mlx_lm_sample.make_sampler = lambda **_kw: object() + mlx_lm_sample.make_logits_processors = lambda **_kw: None + + class _Resp: + def __init__(self, text, tok): + self.text = text + self.token = tok + + def _stream_generate(_model, _tokenizer, **_kw): + yield _Resp("reason", 10) + yield _Resp("", 11) + yield _Resp("answer", 12) + + mlx_lm_pkg.stream_generate = _stream_generate + monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg) + monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample) + + backend = MLXInferenceBackend() + backend._model = object() + backend._tokenizer = SimpleNamespace(all_special_tokens = []) + backend._is_vlm = False + + snapshots = list( + backend.generate_chat_response( + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 3, + ) + ) + assert snapshots == [ + "\n", + "\nreason", + "\nreason", + "\nreasonanswer", + ] + assert all(current.startswith(previous) for previous, current in zip(snapshots, snapshots[1:])) + + +def test_mlx_vlm_normalizes_native_reasoning_channels(monkeypatch): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import MLXInferenceBackend + + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda *_args, **_kwargs: "prompt", + raising = True, + ) + + mlx_vlm_pkg = types.ModuleType("mlx_vlm") + + class _Resp: + def __init__(self, text, tok): + self.text = text + self.token = tok + + def _stream_generate(_model, _processor, _prompt, _images, **_kw): + yield _Resp("<|channel>thought\n", 10) + yield _Resp("vision", 11) + yield _Resp("", 12) + yield _Resp(" answer", 13) + + mlx_vlm_pkg.stream_generate = _stream_generate + monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm_pkg) + + backend = MLXInferenceBackend() + backend._model = SimpleNamespace(config = SimpleNamespace()) + backend._processor = SimpleNamespace( + chat_template = "<|channel>thought\n...", + all_special_tokens = [], + apply_chat_template = lambda *_args, **_kwargs: "prompt", + ) + backend._is_vlm = True + + assert list( + backend.generate_chat_response( + messages = [{"role": "user", "content": "describe"}], + image = object(), + max_new_tokens = 4, + ) + ) == [ + "", + "vision", + "vision", + "vision answer", + ] diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 86e528ae67..02230632b6 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -328,6 +328,7 @@ def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch): pick, label, cancel_event = None, + near_path = None, ): captured["pick"] = pick return None @@ -428,6 +429,32 @@ def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch assert got is not None and Path(got).parent.parent.name == "newest" +def test_download_mtp_prefers_main_snapshot_offline(tmp_path, monkeypatch): + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snapshots = tmp_path / "models--unsloth--gemma" / "snapshots" + old = snapshots / "old" + new = snapshots / "new" + old.mkdir(parents = True) + new.mkdir(parents = True) + main = old / "gemma-UD-Q4_K_XL.gguf" + old_drafter = old / "mtp-gemma.gguf" + new_drafter = new / "mtp-gemma.gguf" + main.write_bytes(b"main") + old_drafter.write_bytes(b"old") + new_drafter.write_bytes(b"new") + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda _repo: [new, old]) + + got = LlamaCppBackend()._download_mtp( + hf_repo = "unsloth/gemma-GGUF", + near_path = str(main), + ) + + assert got == str(old_drafter) + + def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch): # Online, do not reuse a cached copy: go to the download path so a changed # drafter is refetched (hf_hub_download checks the current revision). @@ -447,6 +474,7 @@ def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch): pick, label, cancel_event = None, + near_path = None, ): reached["hit"] = True return None diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index e24e2ca451..0b9a1e704f 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -244,9 +244,7 @@ class TestGgufVariantFileResolution: def test_download_reuses_older_snapshot_when_current_ref_snapshot_is_partial( self, monkeypatch, hf_cache ): - # Cross-snapshot reuse is an offline-resilience path: online, hf_hub_download - # resumes the partial current-ref download and revalidates the revision instead - # of serving an older snapshot's same-name blob. + # Keep coverage for offline reuse; online reuse is tested separately. monkeypatch.setenv("HF_HUB_OFFLINE", "1") backend = LlamaCppBackend() repo = "unsloth/vision-GGUF" @@ -292,8 +290,7 @@ class TestGgufVariantFileResolution: def test_download_reuses_cached_gguf_when_lowercase_partial_cache_shadows_it( self, monkeypatch, hf_cache ): - # Case-variant cross-dir reuse is offline-only; online the canonical repo id - # resolves up front and hf_hub_download fetches the current revision. + # Keep coverage for case-insensitive offline cache lookup. monkeypatch.setenv("HF_HUB_OFFLINE", "1") backend = LlamaCppBackend() canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" @@ -348,45 +345,26 @@ class TestGgufVariantFileResolution: assert out == str(snap / gguf_file) assert seen_repos - def test_download_online_does_not_reuse_old_snapshot(self, monkeypatch, hf_cache): - # Online, an older same-name snapshot must not be served (it may be a stale - # revision); hf_hub_download is called so the current revision is fetched and - # its etag revalidated. + def test_download_online_reuses_complete_cached_snapshot(self, monkeypatch, hf_cache): + # Loads reuse complete cached models across repo revisions. monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) backend = LlamaCppBackend() repo = "unsloth/vision-GGUF" - _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) - downloaded: list[str] = [] + snap = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) - def fake_get_paths_info( - _repo_id, - paths, - token = None, - ): - return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] - - def fake_download( - repo_id, - filename, - token = None, - **kwargs, - ): - downloaded.append(filename) - return f"/fresh/{filename}" + def fail_download(*_args, **_kwargs): + raise AssertionError("must reuse the cached GGUF instead of downloading") with ( patch( "huggingface_hub.list_repo_files", lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"], ), - patch("huggingface_hub.get_paths_info", fake_get_paths_info), - patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), ): out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") - assert downloaded == ["model-UD-Q4_K_XL.gguf"] - assert out == "/fresh/model-UD-Q4_K_XL.gguf" + assert out == str(snap / "model-UD-Q4_K_XL.gguf") def test_download_reuses_older_snapshot_when_offline_env_is_true(self, monkeypatch, hf_cache): # HF_HUB_OFFLINE accepts truthy spellings beyond "1" (true/yes/on); the offline diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index d02a2a4f7e..8742b84ae7 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -1609,11 +1609,11 @@ def test_env_idle_ttl_standalone_when_no_stored_value(monkeypatch): def test_stored_idle_value_overrides_env_and_stays_gated(monkeypatch): # An explicit stored value wins over the env default and remains gated on the # auto-switch toggle. - store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 30} + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 90} monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) - assert settings.get_auto_unload_idle_seconds() == 30 # stored wins, not env + assert settings.get_auto_unload_idle_seconds() == 90 # stored wins, not env monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) assert settings.get_auto_unload_idle_seconds() == 0 # explicit value still gated off @@ -3094,3 +3094,53 @@ def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeyp monkeypatch, enabled = False, active_model_name = "unsloth/Llama-3.2-1B-Instruct" ) assert "Model auto-switch" in non_gguf_loaded + + +def test_setter_rejects_idle_below_floor(monkeypatch): + import storage.studio_db as db + + writes = [] + monkeypatch.setattr(db, "upsert_app_settings", lambda m: writes.append(dict(m))) + settings._cache.clear() + + with pytest.raises(ValueError, match = "at least 60"): + settings.set_openai_auto_switch(True, 30) + assert writes == [] # rejected before any persist + # 0 (off) and >= 60 pass through unchanged. + assert settings.set_openai_auto_switch(True, 0)[1] == 0 + assert settings.set_openai_auto_switch(True, 60)[1] == 60 + assert settings.set_openai_auto_switch(True, 3600)[1] == 3600 + + +def test_put_route_rejects_idle_below_floor(): + import routes.settings as settings_route + from fastapi import HTTPException + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_idle_seconds = 30) + with pytest.raises(HTTPException) as excinfo: + settings_route.update_openai_auto_switch(payload, "tester") + assert excinfo.value.status_code == 400 + + +def test_stored_legacy_idle_below_floor_is_clamped(monkeypatch): + # Values persisted before the floor existed are raised to it on read, for + # both the effective TTL and the value the settings UI displays. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 5} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 60 + assert settings.get_stored_auto_unload_idle_seconds() == 60 + store[settings.AUTO_UNLOAD_IDLE_SETTING_KEY] = 90 + assert settings.get_auto_unload_idle_seconds() == 90 + + +def test_env_idle_below_floor_is_clamped(monkeypatch): + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) + monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "5") + assert settings.get_auto_unload_idle_seconds() == 60 + monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "0") + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600") + assert settings.get_auto_unload_idle_seconds() == 600 + monkeypatch.delenv(settings.MODEL_IDLE_TTL_ENV_VAR) + assert settings.get_auto_unload_idle_seconds() == 0 diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 5300a48557..8725b28ac8 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1054,7 +1054,7 @@ class TestChatCompletionRequestToolFields: monkeypatch.setattr( inference_route, "_detect_safetensors_features", - lambda backend, chat_template: {"supports_tools": True}, + lambda backend, chat_template, tools = None: {"supports_tools": True}, ) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inference_route, "api_monitor", monitor) @@ -6620,6 +6620,29 @@ class TestApiMonitorAudioInput: assert entry["reply"] == "hello world" assert monitor.active_count() == 0 + def failing_chunks(): + yield "partial" + raise RuntimeError("generation failed") + + self._patch_audio_backend(monkeypatch, failing_chunks()) + error_monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", error_monitor) + error_response = await openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + error_chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in error_response.body_iterator + ] + + assert '"type": "server_error"' in error_chunks[-1] + assert error_chunks[-1].endswith("data: [DONE]\n\n") + [error_entry] = error_monitor.snapshot() + assert error_entry["status"] == "error" + assert error_monitor.active_count() == 0 + asyncio.run(_run()) def test_non_gguf_tts_auto_route_records_monitor(self, monkeypatch): diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py index fb80b6d061..3a36500aee 100644 --- a/studio/backend/tests/test_orchestrator_unload_cancel.py +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -34,6 +34,70 @@ def _bare_orchestrator(): return o +def test_adapter_control_raises_stream_errors(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr( + o, + "_generate_dispatched", + lambda **_kwargs: iter([orch_mod.GenStreamError("Error: adapter failed")]), + ) + + with pytest.raises(RuntimeError, match = "adapter failed"): + list(o.generate_with_adapter_control(use_adapter = False)) + + closed = [] + + def _stream(**_kwargs): + try: + yield "token" + yield "late token" + finally: + closed.append(True) + + monkeypatch.setattr(o, "_generate_dispatched", _stream) + generator = o.generate_with_adapter_control(use_adapter = False) + assert next(generator) == "token" + generator.close() + assert closed == [True] + + +def test_worker_closes_cancelled_generator_before_gen_done(): + from core.inference.worker import _handle_generate + + events = [] + + class _Backend: + last_generation_stats = None + + def generate_with_adapter_control(self, **_kwargs): + try: + yield "token" + yield "late token" + finally: + events.append("closed") + + class _Responses: + def __init__(self): + self.items = [] + + def put(self, item): + if item["type"] == "gen_done": + assert events == ["closed"] + self.items.append(item) + + responses = _Responses() + cancel = threading.Event() + cancel.set() + _handle_generate( + _Backend(), + {"request_id": "r1", "messages": [], "use_adapter": False}, + responses, + cancel, + ) + + assert [item["type"] for item in responses.items] == ["gen_done"] + + def test_unload_cancels_inflight_generation_then_unloads(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py index 14ab0efe2e..3e259f6bd0 100644 --- a/studio/backend/tests/test_rag_parsing.py +++ b/studio/backend/tests/test_rag_parsing.py @@ -54,6 +54,56 @@ def test_pdf_markdown_off_uses_plain_text(tmp_path, monkeypatch): assert "#" not in text and "|" not in text # plain text path emits no Markdown markup +def test_pdf_bytes_use_same_extraction_path(tmp_path, monkeypatch): + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", False) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + from_file = parsers.parse(str(pdf)) + from_bytes, total_pages = parsers.parse_pdf_bytes(pdf.read_bytes()) + assert [page.text for page in from_bytes] == [page.text for page in from_file] + assert total_pages == len(from_file) + + +def test_pdf_bytes_limit_pages_before_extraction(monkeypatch): + import pymupdf + + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", False) + doc = pymupdf.open() + for marker in ("page one", "page two", "page three"): + page = doc.new_page() + page.insert_text((40, 40), marker) + data = doc.tobytes() + doc.close() + + pages, total_pages = parsers.parse_pdf_bytes(data, max_pages = 2) + assert len(pages) == 2 + assert "page two" in pages[-1].text + assert total_pages == 3 # full count, not the 2 extracted + + +def test_pdf_markdown_receives_page_limit(monkeypatch): + from core.rag import parsers + + captured = {} + + class _FakePymupdf4llm: + @staticmethod + def to_markdown(doc, **kwargs): + captured.update(kwargs) + return [{"text": "page"} for _ in kwargs["pages"]] + + class _Doc: + page_count = 100 + + monkeypatch.setitem(__import__("sys").modules, "pymupdf4llm", _FakePymupdf4llm) + assert parsers._pdf_markdown(_Doc(), range(2)) == ["page", "page"] + assert captured == {"page_chunks": True, "show_progress": False, "pages": [0, 1]} + + def test_pdf_markdown_passes_only_supported_legacy_kwargs(monkeypatch): # The pinned PyMuPDF4LLM legacy path ignores unknown kwargs; do not pass the # newer layout-only OCR knobs or Markdown extraction silently loses policy control. diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 0ed670ac01..bd3d8d16b9 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -417,6 +417,59 @@ def test_detect_safetensors_features_gemma_native_tool_call_keeps_tools_on(): assert flags["supports_tools"] is True +def test_detect_safetensors_features_gemma_native_reasoning_is_parseable_not_prefilled(): + """Native Gemma channels are normalized to , then split by the route.""" + from routes.inference import _detect_safetensors_features, _sf_reasoning_prefill_mode + + tpl_with_gemma_native = "{% if add_generation_prompt %}<|channel>thought\n{% endif %}" + backend = SimpleNamespace( + active_model_name = "unsloth/gemma-4-E2B-it", + models = { + "unsloth/gemma-4-E2B-it": { + "native_chat_template": tpl_with_gemma_native, + "chat_template_info": {"template": "override has no native markers"}, + } + }, + ) + flags = _detect_safetensors_features(backend, "override has no native markers") + missing_arg_flags = _detect_safetensors_features(backend, None) + + assert flags["supports_reasoning"] is True + assert flags["reasoning_always_on"] is True + assert missing_arg_flags["supports_reasoning"] is True + assert _sf_reasoning_prefill_mode(flags, None, tpl_with_gemma_native) is False + + +def test_detect_safetensors_features_selects_native_reasoning_from_tool_template(): + """Request tools select a marker-bearing named template without affecting default chat.""" + from routes.inference import _detect_safetensors_features + + named_template = { + "default": "plain default template", + "tool_use": "{% if tools %}<|channel>thought\n{% endif %}", + } + backend = SimpleNamespace( + active_model_name = "custom/named-native-reasoning", + models = { + "custom/named-native-reasoning": { + "native_chat_template": named_template, + "chat_template_info": {"template": "{% if tools %}{% endif %}"}, + } + }, + ) + + default_flags = _detect_safetensors_features(backend, "plain override") + tool_flags = _detect_safetensors_features( + backend, + "plain override", + tools = [{"type": "function"}], + ) + + assert default_flags["supports_reasoning"] is False + assert tool_flags["supports_reasoning"] is True + assert tool_flags["reasoning_always_on"] is True + + # Qwen3.5 family pin: the live GGUF + safetensors templates both wrap tool # calls as ``\n...``. Faithful slice so the # classifier never silently regresses for this family. diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py index 4e708139b7..af5a05d266 100644 --- a/studio/backend/tests/test_safetensors_reasoning_stream.py +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -215,3 +215,135 @@ def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): swallowed = _replay_sf_reasoning_stream(events, prefilled = True) assert swallowed["visible"] == "" assert swallowed["reasoning"] == "The capital of France is Paris." + + +def test_native_reasoning_streamer_selected_and_errors_raise(): + import threading + import pytest + + torch = pytest.importorskip("torch") + inf = pytest.importorskip("core.inference.inference") + + class Batch(dict): + def to(self, _device): + return self + + class Tok: + chat_template = "<|channel>thought\n..." + all_special_tokens = [] + eos_token_id = 1 + pad_token_id = None + pieces = {10: "<|channel>thought\n", 11: "r", 12: "", 13: "a"} + + def __call__(self, *_args, **_kwargs): + return Batch({"input_ids": torch.zeros((1, 1), dtype = torch.long)}) + + def decode(self, ids, **_kwargs): + return "".join(self.pieces.get(int(token_id), "") for token_id in ids) + + class Model: + device = "cpu" + generation_config = type("Cfg", (), {"eos_token_id": 1})() + config = generation_config + + def __init__(self, fail = False): + self.fail = fail + self.kwargs = None + + def generate(self, **kwargs): + self.kwargs = kwargs + streamer = kwargs["streamer"] + streamer.put(torch.zeros((1, 1), dtype = torch.long)) + for token_id in [10, 11, 12, 13]: + streamer.put(torch.tensor([token_id])) + if self.fail: + raise RuntimeError("boom") + + backend = inf.InferenceBackend.__new__(inf.InferenceBackend) + backend.active_model_name = "gemma-test" + backend._generation_lock = threading.Lock() + backend.models = {"gemma-test": {"model": Model(), "tokenizer": Tok()}} + + assert list(backend.generate_stream("prompt", max_new_tokens = 4))[-1] == "ra" + + backend.models["gemma-test"]["model"] = Model(fail = True) + + with pytest.raises(inf._GenerationThreadError, match = "boom"): + list(backend.generate_stream("prompt", max_new_tokens = 4)) + + +def test_text_only_vlm_fallback_resolves_native_markers_off(): + import threading + import pytest + + torch = pytest.importorskip("torch") + inf = pytest.importorskip("core.inference.inference") + + class Batch(dict): + def to(self, _device): + return self + + class Tokenizer: + all_special_tokens = [] + eos_token_id = 1 + pad_token_id = None + + def __call__(self, *_args, **_kwargs): + return Batch({"input_ids": torch.zeros((1, 1), dtype = torch.long)}) + + class Processor: + chat_template = "<|channel>thought\n..." + tokenizer = Tokenizer() + + class Model: + device = "cpu" + generation_config = type("Cfg", (), {"eos_token_id": 1})() + config = generation_config + + def generate(self, **_kwargs): + return None + + class EmptyStreamer: + def __next__(self): + raise StopIteration + + def end(self): + return None + + captured = {} + backend = inf.InferenceBackend.__new__(inf.InferenceBackend) + backend.active_model_name = "vision-test" + backend._generation_lock = threading.Lock() + backend.models = { + "vision-test": { + "model": Model(), + "processor": Processor(), + "tokenizer": Processor(), + } + } + backend.format_chat_prompt = lambda *_args, **_kwargs: "manual text-only prompt" + + def make_streamer(*_args, **kwargs): + captured.update(kwargs) + return EmptyStreamer() + + backend._make_text_streamer = make_streamer + + assert ( + list( + backend._generate_vision_response( + messages = [{"role": "user", "content": "hello"}], + system_prompt = "", + image = None, + temperature = 0.7, + top_p = 0.9, + top_k = 40, + min_p = 0.0, + max_new_tokens = 1, + repetition_penalty = 1.0, + ) + ) + == [] + ) + assert captured["reasoning_channel_markers"] is None + assert captured["reasoning_channel_markers_resolved"] is True diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index eae1a75161..915f82ac8e 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -2843,6 +2843,61 @@ class TestLoopBehaviour: ] assert len(duplicate_nudges) == 1 + def test_same_turn_duplicate_does_not_drop_later_parallel_call(self): + # Turn 1 runs search(x). Turn 2's batch is [search(x) duplicate, python]: + # the duplicate is a no-op, but python after it must still run, and the + # no-op nudge must land after python's result rather than splitting it. + captured_messages: list[list[dict]] = [] + turns = iter( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + '{"name":"python","arguments":{"code":"print(1)"}}' + ], + ["final"], + ] + ) + + def fake_single_turn(messages, active_tools = None): + captured_messages.append([dict(m) for m in messages]) + chunks = next(turns) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["search-x", "py-result"]) + _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 4, + ) + ) + + # Turn-1 search and turn-2 python both ran; the turn-2 duplicate search did not. + assert exec_fn.calls == [ + ("web_search", {"query": "x"}), + ("python", {"code": "print(1)"}), + ] + + conv = captured_messages[-1] + turn2 = [m for m in conv if m.get("role") == "assistant" and m.get("tool_calls")][-1] + assert [tc["function"]["name"] for tc in turn2["tool_calls"]] == ["python"] + after = conv[conv.index(turn2) + 1 :] + assert after[0]["role"] == "tool" and after[0]["content"] == "py-result" + assert after[1]["role"] == "user" # deferred duplicate nudge, after the result + assert after[1]["content"].startswith( + "One earlier request to call tool 'web_search' in this batch was not executed" + ) + assert "previous tool request" not in after[1]["content"].lower() + def test_duplicate_tool_call_internal_noop_allows_distinct_followup_tool(self): captured_messages: list[list[dict]] = [] captured_tool_names: list[list[str]] = [] @@ -3150,6 +3205,196 @@ class TestLoopBehaviour: class TestLoopRePrompt: """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Studio always nudges, so these drive the loop with ``nudge_tool_calls=True``.""" + def test_reasoning_intent_does_not_reprompt_a_visible_answer(self): + generations = 0 + + def _gen(_messages, active_tools = None): + nonlocal generations + generations += 1 + yield ( + "Let me prepare the requested summary carefully." + "This is the final visible answer." + ) + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "summarize this"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + ) + ) + + assert generations == 1 + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1].endswith("This is the final visible answer.") + + def test_prefilled_reasoning_intent_does_not_reprompt_a_visible_answer(self): + generations = 0 + + def _gen(_messages, active_tools = None): + nonlocal generations + generations += 1 + yield "Let me prepare the requested summary carefully.This is the final visible answer." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "summarize this"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + ) + + assert generations == 1 + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1].endswith("This is the final visible answer.") + + def test_prefilled_reasoning_with_reemitted_think_does_not_reprompt(self): + generations = 0 + + def _gen(_messages, active_tools = None): + nonlocal generations + generations += 1 + yield ( + "Let me prepare the requested summary carefully." + "more private planningThis is the final visible answer." + ) + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "summarize this"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + ) + + assert generations == 1 + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1].endswith("This is the final visible answer.") + + def test_prefilled_reasoning_with_later_think_does_not_reprompt(self): + generations = 0 + + def _gen(_messages, active_tools = None): + nonlocal generations + generations += 1 + yield ( + "private prefilled planning" + "Let me prepare the requested summary carefully." + "This is the final visible answer." + ) + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "summarize this"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + ) + + assert generations == 1 + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1].endswith("This is the final visible answer.") + + def test_reasoning_only_intent_still_reprompts_and_uses_a_tool(self): + loop, exec_fn = _make_loop( + turns = [ + ["Let me search for that."], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the answer."], + ], + exec_results = ["result"], + nudge_tool_calls = True, + ) + + events = _collect_events(loop) + + assert exec_fn.calls == [("web_search", {"query": "cats"})] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1] == "Here is the answer." + + def test_prefilled_no_close_reasoning_intent_still_reprompts(self): + loop, exec_fn = _make_loop( + turns = [ + ["I need more context.Let me search for that."], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the answer."], + ], + exec_results = ["result"], + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + + events = _collect_events(loop) + + assert exec_fn.calls == [("web_search", {"query": "cats"})] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1] == "Here is the answer." + + def test_prefilled_reasoning_prefix_is_kept_for_reasoning_only_reprompt(self): + loop, exec_fn = _make_loop( + turns = [ + ["Let me search for that.checking details"], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the answer."], + ], + exec_results = ["result"], + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + + events = _collect_events(loop) + + assert exec_fn.calls == [("web_search", {"query": "cats"})] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1] == "Here is the answer." + + def test_reprompt_history_uses_visible_intent_text(self): + captured: list[list[dict]] = [] + + def _gen(messages, active_tools = None): + captured.append([dict(message) for message in messages]) + if len(captured) == 1: + yield "private planning detailsLet me search for that." + elif len(captured) == 2: + yield '{"name":"web_search","arguments":{"query":"cats"}}' + else: + yield "Here is the answer." + + exec_fn = FakeExecuteTool(["result"]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "find cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + ) + ) + + assert exec_fn.calls == [("web_search", {"query": "cats"})] + assert captured[1][1] == {"role": "assistant", "content": "Let me search for that."} + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1] == "Here is the answer." + def test_intent_signal_triggers_reprompt(self): # Turn 1: intent signal, no tool call. # Turn 2 (re-prompt): proper tool call -> executes. diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index 01905b712c..f91eec9817 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -177,6 +177,15 @@ def _sse_objects(chunks): # ── Non-streaming ───────────────────────────────────────────────── +def test_non_reasoning_backend_keeps_literal_think_tags(monkeypatch): + backend = _ScriptedBackend(_fixed("show example tags")) + response = _call(_request(stream = False), monkeypatch, backend, supports_tools = False) + + message = _json_body(response)["choices"][0]["message"] + assert message["content"] == "show example tags" + assert message["reasoning_content"] is None + + def test_xml_healed_to_tool_calls_non_streaming(monkeypatch): backend = _ScriptedBackend(_fixed(_CALL_XML)) payload = _request(tools = [LOOKUP_TOOL], stream = False) @@ -485,6 +494,52 @@ def test_streaming_no_tools_verbatim(monkeypatch): assert finishes == ["stop"] +def test_streaming_gen_stream_error_is_not_model_text(monkeypatch): + from core.inference.orchestrator import GenStreamError + + class _ErrorAfterPartial(_ScriptedBackend): + def __init__(self): + super().__init__(_fixed()) + + def generate_chat_response(self, **_kwargs): + yield "partial" + yield GenStreamError("Error: /tmp/secret traceback") + + backend = _ErrorAfterPartial() + payload = _request(stream = True) + response = _call(payload, monkeypatch, backend, supports_tools = False) + chunks = _collect_sse(response) + objs = _sse_objects(chunks) + + deltas = [o.get("choices", [{}])[0].get("delta", {}) for o in objs if o.get("choices")] + assert any("partial" in json.dumps(delta) for delta in deltas) + assert not any("/tmp/secret" in json.dumps(delta) for delta in deltas) + errors = [o["error"]["message"] for o in objs if "error" in o] + assert errors == ["An internal error occurred."] + assert any( + "data: [DONE]" in (chunk.decode() if isinstance(chunk, bytes) else chunk) + for chunk in chunks + ) + + +def test_server_tool_streaming_invalid_event_is_error(monkeypatch): + class _InvalidEventBackend(_ScriptedBackend): + def __init__(self): + super().__init__(_fixed()) + + def generate_chat_completion_with_tools(self, **_kwargs): + yield {"type": "content", "text": "partial"} + yield "not-an-event" + + backend = _InvalidEventBackend() + payload = _request(tools = [LOOKUP_TOOL], enable_tools = True, stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + + errors = [o["error"]["message"] for o in objs if "error" in o] + assert errors == ["An internal error occurred."] + + def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch): # Repeated then shrunk cumulative snapshots must not double-heal. backend = _ScriptedBackend(_fixed(_CALL_XML, _CALL_XML, _CALL_XML[:5], _CALL_XML)) diff --git a/studio/backend/tests/test_think_prefill_reemit.py b/studio/backend/tests/test_think_prefill_reemit.py index 300ff92776..346399c3b2 100644 --- a/studio/backend/tests/test_think_prefill_reemit.py +++ b/studio/backend/tests/test_think_prefill_reemit.py @@ -2,7 +2,7 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Unit tests for detect_think_prefill. +Unit tests for local reasoning-stream helpers. Reasoning templates (Qwen3.6-style) end the generation prompt with an open ``\\n`` so the model starts reasoning immediately. skip_prompt @@ -16,7 +16,13 @@ import sys _backend = os.path.join(os.path.dirname(__file__), "..") sys.path.insert(0, _backend) -from core.inference.chat_template_helpers import detect_think_prefill +from core.inference.chat_template_helpers import ( + ReasoningChannelNormalizer, + detect_reasoning_channel_markers, + detect_reasoning_channel_markers_from_model_info, + detect_think_prefill, + render_with_native_template_fallback, +) QWEN_PROMPT = "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n" @@ -87,3 +93,170 @@ def test_guard_emits_when_think_not_special(): def test_guard_default_and_empty_keep_emitting(): assert detect_think_prefill(QWEN_PROMPT + "\n", None) == "\n" assert detect_think_prefill(QWEN_PROMPT + "\n", []) == "\n" + + +def test_gemma_channel_detection_uses_active_template_not_token_metadata(): + class TemplateTokenizer: + chat_template = {"default": "...<|channel>thought\\n{{ eoc_token }}"} + + class NamedTemplateTokenizer: + chat_template = { + "default": "plain assistant template", + "tool_use": "...<|channel>thought\\n{{ eoc_token }}", + } + + class TokenMetadataOnly: + chat_template = None + soc_token = "<|channel>" + eoc_token = "" + + class NamedTemplateProcessor: + chat_template = { + "default": "plain processor default", + "tool_use": "<|channel>thought\nprocessor tool template", + } + tokenizer = TokenMetadataOnly() + + def apply_chat_template(self, *_args, **_kwargs): + raise NotImplementedError + + expected = ("<|channel>thought", "") + assert detect_reasoning_channel_markers(TemplateTokenizer()) == expected + assert detect_reasoning_channel_markers(NamedTemplateTokenizer()) is None + assert ( + detect_reasoning_channel_markers( + NamedTemplateTokenizer(), tools = [{"function": {"name": "web_search"}}] + ) + == expected + ) + assert detect_reasoning_channel_markers(NamedTemplateTokenizer(), tools = []) is None + assert ( + detect_reasoning_channel_markers( + NamedTemplateProcessor(), tools = [{"function": {"name": "web_search"}}] + ) + is None + ) + assert detect_reasoning_channel_markers(TokenMetadataOnly()) is None + + +def test_gemma_channel_detection_tries_no_argument_getter_fallback(): + class FallbackTokenizer: + chat_template = "plain fallback template" + + def get_chat_template(self, **kwargs): + if kwargs: + raise ValueError("tools are not supported") + return "...<|channel>thought\n" + + assert detect_reasoning_channel_markers( + FallbackTokenizer(), tools = [{"function": {"name": "web_search"}}] + ) == ("<|channel>thought", "") + + +def test_native_template_fallback_returns_selected_reasoning_metadata(): + from types import SimpleNamespace + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def render(tokenizer, msgs, *, tools, **_kw): + body = "".join(message["content"] for message in msgs) + suffix = "|TOOLS" if tools else "" + return body + suffix if tokenizer.chat_template == "NATIVE <|channel>thought\n" else body + + result = render_with_native_template_fallback( + formatted_prompt = "hi", + tokenizer = SimpleNamespace(chat_template = "OVERRIDE"), + model_info = { + "native_chat_template": "NATIVE <|channel>thought\n", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + }, + active_model_name = "gemma-test", + messages = messages, + tools = tools, + apply_fn = render, + return_metadata = True, + ) + + assert result.prompt == "hi|TOOLS" + assert result.reasoning_channel_markers == ("<|channel>thought", "") + + +def test_cached_native_template_metadata_recovers_reasoning_markers_without_tools(): + from types import SimpleNamespace + + model_info = {"chat_template_info": {"template": "native <|channel>thought\n"}} + + assert detect_reasoning_channel_markers_from_model_info( + SimpleNamespace(chat_template = "override has no native markers"), + model_info, + tools = None, + ) == ("<|channel>thought", "") + result = render_with_native_template_fallback( + formatted_prompt = "prompt from override", + tokenizer = SimpleNamespace(chat_template = "override has no native markers"), + model_info = model_info, + active_model_name = "gemma-test", + messages = [{"role": "user", "content": "hi"}], + tools = None, + return_metadata = True, + ) + assert result.prompt == "prompt from override" + assert result.reasoning_channel_markers == ("<|channel>thought", "") + + +def test_cached_native_markers_do_not_describe_live_tool_template(): + from types import SimpleNamespace + + tools = [{"type": "function", "function": {"name": "web_search"}}] + + class LiveTokenizer: + chat_template = "live tool template without native markers" + + def render(_tokenizer, _messages, *, tools, **_kwargs): + return "prompt with tools" if tools else "prompt without tools" + + result = render_with_native_template_fallback( + formatted_prompt = "prompt with tools", + tokenizer = LiveTokenizer(), + model_info = { + "chat_template_info": {"template": "native <|channel>thought\n"}, + "tokenizer": SimpleNamespace(), + }, + active_model_name = "gemma-test", + messages = [{"role": "user", "content": "hi"}], + tools = tools, + apply_fn = render, + return_metadata = True, + ) + + assert result.prompt == "prompt with tools" + assert result.reasoning_channel_markers is None + + +def test_gemma_channel_normalization_is_prefix_monotonic_and_preserves_tools(): + parser = ReasoningChannelNormalizer("<|channel>thought", "") + output = "" + snapshots = [] + for chunk in ( + "<|chan", + "nel>thought", + "\nReason", + "<|tool_call>web_search", + ): + delta = parser.feed(chunk) + if delta: + output += delta + snapshots.append(output) + + assert snapshots == [ + "", + "Reason", + "Reason<|tool_call>web_search", + ] + assert snapshots[1].startswith(snapshots[0]) + compact = ReasoningChannelNormalizer("<|channel>thought", "") + assert compact.feed("<|channel>thoughtanswer") + compact.finish() == ( + "answer" + ) diff --git a/studio/backend/tests/test_tool_loop_controller.py b/studio/backend/tests/test_tool_loop_controller.py index 0e8ae798af..496c30ac13 100644 --- a/studio/backend/tests/test_tool_loop_controller.py +++ b/studio/backend/tests/test_tool_loop_controller.py @@ -13,6 +13,7 @@ if _BACKEND_DIR not in sys.path: from core.inference.tool_loop_controller import ( ToolLoopController, + append_deferred_nudges, canonical_tool_call_key, coerce_tool_arguments, status_for_tool, @@ -21,6 +22,22 @@ from core.inference.tool_loop_controller import ( ) +def test_append_deferred_nudges_merges_deduped_into_one_message(): + conversation = [{"role": "assistant", "tool_calls": [1]}, {"role": "tool", "content": "r"}] + nudges = [ + {"role": "user", "content": "duplicate"}, + {"role": "user", "content": "duplicate"}, # dropped: same content + {"role": "user", "content": "disabled foo"}, + ] + append_deferred_nudges(conversation, nudges) + # One user message, after the results, with distinct contents joined. + assert conversation[2:] == [{"role": "user", "content": "duplicate\n\ndisabled foo"}] + # Empty is a no-op. + before = list(conversation) + append_deferred_nudges(conversation, []) + assert conversation == before + + def _tool(name: str) -> dict: return {"type": "function", "function": {"name": name}} @@ -111,6 +128,10 @@ def test_successful_duplicate_is_internal_noop_and_keeps_remaining_tools(): assert not duplicate.should_execute assert not duplicate.emit_visible_events duplicate_nudge = completion.model_message()["content"] + assert duplicate_nudge.startswith( + "One earlier request to call tool 'web_search' in this batch was not executed" + ) + assert "previous tool request" not in duplicate_nudge.lower() assert "already completed successfully" in duplicate_nudge assert "different enabled tool" in duplicate_nudge assert completion.model_message()["role"] == "user" @@ -165,7 +186,12 @@ def test_empty_enabled_tool_list_blocks_all_tool_calls(): assert decision.action == "disabled" assert not decision.emit_visible_events assert completion.model_message()["role"] == "user" - assert "not enabled" in completion.model_message()["content"] + disabled_nudge = completion.model_message()["content"] + assert disabled_nudge.startswith( + "One earlier request to call tool 'web_search' in this batch was not executed" + ) + assert "previous tool request" not in disabled_nudge.lower() + assert "not enabled" in disabled_nudge assert controller.force_final_answer assert controller.active_tools() == [] diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py index 457dfc8ea2..0cd702bce2 100644 --- a/studio/backend/tests/test_training_stop_watchdog.py +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -258,15 +258,27 @@ def test_watchdog_no_op_when_worker_superseded(monkeypatch): def test_new_run_gets_its_own_watchdog(monkeypatch): # A stale watchdog sleeping on an old proc must not stop a new run's stop from # creating its own watcher. - monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) - monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) b = TrainingBackend() - _record_force_terminate(monkeypatch, b) + started = [] + release = threading.Event() + + def _blocked_watchdog( + target_proc, + cancel, + watched_job_id = None, + ): + started.append(target_proc) + # No timeout: the finally always releases this, so a superseded watchdog stays + # alive through the assertions regardless of load; as a daemon it can't hang exit. + release.wait() + + monkeypatch.setattr(b, "_stop_watchdog_loop", _blocked_watchdog) old_proc = _FakeProc(alive = True) b._proc = old_proc b._start_stop_watchdog(cancel = False) first_wd = b._stop_watchdog + assert _wait_until(lambda: started == [old_proc]) # New run: fresh worker replaces the handle; its stop must get a new watcher # even though the old (superseded) watchdog is still alive. @@ -276,12 +288,12 @@ def test_new_run_gets_its_own_watchdog(monkeypatch): second_wd = b._stop_watchdog try: + assert _wait_until(lambda: started == [old_proc, new_proc]) assert first_wd.is_alive() assert second_wd is not first_wd, "a new run must get its own watchdog" assert b._stop_watchdog_proc is new_proc finally: - old_proc._alive = False - new_proc._alive = False + release.set() first_wd.join(timeout = 5) second_wd.join(timeout = 5) diff --git a/studio/backend/tests/test_web_fetch_binary_guard.py b/studio/backend/tests/test_web_fetch_binary_guard.py index 3041ed5c34..10db953913 100644 --- a/studio/backend/tests/test_web_fetch_binary_guard.py +++ b/studio/backend/tests/test_web_fetch_binary_guard.py @@ -59,6 +59,18 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str: return tools._fetch_page_text("https://example.com/thing", timeout = 5) +def _pdf_bytes(*page_texts: str) -> bytes: + pymupdf = pytest.importorskip("pymupdf") + doc = pymupdf.open() + for text in page_texts: + page = doc.new_page() + if text: + page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), text, fontsize = 11) + data = doc.tobytes() + doc.close() + return data + + @pytest.mark.parametrize( "content_type,expected", [ @@ -90,10 +102,119 @@ def test_is_text_candidate_content_type(content_type, expected): assert tools._is_text_candidate_content_type(content_type) is expected -def test_pdf_rejected_by_content_type(monkeypatch): - out = _fetch_with(monkeypatch, b"%PDF-1.7\n\xff\xd8\xff\x00\x89PNG" * 200, "application/pdf") - assert "�" not in out - assert "non-text content" in out and "application/pdf" in out +@pytest.mark.parametrize( + "content_type", + ["application/pdf", "application/octet-stream", "text/html", "text/plain", None], +) +def test_pdf_text_extracted(monkeypatch, content_type): + out = _fetch_with( + monkeypatch, + _pdf_bytes("First page marker", "Second page marker"), + content_type, + ) + assert "## Page 1\n\nFirst page marker" in out + assert "## Page 2" in out and "Second page marker" in out + assert "binary content" not in out and "non-text content" not in out + + +@pytest.mark.parametrize("content_type", ["application/pdf", "text/plain"]) +def test_malformed_pdf_returns_safe_placeholder(monkeypatch, content_type): + out = _fetch_with(monkeypatch, b"%PDF-1.7\nnot a complete PDF", content_type) + assert out == "(PDF content could not be read as text)" + + +def test_pdf_without_text_layer_reported(monkeypatch): + out = _fetch_with(monkeypatch, _pdf_bytes(""), "application/pdf") + assert out == "(PDF contains no extractable text)" + + +def test_encrypted_pdf_returns_safe_placeholder(monkeypatch): + pymupdf = pytest.importorskip("pymupdf") + doc = pymupdf.open() + doc.new_page().insert_text((40, 40), "private text") + data = doc.tobytes( + encryption = pymupdf.PDF_ENCRYPT_AES_256, + owner_pw = "owner", + user_pw = "secret", + ) + doc.close() + out = _fetch_with(monkeypatch, data, "application/pdf") + assert out == "(PDF content could not be read as text)" + + +def test_pdf_download_limit_enforced(monkeypatch): + monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", 256) + out = _fetch_with(monkeypatch, _pdf_bytes("Readable but oversized"), "application/pdf") + assert out == "(PDF content exceeds the download limit; not readable as text)" + + +def test_mislabeled_pdf_is_read_past_text_download_cap(monkeypatch): + body = _pdf_bytes("Cross-reference data was fetched") + monkeypatch.setattr(tools, "_MAX_FETCH_BYTES", 128) + monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", len(body) + 100) + out = _fetch_with(monkeypatch, body, "text/plain") + assert "Cross-reference data was fetched" in out + + +def test_pdf_extraction_caps_pages_and_intermediate_text(monkeypatch): + from core.rag.parsers import Page + + seen = {} + + def fake_parse(data, *, max_pages = None): + seen["max_pages"] = max_pages + pages = [Page(text = "x" * 1000, page_number = i, char_count = 1000) for i in range(1, 51)] + return pages, 60 # document actually has more pages than the cap + + monkeypatch.setattr("core.rag.parsers.parse_pdf_bytes", fake_parse) + text = tools._extract_pdf_text(b"unused") + assert seen["max_pages"] == tools._MAX_WEB_PDF_PAGES + assert len(text) <= tools._MAX_PAGE_CHARS + assert "text limited to 16,000 characters" in text + assert "page processing capped at 50 pages" in text + + +def test_pdf_exactly_at_page_cap_not_marked_capped(monkeypatch): + from core.rag.parsers import Page + + # Exactly _MAX_WEB_PDF_PAGES pages are fully read, so no "capped" marker. + monkeypatch.setattr( + "core.rag.parsers.parse_pdf_bytes", + lambda data, *, max_pages = None: ( + [Page(text = "short", page_number = i, char_count = 5) for i in range(1, 51)], + 50, + ), + ) + text = tools._extract_pdf_text(b"unused") + assert "page processing capped" not in text + assert "## Page 50\n\nshort" in text + + +def test_pdf_page_cap_does_not_claim_later_pages_are_textless(monkeypatch): + from core.rag.parsers import Page + monkeypatch.setattr( + "core.rag.parsers.parse_pdf_bytes", + lambda data, *, max_pages = None: ( + [Page(text = "", page_number = i, char_count = 0) for i in range(1, 51)], + 60, + ), + ) + assert tools._extract_pdf_text(b"unused") == ( + "(PDF contains no extractable text in the first 50 pages)" + ) + + +def test_pdf_result_discarded_after_fetch_deadline(monkeypatch): + clock = {"time": 1000.0} + monkeypatch.setattr(tools.time, "monotonic", lambda: clock["time"]) + + def slow_extract(data): + clock["time"] += 10.0 + return "late PDF text" + + monkeypatch.setattr(tools, "_extract_pdf_text", slow_extract) + out = _fetch_with(monkeypatch, _pdf_bytes("Readable text"), "application/pdf") + assert out == "Failed to fetch URL: timed out." def test_text_octet_stream_kept_after_sniffing(monkeypatch): @@ -154,7 +275,6 @@ def test_valid_utf8_binary_caught_by_control_chars(monkeypatch): @pytest.mark.parametrize( "magic", [ - b"%PDF-", b"PK\x03\x04", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", b"\x1f\x8b", @@ -180,8 +300,8 @@ def test_text_labeled_binary_caught_by_magic(monkeypatch, magic): b"\t\xef\xbb\xbf ", ], ) -def test_pdf_magic_after_harmless_prefix(monkeypatch, prefix): - body = prefix + b"%PDF-1.7\n" + b"1 0 obj<>endobj\n" * 100 +def test_binary_magic_after_harmless_prefix(monkeypatch, prefix): + body = prefix + b"\x1f\x8b" + b" printable text-heavy body" * 100 out = _fetch_with(monkeypatch, body, "text/plain") assert "binary content" in out @@ -243,10 +363,10 @@ def test_html_page_unaffected(monkeypatch): def test_content_type_sanitized_in_message(monkeypatch): # Do not echo obs-folded header content into the model response. - out = _fetch_with(monkeypatch, b"\x00\x01\x02" * 500, "application/pdf\r\n data: injected") + out = _fetch_with(monkeypatch, b"PK\x03\x04" * 500, "application/zip\r\n data: injected") assert "\n" not in out and "\r" not in out assert "injected" not in out - assert "application/pdf" in out + assert "application/zip" in out @pytest.mark.parametrize( diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 8d6c919ebd..117ad7b780 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -49,6 +49,11 @@ logger = get_logger(__name__) # and spawn workers copy os.environ. setdefault so an explicit user override wins. os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") +# Studio workers can import MLX without importing unsloth first, so mirror the +# package bootstrap here. Keep an explicit user value authoritative. +if platform.system() == "Darwin" and platform.machine() == "arm64": + os.environ.setdefault("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1") + # ========== Device Enum ========== diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py index 1689395f40..462435e5d5 100644 --- a/studio/backend/utils/openai_auto_switch_settings.py +++ b/studio/backend/utils/openai_auto_switch_settings.py @@ -8,7 +8,9 @@ Two settings, both off by default so existing API behavior is unchanged: names a downloaded local GGUF different from the loaded one transparently loads it before serving (llama-swap-style). Unknown names pass through. - ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is - unloaded after this many idle seconds to free VRAM. + unloaded after this many idle seconds to free VRAM. Enabled values have a + 60s floor (0 stays "off"): a tiny TTL tears the model down between turns of + an active chat, forcing a full weight reload + prompt re-prefill per turn. The idle TTL can also be set at startup via the ``UNSLOTH_MODEL_IDLE_TTL`` env var. Unlike the stored setting (which stays gated on auto-switch), the env value @@ -33,6 +35,7 @@ MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL" DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0 +MIN_AUTO_UNLOAD_IDLE_SECONDS = 60 _CACHE_TTL_S = 2.0 _cache_lock = threading.Lock() @@ -58,6 +61,10 @@ def _coerce_int(value: Any) -> int | None: return None +def _apply_idle_floor(seconds: int) -> int: + return 0 if seconds <= 0 else max(MIN_AUTO_UNLOAD_IDLE_SECONDS, seconds) + + def _cached_setting(key: str, default: Any) -> Any: """Read an app setting, memoized for _CACHE_TTL_S to spare the hot path.""" now = time.monotonic() @@ -91,12 +98,34 @@ def _stored_idle_seconds() -> Optional[int]: return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None)) +_env_floor_warned = False + + def _env_idle_seconds() -> Optional[int]: - """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid.""" + """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid. + + Floored to MIN_AUTO_UNLOAD_IDLE_SECONDS here (with a one-time warning) since + headless/container deploys have no UI to surface a validation error.""" raw = os.environ.get(MODEL_IDLE_TTL_ENV_VAR) if raw is None or not raw.strip(): return None - return _coerce_int(raw) + parsed = _coerce_int(raw) + if parsed is None: + return None + floored = _apply_idle_floor(parsed) + if floored != parsed: + global _env_floor_warned + if not _env_floor_warned: + _env_floor_warned = True + from loggers import get_logger + get_logger(__name__).warning( + "%s=%s is below the %ss minimum; using %ss", + MODEL_IDLE_TTL_ENV_VAR, + parsed, + MIN_AUTO_UNLOAD_IDLE_SECONDS, + floored, + ) + return floored def get_stored_auto_unload_idle_seconds() -> int: @@ -108,7 +137,9 @@ def get_stored_auto_unload_idle_seconds() -> int: """ stored = _stored_idle_seconds() if stored is not None: - return stored + # Floor legacy values persisted before the minimum existed, so the UI + # displays the effective TTL and round-trips it cleanly. + return _apply_idle_floor(stored) env = _env_idle_seconds() return env if env is not None else DEFAULT_AUTO_UNLOAD_IDLE_SECONDS @@ -118,8 +149,9 @@ def get_auto_unload_idle_seconds() -> int: stored = _stored_idle_seconds() if stored is not None: # An explicit UI/API value stays gated on auto-switch: off reports 0 so the - # off state is identical to pre-feature. - return stored if get_openai_auto_switch_enabled() else 0 + # off state is identical to pre-feature. Floored to cover values persisted + # before the minimum existed. + return _apply_idle_floor(stored) if get_openai_auto_switch_enabled() else 0 # No stored value: UNSLOTH_MODEL_IDLE_TTL is a standalone startup default that # enables idle-unload even with auto-switch off (headless/container deploys). env = _env_idle_seconds() @@ -136,6 +168,11 @@ def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]: parsed_idle = _coerce_int(idle_seconds) if parsed_idle is None: raise ValueError("Auto-unload idle seconds must be a non-negative integer.") + if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS: + raise ValueError( + f"Auto-unload idle seconds must be 0 (off) or at least " + f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}." + ) from storage.studio_db import upsert_app_settings upsert_app_settings( diff --git a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx index ee8c26abf1..0345dc6e2a 100644 --- a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx +++ b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx @@ -8,7 +8,7 @@ import { cn } from "@/lib/utils"; import { useAuiState } from "@assistant-ui/react"; import { LayoutTwoColumnIcon as Layout2ColumnIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useLayoutEffect, useMemo } from "react"; +import { useLayoutEffect, useMemo, useRef } from "react"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; import type { ArtifactViewMode } from "./html-frame"; import { @@ -83,15 +83,18 @@ export function ArtifactCard({ ], ); const surface = artifactThreadId ? "panel" : "overlay"; + // Once per mount, so a view-change cleanup can't re-trigger a stale open. + const autoOpenAttemptedRef = useRef(false); useLayoutEffect(() => { if (selectedArtifactId === artifact.id) { updateArtifact(artifact); } - if (!autoOpen) { + if (!autoOpen || autoOpenAttemptedRef.current) { return; } + autoOpenAttemptedRef.current = true; if (hasAutoOpenedArtifact(artifact.id)) { return; } diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 380ce0e0ab..ec0ad977bf 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -247,13 +247,13 @@ const SingleContent = memo(function SingleContent({ useState(false); const [isArtifactSurfaceVisible, setIsArtifactSurfaceVisible] = useState(false); + // Without a URL threadId the artifact must belong to the active thread. const showArtifactPanel = Boolean( artifact && artifactSurface === "panel" && (threadId ? !artifact.threadId || artifact.threadId === threadId - : Boolean(newThreadNonce) || - Boolean(artifact.threadId && artifact.threadId === activeThreadId)), + : Boolean(artifact.threadId && artifact.threadId === activeThreadId)), ); const artifactLayoutActive = showArtifactPanel || isArtifactPanelLayoutActive; @@ -1771,10 +1771,8 @@ export function ChatPage({ useEffect(() => { if (view.mode !== "single") return; - if (view.threadId || view.newThreadNonce || !selectedArtifact) return; - // view excludes __LOCALID_ threads (they fall through to mode:"single" - // with no threadId/nonce). Don't close a canvas whose thread is the - // active local thread. + if (view.threadId || !selectedArtifact) return; + // Close any canvas that doesn't belong to the active thread. if ( selectedArtifact.threadId && selectedArtifact.threadId === activeThreadId diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index 95e5cfbd79..e39955e576 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -244,8 +244,8 @@ export function ChatProvidersSettings({ (s) => s.setConnectionsEnabled, ); const isCustomProvider = isCustomProviderType(providerType); - // Local presets (Ollama, llama.cpp) never use API keys — hide the field. - // vLLM may optionally use a bearer token on secured deployments. + // llama.cpp hides the key field. Ollama and vLLM show an optional key: + // Ollama cloud and secured vLLM need one; local servers leave it empty. const showApiKeyField = !customPresetSkipsApiKeyField(providerType); const showReasoningToggle = supportsProviderReasoningToggle(providerType); diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index bc718abbba..eb9e4656b0 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -184,11 +184,12 @@ export function supportsRemoteModelCatalog( ); } -/** Presets that skip the API-key field (local servers with no auth by default). */ +/** Presets that hide the API-key field. Ollama is not skipped: Ollama cloud + * requires a key; local servers leave the optional field empty. */ export function customPresetSkipsApiKeyField( providerType: string | null | undefined, ): boolean { - return providerType === "ollama" || providerType === "llama_cpp"; + return providerType === "llama_cpp"; } /** Catalog load plus optional manual model IDs. */ diff --git a/studio/frontend/src/features/chat/permission-mode-select.tsx b/studio/frontend/src/features/chat/permission-mode-select.tsx index 4277c1bfcf..a9cb8ce5d1 100644 --- a/studio/frontend/src/features/chat/permission-mode-select.tsx +++ b/studio/frontend/src/features/chat/permission-mode-select.tsx @@ -278,28 +278,23 @@ export function PermissionModeComposerPill({ data-pill-label={active.label} data-active={fullAccess ? "true" : "false"} data-variant={fullAccess ? "danger" : undefined} + data-keep-label="true" aria-label="Permission level for tool calls" title={`${active.label}: ${active.description}`} > {/* The icon doubles as an off switch (mirrors the MCP pill): hover swaps it to an X; clicking it turns bypass permissions Off (no - prompts, sandbox on) without opening the menu. In compact - icon-only mode the glyph is the whole button, so clicks fall - through and open the menu instead. */} + prompts, sandbox on) without opening the menu. data-keep-label + exempts this pill from compact icon-only mode, so the off switch + stays clickable even while the other pills are collapsed. */} { - if (e.currentTarget.closest('[data-pill-compact="true"]')) { - return; - } e.stopPropagation(); }} onClick={(e) => { - if (e.currentTarget.closest('[data-pill-compact="true"]')) { - return; - } e.stopPropagation(); setPermissionMode("off"); }} diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index abdaa871c0..af99458349 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -2,7 +2,11 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import type { RememberedLoadSettings } from "@/components/assistant-ui/model-selector/remembered-load-settings"; -import { cancelStagedModelDownload } from "@/features/hub"; +import { + cancelStagedModelDownload, + mirrorHfTokenInto, + useHfTokenStore, +} from "@/features/hub"; import { toast } from "@/lib/toast"; import { create } from "zustand"; import { isExternalModelId, parseExternalModelId } from "../external-providers"; @@ -25,8 +29,6 @@ import { import { useExternalProvidersStore } from "./external-providers-store"; import { PLUS_MENU_PINS_STORAGE_KEY } from "./plus-menu-prefs-store"; -const HF_TOKEN_KEY = "unsloth_hf_token"; -const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed"; export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled"; export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled"; export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled"; @@ -495,17 +497,6 @@ export function saveSpeculativeType(value: string | null): void { } } -function notifyHfTokenChanged(value: string): void { - if (!canUseStorage()) return; - try { - window.dispatchEvent( - new CustomEvent(HF_TOKEN_CHANGED_EVENT, { detail: value }), - ); - } catch { - // ignore - } -} - /** A local model staged for a deferred load (see `pendingSelection`). Shape is * a subset of the load hook's `SelectedModelInput`, structurally assignable. */ export type PendingModelSelection = { @@ -1144,7 +1135,7 @@ export const useChatRuntimeStore = create((set, get) => ({ runningByThreadId: {}, cancelByThreadId: {}, autoTitle: false, - hfToken: loadString(HF_TOKEN_KEY, ""), + hfToken: useHfTokenStore.getState().token, modelsError: null, lastModelLoadError: null, activeGgufVariant: null, @@ -1349,11 +1340,7 @@ export const useChatRuntimeStore = create((set, get) => ({ setScalarSettingVersion("autoTitle", autoTitle, state.autoTitle); return { autoTitle }; }), - setHfToken: (hfToken) => { - saveString(HF_TOKEN_KEY, hfToken); - set({ hfToken }); - notifyHfTokenChanged(hfToken); - }, + setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken), setModelsError: (modelsError) => set({ modelsError }), setLastModelLoadError: (lastModelLoadError) => set({ lastModelLoadError }), setCheckpoint: (modelId, ggufVariant) => @@ -1837,6 +1824,12 @@ export const useChatRuntimeStore = create((set, get) => ({ setContextUsage: (contextUsage) => set({ contextUsage }), })); +// Mirror token edits made through the shared store (e.g. Studio's field). +const unsubscribeHfTokenMirror = mirrorHfTokenInto(useChatRuntimeStore); +if (import.meta.hot) { + import.meta.hot.dispose(unsubscribeHfTokenMirror); +} + export function resolveSpeculativeSettingsForLoad({ usePersistedPreference = false, }: { diff --git a/studio/frontend/src/features/hub/index.ts b/studio/frontend/src/features/hub/index.ts index ddcef146b0..3515f6ca76 100644 --- a/studio/frontend/src/features/hub/index.ts +++ b/studio/frontend/src/features/hub/index.ts @@ -2,3 +2,8 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 export { cancelStagedModelDownload } from "./download-manager"; +export { + getHfToken, + mirrorHfTokenInto, + useHfTokenStore, +} from "./stores/hf-token-store"; diff --git a/studio/frontend/src/features/hub/stores/hf-token-store.ts b/studio/frontend/src/features/hub/stores/hf-token-store.ts index b1e2560f02..499ba9f644 100644 --- a/studio/frontend/src/features/hub/stores/hf-token-store.ts +++ b/studio/frontend/src/features/hub/stores/hf-token-store.ts @@ -5,11 +5,9 @@ import { create } from "zustand"; import { bumpInventoryVersion } from "./inventory-events"; const HF_TOKEN_KEY = "unsloth_hf_token"; -const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed"; const LEGACY_TRAINING_KEY = "unsloth_training_config_v1"; let storageSyncStarted = false; let storageSyncListener: ((event: StorageEvent) => void) | null = null; -let tokenChangedListener: ((event: Event) => void) | null = null; function canUseStorage(): boolean { return typeof window !== "undefined"; @@ -63,10 +61,6 @@ function stopStorageSync(): void { window.removeEventListener("storage", storageSyncListener); storageSyncListener = null; } - if (tokenChangedListener !== null) { - window.removeEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener); - tokenChangedListener = null; - } storageSyncStarted = false; } @@ -100,10 +94,6 @@ export const useHfTokenStore = create((set) => { applyToken(event.newValue ?? "", false); }; window.addEventListener("storage", storageSyncListener); - tokenChangedListener = (event) => { - applyToken((event as CustomEvent).detail ?? "", false); - }; - window.addEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener); } return { @@ -117,6 +107,21 @@ export function getHfToken(): string { return useHfTokenStore.getState().token; } +// Keep a plain zustand store's `hfToken` field in sync with the shared token: +// seed the current value, then mirror later edits. Returns the unsubscribe so +// callers can wire it to HMR disposal. +export function mirrorHfTokenInto(store: { + getState: () => T; + setState: (partial: Partial) => void; +}): () => void { + store.setState({ hfToken: getHfToken() } as Partial); + return useHfTokenStore.subscribe((state) => { + if (store.getState().hfToken !== state.token) { + store.setState({ hfToken: state.token } as Partial); + } + }); +} + // HF's JS client throws on a non-empty token that isn't `hf_...` instead of // browsing anonymously, so treat anything malformed as no token. export function hfApiToken( diff --git a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx index 5bebefa84c..32b3e53a2c 100644 --- a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx +++ b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx @@ -14,6 +14,9 @@ import { import { SettingsRow } from "./settings-row"; import { SettingsSection } from "./settings-section"; +// Mirrors MIN_AUTO_UNLOAD_IDLE_SECONDS in the backend settings store. +const MIN_IDLE_SECONDS = 60; + export function ModelAutoSwitchSection() { const t = useT(); const [settings, setSettings] = useState( @@ -45,13 +48,16 @@ export function ModelAutoSwitchSection() { }; }, [t]); - // Parse the idle-seconds draft to a non-negative integer; empty/invalid -> null. + // Parse the idle-seconds draft: 0 (off) or >= MIN_IDLE_SECONDS; else null. const parseIdleSeconds = (): number | null => { if (!draftIdleSeconds.trim()) { return null; } const parsed = Number(draftIdleSeconds); - return Number.isInteger(parsed) && parsed >= 0 ? parsed : null; + if (!Number.isInteger(parsed)) { + return null; + } + return parsed === 0 || parsed >= MIN_IDLE_SECONDS ? parsed : null; }; const persist = async ( diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index d00d192bb2..6aa9329609 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -54,7 +54,7 @@ import { // Imported directly from the store module rather than the "@/features/training" // barrel to avoid an import cycle (the barrel re-exports this section's siblings). import { hasSeparateStreamingEvalSplit } from "@/features/training/stores/training-config-store"; -import { useDebouncedValue, useHfTokenValidation } from "@/hooks"; +import { useDebouncedValue } from "@/hooks"; import { translate, useT } from "@/i18n"; import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; import { toast } from "@/lib/toast"; @@ -70,11 +70,13 @@ import { } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useNavigate } from "@tanstack/react-router"; +import { motion, useReducedMotion } from "motion/react"; import { type ChangeEvent, type DragEvent, useCallback, useEffect, + useId, useMemo, useRef, useState, @@ -153,6 +155,9 @@ function normalizeSliceInput(value: string): string | null { export function DatasetSection() { const t = useT(); const navigate = useNavigate(); + const reducedMotion = useReducedMotion(); + // Scopes the pill layoutId so multiple instances never share one. + const sourcePillLayoutId = useId(); const { dataset, datasetSource, @@ -398,9 +403,6 @@ export function DatasetSection() { enabled: pickerTab === "huggingface", }); - const { error: tokenValidationError, isChecking: isCheckingToken } = - useHfTokenValidation(hfToken); - const hfResultIds = useMemo(() => { const ids = hfResults.map((r) => r.id); if (dataset && !ids.includes(dataset)) { @@ -683,16 +685,15 @@ export function DatasetSection() { title={t("studio.dataset.title")} description={t("studio.dataset.description")} accent="indigo" - className={`dark:shadow-border ${ - advancedOpen || (datasetSource === "upload" && uploadedFile) - ? "min-h-studio-config-column" - : "h-studio-config-column" - }`} + className="dark:shadow-border min-h-studio-config-column" >
{(() => { // Hub-style sliding-pill segmented control, matching the Hub tabs // via the shared .hub-tab-toggle / .hub-tab-toggle-pill classes. + // flex-auto buttons share leftover space equally so padding stays + // equal for all labels; the pill sits inside the active button so + // it always matches its bounds. const sourceTabs: { value: "huggingface" | "upload" | "s3"; label: string; @@ -703,24 +704,12 @@ export function DatasetSection() { ? [] : [{ value: "s3" as const, label: "Amazon S3" }]), ]; - const activeIndex = Math.max( - 0, - sourceTabs.findIndex((item) => item.value === datasetSource), - ); return (
-
@@ -1009,9 +1015,9 @@ export function DatasetSection() {
- {(tokenValidationError ?? hfSearchError) && ( + {hfSearchError && (

- {tokenValidationError ?? hfSearchError} + {hfSearchError} {" — "}

)} - {isCheckingToken && ( -

- {t("studio.dataset.checkingToken")} -

- )} {pickerTab !== activeSourceTab && (

{t("studio.dataset.browsingSource", { diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 8d6e18b83f..ad438f63f1 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -192,7 +192,6 @@ export function ParamsSection(): ReactElement { const showVisionImageSize = showVisionLora && !isDeepseekOcr; const [loraOpen, setLoraOpen] = useState(false); const [hyperOpen, setHyperOpen] = useState(false); - const needsExpandedHeight = isCpt || (isLora && loraOpen) || hyperOpen; const [ctxInput, setCtxInput] = useState(String(store.contextLength)); const ctxAnchorRef = useRef(null); const ctxItems = CONTEXT_LENGTHS.map(String); @@ -233,11 +232,7 @@ export function ParamsSection(): ReactElement { title={t("studio.params.title")} description={t("studio.params.description")} accent="orange" - className={`${ - needsExpandedHeight - ? "min-h-studio-config-column" - : "h-studio-config-column" - } duration-150`} + className="min-h-studio-config-column" >

diff --git a/studio/frontend/src/features/studio/sections/training-section.tsx b/studio/frontend/src/features/studio/sections/training-section.tsx index 6d1ee9411b..5650b9c145 100644 --- a/studio/frontend/src/features/studio/sections/training-section.tsx +++ b/studio/frontend/src/features/studio/sections/training-section.tsx @@ -53,7 +53,6 @@ export function TrainingSection() { ((!store.isVisionModel && store.isDatasetImage === true) || (!store.isAudioModel && store.isDatasetAudio === true)); const configValidation = validateTrainingConfig(store); - const hasMessage = !!(startError || isIncompatible || (!configValidation.ok && configValidation.message)); const fileInputRef = useRef(null); const handleFileUpload = (e: React.ChangeEvent) => { @@ -125,7 +124,7 @@ export function TrainingSection() { title={t("studio.training.title")} description={t("studio.training.description")} accent="blue" - className={hasMessage ? "min-h-studio-config-column" : "h-studio-config-column"} + className="min-h-studio-config-column" >
{/* Loss chart */} diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 306c02475c..a927f83fd8 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -3,6 +3,7 @@ import { CPT_TARGET_MODULES, DEFAULT_HYPERPARAMS, LR_DEFAULT_CPT, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS, TARGET_MODULES } from "@/config/training"; import { authFetch } from "@/features/auth"; +import { getHfToken, mirrorHfTokenInto, useHfTokenStore } from "@/features/hub"; import { isAdapterMethod } from "@/types/training"; import type { DatasetFormat } from "@/types/training"; import type { ModelType, StepNumber, TrainingMethod } from "@/types/training"; @@ -117,7 +118,9 @@ let _datasetFormatAutoForcedByCpt = false; // modelType / isVisionModel / isAudioModel persist so multimodal-only UI // paints right on reload; the model-config fetch still re-derives them. +// hfToken mirrors the shared hf-token-store and is persisted there instead. const NON_PERSISTED_STATE_KEYS: ReadonlySet = new Set([ + "hfToken", "isCheckingVision", "isEmbeddingModel", "isLoadingModelDefaults", @@ -632,8 +635,7 @@ export const useTrainingConfigStore = create()( ), ); }, - setHfToken: (hfToken) => - set({ hfToken: hfToken.trim().replace(/^["']+|["']+$/g, "") }), + setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken), setDatasetSource: (datasetSource) => set({ datasetSource }), selectHfDataset: (dataset) => { _datasetCheckController?.abort(); @@ -923,7 +925,7 @@ export const useTrainingConfigStore = create()( _learningRateManuallySet = false; _yamlLearningRate = undefined; clearCptDatasetFormatTracking(); - set(initialState); + set({ ...initialState, hfToken: getHfToken() }); }, resetToModelDefaults: () => { const { selectedModel } = get(); @@ -947,7 +949,7 @@ export const useTrainingConfigStore = create()( }, { name: "unsloth_training_config_v1", - version: 11, + version: 12, migrate: (persisted, version) => { const s = persisted as Record; if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) { @@ -1000,6 +1002,15 @@ export const useTrainingConfigStore = create()( // own version guard. s.datasetStreaming ??= false; } + if (version < 12) { + // hfToken moved to the shared hf-token-store; seed it once so an + // existing Studio-only token isn't lost. + const legacyToken = typeof s.hfToken === "string" ? s.hfToken.trim() : ""; + if (legacyToken && !getHfToken()) { + useHfTokenStore.getState().setToken(legacyToken); + } + delete s.hfToken; + } return s as unknown as TrainingConfigStore; }, partialize: partializePersistedState, @@ -1022,3 +1033,8 @@ export const useTrainingConfigStore = create()( }, ), ); + +const unsubscribeHfTokenMirror = mirrorHfTokenInto(useTrainingConfigStore); +if (import.meta.hot) { + import.meta.hot.dispose(unsubscribeHfTokenMirror); +} diff --git a/studio/frontend/src/i18n/locales/ar.ts b/studio/frontend/src/i18n/locales/ar.ts index 28f404a384..744a2002c9 100644 --- a/studio/frontend/src/i18n/locales/ar.ts +++ b/studio/frontend/src/i18n/locales/ar.ts @@ -155,14 +155,14 @@ export const ar = { "عندما يسمّي طلب متوافق مع OpenAI ملف GGUF مُنزّلاً مختلفًا، يتم تحميله قبل الخدمة. مُعطّل افتراضيًا؛ الأسماء غير المعروفة تُبقي على النموذج المُحمَّل.", idleUnload: "الإلغاء التلقائي عند الخمول", idleUnloadDescription: - "إلغاء تحميل النموذج بعد هذا العدد من ثواني الخمول لتحرير الـ VRAM؛ الطلب التالي يعيد تحميله. القيمة 0 تُبقيه محمَّلاً.", + "إلغاء تحميل النموذج بعد هذا العدد من ثواني الخمول لتحرير الـ VRAM؛ الطلب التالي يعيد تحميله. القيمة 0 تُبقيه محمَّلاً. الحد الأدنى 60 ثانية.", idleNeedsEnable: "فعّل تبديل النموذج حسب الطلب حتى يعاد تحميل النموذج غير المحمَّل عند الاستخدام التالي.", idleActiveViaEnv: "الإلغاء التلقائي عند الخمول مُفعَّل عبر متغير البيئة UNSLOTH_MODEL_IDLE_TTL.", loadError: "فشل تحميل إعدادات التبديل التلقائي للنموذج.", saveError: "فشل حفظ إعدادات التبديل التلقائي للنموذج.", - idleError: "أدخل عددًا صحيحًا من الثواني (0 أو أكثر).", + idleError: "أدخل 0 لإبقاء النموذج محمَّلاً، أو 60 ثانية على الأقل.", }, previewSharing: { sectionTitle: "مشاركة المعاينة", diff --git a/studio/frontend/src/i18n/locales/de.ts b/studio/frontend/src/i18n/locales/de.ts index e38cdbfa0e..7d94e7656e 100644 --- a/studio/frontend/src/i18n/locales/de.ts +++ b/studio/frontend/src/i18n/locales/de.ts @@ -158,7 +158,7 @@ export const de = { "Wenn eine OpenAI-kompatible Anfrage ein anderes heruntergeladenes GGUF nennt, wird dieses vor der Auslieferung geladen. Standardmäßig aus; unbekannte Namen liefern weiterhin das geladene Modell aus.", idleUnload: "Automatisches Entladen bei Inaktivität", idleUnloadDescription: - "Entlädt das Modell nach dieser Anzahl inaktiver Sekunden, um VRAM freizugeben; die nächste Anfrage lädt es erneut. 0 hält es geladen.", + "Entlädt das Modell nach dieser Anzahl inaktiver Sekunden, um VRAM freizugeben; die nächste Anfrage lädt es erneut. 0 hält es geladen. Minimum 60 Sekunden.", idleNeedsEnable: "Aktivieren Sie \"Modell je Anfrage wechseln\", damit ein entladenes Modell bei der nächsten Nutzung erneut geladen wird.", idleActiveViaEnv: @@ -167,7 +167,7 @@ export const de = { "Einstellungen für automatischen Modellwechsel konnten nicht geladen werden.", saveError: "Einstellungen für automatischen Modellwechsel konnten nicht gespeichert werden.", - idleError: "Geben Sie eine ganze Anzahl an Sekunden ein (0 oder mehr).", + idleError: "Geben Sie 0 ein, um das Modell geladen zu halten, oder mindestens 60 Sekunden.", }, previewSharing: { sectionTitle: "Vorschau-Freigabe", diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index fe3a6f8542..de8ac17c29 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -224,14 +224,14 @@ export const en = { "When an OpenAI-compatible request names a different downloaded GGUF, load it before serving. Off by default; unknown names keep serving the loaded model.", idleUnload: "Idle auto-unload", idleUnloadDescription: - "Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded.", + "Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded. Minimum 60 seconds.", idleNeedsEnable: "Turn on Switch model by request so an unloaded model reloads on next use.", idleActiveViaEnv: "Idle auto-unload is active via the UNSLOTH_MODEL_IDLE_TTL environment variable.", loadError: "Failed to load model auto-switch settings.", saveError: "Failed to save model auto-switch settings.", - idleError: "Enter a whole number of seconds (0 or more).", + idleError: "Enter 0 to keep the model loaded, or at least 60 seconds.", }, previewSharing: { sectionTitle: "Preview sharing", diff --git a/studio/frontend/src/i18n/locales/es.ts b/studio/frontend/src/i18n/locales/es.ts index e5f9650bef..988c109a3f 100644 --- a/studio/frontend/src/i18n/locales/es.ts +++ b/studio/frontend/src/i18n/locales/es.ts @@ -157,7 +157,7 @@ export const es = { "Cuando una solicitud compatible con OpenAI nombra un GGUF descargado distinto, se carga antes de responder. Desactivado por defecto; los nombres desconocidos siguen usando el modelo cargado.", idleUnload: "Descarga automática por inactividad", idleUnloadDescription: - "Descarga el modelo tras este número de segundos inactivo para liberar VRAM; la siguiente solicitud lo recarga. 0 lo mantiene cargado.", + "Descarga el modelo tras este número de segundos inactivo para liberar VRAM; la siguiente solicitud lo recarga. 0 lo mantiene cargado. Mínimo 60 segundos.", idleNeedsEnable: "Activa Cambiar de modelo según la solicitud para que un modelo descargado se recargue en el próximo uso.", idleActiveViaEnv: @@ -166,7 +166,7 @@ export const es = { "No se pudo cargar la configuración de cambio automático de modelo.", saveError: "No se pudo guardar la configuración de cambio automático de modelo.", - idleError: "Introduce un número entero de segundos (0 o más).", + idleError: "Introduce 0 para mantener el modelo cargado, o al menos 60 segundos.", }, previewSharing: { sectionTitle: "Compartir vista previa", diff --git a/studio/frontend/src/i18n/locales/fr.ts b/studio/frontend/src/i18n/locales/fr.ts index 190284175d..e1f2a0c5ec 100644 --- a/studio/frontend/src/i18n/locales/fr.ts +++ b/studio/frontend/src/i18n/locales/fr.ts @@ -157,7 +157,7 @@ export const fr = { "Lorsqu'une requête compatible OpenAI nomme un autre GGUF téléchargé, le charger avant de répondre. Désactivé par défaut ; les noms inconnus continuent de servir le modèle chargé.", idleUnload: "Déchargement automatique en cas d'inactivité", idleUnloadDescription: - "Décharger le modèle après ce nombre de secondes d'inactivité pour libérer la VRAM ; la requête suivante le recharge. 0 le maintient chargé.", + "Décharger le modèle après ce nombre de secondes d'inactivité pour libérer la VRAM ; la requête suivante le recharge. 0 le maintient chargé. Minimum 60 secondes.", idleNeedsEnable: "Activez Changer de modèle par requête pour qu'un modèle déchargé se recharge à la prochaine utilisation.", idleActiveViaEnv: @@ -166,7 +166,7 @@ export const fr = { "Échec du chargement des paramètres de changement automatique de modèle.", saveError: "Échec de l'enregistrement des paramètres de changement automatique de modèle.", - idleError: "Saisissez un nombre entier de secondes (0 ou plus).", + idleError: "Saisissez 0 pour garder le modèle chargé, ou au moins 60 secondes.", }, previewSharing: { sectionTitle: "Partage de l'aperçu", diff --git a/studio/frontend/src/i18n/locales/hi.ts b/studio/frontend/src/i18n/locales/hi.ts index 97f55251c5..77b6265e7b 100644 --- a/studio/frontend/src/i18n/locales/hi.ts +++ b/studio/frontend/src/i18n/locales/hi.ts @@ -154,14 +154,14 @@ export const hi = { "जब कोई OpenAI-संगत अनुरोध किसी अन्य डाउनलोड किए गए GGUF का नाम लेता है, तो सर्व करने से पहले उसे लोड करें। डिफ़ॉल्ट रूप से बंद; अज्ञात नाम लोड किए गए मॉडल को सर्व करते रहते हैं।", idleUnload: "निष्क्रिय ऑटो-अनलोड", idleUnloadDescription: - "VRAM मुक्त करने के लिए इतने निष्क्रिय सेकंड के बाद मॉडल को अनलोड करें; अगला अनुरोध इसे फिर से लोड करता है। 0 इसे लोड रखता है।", + "VRAM मुक्त करने के लिए इतने निष्क्रिय सेकंड के बाद मॉडल को अनलोड करें; अगला अनुरोध इसे फिर से लोड करता है। 0 इसे लोड रखता है। न्यूनतम 60 सेकंड।", idleNeedsEnable: "अनुरोध के अनुसार मॉडल बदलें चालू करें ताकि अनलोड किया गया मॉडल अगले उपयोग पर फिर से लोड हो।", idleActiveViaEnv: "निष्क्रिय ऑटो-अनलोड UNSLOTH_MODEL_IDLE_TTL एनवायरनमेंट वेरिएबल के माध्यम से सक्रिय है।", loadError: "मॉडल ऑटो-स्विच सेटिंग्स लोड करने में विफल।", saveError: "मॉडल ऑटो-स्विच सेटिंग्स सहेजने में विफल।", - idleError: "सेकंड की पूरी संख्या दर्ज करें (0 या अधिक)।", + idleError: "मॉडल को लोड रखने के लिए 0 दर्ज करें, या कम से कम 60 सेकंड।", }, previewSharing: { sectionTitle: "पूर्वावलोकन साझाकरण", diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts index 23752b6c26..a261994f03 100644 --- a/studio/frontend/src/i18n/locales/ja.ts +++ b/studio/frontend/src/i18n/locales/ja.ts @@ -149,12 +149,12 @@ export const ja = { enable: "リクエストごとにモデルを切り替え", enableDescription: "OpenAI互換のリクエストが別のダウンロード済み GGUF を指定した場合、応答する前にそのモデルを読み込みます。デフォルトはオフです。不明な名前の場合は、読み込み済みのモデルで応答を続けます。", idleUnload: "アイドル時の自動アンロード", - idleUnloadDescription: "指定した秒数だけアイドル状態が続くとモデルをアンロードして VRAM を解放します。次のリクエストで再読み込みされます。0 にすると読み込んだままにします。", + idleUnloadDescription: "指定した秒数だけアイドル状態が続くとモデルをアンロードして VRAM を解放します。次のリクエストで再読み込みされます。0 にすると読み込んだままにします。最小 60 秒。", idleNeedsEnable: "アンロードされたモデルが次回使用時に再読み込みされるように、「リクエストごとにモデルを切り替え」をオンにしてください。", idleActiveViaEnv: "アイドル時の自動アンロードは UNSLOTH_MODEL_IDLE_TTL 環境変数によって有効になっています。", loadError: "モデル自動切り替え設定の読み込みに失敗しました。", saveError: "モデル自動切り替え設定の保存に失敗しました。", - idleError: "秒数を整数（0 以上）で入力してください。", + idleError: "モデルを読み込んだままにするには 0 を、それ以外は 60 秒以上を入力してください。", }, previewSharing: { sectionTitle: "プレビュー共有", diff --git a/studio/frontend/src/i18n/locales/ko.ts b/studio/frontend/src/i18n/locales/ko.ts index a11a94faf2..7c5691925e 100644 --- a/studio/frontend/src/i18n/locales/ko.ts +++ b/studio/frontend/src/i18n/locales/ko.ts @@ -153,14 +153,14 @@ export const ko = { "OpenAI 호환 요청이 다운로드된 다른 GGUF를 지정하면, 응답하기 전에 해당 모델을 불러옵니다. 기본값은 꺼짐이며, 알 수 없는 이름은 불러온 모델을 계속 제공합니다.", idleUnload: "유휴 시 자동 해제", idleUnloadDescription: - "지정한 유휴 시간(초)이 지나면 모델을 해제하여 VRAM을 확보합니다. 다음 요청 시 다시 불러옵니다. 0으로 설정하면 계속 로드된 상태로 유지됩니다.", + "지정한 유휴 시간(초)이 지나면 모델을 해제하여 VRAM을 확보합니다. 다음 요청 시 다시 불러옵니다. 0으로 설정하면 계속 로드된 상태로 유지됩니다. 최소 60초입니다.", idleNeedsEnable: "해제된 모델이 다음 사용 시 다시 로드되도록 하려면 요청에 따라 모델 전환을 켜세요.", idleActiveViaEnv: "유휴 시 자동 해제가 UNSLOTH_MODEL_IDLE_TTL 환경 변수를 통해 활성화되어 있습니다.", loadError: "모델 자동 전환 설정을 불러오지 못했습니다.", saveError: "모델 자동 전환 설정을 저장하지 못했습니다.", - idleError: "정수(초)를 입력하세요(0 이상).", + idleError: "모델을 로드 상태로 유지하려면 0을, 그렇지 않으면 60초 이상을 입력하세요.", }, previewSharing: { sectionTitle: "미리보기 공유", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts index 07832ef1d0..e6d2347c10 100644 --- a/studio/frontend/src/i18n/locales/pt-br.ts +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -157,14 +157,14 @@ export const ptBR = { "Quando uma requisição compatível com OpenAI nomear um GGUF baixado diferente, carrega-o antes de responder. Desativado por padrão; nomes desconhecidos continuam usando o modelo carregado.", idleUnload: "Descarregamento automático por inatividade", idleUnloadDescription: - "Descarrega o modelo após esta quantidade de segundos de inatividade para liberar VRAM; a próxima requisição o recarrega. 0 mantém o modelo carregado.", + "Descarrega o modelo após esta quantidade de segundos de inatividade para liberar VRAM; a próxima requisição o recarrega. 0 mantém o modelo carregado. Mínimo de 60 segundos.", idleNeedsEnable: "Ative Trocar de modelo por requisição para que um modelo descarregado seja recarregado no próximo uso.", idleActiveViaEnv: "O descarregamento automático por inatividade está ativo por meio da variável de ambiente UNSLOTH_MODEL_IDLE_TTL.", loadError: "Falha ao carregar as configurações de troca automática de modelo.", saveError: "Falha ao salvar as configurações de troca automática de modelo.", - idleError: "Insira um número inteiro de segundos (0 ou mais).", + idleError: "Insira 0 para manter o modelo carregado, ou pelo menos 60 segundos.", }, previewSharing: { sectionTitle: "Compartilhamento de pré-visualização", diff --git a/studio/frontend/src/i18n/locales/ru.ts b/studio/frontend/src/i18n/locales/ru.ts index 81d20cc2ea..c7464a3b44 100644 --- a/studio/frontend/src/i18n/locales/ru.ts +++ b/studio/frontend/src/i18n/locales/ru.ts @@ -154,14 +154,14 @@ export const ru = { "Когда OpenAI-совместимый запрос указывает другую загруженную GGUF, загружать её перед обслуживанием. По умолчанию выключено; неизвестные имена продолжают обслуживать загруженную модель.", idleUnload: "Автовыгрузка при простое", idleUnloadDescription: - "Выгружать модель после указанного числа секунд простоя, чтобы освободить VRAM; следующий запрос загрузит её снова. 0 оставляет модель загруженной.", + "Выгружать модель после указанного числа секунд простоя, чтобы освободить VRAM; следующий запрос загрузит её снова. 0 оставляет модель загруженной. Минимум 60 секунд.", idleNeedsEnable: "Включите «Переключать модель по запросу», чтобы выгруженная модель загружалась при следующем использовании.", idleActiveViaEnv: "Автовыгрузка при простое активна через переменную окружения UNSLOTH_MODEL_IDLE_TTL.", loadError: "Не удалось загрузить настройки автопереключения модели.", saveError: "Не удалось сохранить настройки автопереключения модели.", - idleError: "Введите целое число секунд (0 или больше).", + idleError: "Введите 0, чтобы модель оставалась загруженной, или не менее 60 секунд.", }, previewSharing: { sectionTitle: "Публикация предпросмотра", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index 4c51755244..ff218adad2 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -152,14 +152,14 @@ export const zhCN = { "当兼容 OpenAI 的请求指定了另一个已下载的 GGUF 时，先加载它再提供服务。默认关闭；未知名称将继续使用已加载的模型。", idleUnload: "空闲自动卸载", idleUnloadDescription: - "空闲达到该秒数后卸载模型以释放 VRAM；下次请求会重新加载。设为 0 则保持加载。", + "空闲达到该秒数后卸载模型以释放 VRAM；下次请求会重新加载。设为 0 则保持加载。最小 60 秒。", idleNeedsEnable: "开启“按请求切换模型”，以便已卸载的模型在下次使用时重新加载。", idleActiveViaEnv: "空闲自动卸载已通过 UNSLOTH_MODEL_IDLE_TTL 环境变量启用。", loadError: "加载模型自动切换设置失败。", saveError: "保存模型自动切换设置失败。", - idleError: "请输入整数秒数（0 或以上）。", + idleError: "输入 0 保持模型加载，或输入至少 60 秒。", }, previewSharing: { sectionTitle: "预览分享", diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css index 159eba743f..6d4c21eec8 100644 --- a/studio/frontend/src/index.css +++ b/studio/frontend/src/index.css @@ -1467,7 +1467,9 @@ html[data-chat-font] .aui-root { } /* With more than 4 tools on, drop pill labels to icons only to cut clutter. - Compare keeps its label via data-keep-label. */ + Compare and the bypass-permissions pill keep their labels via + data-keep-label; the permission pill sits before the collapsed icons, so + they line up to its right. */ [data-pill-compact="true"] .composer-pill-btn:not([data-keep-label]) > span:not(.composer-pill-glyph) { @@ -2048,15 +2050,11 @@ html[data-chat-font] .aui-root { } } - /* Fine-tuning Studio: equal default height, expandable when needed (md+) */ + /* Fine-tuning Studio: equal minimum card height, grows with content (md+) */ .min-h-studio-config-column { @apply md:min-h-[520px]; } - .h-studio-config-column { - @apply md:h-[520px]; - } - [data-streamdown="unordered-list"] { list-style-type: disc; list-style-position: outside; diff --git a/tests/python/test_get_lora_parameters_bias_fp8_block_size.py b/tests/python/test_get_lora_parameters_bias_fp8_block_size.py new file mode 100644 index 0000000000..835ad19542 --- /dev/null +++ b/tests/python/test_get_lora_parameters_bias_fp8_block_size.py @@ -0,0 +1,82 @@ +import ast +from pathlib import Path + + +def _load_function(name): + # Extract a function from kernels/utils.py without importing unsloth (which + # needs a GPU / torch / bitsandbytes). The functions exercised here only use + # getattr and the _FP8_WEIGHT_DTYPES name on the paths under test. + source = Path(__file__).parents[2] / "unsloth" / "kernels" / "utils.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + funcs = [ + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name + ] + assert len(funcs) == 1, (name, funcs) + namespace = {"getattr": getattr, "_FP8_WEIGHT_DTYPES": ()} + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace[name] + + +class _Obj: + pass + + +def _make_disabled_block_fp8_proj(block_size): + # A merged/disabled projection whose base layer is a block-fp8 weight that + # ships a non-default block size on its checkpoint. + weight = _Obj() + weight.quant_state = _Obj() + base_layer = _Obj() + base_layer.weight = weight + base_layer.quant_method = "fp8" + base_layer.block_size = block_size + base_layer.bias = None + proj = _Obj() + proj.base_layer = base_layer + proj.merged = True + proj.disable_adapters = True + return proj, weight.quant_state + + +def test_bias_variant_propagates_fp8_block_size_on_disabled_path(): + # Downstream fp8 kernels read getattr(weight_scale, "block_size", [128, 128]), + # so the checkpoint's real block size must survive the merged/disabled path, + # exactly as it does for the non-bias sibling get_lora_parameters. + get_lora_parameters_bias = _load_function("get_lora_parameters_bias") + + proj, weight_scale = _make_disabled_block_fp8_proj([64, 128]) + get_lora_parameters_bias(proj) + + assert getattr(weight_scale, "block_size", [128, 128]) == [64, 128] + + +def _make_decompressed_merged_proj(): + # A merged compressed-tensors layer that was decompressed back to bf16. It keeps + # quant_method == "fp8" from the checkpoint metadata, but the live weight is bf16 + # so there is no quant state to attach a block size to. + weight = _Obj() + weight.dtype = "bfloat16" + base_layer = _Obj() + base_layer.weight = weight + base_layer.quant_method = "fp8" + base_layer.block_size = [128, 128] + base_layer.bias = None + proj = _Obj() + proj.base_layer = base_layer + proj.merged = True + proj.disable_adapters = True + return proj + + +def test_bias_variant_keeps_none_quant_state_for_decompressed_layer(): + # Such a layer has no quant state, and fast_linear_forward relies on getting + # W_quant None back so it can fall back to a plain matmul, so setting the block + # size must not assume a quant state is present. + get_lora_parameters_bias = _load_function("get_lora_parameters_bias") + + W, W_quant = get_lora_parameters_bias(_make_decompressed_merged_proj())[:2] + + assert W_quant is None + assert getattr(W, "block_size", None) == [128, 128] diff --git a/tests/python/test_get_lora_parameters_fp8_block_size.py b/tests/python/test_get_lora_parameters_fp8_block_size.py new file mode 100644 index 0000000000..f5f1359125 --- /dev/null +++ b/tests/python/test_get_lora_parameters_fp8_block_size.py @@ -0,0 +1,81 @@ +import ast +from pathlib import Path + + +def _load_function(name): + # Extract a function from kernels/utils.py without importing unsloth (which + # needs a GPU / torch / bitsandbytes). get_lora_parameters only uses getattr, + # hasattr and the _FP8_WEIGHT_DTYPES name on the paths under test. + source = Path(__file__).parents[2] / "unsloth" / "kernels" / "utils.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + funcs = [ + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name + ] + assert len(funcs) == 1, (name, funcs) + namespace = {"getattr": getattr, "hasattr": hasattr, "_FP8_WEIGHT_DTYPES": ()} + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace[name] + + +class _Obj: + pass + + +def _make_disabled_block_fp8_proj(block_size): + # A merged/disabled projection whose base layer is a block-fp8 weight that + # ships a non-default block size on its checkpoint. + weight = _Obj() + weight.quant_state = _Obj() + base_layer = _Obj() + base_layer.weight = weight + base_layer.quant_method = "fp8" + base_layer.block_size = block_size + proj = _Obj() + proj.base_layer = base_layer + proj.merged = True + proj.disable_adapters = True + return proj, weight.quant_state + + +def test_propagates_fp8_block_size_on_disabled_path(): + # get_lora_parameters already sets block_size before its early return; downstream + # fp8 kernels read getattr(weight_scale, "block_size", [128, 128]), so the + # checkpoint's real block size must survive the merged/disabled path. + get_lora_parameters = _load_function("get_lora_parameters") + + proj, weight_scale = _make_disabled_block_fp8_proj([64, 128]) + get_lora_parameters(proj) + + assert getattr(weight_scale, "block_size", [128, 128]) == [64, 128] + + +def _make_decompressed_merged_proj(): + # A merged compressed-tensors layer that was decompressed back to bf16. It keeps + # quant_method == "fp8" from the checkpoint metadata, but the live weight is bf16 + # so there is no quant state to attach a block size to. + weight = _Obj() + weight.dtype = "bfloat16" + base_layer = _Obj() + base_layer.weight = weight + base_layer.quant_method = "fp8" + base_layer.block_size = [128, 128] + proj = _Obj() + proj.base_layer = base_layer + proj.merged = True + proj.disable_adapters = True + return proj + + +def test_keeps_none_quant_state_for_decompressed_layer(): + # Mirrors the get_lora_parameters_bias guard: with no quant state, assigning + # W_quant.block_size must not assume one is present, or it raises AttributeError + # on None. fast_lora relies on getting W_quant None back to fall back to a plain + # matmul, so this path must stay crash-free. + get_lora_parameters = _load_function("get_lora_parameters") + + W, W_quant = get_lora_parameters(_make_decompressed_merged_proj())[:2] + + assert W_quant is None + assert getattr(W, "block_size", None) == [128, 128] diff --git a/tests/saving/test_export_dispatch.py b/tests/saving/test_export_dispatch.py index 3870d6269d..ad6d51e0b7 100644 --- a/tests/saving/test_export_dispatch.py +++ b/tests/saving/test_export_dispatch.py @@ -8,6 +8,8 @@ regressions that pure AST checks cannot (e.g. wrong scheme/suffix/outtype passed from __future__ import annotations +import inspect + import pytest import unsloth.save as save_mod @@ -126,6 +128,80 @@ def test_gguf_lora_push_to_hub_is_rejected(tmp_path): ) +# The above rejection points users at push_to_hub_gguf(save_method='lora'), so that path +# has to work; it is only ever exercised here. + + +def test_push_to_hub_gguf_lora_dispatches(monkeypatch): + seen = {} + monkeypatch.setattr( + save_mod, + "_unsloth_save_lora_gguf", + lambda model, tok, sd, **kw: seen.update(kw), + ) + save_mod.unsloth_push_to_hub_gguf( + _FakeModel(), + "repo/id", + tokenizer = object(), + save_method = "lora", + quantization_method = "q8_0", + ) + assert seen.get("outtype") == "q8_0" + assert seen.get("push_to_hub") is True + + +def test_push_to_hub_gguf_lora_skips_non_main_process(monkeypatch): + calls = [] + monkeypatch.setattr( + save_mod, + "_unsloth_save_lora_gguf", + lambda *a, **kw: calls.append(kw), + ) + result = save_mod.unsloth_push_to_hub_gguf( + _FakeModel(), + "repo/id", + tokenizer = object(), + save_method = "lora", + is_main_process = False, + ) + assert result is None + assert calls == [] + + +def test_push_to_hub_gguf_skips_non_main_process_before_merged_conversion(monkeypatch): + calls = [] + monkeypatch.setattr( + save_mod, + "unsloth_save_pretrained_gguf", + lambda **kw: calls.append(kw), + ) + result = save_mod.unsloth_push_to_hub_gguf( + _FakeModel(), + "repo/id", + tokenizer = object(), + is_main_process = False, + ) + assert result is None + assert calls == [] + + +def test_push_to_hub_gguf_preserves_positional_max_shard_size(): + bound = inspect.signature(save_mod.unsloth_push_to_hub_gguf).bind( + _FakeModel(), + "repo/id", + object(), + "q4_k_m", + None, + None, + None, + None, + "token", + "50GB", + ) + assert bound.arguments["max_shard_size"] == "50GB" + assert "is_main_process" not in bound.arguments + + # -- torchao PTQ / QAT dispatch ------------------------------------------------------------ diff --git a/tests/saving/test_fix_sentencepiece_tokenizer_guard.py b/tests/saving/test_fix_sentencepiece_tokenizer_guard.py new file mode 100644 index 0000000000..1ee523d57b --- /dev/null +++ b/tests/saving/test_fix_sentencepiece_tokenizer_guard.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: AGPL-3.0-only +import gc +import os + +os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python") + +import transformers +from transformers.utils import sentencepiece_model_pb2 + +from unsloth.tokenizer_utils import fix_sentencepiece_tokenizer + + +NORMAL, CONTROL = 1, 3 + + +def _spm_bytes(pieces): + m = sentencepiece_model_pb2.ModelProto() + for piece, score, typ in pieces: + p = m.pieces.add() + p.piece = piece + p.score = score + p.type = typ + return m.SerializeToString() + + +def _read_pieces(path): + m = sentencepiece_model_pb2.ModelProto() + with open(path, "rb") as f: + m.ParseFromString(f.read()) + return [p.piece for p in m.pieces] + + +class _FakeTokenizer: + """Minimal stand-in for a sentencepiece-backed slow tokenizer. + + ``save_pretrained`` writes a tokenizer.model, which is what the real slow + tokenizers do and what fix_sentencepiece_tokenizer reads back. + """ + + def __init__( + self, + name, + spm_bytes = None, + vocab = None, + ): + self.name = name + self.eos_token = "" + self.pad_token = "" + self._spm_bytes = spm_bytes + self._vocab = vocab or {} + self.saved_to = [] + + def save_pretrained(self, location): + self.saved_to.append(location) + os.makedirs(location, exist_ok = True) + if self._spm_bytes is not None: + with open(os.path.join(location, "tokenizer.model"), "wb") as f: + f.write(self._spm_bytes) + + def __call__( + self, + texts, + add_special_tokens = False, + ): + class _Encoded: + pass + + encoded = _Encoded() + encoded.input_ids = [[self._vocab[text]] for text in texts] + return encoded + + +def _tokenizers(): + pieces = [("", 0.0, CONTROL), ("a", -1.0, NORMAL), ("", 0.0, CONTROL)] + old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"": 2}) + new = _FakeTokenizer("new") + return old, new + + +class _ReloadedTokenizer: + """Weakref-able stand-in for the tokenizer AutoTokenizer.from_pretrained returns.""" + + def __init__(self, location): + self.location = location + + +def _stub_auto_tokenizer(monkeypatch): + """fix_sentencepiece_tokenizer reloads the patched directory through + AutoTokenizer at the end; that needs a full tokenizer on disk, which is + out of scope here. Record the reload location and hand back a sentinel. + """ + loaded = [] + + class _StubAutoTokenizer: + @staticmethod + def from_pretrained(location, **kwargs): + loaded.append(location) + return _ReloadedTokenizer(location) + + monkeypatch.setattr(transformers, "AutoTokenizer", _StubAutoTokenizer) + return loaded + + +def test_old_tokenizer_is_saved_so_its_model_can_be_read(tmp_path, monkeypatch): + """The guard must not skip the body on a fresh temporary directory. + + fix_sentencepiece_tokenizer creates its scratch directory itself and then + checks for a tokenizer.model inside it, but that file only appears once + old_tokenizer.save_pretrained() has run. + """ + _stub_auto_tokenizer(monkeypatch) + old, new = _tokenizers() + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location) + + assert old.saved_to, "old tokenizer was never saved: the body did not run" + + +def test_token_mapping_is_applied_to_the_sentencepiece_model(tmp_path, monkeypatch): + loaded = _stub_auto_tokenizer(monkeypatch) + old, new = _tokenizers() + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + # Hold the returned tokenizer so its scratch dir survives until we read it. + tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location) + + assert "<|im_end|>" in _read_pieces(f"{loaded[-1]}/tokenizer.model") + assert tok is not None + + +def test_tokenizer_without_a_sentencepiece_model_is_returned_untouched(tmp_path, monkeypatch): + """A fast-only tokenizer writes no tokenizer.model, so the guard still + short-circuits and the caller gets new_tokenizer back unchanged. Its scratch + dir is unreferenced and reclaimed immediately. + """ + _stub_auto_tokenizer(monkeypatch) + old = _FakeTokenizer("old", spm_bytes = None) + new = _FakeTokenizer("new") + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + result = fix_sentencepiece_tokenizer( + old, new, {"": "<|im_end|>"}, temporary_location = location + ) + + assert result is new + assert not any( + name.startswith("tokenizer_") for name in os.listdir(location) + ), "the fast-only scratch dir was not reclaimed" + + +def test_each_call_uses_a_fresh_isolated_subdirectory(tmp_path, monkeypatch): + """Each call must work in its own unique subdirectory, so concurrent or + repeated calls never share scratch files, stale artifacts never leak into + the reload, and nothing the caller left in the scratch location is deleted. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + location = str(tmp_path / "_unsloth_sentencepiece_temp") + os.makedirs(location, exist_ok = True) + + # A pre-existing artifact in the shared scratch location. + marker = os.path.join(location, "leftover.json") + with open(marker, "w") as f: + f.write("{}") + + old1, new1 = _tokenizers() + old2, new2 = _tokenizers() + # Hold both returned tokenizers so their scratch dirs stay alive. + tok1 = fix_sentencepiece_tokenizer( + old1, new1, {"": "<|im_end|>"}, temporary_location = location + ) + tok2 = fix_sentencepiece_tokenizer( + old2, new2, {"": "<|im_end|>"}, temporary_location = location + ) + + work1, work2 = loaded[0], loaded[1] + assert work1 != work2, "two calls reused the same directory" + assert os.path.dirname(work1) == location and os.path.dirname(work2) == location + assert os.path.isdir(work1) and os.path.isdir(work2) + # Nothing the caller left behind is deleted, and it never leaks into a work dir. + assert os.path.isfile(marker), "a pre-existing scratch file was deleted" + assert not os.path.isfile(os.path.join(work1, "leftover.json")) + assert not os.path.isfile(os.path.join(work2, "leftover.json")) + assert tok1 is not None and tok2 is not None + + +def test_sentencepiece_scratch_dir_is_reclaimed_once_the_tokenizer_is_gone(tmp_path, monkeypatch): + """The scratch dir must live as long as the returned tokenizer (its vocab_file + points there), then be reclaimed when the tokenizer is garbage collected. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + old, new = _tokenizers() + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location) + work = loaded[-1] + assert os.path.isdir(work), "scratch dir vanished while the tokenizer was alive" + + del tok + gc.collect() + assert not os.path.isdir(work), "scratch dir was not reclaimed after the tokenizer was freed" + + +class _CopyFromSubdirTokenizer: + """A slow tokenizer whose sentencepiece source lives elsewhere (like the + tokenizers convert_to_fast_tokenizer produces under {location}/{name}). + save_pretrained copies that source into the destination, as HF slow + tokenizers copy their vocab_file. + """ + + def __init__(self, source_model_path): + self.eos_token = "" + self.pad_token = "" + self._source_model_path = source_model_path + + def save_pretrained(self, location): + os.makedirs(location, exist_ok = True) + if os.path.isfile(self._source_model_path): + with open(self._source_model_path, "rb") as src: + data = src.read() + with open(os.path.join(location, "tokenizer.model"), "wb") as dst: + dst.write(data) + + def __call__( + self, + texts, + add_special_tokens = False, + ): + class _Encoded: + pass + + encoded = _Encoded() + encoded.input_ids = [[2] for _ in texts] + return encoded + + +def test_source_vocab_outside_the_work_directory_is_not_disturbed(tmp_path, monkeypatch): + """A tokenizer whose sentencepiece source lives elsewhere (e.g. the subtree + convert_to_fast_tokenizer created) is copied into the fresh work directory + and patched there; the original source is left untouched. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + location = str(tmp_path / "_unsloth_sentencepiece_temp") + subdir = os.path.join(location, "some_model") + os.makedirs(subdir, exist_ok = True) + + pieces = [("", 0.0, CONTROL), ("a", -1.0, NORMAL), ("", 0.0, CONTROL)] + source_model = os.path.join(subdir, "tokenizer.model") + with open(source_model, "wb") as f: + f.write(_spm_bytes(pieces)) + + old = _CopyFromSubdirTokenizer(source_model) + new = _FakeTokenizer("new") + tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location) + + assert _read_pieces(source_model) == [ + "", + "a", + "", + ], "the original source vocab was modified" + assert "<|im_end|>" in _read_pieces(f"{loaded[-1]}/tokenizer.model") + assert tok is not None + + +def test_swap_mapping_swaps_both_pieces_without_duplicating(tmp_path, monkeypatch): + """When the caller swaps eos and stop_word in the fast JSON it must pass both + directions here; a one-way mapping would leave two stop_word pieces and no eos. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + pieces = [("", 0.0, CONTROL), ("<|im_end|>", -1.0, NORMAL), ("", 0.0, CONTROL)] + old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"": 2, "<|im_end|>": 1}) + new = _FakeTokenizer("new") + + tok = fix_sentencepiece_tokenizer( + old, new, {"": "<|im_end|>", "<|im_end|>": ""}, temporary_location = location + ) + + result = _read_pieces(f"{loaded[-1]}/tokenizer.model") + assert result.count("<|im_end|>") == 1 and result.count("") == 1, result + assert tok is not None + + +def test_only_applied_mappings_are_patched(tmp_path, monkeypatch): + """When the caller skips a mapping whose target already exists, it must not + pass that mapping here, or the skipped source token gets renamed anyway and + duplicates the existing target in the model. + """ + loaded = _stub_auto_tokenizer(monkeypatch) + location = str(tmp_path / "_unsloth_sentencepiece_temp") + + pieces = [ + ("", 0.0, CONTROL), + ("aa", -1.0, NORMAL), + ("bb", -1.0, NORMAL), + ("X", -1.0, NORMAL), + ] + old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"aa": 1, "bb": 2}) + new = _FakeTokenizer("new") + + # Caller skipped aa->X (X already exists) and applied bb->Y, so only bb->Y is passed. + tok = fix_sentencepiece_tokenizer(old, new, {"bb": "Y"}, temporary_location = location) + + result = _read_pieces(f"{loaded[-1]}/tokenizer.model") + assert result.count("X") == 1 and "Y" in result and "aa" in result, result + assert tok is not None diff --git a/tests/studio/test_studio_text_descender_clipping.py b/tests/studio/test_studio_text_descender_clipping.py index 7b9b375374..73d1244ee3 100644 --- a/tests/studio/test_studio_text_descender_clipping.py +++ b/tests/studio/test_studio_text_descender_clipping.py @@ -34,19 +34,22 @@ def test_model_selector_trigger_label_uses_leading_tight(): def test_sidebar_account_block_uses_leading_tight(): src = _read(APP_SIDEBAR) - # Match class membership without assuming utility order. - pattern = re.compile( - r'', - ) - matches = pattern.findall(src) + class_names = re.findall(r'