From bb5cea63450a6e49b4b0459c3af50dbb3dc0b4f5 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:29:18 -0700 Subject: [PATCH] test: split test_server_patches, unify _real_apple_gpu in helpers, drop slow marker, drop nonstandard module aliases --- gmlx/serve/server.py | 10 +- pyproject.toml | 1 - tests/assistant/test_assistant_memory.py | 18 +- tests/conftest.py | 2 + tests/gen/test_batch_parity.py | 2 +- tests/gen/test_hy_v4_long_context.py | 2 +- tests/gen/test_long_context.py | 2 +- tests/helpers.py | 18 + tests/load/test_vlm_integration.py | 2 +- tests/models/test_qwen4exp_ple.py | 13 +- .../models/test_qwen4exp_prefill_dispatch.py | 14 +- tests/models/test_vlm_mtp_integration.py | 2 +- .../test_vlm_plain_image_integration.py | 2 +- tests/serve/conftest.py | 106 ++ tests/serve/test_patches_chat_behavior.py | 676 ++++++++ tests/serve/test_patches_routes.py | 344 ++++ tests/serve/test_patches_sampling.py | 488 ++++++ tests/serve/test_serve_apc_engagement.py | 2 +- tests/serve/test_server_patches.py | 1494 ----------------- tests/spec/test_full_prompt_prefill.py | 2 +- tests/test_eval_guard.py | 14 +- tests/test_pinned_invariants.py | 15 +- 22 files changed, 1664 insertions(+), 1565 deletions(-) create mode 100644 tests/helpers.py create mode 100644 tests/serve/conftest.py create mode 100644 tests/serve/test_patches_chat_behavior.py create mode 100644 tests/serve/test_patches_routes.py create mode 100644 tests/serve/test_patches_sampling.py diff --git a/gmlx/serve/server.py b/gmlx/serve/server.py index 34566e84..bc4f6d6d 100644 --- a/gmlx/serve/server.py +++ b/gmlx/serve/server.py @@ -2,7 +2,7 @@ Composes the gmlx loader bridge (:func:`server_bridge_vlm.install_gguf_server_bridge`), the multi-model residency pool (:func:`residency.install_gguf_residency_pool`), -the config-driven HTTP surface (:func:`server_patches.install_server_patches`), +the config-driven HTTP surface (:func:`gmlx.serve.patches.install_server_patches`), and mlx-vlm's FastAPI app + ``BatchGenerator`` continuous-batching engine into one GGUF-only HTTP server - **text**, **VLM** (LLM GGUF + float ``mmproj``), and **speculative/MTP** models. @@ -92,7 +92,7 @@ def _import_serving(): bare ModuleNotFoundError.""" try: from . import residency # noqa: F401 - from . import patches as server_patches # noqa: F401 + from . import patches # noqa: F401 from . import bridge_vlm as server_bridge_vlm # noqa: F401 except ImportError as exc: root = (exc.name or "").split(".")[0] @@ -2052,7 +2052,7 @@ def _on_sighup(_sig, _frame): import uvicorn - from . import patches as server_patches + from . import patches # Preload off the startup path: mlx-vlm's lifespan loads MLX_VLM_PRELOAD_MODEL # synchronously *before* the port accepts connections, so a big model makes the @@ -2068,7 +2068,7 @@ def _on_sighup(_sig, _frame): if m != preload] if preload or extras: os.environ.pop("MLX_VLM_PRELOAD_MODEL", None) - server_patches.spawn_preload_warm(preload, extras) + patches.spawn_preload_warm(preload, extras) if extras: print(f"[server] preload: warming {', '.join(extras)} in background") @@ -2087,7 +2087,7 @@ def _on_sighup(_sig, _frame): # gmlx/mlx_vlm loggers, so the flag governs the whole server # while the timestamped formatters and noise filters stay. log_level=getattr(a, "log_level", None), - log_config=server_patches.uvicorn_log_config( + log_config=patches.uvicorn_log_config( getattr(a, "log_level", None))) return 0 diff --git a/pyproject.toml b/pyproject.toml index ab020013..5f8c4f59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -155,7 +155,6 @@ gmlx = ["_vendor/misaki/LICENSE", "_vendor/misaki/data/*", testpaths = ["tests"] markers = [ "integration: end-to-end test needing a real GGUF (and, for parity, a llama.cpp binary); skipped unless KQUANT_TEST_GGUF_DIR is set", - "slow: long-running (large prompts / long generation)", "needs_kvarn_ops: dispatches kvarn Metal kernels; skipped when the mlx-kquant ops probe fails (CPU device, op-less wheel)", "needs_kvarn_row_ends: dispatches the decode kernels with per-row ends; skipped unless the installed mlx-kquant is 0.4.9 or later", ] diff --git a/tests/assistant/test_assistant_memory.py b/tests/assistant/test_assistant_memory.py index 8d7f547a..32bf7369 100644 --- a/tests/assistant/test_assistant_memory.py +++ b/tests/assistant/test_assistant_memory.py @@ -94,12 +94,12 @@ def test_prune_time_corruption_also_quarantined(tmp_path, monkeypatch): # open-time prune; construction must still self-heal, not abort. import sqlite3 - import gmlx.assistant.memory as talk_memory + from gmlx.assistant import memory def bad_prune(self): raise sqlite3.DatabaseError("database disk image is malformed") - monkeypatch.setattr(talk_memory.MemoryStore, "_prune", bad_prune) + monkeypatch.setattr(memory.MemoryStore, "_prune", bad_prune) warned = [] m = _store(tmp_path, warn=warned.append) m.remember("I like green tea", "Noted.") @@ -296,8 +296,8 @@ def counting_execute(sql, *a): def test_cache_capacity_doubles(tmp_path, monkeypatch): - import gmlx.assistant.memory as talk_memory - monkeypatch.setattr(talk_memory, "_INITIAL_CAP", 1) + from gmlx.assistant import memory + monkeypatch.setattr(memory, "_INITIAL_CAP", 1) m = _store(tmp_path) m.remember("green tea daily", "") m.recall("tea?") # cache: capacity 1, n 1 @@ -311,8 +311,8 @@ def test_cache_capacity_doubles(tmp_path, monkeypatch): def test_recall_correct_after_growth(tmp_path, monkeypatch): # The promise behind the doubling rule: recall stays CORRECT after the # cache matrix has grown - rows appended through growth are recallable. - import gmlx.assistant.memory as talk_memory - monkeypatch.setattr(talk_memory, "_INITIAL_CAP", 1) + from gmlx.assistant import memory + monkeypatch.setattr(memory, "_INITIAL_CAP", 1) m = _store(tmp_path) m.remember("I like green tea", "") m.recall("tea?") # cache built at capacity 1 @@ -433,7 +433,7 @@ def test_old_schema_gains_recalled_column(tmp_path): def test_make_extractor_prompts_and_parses(monkeypatch): - import gmlx.assistant.memory as talk_memory + from gmlx.assistant import memory seen = {} def fake_stream(base_url, *, model, messages, max_tokens, @@ -444,8 +444,8 @@ def fake_stream(base_url, *, model, messages, max_tokens, yield {"content": "NONE\n2) rides a red bike"} yield {"_finish": "stop"} - monkeypatch.setattr(talk_memory.talk_client, "stream_chat", fake_stream) - extract = talk_memory.make_extractor("http://h:1/v1", "m1") + monkeypatch.setattr(memory.talk_client, "stream_chat", fake_stream) + extract = memory.make_extractor("http://h:1/v1", "m1") assert extract("I like tea", "noted") == ["likes green tea", "rides a red bike"] assert seen["model"] == "m1" diff --git a/tests/conftest.py b/tests/conftest.py index aa0b3c05..fd44f20d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -32,6 +32,8 @@ _p = str(_sub) if _p not in sys.path: sys.path.insert(0, _p) +if str(_TESTS_DIR) not in sys.path: + sys.path.insert(0, str(_TESTS_DIR)) def pytest_configure(config): diff --git a/tests/gen/test_batch_parity.py b/tests/gen/test_batch_parity.py index 3b770b62..e3853d5f 100644 --- a/tests/gen/test_batch_parity.py +++ b/tests/gen/test_batch_parity.py @@ -90,7 +90,7 @@ GREEDY = lambda x: mx.argmax(x, axis=-1) -pytestmark = [pytest.mark.integration, pytest.mark.slow] +pytestmark = pytest.mark.integration # helpers diff --git a/tests/gen/test_hy_v4_long_context.py b/tests/gen/test_hy_v4_long_context.py index b98e94d2..4aa98b26 100644 --- a/tests/gen/test_hy_v4_long_context.py +++ b/tests/gen/test_hy_v4_long_context.py @@ -51,7 +51,7 @@ # not the total depth. ABOVE_BOUNDARY = int(os.environ.get("GMLX_HY4_LONGCTX_TOKENS", "4096")) -pytestmark = [pytest.mark.integration, pytest.mark.slow] +pytestmark = pytest.mark.integration def _greedy_text(model, tok, ids, n): diff --git a/tests/gen/test_long_context.py b/tests/gen/test_long_context.py index 9bf84d10..56a75d03 100644 --- a/tests/gen/test_long_context.py +++ b/tests/gen/test_long_context.py @@ -111,7 +111,7 @@ _NEEDLE_QUERY = ("\nAs noted near the beginning of this document, the vault " "access code written in the expedition logbook is") -pytestmark = [pytest.mark.integration, pytest.mark.slow] +pytestmark = pytest.mark.integration # helpers diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 00000000..e9625779 --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,18 @@ +"""Shared test helpers importable from any test directory.""" + +import os + + +def _real_apple_gpu() -> bool: + """True on a real Apple GPU; False under KQUANT_FORCE_CPU, in a VM + (Paravirtual device), or when mlx cannot report a device. Import-time + callable so module-level pytestmark can use it (not a fixture).""" + if os.environ.get("KQUANT_FORCE_CPU"): + return False + try: + import mlx.core as mx + + name = str(mx.device_info().get("device_name", "")) + except Exception: + return False + return "Apple" in name and "Paravirtual" not in name diff --git a/tests/load/test_vlm_integration.py b/tests/load/test_vlm_integration.py index cb051ceb..bfe06956 100644 --- a/tests/load/test_vlm_integration.py +++ b/tests/load/test_vlm_integration.py @@ -21,7 +21,7 @@ pytest.importorskip("mlx_vlm") -pytestmark = [pytest.mark.integration, pytest.mark.slow] +pytestmark = pytest.mark.integration # A 30B-class VLM forward is a GPU workload; the forced-CPU stream keeps to # logic tests. diff --git a/tests/models/test_qwen4exp_ple.py b/tests/models/test_qwen4exp_ple.py index a2bc4121..701db305 100644 --- a/tests/models/test_qwen4exp_ple.py +++ b/tests/models/test_qwen4exp_ple.py @@ -4,12 +4,13 @@ import dataclasses import os -import re import mlx.core as mx import numpy as np import pytest +from helpers import _real_apple_gpu + import gmlx.models.qwen4_exp.model as q4 from gmlx.models.qwen4_exp.model import ModelArgs, PLEEmbedding @@ -33,16 +34,6 @@ ) -def _real_apple_gpu() -> bool: - if os.environ.get("KQUANT_FORCE_CPU"): - return False - try: - name = str(mx.device_info().get("device_name", "")) - except Exception: - return False - return bool(re.search(r"Apple M\d", name)) - - def _ple(): flds = {f.name for f in dataclasses.fields(ModelArgs)} return PLEEmbedding(ModelArgs( diff --git a/tests/models/test_qwen4exp_prefill_dispatch.py b/tests/models/test_qwen4exp_prefill_dispatch.py index 35c01619..0ba1b9a8 100644 --- a/tests/models/test_qwen4exp_prefill_dispatch.py +++ b/tests/models/test_qwen4exp_prefill_dispatch.py @@ -5,12 +5,12 @@ from __future__ import annotations -import os -import re import mlx.core as mx import pytest +from helpers import _real_apple_gpu + import gmlx.models.qwen4_exp.model as q4 from gmlx.models.qwen4_exp.model import ( Attention, @@ -21,16 +21,6 @@ ) -def _real_apple_gpu() -> bool: - if os.environ.get("KQUANT_FORCE_CPU"): - return False - try: - name = str(mx.device_info().get("device_name", "")) - except Exception: - return False - return bool(re.search(r"Apple M\d", name)) - - gpu_only = pytest.mark.skipif( not _real_apple_gpu(), reason="needs a real Apple GPU") diff --git a/tests/models/test_vlm_mtp_integration.py b/tests/models/test_vlm_mtp_integration.py index 2f697acf..864768a6 100644 --- a/tests/models/test_vlm_mtp_integration.py +++ b/tests/models/test_vlm_mtp_integration.py @@ -22,7 +22,7 @@ pytest.importorskip("mlx_vlm") -pytestmark = [pytest.mark.integration, pytest.mark.slow] +pytestmark = pytest.mark.integration _NEEDS_GPU = pytest.mark.skipif( bool(os.environ.get("KQUANT_FORCE_CPU")), diff --git a/tests/models/test_vlm_plain_image_integration.py b/tests/models/test_vlm_plain_image_integration.py index d723fb99..1a9168a4 100644 --- a/tests/models/test_vlm_plain_image_integration.py +++ b/tests/models/test_vlm_plain_image_integration.py @@ -22,7 +22,7 @@ pytest.importorskip("mlx_vlm") -pytestmark = [pytest.mark.integration, pytest.mark.slow] +pytestmark = pytest.mark.integration _NEEDS_GPU = pytest.mark.skipif( bool(os.environ.get("KQUANT_FORCE_CPU")), diff --git a/tests/serve/conftest.py b/tests/serve/conftest.py new file mode 100644 index 00000000..5ab5d429 --- /dev/null +++ b/tests/serve/conftest.py @@ -0,0 +1,106 @@ +"""Shared serve-suite fixtures. + +The seam snapshot/restore fixture guards every test in this directory; +it is a no-op when mlx-vlm is not installed.""" +from __future__ import annotations + +import importlib +import sys + +import pytest + +try: + import gmlx.serve.bridge_vlm as serving + from gmlx.serve.patches import _common as sp_common + from gmlx.serve.patches import hardening as sp_hardening + _APP = importlib.import_module("mlx_vlm.server.app") + _UTILS = importlib.import_module("mlx_vlm.utils") + _PKG = importlib.import_module("mlx_vlm.server") +except ImportError: + _APP = None + + +_PATCH_MODULES = { + "test_server_patches", + "test_patches_chat_behavior", + "test_patches_routes", + "test_patches_sampling", +} + + +@pytest.fixture(autouse=True) +def _restore_mlxvlm(request): + if _APP is None or request.module.__name__ not in _PATCH_MODULES: + yield + return + """Snapshot every mlx-vlm seam these patches mutate, restore after each test.""" + fastapi_app = _APP.app + saved = { + "build_gen_args": _APP._build_gen_args, + "snapshot": _APP._server_runtime_snapshot, + "get_model_path": _UTILS.get_model_path, + "routes": sp_common._snapshot_routes(fastapi_app), + "handlers": dict(fastapi_app.exception_handlers), + "deps": getattr(getattr(_APP, "_protocol_deps", None), "build_gen_args", None), + "pool": getattr(_PKG, "_kq_residency_pool", None), + # Middleware installs (auth / host guard) append to user_middleware and + # set app.state flags; CORS hardening mutates a Middleware's kwargs in + # place - snapshot all three or one test's auth poisons the rest. + "middleware": list(fastapi_app.user_middleware), + "mw_kwargs": [(m, dict(getattr(m, "kwargs", {}) or {})) + for m in fastapi_app.user_middleware], + } + openai = sys.modules.get("mlx_vlm.server.openai") + anthropic = sys.modules.get("mlx_vlm.server.anthropic") + saved["openai_bga"] = getattr(openai, "_build_gen_args", None) + saved["anthropic_bga"] = getattr(anthropic, "_build_gen_args", None) + apc = sys.modules.get("mlx_vlm.apc") or importlib.import_module("mlx_vlm.apc") + saved["apc_harvest"] = apc.harvest_blocks_from_batch_cache + saved["apc_lone_flag"] = getattr(apc, "_kq_lone_harvest", False) + gen = importlib.import_module("mlx_vlm.server.generation") + saved["to_template_kwargs"] = gen.GenerationArguments.to_template_kwargs + pu = importlib.import_module("mlx_vlm.prompt_utils") + saved["get_chat_template"] = pu.get_chat_template + saved["make_sampler"] = gen.ResponseGenerator._make_sampler + saved["make_tb_criteria"] = gen.ResponseGenerator._make_thinking_budget_criteria + schemas = importlib.import_module("mlx_vlm.server.schemas") + saved["stream_chunk_dump"] = schemas.ChatStreamChunk.model_dump_json + saved["stopping_call"] = _UTILS.StoppingCriteria.__call__ + serving.clear_resolved_models() + yield + _UTILS.StoppingCriteria.__call__ = saved["stopping_call"] + apc.harvest_blocks_from_batch_cache = saved["apc_harvest"] + apc._kq_lone_harvest = saved["apc_lone_flag"] + gen.GenerationArguments.to_template_kwargs = saved["to_template_kwargs"] + pu.get_chat_template = saved["get_chat_template"] + gen.ResponseGenerator._make_sampler = saved["make_sampler"] + gen.ResponseGenerator._make_thinking_budget_criteria = saved["make_tb_criteria"] + schemas.ChatStreamChunk.model_dump_json = saved["stream_chunk_dump"] + _APP._build_gen_args = saved["build_gen_args"] + _APP._server_runtime_snapshot = saved["snapshot"] + _UTILS.get_model_path = saved["get_model_path"] + sp_common._restore_routes(fastapi_app, saved["routes"]) + fastapi_app.exception_handlers.clear() + fastapi_app.exception_handlers.update(saved["handlers"]) + if getattr(_APP, "_protocol_deps", None) is not None and saved["deps"] is not None: + _APP._protocol_deps.build_gen_args = saved["deps"] + if openai is not None: + openai._build_gen_args = saved["openai_bga"] + if anthropic is not None: + anthropic._build_gen_args = saved["anthropic_bga"] + if saved["pool"] is None: + if hasattr(_PKG, "_kq_residency_pool"): + delattr(_PKG, "_kq_residency_pool") + else: + _PKG._kq_residency_pool = saved["pool"] + fastapi_app.user_middleware[:] = saved["middleware"] + for m, kw in saved["mw_kwargs"]: + if getattr(m, "kwargs", None) is not None: + m.kwargs.clear() + m.kwargs.update(kw) + fastapi_app.middleware_stack = None # force a rebuild from the restored list + for flag in (sp_hardening._AUTH_FLAG, sp_hardening._HOST_GUARD_FLAG, + sp_hardening._JSON_CT_FLAG): + if hasattr(fastapi_app.state, flag): + delattr(fastapi_app.state, flag) + serving.clear_resolved_models() diff --git a/tests/serve/test_patches_chat_behavior.py b/tests/serve/test_patches_chat_behavior.py new file mode 100644 index 00000000..a882564a --- /dev/null +++ b/tests/serve/test_patches_chat_behavior.py @@ -0,0 +1,676 @@ +#!/usr/bin/env python3 +"""Chat-behavior patches: thinking budget, template kwargs, role +normalization, harmony split - carved from test_server_patches.py.""" +from __future__ import annotations + +import importlib +import types + +import pytest + +pytest.importorskip("mlx_vlm") + +import gmlx.serve.patches as sp # noqa: E402 +from gmlx.serve.patches import _common as sp_common # noqa: E402 +from gmlx.serve.patches import chat_behavior as sp_chat # noqa: E402 +import gmlx.serve.bridge_vlm as serving # noqa: E402 +from gmlx.config import ResolvedModel # noqa: E402 + +_APP = importlib.import_module("mlx_vlm.server.app") +_UTILS = importlib.import_module("mlx_vlm.utils") +_PKG = importlib.import_module("mlx_vlm.server") + +from test_server_patches import _FakeThinkTok # noqa: E402 + + +# 1b. thinking_budget enforcement fix (generate- models) +def _drive(criteria, tokens): + """Feed token ids; return the forced id (or None) the criteria emits each step.""" + return [criteria(t) for t in tokens] + + +def _armed(budget, prompt_open): + cls = sp_chat._armed_thinking_budget_criteria_cls() + return cls(tokenizer=_FakeThinkTok(), thinking_budget=budget, + thinking_start_token="", thinking_end_token="", + enable_thinking=True, prompt_open_thinking=prompt_open) + + +def test_armed_criteria_caps_generated_think(): + # prompt did NOT pre-fill (Qwen3 case); the model generates it. + c = _armed(3, prompt_open=False) + assert c.in_thinking is False # not started in a block + forced = _drive(c, [99, 1, 2, 3, 4, 5, 6]) # then 6 words + assert 10 in forced and 100 in forced # forced \n then + assert forced.index(10) < forced.index(100) + + +def test_armed_criteria_caps_prefilled_think(): + # prompt pre-filled (GLM-5.2 case): counting starts immediately. + c = _armed(2, prompt_open=True) + assert c.in_thinking is True + forced = _drive(c, [1, 2, 3, 4, 5]) + assert 100 in forced # forced close + + +def test_armed_criteria_never_forces_non_thinking_answer(): + # thinking enabled + budget set, but the model never opens a -> + # no token is ever counted, so nothing is force-closed (no corruption). + c = _armed(2, prompt_open=False) + forced = _drive(c, [1, 2, 3, 4, 5, 6, 7, 8]) + assert all(f is None for f in forced) + + +def test_armed_criteria_reset_restores_prompt_seed(): + c = _armed(2, prompt_open=False) + c.in_thinking = True + c.reset_thinking_state() + assert c.in_thinking is False # back to the prompt seed + + +def test_prompt_tail_opens_thinking_cases(): + pairs = (("", ""),) + f = sp_chat._prompt_tail_opens_thinking + assert f("", pairs) is False + assert f("plain prompt, no markers", pairs) is False + assert f("<|im_start|>assistant\n\n", pairs) is True # Qwen3.6 pre-fill + assert f("x\n\n\n\n", pairs) is False # thinking off + assert f("axb", pairs) is True # last pair open + assert f(None, pairs) is False + + +def test_stream_thinking_seed_reseeds_from_prompt(): + rs = importlib.import_module("mlx_vlm.server.responses_state") + cls = rs.ThinkingStreamState + original_init = cls.__init__ + try: + sp.install_stream_thinking_seed() + assert getattr(cls.__init__, sp_chat._STREAM_SEED_FLAG, False) + patched = cls.__init__ + sp.install_stream_thinking_seed() # idempotent + assert cls.__init__ is patched + tok = sp_chat._LAST_RENDERED_PROMPT.set("rendered, no thinking scaffold") + try: + st = cls(True) # enable_thinking forced True (7b) + assert st.in_thinking is False # gemma-4 default-off: content mode + sp_chat._LAST_RENDERED_PROMPT.set("<|im_start|>assistant\n\n") + st = cls(True) + assert st.in_thinking is True # Qwen3.6 pre-fill: reasoning first + st = cls(False) + assert st.in_thinking is True # prompt truth beats the flag + sp_chat._LAST_RENDERED_PROMPT.set(None) + st = cls(True) + assert st.in_thinking is True # no render seen -> stock seed + finally: + sp_chat._LAST_RENDERED_PROMPT.reset(tok) + finally: + cls.__init__ = original_init + + +def test_stream_thinking_seed_wraps_render_binding(): + openai_mod = importlib.import_module("mlx_vlm.server.openai") + rs = importlib.import_module("mlx_vlm.server.responses_state") + original_fn = openai_mod.apply_chat_template + original_init = rs.ThinkingStreamState.__init__ + try: + openai_mod.apply_chat_template = lambda *a, **kw: "tail \n" + sp.install_stream_thinking_seed() + wrapped = openai_mod.apply_chat_template + assert getattr(wrapped, sp_chat._STREAM_SEED_FLAG, False) + out = wrapped("processor", "config", []) + assert out == "tail \n" + assert sp_chat._LAST_RENDERED_PROMPT.get() == out # stashed + sp_chat._LAST_RENDERED_PROMPT.set(None) + finally: + openai_mod.apply_chat_template = original_fn + rs.ThinkingStreamState.__init__ = original_init + + +def test_nonstream_split_truncated_thinking_seeded_from_prompt(): + app_mod = importlib.import_module("mlx_vlm.server.app") + sp.install_stream_thinking_seed() + split = app_mod._split_thinking_text + assert getattr(split, sp_chat._STREAM_SEED_FLAG, False) + tok = sp_chat._LAST_RENDERED_PROMPT.set(None) + try: + # No stashed render (off-request callers): stock classification. + assert split("half a plan, cut off") == (None, "half a plan, cut off") + sp_chat._LAST_RENDERED_PROMPT.set("<|im_start|>assistant\n\n") + # Prompt-opened block, no marker in the truncated text: reasoning. + assert split("half a plan, cut off") == ("half a plan, cut off", "") + # A close marker in the text keeps the stock split untouched. + r, c = split("plan\n\n\nanswer") + assert r == "plan" and c == "answer" + # Prompt did not open a block: stock classification. + sp_chat._LAST_RENDERED_PROMPT.set("<|im_start|>assistant\n") + assert split("just an answer") == (None, "just an answer") + finally: + sp_chat._LAST_RENDERED_PROMPT.reset(tok) + + +def test_nonstream_split_strips_xtml_section_markers(): + """Kimi-K3: the response/message closers and turn terminator must not + leak into chat content; other think spellings stay untouched.""" + app_mod = importlib.import_module("mlx_vlm.server.app") + sp.install_stream_thinking_seed() + split = app_mod._split_thinking_text + tok = sp_chat._LAST_RENDERED_PROMPT.set(None) + try: + text = ("plan<|close|>think<|sep|><|open|>response<|sep|>" + "2+2 equals 4.<|close|>response<|sep|>" + "<|close|>message<|sep|><|end_of_msg|>") + r, c = split(text, "<|open|>think<|sep|>", "<|close|>think<|sep|>") + assert r == "plan" + assert c == "2+2 equals 4." + # Gate: a different think spelling passes content through unchanged. + r, c = split("xkeep <|close|>message<|sep|> text", + "", "") + assert c == "keep <|close|>message<|sep|> text" + finally: + sp_chat._LAST_RENDERED_PROMPT.reset(tok) + + +def test_nonstream_split_keeps_first_line_code_indent(): + """The stock splitter .strip()s content, which deletes the first line's + leading indent from verbatim-code replies. The seeded splitter trims + newlines only, for every marker branch and the no-marker fallthrough.""" + app_mod = importlib.import_module("mlx_vlm.server.app") + rs = importlib.import_module("mlx_vlm.server.responses_state") + sp.install_stream_thinking_seed() + split = app_mod._split_thinking_text + tok = sp_chat._LAST_RENDERED_PROMPT.set(None) + try: + # open+close pair + r, c = split("plan\n\n if x:\n y()") + assert r == "plan" + assert c == " if x:\n y()" + # close-only (prompt-opened block) + r, c = split("plan\n\n\n\tdoc = doc || document;") + assert r == "plan" + assert c == "\tdoc = doc || document;" + # no marker at all + assert split(" indented") == (None, " indented") + # whitespace-only content collapses to empty + assert split("plan\n \n") == ("plan", "") + # XTML section strip keeps the indent too + text = ("plan<|close|>think<|sep|><|open|>response<|sep|>" + " return 4;<|close|>response<|sep|><|end_of_msg|>") + r, c = split(text, "<|open|>think<|sep|>", "<|close|>think<|sep|>") + assert c == " return 4;" + # the /v1/responses module global got the same splitter + assert rs._split_thinking is split + finally: + sp_chat._LAST_RENDERED_PROMPT.reset(tok) + + +def test_install_thinking_budget_fix_applies_and_idempotent(): + # Fail-loud guard: asserts the seam bound to the REAL mlx-vlm symbol, so a + # rename of ResponseGenerator._make_thinking_budget_criteria turns into a CI + # failure instead of a silent no-op. + gen = importlib.import_module("mlx_vlm.server.generation") + cls = gen.ResponseGenerator + original = cls._make_thinking_budget_criteria + try: + sp.install_thinking_budget_fix() + patched = cls._make_thinking_budget_criteria + assert patched is not original # actually bound + assert getattr(patched, sp_chat._TBUDGET_FLAG, False) + sp.install_thinking_budget_fix() # idempotent + assert cls._make_thinking_budget_criteria is patched + finally: + cls._make_thinking_budget_criteria = original + + +def test_make_criteria_honors_budget_when_enable_thinking_false(): + # A configured thinking_budget must arm even when enable_thinking is False: + # a group/profile config may disable thinking, but the model can still emit + # , and an explicit budget must cap it. None budget still opts out. + gen = importlib.import_module("mlx_vlm.server.generation") + cls = gen.ResponseGenerator + original = cls._make_thinking_budget_criteria + try: + sp.install_thinking_budget_fix() + make = cls._make_thinking_budget_criteria + me = types.SimpleNamespace( + tokenizer=_FakeThinkTok(), + _thinking_token_ids=lambda args: (99, 100), # =99, =100 + ) + args = types.SimpleNamespace( + thinking_budget=8, enable_thinking=False, + thinking_start_token=None, thinking_end_token=None) + # generate-style prompt (no open ): armed but not seeded in-block + criteria = make(me, args, [1, 2, 3]) + assert criteria is not None # armed despite enable_thinking=False + assert criteria.in_thinking is False + # pre-fill prompt ending with an open seeds in_thinking True, even + # though enable_thinking is False - the cap must still fire (GLM-style). + criteria2 = make(me, args, [1, 2, 99]) + assert criteria2 is not None and criteria2.in_thinking is True + args.thinking_budget = None + assert make(me, args, [1, 2, 3]) is None # no budget -> still opts out + finally: + cls._make_thinking_budget_criteria = original + + +def test_seed_wrapper_survives_thinking_budget_fix(): + # Regression: the install order used to be seed -> tbfix, and tbfix rebinds + # the criteria seam without delegating, so it clobbered the seed wrapper and + # per-request seeds were dead on the serve path. The order is now tbfix -> + # seed: one call must stash the seed AND run tbfix's construction, and a + # tbfix re-install must see its flag through the wrapper and no-op. + import gmlx.serve.seed_rows as sr + gen = importlib.import_module("mlx_vlm.server.generation") + ar = importlib.import_module("mlx_vlm.generate.ar") + cls = gen.ResponseGenerator + saved = (cls._make_thinking_budget_criteria, ar.BatchGenerator.insert, + ar.GenerationBatch._step, ar.PromptProcessingBatch.generate, + ar.SpeculativeGenerationBatch.next) + sr._PENDING.clear() + try: + sp.install_thinking_budget_fix() + sr.install_per_request_seed() + crit = cls._make_thinking_budget_criteria + assert getattr(crit, sr._INSTALLED_FLAG, False) # seed outermost + assert getattr(crit, sp_chat._TBUDGET_FLAG, False) # tbfix flag carried + sp.install_thinking_budget_fix() + assert cls._make_thinking_budget_criteria is crit # re-install no-ops + me = types.SimpleNamespace( + tokenizer=_FakeThinkTok(), + _thinking_token_ids=lambda args: (99, 100)) + args = types.SimpleNamespace( + seed=7, temperature=1.0, thinking_budget=8, enable_thinking=False, + thinking_start_token=None, thinking_end_token=None) + out = crit(me, args, [1, 2, 3]) + assert out is not None and out.in_thinking is False # tbfix ran + assert sr._PENDING == [7] # seed stashed + finally: + (cls._make_thinking_budget_criteria, ar.BatchGenerator.insert, + ar.GenerationBatch._step, ar.PromptProcessingBatch.generate, + ar.SpeculativeGenerationBatch.next) = saved + sr._PENDING.clear() + + +# 7b. chat_template_kwargs passthrough +def _spec_ctkw(**ctkw): + return ResolvedModel(id="m", path="/p", sampling={}, load={}, cache={}, + system=None, speculative=False, mmproj=None, + draft_gguf=None, pin=False, ttl_s=None, + chat_template_kwargs=ctkw) + + +def test_merged_template_kwargs_request_wins_over_profile(): + spec = _spec_ctkw(preserve_thinking=True, foo="profile") + request = types.SimpleNamespace(chat_template_kwargs={"foo": "request"}) + merged = sp_chat._merged_template_kwargs(request, spec) + assert merged == {"preserve_thinking": True, "foo": "request"} + + +def test_merged_template_kwargs_each_side_alone_and_empty(): + # request only (single-model mode: no active spec) + req = types.SimpleNamespace(chat_template_kwargs={"preserve_thinking": True}) + assert sp_chat._merged_template_kwargs(req, None) == {"preserve_thinking": True} + # profile only (request carries nothing) + spec = _spec_ctkw(preserve_thinking=False) + assert sp_chat._merged_template_kwargs(types.SimpleNamespace(), spec) == { + "preserve_thinking": False} + # neither => {} + assert sp_chat._merged_template_kwargs(types.SimpleNamespace(), None) == {} + + +def test_merged_template_kwargs_spec_thinking_controls_mapped(): + """Profile-level thinking/reasoning_effort are dedicated controls: mapped + onto whatever switch the serving model's template reads.""" + spec = _spec_ctkw() + spec.thinking = "off" + req = types.SimpleNamespace() + assert sp_chat._merged_template_kwargs( + req, spec, "{% if enable_thinking %}...{% endif %}") == \ + {"enable_thinking": False} + assert sp_chat._merged_template_kwargs( + req, spec, "reasoning_effort in ['low','high','no_think']") == \ + {"reasoning_effort": "no_think"} + spec.thinking = "adaptive" + assert sp_chat._merged_template_kwargs( + req, spec, 'thinking_mode == "adaptive"') == \ + {"thinking_mode": "adaptive"} + spec.thinking = None + spec.reasoning_effort = "high" + assert sp_chat._merged_template_kwargs( + req, spec, 'set reasoning_effort = "medium"') == \ + {"reasoning_effort": "high"} + + +def test_merged_template_kwargs_request_kwargs_beat_spec_controls(): + """A request's explicit chat_template_kwargs pass through verbatim and win + over the profile's mapped controls.""" + spec = _spec_ctkw() + spec.thinking = "off" + spec.reasoning_effort = "low" + req = types.SimpleNamespace( + chat_template_kwargs={"enable_thinking": True}) + merged = sp_chat._merged_template_kwargs( + req, spec, "{% if enable_thinking %}{% endif %} reasoning_effort") + assert merged["enable_thinking"] is True + assert merged["reasoning_effort"] == "low" + + +def test_install_chat_template_kwargs_forwards_into_to_template_kwargs(): + """End-to-end seam: the gen-args wrapper stashes the merged dict and the + patched to_template_kwargs folds it into what mlx-vlm hands the template.""" + gen = importlib.import_module("mlx_vlm.server.generation") + + def stub(request, processor=None, tenant_id=None): + return gen.GenerationArguments() + + _APP._build_gen_args = stub + sp.install_gen_args_profile_injection() + sp.install_chat_template_kwargs() + fn = _APP._build_gen_args + assert getattr(fn, sp_chat._CTKW_FLAG, False) # stash carried on the chain + + spec = _spec_ctkw(preserve_thinking=True) + tok = serving.set_active_spec(spec) + try: + req = types.SimpleNamespace(model_fields_set=set(), + chat_template_kwargs={"foo": "bar"}) + args = _APP._build_gen_args(req) + finally: + serving.reset_active_spec(tok) + kw = args.to_template_kwargs() + assert kw["preserve_thinking"] is True # from the profile + assert kw["foo"] == "bar" # from the request + # enable_thinking was not explicit (request/spec/env) -> dropped from the + # template kwargs so the chat template's own default governs (b90aa60), + # while the args flag stays True for the generation path. + assert "enable_thinking" not in kw + assert args.enable_thinking is True + + # explicitly set on the request -> preserved verbatim + spec = _spec_ctkw() + tok = serving.set_active_spec(spec) + try: + req = types.SimpleNamespace(model_fields_set={"enable_thinking"}, + chat_template_kwargs=None, + enable_thinking=False) + args = _APP._build_gen_args(req) + finally: + serving.reset_active_spec(tok) + assert "enable_thinking" in args.to_template_kwargs() + + +def test_install_chat_template_kwargs_idempotent_and_noop_default(): + sp.install_chat_template_kwargs() + gen = importlib.import_module("mlx_vlm.server.generation") + first = gen.GenerationArguments.to_template_kwargs + sp.install_chat_template_kwargs() + assert gen.GenerationArguments.to_template_kwargs is first + # a request/spec with no kwargs leaves to_template_kwargs untouched (stock keys) + assert gen.GenerationArguments().to_template_kwargs() == { + "enable_thinking": gen.GenerationArguments().enable_thinking} + + +_KIMI_TAIL = ( + "{%- if thinking is defined and thinking is false -%}" + "{%- else -%}{%- endif -%}") + + +def test_request_thinking_off_maps_onto_kimi_bare_switch(): + """A plain --thinking value forwarded by the chat client (or any client's + `thinking: "off"`) must reach the template as the model's own switch + spelling - Kimi K2.x reads a bare `thinking` variable.""" + gen = importlib.import_module("mlx_vlm.server.generation") + proc = types.SimpleNamespace(chat_template=_KIMI_TAIL) + req = types.SimpleNamespace(model_fields_set=set(), + chat_template_kwargs=None, thinking="off") + out = sp_chat._stash_template_kwargs(gen.GenerationArguments(), req, proc) + assert out._kq_template_kwargs == {"thinking": False} + assert out.enable_thinking is False + assert out._kq_thinking_explicit is True + + req = types.SimpleNamespace(model_fields_set=set(), + chat_template_kwargs=None, thinking="on") + out = sp_chat._stash_template_kwargs(gen.GenerationArguments(), req, proc) + assert out._kq_template_kwargs == {"thinking": True} + assert out.enable_thinking is True + + +def test_request_reasoning_effort_field_maps_onto_template(): + gen = importlib.import_module("mlx_vlm.server.generation") + proc = types.SimpleNamespace(chat_template="reads reasoning_effort") + req = types.SimpleNamespace(model_fields_set=set(), + chat_template_kwargs=None, + reasoning_effort="low") + out = sp_chat._stash_template_kwargs(gen.GenerationArguments(), req, proc) + assert out._kq_template_kwargs == {"reasoning_effort": "low"} + assert out.enable_thinking is True # effort alone: not a switch + assert out._kq_thinking_explicit is False + + +class _RecordingTok: + """A tokenizer stand-in whose **kwargs signature makes mlx-vlm's + enable_thinking capability probe say yes (the 0.6.15 injection path).""" + chat_template = "{{ messages }}" + + def __init__(self): + self.kwargs = {} + + def apply_chat_template(self, messages, **kwargs): + self.kwargs = kwargs + return "rendered" + + +def test_template_default_thinking_blocks_0615_false_injection(): + """Regression: mlx-vlm >= 0.6.15 get_chat_template injects + enable_thinking=False when the kwarg is absent, so served models with + default-on reasoning templates rendered the dead think prefill and + stopped thinking. The seeded Jinja Undefined must reach the tokenizer + instead of False, and an explicit value must pass through verbatim.""" + import jinja2 + + import gmlx.tui.reasoning as reasoning + + sp.install_chat_template_kwargs() # installs the render guard too + pu = importlib.import_module("mlx_vlm.prompt_utils") + assert getattr(pu.get_chat_template, reasoning._TEMPLATE_DEFAULT_FLAG, False) + + tok = _RecordingTok() + msgs = [{"role": "user", "content": "hi"}] + assert pu.get_chat_template(tok, msgs, True) == "rendered" + assert isinstance(tok.kwargs["enable_thinking"], jinja2.Undefined) + + pu.get_chat_template(tok, msgs, True, enable_thinking=False) + assert tok.kwargs["enable_thinking"] is False + pu.get_chat_template(tok, msgs, True, enable_thinking=True) + assert tok.kwargs["enable_thinking"] is True + + # idempotent: a second install keeps the same wrapper + fn = pu.get_chat_template + reasoning.install_template_default_thinking() + assert pu.get_chat_template is fn + + +# The old-style role dispatch from issue #66: no developer alias, else-raise. +_ROLE_RAISE_TMPL = ("{% if m.role == 'system' %}{% elif m.role == 'user' %}" + "{% else %}{{ raise_exception('Unexpected message role.') }}" + "{% endif %}") + + +def test_normalize_developer_roles(): + msgs = [{"role": "developer", "content": "terse"}, + {"role": "user", "content": "hi"}, "not-a-dict"] + out = sp_chat._normalize_developer_roles(msgs, _ROLE_RAISE_TMPL) + assert out[0] == {"role": "system", "content": "terse"} + assert out[1]["role"] == "user" and out[2] == "not-a-dict" + assert msgs[0]["role"] == "developer" # input not mutated + # a template that handles developer gets the messages verbatim + assert sp_chat._normalize_developer_roles( + msgs, "role == 'developer'") is msgs + # nothing to rewrite -> same object + plain = [{"role": "user", "content": "hi"}] + assert sp_chat._normalize_developer_roles(plain, _ROLE_RAISE_TMPL) is plain + assert sp_chat._normalize_developer_roles("prompt", _ROLE_RAISE_TMPL) == \ + "prompt" + + +def test_install_role_normalization_rewrites_before_render(): + """Issue #66: a developer-role request against a template without the + alias must render as system instead of raising in the template.""" + class _Recorder: + chat_template = _ROLE_RAISE_TMPL + + def apply_chat_template(self, messages, **kwargs): + self.messages = messages + return "rendered" + + sp.install_role_normalization() + openai = importlib.import_module("mlx_vlm.server.openai") + assert getattr(openai.apply_chat_template, + sp_chat._ROLE_NORM_FLAG, False) + + tok = _Recorder() + out = openai.apply_chat_template( + tok, {"model_type": "gguf-llama"}, + [{"role": "developer", "content": "terse"}, + {"role": "user", "content": "hi"}]) + assert out == "rendered" + assert [m["role"] for m in tok.messages + if isinstance(m, dict) and "role" in m][:2] == ["system", "user"] + + # idempotent + fn = openai.apply_chat_template + sp.install_role_normalization() + assert openai.apply_chat_template is fn + + +def test_template_error_becomes_clean_400(): + """A raise_exception from the chat template answers 400 with the + template's message; template bugs (subclasses) stay 500, clean body.""" + import jinja2 + from fastapi.testclient import TestClient + + app = _APP.app + if not any(getattr(r, "path", None) == "/test/raise-template" + for r in app.router.routes): + @app.get("/test/raise-template") + async def _raise_template(): + raise jinja2.exceptions.TemplateError("Unexpected message role.") + + @app.get("/test/raise-template-bug") + async def _raise_template_bug(): + raise jinja2.exceptions.TemplateSyntaxError("bad", 1) + + sp.install_resolver_error_handlers() + client = TestClient(app, raise_server_exceptions=False) + r = client.get("/test/raise-template") + assert r.status_code == 400 + assert r.json()["error"] == { + "type": "invalid_request_error", + "message": "chat template rejected the conversation: " + "Unexpected message role."} + r2 = client.get("/test/raise-template-bug") + assert r2.status_code == 500 + assert r2.json()["error"]["type"] == "server_error" + assert "chat template failed to render" in r2.json()["error"]["message"] + + +# harmony (gpt-oss) serve-side split +_HARMONY_PROMPT = ("<|start|>system<|message|>You are helpful.<|end|>" + "<|start|>user<|message|>hi<|end|><|start|>assistant") +_HARMONY_REPLY = ('<|channel|>analysis<|message|>User greets; keep it short.' + "<|end|><|start|>assistant<|channel|>final<|message|>" + "Hello! How can I help?") + + +def test_nonstream_split_harmony_reply(): + app_mod = importlib.import_module("mlx_vlm.server.app") + sp.install_stream_thinking_seed() + split = app_mod._split_thinking_text + tok = sp_chat._LAST_RENDERED_PROMPT.set(None) + try: + r, c = split(_HARMONY_REPLY) + assert r == "User greets; keep it short." + assert c == "Hello! How can I help?" + assert "<|" not in c and "<|" not in r + # Length-capped inside analysis: all reasoning, empty content + # (the truncated-thinking convention). + r, c = split("<|channel|>analysis<|message|>Entry 55 reads 4") + assert r == "Entry 55 reads 4" + assert c == "" + # Gemma's lopsided spelling must not take the harmony branch. + r, c = split("<|channel>thought\nplan\nHi there.") + assert c and "<|channel|>" not in c + finally: + sp_chat._LAST_RENDERED_PROMPT.reset(tok) + + +def test_stream_harmony_filter_routes_channels(): + rs = importlib.import_module("mlx_vlm.server.responses_state") + cls = rs.ThinkingStreamState + original_init = cls.__init__ + original_feed = cls.feed + try: + sp.install_stream_thinking_seed() + tok = sp_chat._LAST_RENDERED_PROMPT.set(_HARMONY_PROMPT) + try: + st = cls(True) + assert getattr(st, "_kq_harmony", None) is not None + reasoning, content, closes = [], [], 0 + for i in range(0, len(_HARMONY_REPLY), 7): + d = st.feed(_HARMONY_REPLY[i:i + 7]) + if d.reasoning: + reasoning.append(d.reasoning) + if d.content: + content.append(d.content) + closes += bool(d.thinking_closed) + assert "".join(reasoning) == "User greets; keep it short." + assert "".join(content) == "Hello! How can I help?" + assert closes == 1 + # Non-harmony prompt: the stock state machine still drives. + sp_chat._LAST_RENDERED_PROMPT.set("<|im_start|>assistant\n\n") + st = cls(True) + assert getattr(st, "_kq_harmony", None) is None + d = st.feed("plananswer") + assert d.reasoning == "plan" and d.content == "answer" + finally: + sp_chat._LAST_RENDERED_PROMPT.reset(tok) + finally: + cls.__init__ = original_init + cls.feed = original_feed + + +def test_faithful_history_aliases_gpt_oss_thinking(): + from gmlx.serve.patches import render as sp_render + + def fake(processor, config, prompt, add_generation_prompt=True, + return_messages=False, num_images=0, num_audios=0, **kwargs): + return [dict(m) for m in prompt] + + # Exercise through the installer against a stub target module. + target = types.SimpleNamespace(apply_chat_template=fake) + orig_targets = sp_common._render_target_modules + sp_common._render_target_modules = lambda: [target] + try: + sp_render.install_faithful_history() + finally: + sp_common._render_target_modules = orig_targets + wrapped = target.apply_chat_template + assert wrapped is not fake + msgs = wrapped( + "processor", {"model_type": "gpt_oss"}, + [{"role": "assistant", "content": "Hi.", + "reasoning_content": "Short greeting."}], + return_messages=True) + assert msgs[0]["thinking"] == "Short greeting." + # Explicit thinking key wins; non-gpt-oss untouched. + msgs = wrapped( + "processor", {"model_type": "gpt_oss"}, + [{"role": "assistant", "content": "Hi.", "thinking": "keep", + "reasoning_content": "drop"}], + return_messages=True) + assert msgs[0]["thinking"] == "keep" + msgs = wrapped( + "processor", {"model_type": "qwen3"}, + [{"role": "assistant", "content": "Hi.", + "reasoning_content": "r"}], + return_messages=True) + assert "thinking" not in msgs[0] diff --git a/tests/serve/test_patches_routes.py b/tests/serve/test_patches_routes.py new file mode 100644 index 00000000..bbdea379 --- /dev/null +++ b/tests/serve/test_patches_routes.py @@ -0,0 +1,344 @@ +#!/usr/bin/env python3 +"""/v1/models payload, HF gate, snapshot enrichment, unload and keep +routes - carved from test_server_patches.py. CPU-only.""" +from __future__ import annotations + +import importlib +import os + +import pytest + +pytest.importorskip("mlx_vlm") + +import gmlx.serve.patches as sp # noqa: E402 +from gmlx.serve.patches import _common as sp_common # noqa: E402 +from gmlx.serve.patches import routes as sp_routes # noqa: E402 +import gmlx.serve.bridge_vlm as serving # noqa: E402 + +_APP = importlib.import_module("mlx_vlm.server.app") +_UTILS = importlib.import_module("mlx_vlm.utils") +_PKG = importlib.import_module("mlx_vlm.server") + +from test_server_patches import _FakeKeepPool, _register # noqa: E402 + + +def test_models_payload_lists_configured_ids_not_hf(): + _register({"models": { + "qwen": {"path": "/abs/qwen.gguf"}, + "gemma-vlm": {"path": "/abs/g.gguf", "mmproj": "/abs/mm.gguf"}, + }}) + payload = sp_routes._models_payload() + ids = {m["id"] for m in payload["data"]} + assert ids == {"qwen", "gemma-vlm"} + vlm = next(m for m in payload["data"] if m["id"] == "gemma-vlm") + assert vlm["vlm"] is True + assert all(m["resident"] is False for m in payload["data"]) # no pool + + +def test_models_payload_marks_resident_from_pool(): + _register({"models": {"qwen": {"path": "/abs/qwen.gguf", "pin": True}}}) + + class _FakePool: + def stats(self): + return {"resident": [{"model_path": "/abs/qwen.gguf", "pinned": True, + "footprint_bytes": 10, "idle_s": 3.0, + "ttl_s": 900}]} + + _PKG._kq_residency_pool = _FakePool() + m = sp_routes._models_payload()["data"][0] + assert m["resident"] is True and m["pinned"] is True + + +def test_models_payload_lists_aliases_as_pickable_entries(): + _register({ + "profiles": {"coder": {"sampling": {"temperature": 0.2}}}, + "models": {"qwen": {"path": "/abs/qwen.gguf", "speculative": False}}, + "aliases": {"big": "qwen", "coder-preset": "qwen@coder"}, + }) + payload = sp_routes._models_payload() + by_id = {m["id"]: m for m in payload["data"]} + assert set(by_id) == {"qwen", "big", "coder-preset"} # aliases listed + assert by_id["big"]["alias_of"] == "qwen" + assert by_id["coder-preset"]["alias_of"] == "qwen" + assert by_id["coder-preset"]["profile"] == "coder" # baked profile shown + assert "alias_of" not in by_id["qwen"] # real model unmarked + + +def test_models_payload_marks_default(): + _register({ + "server": {"defaults": {"model": "qwen"}}, + "models": {"qwen": {"path": "/abs/qwen.gguf"}, + "gemma": {"path": "/abs/g.gguf"}}, + }) + by_id = {m["id"]: m for m in sp_routes._models_payload()["data"]} + assert by_id["qwen"]["default"] is True + assert by_id["gemma"]["default"] is False + + +def test_models_override_registers_single_route(): + sp.install_models_endpoint_override() + paths = [getattr(r, "path", None) for r in _APP.app.router.routes] + assert paths.count("/v1/models") == 1 + sp.install_models_endpoint_override() # idempotent-ish + paths = [getattr(r, "path", None) for r in _APP.app.router.routes] + assert paths.count("/v1/models") == 1 + + +# 4. HF gate +def test_gate_allows_local_and_gguf(tmp_path): + calls = [] + orig = lambda p, *a, **k: calls.append(p) or "OK" + local = tmp_path / "f" + local.write_text("x") + assert sp_routes._gate_model_path(str(local), False, orig) == "OK" + assert sp_routes._gate_model_path("/x/model.gguf", False, orig) == "OK" + assert len(calls) == 2 + + +def test_gate_blocks_hf_id_when_disabled(): + orig = lambda p, *a, **k: "OK" + with pytest.raises(sp.HFAccessDisabled): + sp_routes._gate_model_path("org/model", False, orig) + + +def test_gate_allows_hf_id_when_cache_on(): + calls = [] + orig = lambda p, *a, **k: calls.append(p) or "OK" + assert sp_routes._gate_model_path("org/model", True, orig) == "OK" + assert calls == ["org/model"] + + +def test_install_hf_gate_sets_offline_env(monkeypatch): + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + sp.install_hf_download_gate(hf_cache=True) + assert os.environ.get("HF_HUB_OFFLINE") == "1" + + +# 5. runtime-snapshot enrichment +def test_snapshot_enrichment_adds_resident_models(): + _APP._server_runtime_snapshot = lambda: {"loaded_model": "x"} + + class _FakePool: + def stats(self): + return {"resident": [{"model_path": "/abs/qwen.gguf", "pinned": False, + "busy": 3, "footprint_bytes": 99, + "idle_s": 1.234, "ttl_s": 900}]} + + _PKG._kq_residency_pool = _FakePool() + serving._PATH_TO_IDS["/abs/qwen.gguf"] = ["qwen"] + try: + sp.install_runtime_snapshot_enrichment() + snap = _APP._server_runtime_snapshot() + finally: + serving._PATH_TO_IDS.pop("/abs/qwen.gguf", None) + assert snap["loaded_model"] == "x" # base preserved + assert snap["resident_models"][0]["ids"] == ["qwen"] + assert snap["resident_models"][0]["idle_s"] == 1.2 # rounded + assert snap["resident_models"][0]["busy"] == 3 # in-flight count + + +# 6. error handlers + reload + unload route +def test_error_content_dialect_shapes(): + # One condition, two envelopes: OpenAI-style everywhere, Anthropic's + # {"type": "error", ...} with its fixed taxonomy on /v1/messages. + openai = sp_common._error_content( + "/v1/chat/completions", 404, "model_not_found", "no such model", + available_models=["a"]) + assert openai == {"error": {"type": "model_not_found", + "message": "no such model", + "available_models": ["a"]}} + anthropic = sp_common._error_content( + "/v1/messages", 404, "model_not_found", "no such model") + assert anthropic["type"] == "error" + assert anthropic["error"] == {"type": "not_found_error", + "message": "no such model"} + assert sp_common._error_content("/v1/messages", 500, "server_error", + "x")["error"]["type"] == "api_error" + + +def test_http_exception_envelope_unwrapped(): + # The residency resolver path raises HTTPException carrying the unified + # {"error": {...}} detail; the app-level handler must serve that body + # directly (no {"detail": ...} wrapper) and wrap plain-string details. + from fastapi import HTTPException + from fastapi.testclient import TestClient + + app = _APP.app + if not any(getattr(r, "path", None) == "/test/raise-envelope" + for r in app.router.routes): + @app.get("/test/raise-envelope") + async def _raise_envelope(): + raise HTTPException(status_code=404, detail={"error": { + "type": "model_not_found", "message": "no such model", + "available_models": ["a"]}}) + + @app.get("/test/raise-string") + async def _raise_string(): + raise HTTPException(status_code=500, detail="it broke") + + sp.install_resolver_error_handlers() + client = TestClient(app) + r = client.get("/test/raise-envelope") + assert r.status_code == 404 + assert r.json() == {"error": {"type": "model_not_found", + "message": "no such model", + "available_models": ["a"]}} + r2 = client.get("/test/raise-string") + assert r2.status_code == 500 + assert r2.json() == {"error": {"type": "server_error", + "message": "it broke"}} + + +def test_resolver_error_handlers_registered(): + sp.install_resolver_error_handlers() + handlers = _APP.app.exception_handlers + assert serving.ModelNotFound in handlers + assert serving.ModelFileMissing in handlers + assert serving.UnknownProfile in handlers + assert sp.HFAccessDisabled in handlers + + +def test_unload_and_reload_routes_register(): + sp.install_pool_aware_unload() + sp.install_reload_route(lambda: {"reloaded": 1}) + paths = [getattr(r, "path", None) for r in _APP.app.router.routes] + assert paths.count("/unload") == 1 + assert paths.count("/v1/reload") == 1 + + +def test_unload_accepts_body_and_empty_post(): + """Regression: the ``request: Request`` annotation must resolve at module level. + Under ``from __future__ import annotations`` a locally-imported ``Request`` left + FastAPI treating ``request`` as a required query param, so every POST 422'd before + the body was read (caught only at e2e). A route-count check can't see this - POST + it for real and assert the body is actually consumed.""" + from fastapi.testclient import TestClient + + sp.install_pool_aware_unload() + client = TestClient(_APP.app) + # no pool registered -> handler runs (no 422) and reports the absence + # cleanly, as a 503: the unload cannot be honored without a pool + r_body = client.post("/unload", json={"model": "m"}) + assert r_body.status_code == 503, r_body.text + assert r_body.json() == {"status": "error", "message": "no residency pool"} + r_empty = client.post("/unload") + assert r_empty.status_code == 200, r_empty.text + assert r_empty.json()["status"] == "no_model_loaded" + + +def test_json_content_type_tolerance(): + # `curl -d '{...}'` (every doc example) sends form-encoded; the middleware + # must rewrite it to application/json so pydantic parses the body instead + # of 422ing. Multipart (audio uploads) must pass through untouched. + from fastapi.testclient import TestClient + + app = _APP.app + if not any(getattr(r, "path", None) == "/test/echo-ct" + for r in app.router.routes): + # `dict` (a builtin) survives this module's stringized annotations; a + # test-local pydantic class would resolve as a query param instead. + @app.post("/test/echo-ct") + async def _echo_ct(body: dict): + return {"model": body.get("model")} + + sp.install_json_content_type_tolerance() + client = TestClient(app) + r = client.post("/test/echo-ct", content=b'{"model": "m1"}', + headers={"Content-Type": "application/x-www-form-urlencoded"}) + assert r.status_code == 200 and r.json() == {"model": "m1"} + r = client.post("/test/echo-ct", content=b'{"model": "m2"}', + headers={"Content-Type": "text/plain"}) + assert r.status_code == 200 and r.json() == {"model": "m2"} + r = client.post("/test/echo-ct", json={"model": "m3"}) # normal path intact + assert r.status_code == 200 + r = client.post("/test/echo-ct", files={"file": ("a.txt", b"x")}) + assert r.status_code == 422 # multipart not rewritten + + +def test_keep_route_registers(): + sp.install_keep_route() + paths = [getattr(r, "path", None) for r in _APP.app.router.routes] + assert paths.count("/v1/keep") == 1 + + +def test_keep_no_pool_reports_error(): + from fastapi.testclient import TestClient + + sp.install_keep_route() + client = TestClient(_APP.app) + r = client.post("/v1/keep", json={"model": "m"}) + assert r.status_code == 503, r.text + assert r.json() == {"status": "error", "message": "no residency pool"} + + +def test_keep_marks_resolved_model(monkeypatch): + from fastapi.testclient import TestClient + + monkeypatch.setattr(sp_routes, "_spawn_keep_warm", lambda model_id: None) + _register({"models": {"qwen": {"path": "/abs/qwen.gguf"}}}) + pool = _FakeKeepPool() + _PKG._kq_residency_pool = pool + sp.install_keep_route() + client = TestClient(_APP.app) + r = client.post("/v1/keep", json={"model": "qwen", "warm": False}) + assert r.status_code == 200, r.text + assert r.json() == {"status": "kept", "model": "qwen", "warming": False} + assert pool.kept == [("/abs/qwen.gguf", True)] + + +def test_keep_warm_default_spawns_warm(monkeypatch): + from fastapi.testclient import TestClient + + warmed = [] + monkeypatch.setattr(sp_routes, "_spawn_keep_warm", lambda model_id: warmed.append(model_id)) + _register({"models": {"qwen": {"path": "/abs/qwen.gguf"}}}) + _PKG._kq_residency_pool = _FakeKeepPool() + sp.install_keep_route() + client = TestClient(_APP.app) + r = client.post("/v1/keep", json={"model": "qwen"}) # warm omitted -> default True + assert r.json() == {"status": "kept", "model": "qwen", "warming": True} + assert warmed == ["qwen"] + + +def test_keep_false_releases_without_evicting(monkeypatch): + # A voice session ending releases its hold; the model stays resident + # under normal LRU/TTL rather than being dumped. + from fastapi.testclient import TestClient + + warmed = [] + monkeypatch.setattr(sp_routes, "_spawn_keep_warm", lambda model_id: warmed.append(model_id)) + _register({"models": {"qwen": {"path": "/abs/qwen.gguf"}}}) + pool = _FakeKeepPool() + _PKG._kq_residency_pool = pool + sp.install_keep_route() + client = TestClient(_APP.app) + r = client.post("/v1/keep", json={"model": "qwen", "keep": False}) + assert r.status_code == 200, r.text + assert r.json() == {"status": "released", "model": "qwen"} + assert pool.kept == [("/abs/qwen.gguf", False)] + assert warmed == [] # release never warms + + +def test_keep_unknown_model_graceful(): + from fastapi.testclient import TestClient + + _register({"models": {"qwen": {"path": "/abs/qwen.gguf"}}}) + _PKG._kq_residency_pool = _FakeKeepPool() + sp.install_keep_route() + client = TestClient(_APP.app) + r = client.post("/v1/keep", json={"model": "nope"}) + # 404, not 200: a typo'd keep must not read as success (launch checks the + # status code); the body still names the id for older/other clients. + assert r.status_code == 404, r.text + assert r.json() == {"status": "unknown_model", "model": "nope"} + + +def test_keep_missing_model_field(): + from fastapi.testclient import TestClient + + _PKG._kq_residency_pool = _FakeKeepPool() + sp.install_keep_route() + client = TestClient(_APP.app) + r = client.post("/v1/keep", json={}) + assert r.status_code == 400, r.text + assert r.json() == {"status": "error", "message": "missing 'model'"} diff --git a/tests/serve/test_patches_sampling.py b/tests/serve/test_patches_sampling.py new file mode 100644 index 00000000..dc82c54c --- /dev/null +++ b/tests/serve/test_patches_sampling.py @@ -0,0 +1,488 @@ +#!/usr/bin/env python3 +"""Sampling-profile injection, MTP thinking-budget transport, XTC, and the +batch sampler - carved from test_server_patches.py. CPU-only.""" +from __future__ import annotations + +import importlib +import types + +import pytest + +pytest.importorskip("mlx_vlm") + +import gmlx.serve.patches as sp # noqa: E402 +from gmlx.serve.patches import _common as sp_common # noqa: E402 +from gmlx.serve.patches import sampling as sp_sampling # noqa: E402 +import gmlx.serve.bridge_vlm as serving # noqa: E402 + +_APP = importlib.import_module("mlx_vlm.server.app") +_UTILS = importlib.import_module("mlx_vlm.utils") +_PKG = importlib.import_module("mlx_vlm.server") + +from test_server_patches import _FakeThinkTok, _spec # noqa: E402 + + +# 1. sampling injection (pure) +def test_inject_overrides_unset_keeps_client_set(): + args = types.SimpleNamespace(temperature=0.7, top_p=0.95, top_k=5, max_tokens=512) + request = types.SimpleNamespace(model_fields_set={"top_k"}) # client set top_k + spec = _spec(temperature=0.2, top_p=0.9, top_k=99, max_tokens=2048) + sp_sampling._inject_profile_sampling(args, request, spec) + assert args.temperature == 0.2 # injected (unset) + assert args.top_p == 0.9 # injected (unset) + assert args.top_k == 5 # kept (client set) + assert args.max_tokens == 2048 # injected (unset) + + +def test_inject_max_tokens_alias_respects_max_output_tokens(): + args = types.SimpleNamespace(max_tokens=512) + request = types.SimpleNamespace(model_fields_set={"max_output_tokens"}) + sp_sampling._inject_profile_sampling(args, request, _spec(max_tokens=2048)) + assert args.max_tokens == 512 # responses API set it -> not overridden + + +# 1b. ignore-eos: forced-length decode +def test_install_ignore_eos_suppresses_stop(): + crit = _UTILS.StoppingCriteria([7, 8]) + assert crit(7) is True # baseline: 7 is an eos id -> stop + sp.install_ignore_eos() + assert crit(7) is False # patched: EOS never stops decode + assert crit(8) is False + assert crit(123) is False + sp.install_ignore_eos() # idempotent + assert crit(7) is False + + +def test_inject_noop_without_spec_or_sampling(): + args = types.SimpleNamespace(temperature=0.7) + request = types.SimpleNamespace(model_fields_set=set()) + sp_sampling._inject_profile_sampling(args, request, None) + sp_sampling._inject_profile_sampling(args, request, _spec()) # empty sampling + assert args.temperature == 0.7 + + +def test_inject_skips_unknown_arg_attr(): + args = types.SimpleNamespace(temperature=0.7) # no top_p attr + request = types.SimpleNamespace(model_fields_set=set()) + sp_sampling._inject_profile_sampling(args, request, _spec(top_p=0.5)) + assert not hasattr(args, "top_p") + + +def test_inject_thinking_budget_from_profile(): + # off by default: GenerationArguments.thinking_budget is None; a profile/model + # value seeds it when the client didn't ask. + args = types.SimpleNamespace(thinking_budget=None) + request = types.SimpleNamespace(model_fields_set=set()) + sp_sampling._inject_profile_sampling(args, request, _spec(thinking_budget=1024)) + assert args.thinking_budget == 1024 + + +def test_inject_thinking_budget_request_wins(): + args = types.SimpleNamespace(thinking_budget=256) # client sent 256 + request = types.SimpleNamespace(model_fields_set={"thinking_budget"}) + sp_sampling._inject_profile_sampling(args, request, _spec(thinking_budget=1024)) + assert args.thinking_budget == 256 # not clobbered + + +def test_inject_thinking_budget_off_by_default(): + args = types.SimpleNamespace(thinking_budget=None) + request = types.SimpleNamespace(model_fields_set=set()) + sp_sampling._inject_profile_sampling(args, request, _spec(temperature=0.2)) # no budget + assert args.thinking_budget is None # stays off + + +# 1c. server thinking_budget on MTP models (mtp_thinking) +from gmlx.serve.patches import mtp_thinking as sp_mtp # noqa: E402 + + +@pytest.fixture +def _mtp_seams(): + """Force the owned-prefill class flag on (install-order precondition) and + snapshot the three methods mtp_thinking wraps.""" + from gmlx.spec.engine import _FULL_PREFILL_FLAG + gen = importlib.import_module("mlx_vlm.server.generation") + ar = importlib.import_module("mlx_vlm.generate.ar") + cls = gen.ResponseGenerator + had_flag = getattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG, False) + setattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG, True) + saved = (cls.generate, cls._make_thinking_budget_criteria, + ar.PromptProcessingBatch.generate) + yield gen, ar + (cls.generate, cls._make_thinking_budget_criteria, + ar.PromptProcessingBatch.generate) = saved + if not had_flag: + delattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG) + + +def _mtp_self(): + return types.SimpleNamespace( + draft_model=object(), draft_kind="mtp", + tokenizer=_FakeThinkTok(), + _thinking_token_ids=lambda args: (99, 100)) + + +def _budget_args(**kw): + base = dict(thinking_budget=6, enable_thinking=False, seed=None, + temperature=1.0, thinking_start_token=None, + thinking_end_token=None) + base.update(kw) + return types.SimpleNamespace(**base) + + +def test_mtp_thinking_install_refuses_without_owned_prefill(): + from gmlx.spec.engine import _FULL_PREFILL_FLAG + gen = importlib.import_module("mlx_vlm.server.generation") + ar = importlib.import_module("mlx_vlm.generate.ar") + had_flag = getattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG, False) + if had_flag: + delattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG) + before = gen.ResponseGenerator.generate + try: + sp_mtp.install_mtp_thinking_budget() + assert gen.ResponseGenerator.generate is before # refused, unbound + finally: + if had_flag: + setattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG, True) + + +def test_mtp_thinking_defers_budget_only_for_mtp(_mtp_seams): + gen, _ar = _mtp_seams + cls = gen.ResponseGenerator + seen = [] + + def stub(self, prompt, images=None, audio=None, args=None, videos=None): + seen.append(args.thinking_budget if args is not None else None) + return "gen" + + cls.generate = stub + sp_mtp.install_mtp_thinking_budget() + args = _budget_args() + assert cls.generate(_mtp_self(), "p", args=args) == "gen" + assert seen[-1] is None # moved aside + assert getattr(args, sp_mtp._DEFERRED_ATTR) == 6 + # Non-MTP drafter: untouched, so the upstream raise still fires there. + eagle = types.SimpleNamespace(draft_model=object(), draft_kind="eagle") + args2 = _budget_args() + cls.generate(eagle, "p", args=args2) + assert seen[-1] == 6 and not hasattr(args2, sp_mtp._DEFERRED_ATTR) + # Plain model and args=None: untouched. + plain = types.SimpleNamespace(draft_model=None, draft_kind=None) + args3 = _budget_args() + cls.generate(plain, "p", args=args3) + assert seen[-1] == 6 + cls.generate(_mtp_self(), "p") # args=None tolerated + + +def test_mtp_thinking_criteria_restores_even_on_early_out(_mtp_seams): + gen, _ar = _mtp_seams + cls = gen.ResponseGenerator + cls._make_thinking_budget_criteria = lambda self, args, input_ids: None + sp_mtp.install_mtp_thinking_budget() + make = cls._make_thinking_budget_criteria + args = _budget_args(thinking_budget=None) + setattr(args, sp_mtp._DEFERRED_ATTR, 6) + crit = make(_mtp_self(), args, [1, 2, 3]) + assert args.thinking_budget == 6 # restored + assert not hasattr(args, sp_mtp._DEFERRED_ATTR) + # Delegate returned None: the hook rides a duck-shaped carrier that the + # plain batch loop can call without raising. + hook = crit._kq_mtp_hook + assert hook is not None and hook.budget == 6 + assert crit(5) is None and crit.pop_forced_token_id() is None + # Prompt ending inside an open think block seeds the hook in-thinking. + args_open = _budget_args(thinking_budget=None) + setattr(args_open, sp_mtp._DEFERRED_ATTR, 6) + assert make(_mtp_self(), args_open, [1, 2, 99])._kq_mtp_hook.in_thinking + + +def test_mtp_thinking_criteria_restores_on_raise(_mtp_seams): + gen, _ar = _mtp_seams + cls = gen.ResponseGenerator + + def boom(self, args, input_ids): + raise RuntimeError("delegate failed") + + cls._make_thinking_budget_criteria = boom + sp_mtp.install_mtp_thinking_budget() + args = _budget_args(thinking_budget=None) + setattr(args, sp_mtp._DEFERRED_ATTR, 6) + with pytest.raises(RuntimeError): + cls._make_thinking_budget_criteria(_mtp_self(), args, [1]) + assert args.thinking_budget == 6 # not stranded + + +def test_mtp_thinking_full_chain_with_seed_and_tbfix(_mtp_seams): + # Runtime chain mtp -> seed -> tbfix: one call restores the deferred + # budget, stashes the seed, builds the armed criteria, and attaches the + # rounds hook to it. + import gmlx.serve.seed_rows as sr + gen, ar = _mtp_seams + cls = gen.ResponseGenerator + saved_insert = (ar.BatchGenerator.insert, ar.GenerationBatch._step, + ar.SpeculativeGenerationBatch.next) + sr._PENDING.clear() + try: + sp.install_thinking_budget_fix() + sr.install_per_request_seed() + sp_mtp.install_mtp_thinking_budget() + args = _budget_args(thinking_budget=None, seed=11) + setattr(args, sp_mtp._DEFERRED_ATTR, 6) + crit = cls._make_thinking_budget_criteria(_mtp_self(), args, [1, 2]) + assert args.thinking_budget == 6 + assert sr._PENDING == [11] + assert crit is not None and crit.in_thinking is False # tbfix armed + assert crit._kq_mtp_hook is not None and crit._kq_mtp_hook.budget == 6 + finally: + (ar.BatchGenerator.insert, ar.GenerationBatch._step, + ar.SpeculativeGenerationBatch.next) = saved_insert + sr._PENDING.clear() + + +def test_mtp_thinking_transport_stash_and_batch_drop(_mtp_seams): + _gen, ar = _mtp_seams + ar.PromptProcessingBatch.generate = \ + lambda self, sampler, *a, **k: self._out + sp_mtp.install_mtp_thinking_budget() + wrapper = ar.PromptProcessingBatch.generate + hook = object() + crit = types.SimpleNamespace(_kq_mtp_hook=hook) + cache_entry = types.SimpleNamespace() + batch = types.SimpleNamespace(prompt_cache=[cache_entry], uids=["u"]) + me = types.SimpleNamespace( + draft_model=object(), draft_kind="mtp", + thinking_budget_criteria=[crit], _out=batch) + assert wrapper(me, None) is batch + assert cache_entry._kq_mtp_thinking_hook is hook # B==1 stash + # B>1: dropped, nothing stashed. + c2 = types.SimpleNamespace() + batch2 = types.SimpleNamespace(prompt_cache=[c2], uids=["u", "v"]) + me2 = types.SimpleNamespace( + draft_model=object(), draft_kind="mtp", + thinking_budget_criteria=[crit, crit], _out=batch2) + wrapper(me2, None) + assert not hasattr(c2, "_kq_mtp_thinking_hook") + # Criteria/rows mismatch: dropped, not indexed blindly. + c3 = types.SimpleNamespace() + batch3 = types.SimpleNamespace(prompt_cache=[c3], uids=["u"]) + me3 = types.SimpleNamespace( + draft_model=object(), draft_kind="mtp", + thinking_budget_criteria=[crit, crit], _out=batch3) + wrapper(me3, None) + assert not hasattr(c3, "_kq_mtp_thinking_hook") + # Non-MTP batch: untouched. + c4 = types.SimpleNamespace() + batch4 = types.SimpleNamespace(prompt_cache=[c4], uids=["u"]) + me4 = types.SimpleNamespace( + draft_model=None, draft_kind=None, + thinking_budget_criteria=[crit], _out=batch4) + wrapper(me4, None) + assert not hasattr(c4, "_kq_mtp_thinking_hook") + + +def test_mtp_thinking_flags_carry_through_preflight(_mtp_seams): + # The defer wrap carries earlier flags forward and stamps its own, so a + # later mem_preflight re-install must see its flag and not double-wrap. + from gmlx.serve import mem_preflight as mp + gen, _ar = _mtp_seams + cls = gen.ResponseGenerator + + def stub(self, prompt, images=None, audio=None, args=None, videos=None): + return "gen" + + stub.__dict__[mp._INSTALLED_FLAG] = True # preflight installed + cls.generate = stub + sp_mtp.install_mtp_thinking_budget() + wrapped = cls.generate + assert wrapped is not stub + assert getattr(wrapped, mp._INSTALLED_FLAG, False) # carried forward + mp.install_memory_preflight() + assert cls.generate is wrapped # no double wrap + sp_mtp._install_defer(cls) + assert cls.generate is wrapped # own re-install no-ops + + +# 7. XTC sampling injection +def test_attach_xtc_noop_without_request_or_profile(): + args = types.SimpleNamespace(logits_processors=None) + request = types.SimpleNamespace(model_fields_set=set()) + sp_sampling._attach_xtc(args, request, None) + assert args.logits_processors is None + + +def test_attach_xtc_appends_processor_from_request_extras(): + args = types.SimpleNamespace(logits_processors=None) + request = types.SimpleNamespace(xtc_probability=1.0, xtc_threshold=0.2) + sp_sampling._attach_xtc(args, request, None) + assert args.logits_processors is not None and len(args.logits_processors) == 1 + # functional: prob=1.0 always triggers; threshold 0.2 with probs ~[.6,.3,.1] + # masks the top token, so argmax moves to the runner-up. + import math + + import mlx.core as mx + logits = mx.log(mx.array([[0.6, 0.3, 0.1]])) + out = args.logits_processors[0](mx.array([0]), logits) + assert int(mx.argmax(out, axis=-1).item()) == 1 + assert math.isinf(float(out[0, 0].item())) + + +def test_attach_xtc_profile_fallback_and_request_precedence(): + spec = _spec(xtc_probability=1.0, xtc_threshold=0.3) + token = serving.set_active_spec(spec) + try: + args = types.SimpleNamespace(logits_processors=None) + sp_sampling._attach_xtc(args, types.SimpleNamespace(), None) + assert args.logits_processors and len(args.logits_processors) == 1 + # an explicit client 0.0 wins over the profile and disables XTC + args2 = types.SimpleNamespace(logits_processors=None) + sp_sampling._attach_xtc(args2, types.SimpleNamespace(xtc_probability=0.0), None) + assert args2.logits_processors is None + finally: + serving.reset_active_spec(token) + + +def test_attach_xtc_string_zero_disables(): + # extra="allow" preserves raw JSON types: a client's "0" (string) is truthy, + # but must still disable XTC after coercion - the live bug this pins down. + for raw in ("0", "0.0", 0, 0.0): + args = types.SimpleNamespace(logits_processors=None) + sp_sampling._attach_xtc(args, types.SimpleNamespace(xtc_probability=raw), None) + assert args.logits_processors is None, f"xtc_probability={raw!r}" + + +def test_attach_xtc_string_prob_attaches(): + args = types.SimpleNamespace(logits_processors=None) + request = types.SimpleNamespace(xtc_probability="0.5", xtc_threshold="0.2") + sp_sampling._attach_xtc(args, request, None) + assert args.logits_processors is not None and len(args.logits_processors) == 1 + + +def test_attach_xtc_garbage_prob_rejects_400(): + # matches the neighboring coercion behavior (_sampling_float): typed 400, + # never a 500 out of the handler, and args stay untouched. + from fastapi import HTTPException + args = types.SimpleNamespace(logits_processors=None) + request = types.SimpleNamespace(xtc_probability="lots") + with pytest.raises(HTTPException) as ei: + sp_sampling._attach_xtc(args, request, None) + assert ei.value.status_code == 400 + assert args.logits_processors is None + + +def test_xtc_special_tokens_dedup_and_defensive(): + class _Tok: + eos_token_id = 7 + + def encode(self, s, add_special_tokens=True): + return [7] + + assert sp_sampling._xtc_special_tokens(types.SimpleNamespace(tokenizer=_Tok())) == [7] + assert sp_sampling._xtc_special_tokens(None) == [] + + class _IntEosTok(_Tok): + # regression (live server): TokenizersBackend exposes eos_token_ids as + # a bare int - iterating it raised "'int' object is not iterable" + eos_token_ids = 9 + + assert sp_sampling._xtc_special_tokens( + types.SimpleNamespace(tokenizer=_IntEosTok())) == [7, 9] + + +def test_install_xtc_wraps_and_stacks_with_profile_injection(): + sp.install_gen_args_profile_injection() + sp.install_xtc_sampling() + fn = _APP._build_gen_args + assert getattr(fn, sp_common._PATCH_FLAG, False) # carried forward + assert getattr(fn, sp_sampling._XTC_FLAG, False) + sp.install_xtc_sampling() # idempotent + assert _APP._build_gen_args is fn + + +# 7a2. top_k / min_p aware batch sampler (the historical dropped-top_k bug class) +def _kept_ids(sampler, probs): + """Vocab ids surviving the sampler's filter for one row of probs, plus the + masked [1, k] logits (sorted desc by prob).""" + import mlx.core as mx + logits = mx.log(mx.array([probs])) + masked, part, order = sampler._filtered(logits) + kept = [] + for j in range(masked.shape[-1]): + if float(masked[0, j].item()) != float("-inf"): + kept.append(int(part[0, int(order[0, j].item())].item())) + return kept, masked + + +def test_fast_sampler_hierarchical_topk_matches_flat(): + # Large vocabs route _filtered's top-k through the hierarchical id + # selector; the surviving id SET must equal the flat argpartition's + # (order within the set is re-sorted downstream either way). + import mlx.core as mx + for v, seed in ((201088, 0), (200005, 1), (131072, 2)): + lp = mx.random.normal((1, v), key=mx.random.key(seed)) + lp = lp.astype(mx.float32) + mx.eval(lp) + hier = set(sp_sampling._topk_ids(lp, 20)[0].tolist()) + flat = set(mx.argpartition(-lp, kth=19, axis=-1)[:, :20][0].tolist()) + assert hier == flat + + +def test_fast_sampler_masking(): + S = sp_sampling._FastPositionedSampler + probs = [0.4, 0.3, 0.2, 0.1] + # top_k=2: exactly the two most probable survive + kept, _ = _kept_ids(S(temperature=1.0, top_k=2), probs) + assert kept == [0, 1] + # top_p=0.5: nucleus keeps ids 0,1 (mass-before 0.0 and 0.4 < 0.5) + kept, _ = _kept_ids(S(temperature=1.0, top_p=0.5), probs) + assert kept == [0, 1] + # min_p=0.6: threshold 0.4*0.6=0.24 -> 0.3 stays, 0.2 pruned + kept, _ = _kept_ids(S(temperature=1.0, min_p=0.6), probs) + assert kept == [0, 1] + # llama.cpp order: top_k FIRST, top_p over the k renormalized survivors. + # [0.36, 0.34, 0.30] @ top_k=2 renorms to [0.514, 0.486]; top_p=0.45 then + # drops the runner-up (mass-before 0.514 > 0.45). Vocab-order top_p would + # have kept it (0.36 < 0.45). + kept, _ = _kept_ids(S(temperature=1.0, top_k=2, top_p=0.45), + [0.36, 0.34, 0.30]) + assert kept == [0] + # the argmax can never be filtered away (_MIN_KEEP) + kept, _ = _kept_ids(S(temperature=1.0, top_p=1e-9), probs) + assert kept == [0] + # temperature scales the surviving logits (applied last) + import math + _, masked = _kept_ids(S(temperature=0.5, top_k=2), probs) + assert math.isclose(float(masked[0, 0].item()), math.log(0.4) / 0.5, + rel_tol=1e-5) + + +def test_fast_sampler_call_shapes_and_determinism(): + import mlx.core as mx + s = sp_sampling._FastPositionedSampler(temperature=0.7, top_k=1) + logits = mx.log(mx.array([[0.1, 0.2, 0.6, 0.1]])) + # top_k=1 leaves a single candidate -> always the argmax id + assert int(s(logits).item()) == 2 + # a drafter's [B, 1, V] block keeps its leading shape + assert s(logits[:, None, :]).shape == (1, 1) + + +def test_fast_sampler_install_lands(): + """Identity check on the REAL upstream class: an mlx-vlm rename of + ResponseGenerator._make_sampler must fail here, not silently no-op.""" + gen = importlib.import_module("mlx_vlm.server.generation") + cls = gen.ResponseGenerator + original = cls._make_sampler + try: + sp.install_fast_sampler() + patched = cls._make_sampler + assert patched is not original # actually swapped + assert getattr(patched, sp_sampling._FAST_SAMPLER_FLAG, False) + sp.install_fast_sampler() # idempotent + assert cls._make_sampler is patched + me = types.SimpleNamespace() + # greedy keeps the batch engine's argmax fast path + assert patched(me, types.SimpleNamespace(temperature=0)) is None + s = patched(me, types.SimpleNamespace(temperature=0.6, top_p=0.9, + top_k=40, min_p=0.05, seed=3)) + assert isinstance(s, sp_sampling._FastPositionedSampler) + assert (s.top_k, s.min_p, s.seed) == (40, 0.05, 3) + finally: + cls._make_sampler = original diff --git a/tests/serve/test_serve_apc_engagement.py b/tests/serve/test_serve_apc_engagement.py index 849e1be5..f5ddb87f 100644 --- a/tests/serve/test_serve_apc_engagement.py +++ b/tests/serve/test_serve_apc_engagement.py @@ -29,7 +29,7 @@ class hides. import mlx.core as mx # noqa: E402 -pytestmark = [pytest.mark.integration, pytest.mark.slow] +pytestmark = pytest.mark.integration GREEDY = lambda x: mx.argmax(x, axis=-1) # noqa: E731 N_DECODE = 8 diff --git a/tests/serve/test_server_patches.py b/tests/serve/test_server_patches.py index 82f0e9c3..7948df37 100644 --- a/tests/serve/test_server_patches.py +++ b/tests/serve/test_server_patches.py @@ -8,7 +8,6 @@ import importlib import os -import sys import time import types @@ -31,76 +30,6 @@ _PKG = importlib.import_module("mlx_vlm.server") -@pytest.fixture(autouse=True) -def _restore_mlxvlm(): - """Snapshot every mlx-vlm seam these patches mutate, restore after each test.""" - fastapi_app = _APP.app - saved = { - "build_gen_args": _APP._build_gen_args, - "snapshot": _APP._server_runtime_snapshot, - "get_model_path": _UTILS.get_model_path, - "routes": sp_common._snapshot_routes(fastapi_app), - "handlers": dict(fastapi_app.exception_handlers), - "deps": getattr(getattr(_APP, "_protocol_deps", None), "build_gen_args", None), - "pool": getattr(_PKG, "_kq_residency_pool", None), - # Middleware installs (auth / host guard) append to user_middleware and - # set app.state flags; CORS hardening mutates a Middleware's kwargs in - # place - snapshot all three or one test's auth poisons the rest. - "middleware": list(fastapi_app.user_middleware), - "mw_kwargs": [(m, dict(getattr(m, "kwargs", {}) or {})) - for m in fastapi_app.user_middleware], - } - openai = sys.modules.get("mlx_vlm.server.openai") - anthropic = sys.modules.get("mlx_vlm.server.anthropic") - saved["openai_bga"] = getattr(openai, "_build_gen_args", None) - saved["anthropic_bga"] = getattr(anthropic, "_build_gen_args", None) - apc = sys.modules.get("mlx_vlm.apc") or importlib.import_module("mlx_vlm.apc") - saved["apc_harvest"] = apc.harvest_blocks_from_batch_cache - saved["apc_lone_flag"] = getattr(apc, "_kq_lone_harvest", False) - gen = importlib.import_module("mlx_vlm.server.generation") - saved["to_template_kwargs"] = gen.GenerationArguments.to_template_kwargs - pu = importlib.import_module("mlx_vlm.prompt_utils") - saved["get_chat_template"] = pu.get_chat_template - saved["make_sampler"] = gen.ResponseGenerator._make_sampler - schemas = importlib.import_module("mlx_vlm.server.schemas") - saved["stream_chunk_dump"] = schemas.ChatStreamChunk.model_dump_json - saved["stopping_call"] = _UTILS.StoppingCriteria.__call__ - serving.clear_resolved_models() - yield - _UTILS.StoppingCriteria.__call__ = saved["stopping_call"] - apc.harvest_blocks_from_batch_cache = saved["apc_harvest"] - apc._kq_lone_harvest = saved["apc_lone_flag"] - gen.GenerationArguments.to_template_kwargs = saved["to_template_kwargs"] - pu.get_chat_template = saved["get_chat_template"] - gen.ResponseGenerator._make_sampler = saved["make_sampler"] - schemas.ChatStreamChunk.model_dump_json = saved["stream_chunk_dump"] - _APP._build_gen_args = saved["build_gen_args"] - _APP._server_runtime_snapshot = saved["snapshot"] - _UTILS.get_model_path = saved["get_model_path"] - sp_common._restore_routes(fastapi_app, saved["routes"]) - fastapi_app.exception_handlers.clear() - fastapi_app.exception_handlers.update(saved["handlers"]) - if getattr(_APP, "_protocol_deps", None) is not None and saved["deps"] is not None: - _APP._protocol_deps.build_gen_args = saved["deps"] - if openai is not None: - openai._build_gen_args = saved["openai_bga"] - if anthropic is not None: - anthropic._build_gen_args = saved["anthropic_bga"] - if saved["pool"] is None: - if hasattr(_PKG, "_kq_residency_pool"): - delattr(_PKG, "_kq_residency_pool") - else: - _PKG._kq_residency_pool = saved["pool"] - fastapi_app.user_middleware[:] = saved["middleware"] - for m, kw in saved["mw_kwargs"]: - if getattr(m, "kwargs", None) is not None: - m.kwargs.clear() - m.kwargs.update(kw) - fastapi_app.middleware_stack = None # force a rebuild from the restored list - for flag in (sp_hardening._AUTH_FLAG, sp_hardening._HOST_GUARD_FLAG): - if hasattr(fastapi_app.state, flag): - delattr(fastapi_app.state, flag) - serving.clear_resolved_models() def _spec(**sampling): @@ -109,76 +38,8 @@ def _spec(**sampling): draft_gguf=None, pin=False, ttl_s=None) -# 1. sampling injection (pure) -def test_inject_overrides_unset_keeps_client_set(): - args = types.SimpleNamespace(temperature=0.7, top_p=0.95, top_k=5, max_tokens=512) - request = types.SimpleNamespace(model_fields_set={"top_k"}) # client set top_k - spec = _spec(temperature=0.2, top_p=0.9, top_k=99, max_tokens=2048) - sp_sampling._inject_profile_sampling(args, request, spec) - assert args.temperature == 0.2 # injected (unset) - assert args.top_p == 0.9 # injected (unset) - assert args.top_k == 5 # kept (client set) - assert args.max_tokens == 2048 # injected (unset) - - -def test_inject_max_tokens_alias_respects_max_output_tokens(): - args = types.SimpleNamespace(max_tokens=512) - request = types.SimpleNamespace(model_fields_set={"max_output_tokens"}) - sp_sampling._inject_profile_sampling(args, request, _spec(max_tokens=2048)) - assert args.max_tokens == 512 # responses API set it -> not overridden - - -# 1b. ignore-eos: forced-length decode -def test_install_ignore_eos_suppresses_stop(): - crit = _UTILS.StoppingCriteria([7, 8]) - assert crit(7) is True # baseline: 7 is an eos id -> stop - sp.install_ignore_eos() - assert crit(7) is False # patched: EOS never stops decode - assert crit(8) is False - assert crit(123) is False - sp.install_ignore_eos() # idempotent - assert crit(7) is False - - -def test_inject_noop_without_spec_or_sampling(): - args = types.SimpleNamespace(temperature=0.7) - request = types.SimpleNamespace(model_fields_set=set()) - sp_sampling._inject_profile_sampling(args, request, None) - sp_sampling._inject_profile_sampling(args, request, _spec()) # empty sampling - assert args.temperature == 0.7 - - -def test_inject_skips_unknown_arg_attr(): - args = types.SimpleNamespace(temperature=0.7) # no top_p attr - request = types.SimpleNamespace(model_fields_set=set()) - sp_sampling._inject_profile_sampling(args, request, _spec(top_p=0.5)) - assert not hasattr(args, "top_p") - - -def test_inject_thinking_budget_from_profile(): - # off by default: GenerationArguments.thinking_budget is None; a profile/model - # value seeds it when the client didn't ask. - args = types.SimpleNamespace(thinking_budget=None) - request = types.SimpleNamespace(model_fields_set=set()) - sp_sampling._inject_profile_sampling(args, request, _spec(thinking_budget=1024)) - assert args.thinking_budget == 1024 - - -def test_inject_thinking_budget_request_wins(): - args = types.SimpleNamespace(thinking_budget=256) # client sent 256 - request = types.SimpleNamespace(model_fields_set={"thinking_budget"}) - sp_sampling._inject_profile_sampling(args, request, _spec(thinking_budget=1024)) - assert args.thinking_budget == 256 # not clobbered - - -def test_inject_thinking_budget_off_by_default(): - args = types.SimpleNamespace(thinking_budget=None) - request = types.SimpleNamespace(model_fields_set=set()) - sp_sampling._inject_profile_sampling(args, request, _spec(temperature=0.2)) # no budget - assert args.thinking_budget is None # stays off -# 1b. thinking_budget enforcement fix (generate- models) class _FakeThinkTok: """Encodes the three control strings the criteria resolves.""" _MAP = {"": [99], "": [100], "\n": [10]} @@ -187,480 +48,8 @@ def encode(self, text, add_special_tokens=True): return self._MAP.get(text, [7]) -def _drive(criteria, tokens): - """Feed token ids; return the forced id (or None) the criteria emits each step.""" - return [criteria(t) for t in tokens] - - -def _armed(budget, prompt_open): - cls = sp_chat._armed_thinking_budget_criteria_cls() - return cls(tokenizer=_FakeThinkTok(), thinking_budget=budget, - thinking_start_token="", thinking_end_token="", - enable_thinking=True, prompt_open_thinking=prompt_open) - - -def test_armed_criteria_caps_generated_think(): - # prompt did NOT pre-fill (Qwen3 case); the model generates it. - c = _armed(3, prompt_open=False) - assert c.in_thinking is False # not started in a block - forced = _drive(c, [99, 1, 2, 3, 4, 5, 6]) # then 6 words - assert 10 in forced and 100 in forced # forced \n then - assert forced.index(10) < forced.index(100) - - -def test_armed_criteria_caps_prefilled_think(): - # prompt pre-filled (GLM-5.2 case): counting starts immediately. - c = _armed(2, prompt_open=True) - assert c.in_thinking is True - forced = _drive(c, [1, 2, 3, 4, 5]) - assert 100 in forced # forced close - - -def test_armed_criteria_never_forces_non_thinking_answer(): - # thinking enabled + budget set, but the model never opens a -> - # no token is ever counted, so nothing is force-closed (no corruption). - c = _armed(2, prompt_open=False) - forced = _drive(c, [1, 2, 3, 4, 5, 6, 7, 8]) - assert all(f is None for f in forced) - - -def test_armed_criteria_reset_restores_prompt_seed(): - c = _armed(2, prompt_open=False) - c.in_thinking = True - c.reset_thinking_state() - assert c.in_thinking is False # back to the prompt seed - - -def test_prompt_tail_opens_thinking_cases(): - pairs = (("", ""),) - f = sp_chat._prompt_tail_opens_thinking - assert f("", pairs) is False - assert f("plain prompt, no markers", pairs) is False - assert f("<|im_start|>assistant\n\n", pairs) is True # Qwen3.6 pre-fill - assert f("x\n\n\n\n", pairs) is False # thinking off - assert f("axb", pairs) is True # last pair open - assert f(None, pairs) is False - - -def test_stream_thinking_seed_reseeds_from_prompt(): - rs = importlib.import_module("mlx_vlm.server.responses_state") - cls = rs.ThinkingStreamState - original_init = cls.__init__ - try: - sp.install_stream_thinking_seed() - assert getattr(cls.__init__, sp_chat._STREAM_SEED_FLAG, False) - patched = cls.__init__ - sp.install_stream_thinking_seed() # idempotent - assert cls.__init__ is patched - tok = sp_chat._LAST_RENDERED_PROMPT.set("rendered, no thinking scaffold") - try: - st = cls(True) # enable_thinking forced True (7b) - assert st.in_thinking is False # gemma-4 default-off: content mode - sp_chat._LAST_RENDERED_PROMPT.set("<|im_start|>assistant\n\n") - st = cls(True) - assert st.in_thinking is True # Qwen3.6 pre-fill: reasoning first - st = cls(False) - assert st.in_thinking is True # prompt truth beats the flag - sp_chat._LAST_RENDERED_PROMPT.set(None) - st = cls(True) - assert st.in_thinking is True # no render seen -> stock seed - finally: - sp_chat._LAST_RENDERED_PROMPT.reset(tok) - finally: - cls.__init__ = original_init - - -def test_stream_thinking_seed_wraps_render_binding(): - openai_mod = importlib.import_module("mlx_vlm.server.openai") - rs = importlib.import_module("mlx_vlm.server.responses_state") - original_fn = openai_mod.apply_chat_template - original_init = rs.ThinkingStreamState.__init__ - try: - openai_mod.apply_chat_template = lambda *a, **kw: "tail \n" - sp.install_stream_thinking_seed() - wrapped = openai_mod.apply_chat_template - assert getattr(wrapped, sp_chat._STREAM_SEED_FLAG, False) - out = wrapped("processor", "config", []) - assert out == "tail \n" - assert sp_chat._LAST_RENDERED_PROMPT.get() == out # stashed - sp_chat._LAST_RENDERED_PROMPT.set(None) - finally: - openai_mod.apply_chat_template = original_fn - rs.ThinkingStreamState.__init__ = original_init - - -def test_nonstream_split_truncated_thinking_seeded_from_prompt(): - app_mod = importlib.import_module("mlx_vlm.server.app") - sp.install_stream_thinking_seed() - split = app_mod._split_thinking_text - assert getattr(split, sp_chat._STREAM_SEED_FLAG, False) - tok = sp_chat._LAST_RENDERED_PROMPT.set(None) - try: - # No stashed render (off-request callers): stock classification. - assert split("half a plan, cut off") == (None, "half a plan, cut off") - sp_chat._LAST_RENDERED_PROMPT.set("<|im_start|>assistant\n\n") - # Prompt-opened block, no marker in the truncated text: reasoning. - assert split("half a plan, cut off") == ("half a plan, cut off", "") - # A close marker in the text keeps the stock split untouched. - r, c = split("plan\n\n\nanswer") - assert r == "plan" and c == "answer" - # Prompt did not open a block: stock classification. - sp_chat._LAST_RENDERED_PROMPT.set("<|im_start|>assistant\n") - assert split("just an answer") == (None, "just an answer") - finally: - sp_chat._LAST_RENDERED_PROMPT.reset(tok) - - -def test_nonstream_split_strips_xtml_section_markers(): - """Kimi-K3: the response/message closers and turn terminator must not - leak into chat content; other think spellings stay untouched.""" - app_mod = importlib.import_module("mlx_vlm.server.app") - sp.install_stream_thinking_seed() - split = app_mod._split_thinking_text - tok = sp_chat._LAST_RENDERED_PROMPT.set(None) - try: - text = ("plan<|close|>think<|sep|><|open|>response<|sep|>" - "2+2 equals 4.<|close|>response<|sep|>" - "<|close|>message<|sep|><|end_of_msg|>") - r, c = split(text, "<|open|>think<|sep|>", "<|close|>think<|sep|>") - assert r == "plan" - assert c == "2+2 equals 4." - # Gate: a different think spelling passes content through unchanged. - r, c = split("xkeep <|close|>message<|sep|> text", - "", "") - assert c == "keep <|close|>message<|sep|> text" - finally: - sp_chat._LAST_RENDERED_PROMPT.reset(tok) - - -def test_nonstream_split_keeps_first_line_code_indent(): - """The stock splitter .strip()s content, which deletes the first line's - leading indent from verbatim-code replies. The seeded splitter trims - newlines only, for every marker branch and the no-marker fallthrough.""" - app_mod = importlib.import_module("mlx_vlm.server.app") - rs = importlib.import_module("mlx_vlm.server.responses_state") - sp.install_stream_thinking_seed() - split = app_mod._split_thinking_text - tok = sp_chat._LAST_RENDERED_PROMPT.set(None) - try: - # open+close pair - r, c = split("plan\n\n if x:\n y()") - assert r == "plan" - assert c == " if x:\n y()" - # close-only (prompt-opened block) - r, c = split("plan\n\n\n\tdoc = doc || document;") - assert r == "plan" - assert c == "\tdoc = doc || document;" - # no marker at all - assert split(" indented") == (None, " indented") - # whitespace-only content collapses to empty - assert split("plan\n \n") == ("plan", "") - # XTML section strip keeps the indent too - text = ("plan<|close|>think<|sep|><|open|>response<|sep|>" - " return 4;<|close|>response<|sep|><|end_of_msg|>") - r, c = split(text, "<|open|>think<|sep|>", "<|close|>think<|sep|>") - assert c == " return 4;" - # the /v1/responses module global got the same splitter - assert rs._split_thinking is split - finally: - sp_chat._LAST_RENDERED_PROMPT.reset(tok) - - -def test_install_thinking_budget_fix_applies_and_idempotent(): - # Fail-loud guard: asserts the seam bound to the REAL mlx-vlm symbol, so a - # rename of ResponseGenerator._make_thinking_budget_criteria turns into a CI - # failure instead of a silent no-op. - gen = importlib.import_module("mlx_vlm.server.generation") - cls = gen.ResponseGenerator - original = cls._make_thinking_budget_criteria - try: - sp.install_thinking_budget_fix() - patched = cls._make_thinking_budget_criteria - assert patched is not original # actually bound - assert getattr(patched, sp_chat._TBUDGET_FLAG, False) - sp.install_thinking_budget_fix() # idempotent - assert cls._make_thinking_budget_criteria is patched - finally: - cls._make_thinking_budget_criteria = original - - -def test_make_criteria_honors_budget_when_enable_thinking_false(): - # A configured thinking_budget must arm even when enable_thinking is False: - # a group/profile config may disable thinking, but the model can still emit - # , and an explicit budget must cap it. None budget still opts out. - gen = importlib.import_module("mlx_vlm.server.generation") - cls = gen.ResponseGenerator - original = cls._make_thinking_budget_criteria - try: - sp.install_thinking_budget_fix() - make = cls._make_thinking_budget_criteria - me = types.SimpleNamespace( - tokenizer=_FakeThinkTok(), - _thinking_token_ids=lambda args: (99, 100), # =99, =100 - ) - args = types.SimpleNamespace( - thinking_budget=8, enable_thinking=False, - thinking_start_token=None, thinking_end_token=None) - # generate-style prompt (no open ): armed but not seeded in-block - criteria = make(me, args, [1, 2, 3]) - assert criteria is not None # armed despite enable_thinking=False - assert criteria.in_thinking is False - # pre-fill prompt ending with an open seeds in_thinking True, even - # though enable_thinking is False - the cap must still fire (GLM-style). - criteria2 = make(me, args, [1, 2, 99]) - assert criteria2 is not None and criteria2.in_thinking is True - args.thinking_budget = None - assert make(me, args, [1, 2, 3]) is None # no budget -> still opts out - finally: - cls._make_thinking_budget_criteria = original - - -def test_seed_wrapper_survives_thinking_budget_fix(): - # Regression: the install order used to be seed -> tbfix, and tbfix rebinds - # the criteria seam without delegating, so it clobbered the seed wrapper and - # per-request seeds were dead on the serve path. The order is now tbfix -> - # seed: one call must stash the seed AND run tbfix's construction, and a - # tbfix re-install must see its flag through the wrapper and no-op. - import gmlx.serve.seed_rows as sr - gen = importlib.import_module("mlx_vlm.server.generation") - ar = importlib.import_module("mlx_vlm.generate.ar") - cls = gen.ResponseGenerator - saved = (cls._make_thinking_budget_criteria, ar.BatchGenerator.insert, - ar.GenerationBatch._step, ar.PromptProcessingBatch.generate, - ar.SpeculativeGenerationBatch.next) - sr._PENDING.clear() - try: - sp.install_thinking_budget_fix() - sr.install_per_request_seed() - crit = cls._make_thinking_budget_criteria - assert getattr(crit, sr._INSTALLED_FLAG, False) # seed outermost - assert getattr(crit, sp_chat._TBUDGET_FLAG, False) # tbfix flag carried - sp.install_thinking_budget_fix() - assert cls._make_thinking_budget_criteria is crit # re-install no-ops - me = types.SimpleNamespace( - tokenizer=_FakeThinkTok(), - _thinking_token_ids=lambda args: (99, 100)) - args = types.SimpleNamespace( - seed=7, temperature=1.0, thinking_budget=8, enable_thinking=False, - thinking_start_token=None, thinking_end_token=None) - out = crit(me, args, [1, 2, 3]) - assert out is not None and out.in_thinking is False # tbfix ran - assert sr._PENDING == [7] # seed stashed - finally: - (cls._make_thinking_budget_criteria, ar.BatchGenerator.insert, - ar.GenerationBatch._step, ar.PromptProcessingBatch.generate, - ar.SpeculativeGenerationBatch.next) = saved - sr._PENDING.clear() -# 1c. server thinking_budget on MTP models (mtp_thinking) -from gmlx.serve.patches import mtp_thinking as sp_mtp # noqa: E402 - - -@pytest.fixture -def _mtp_seams(): - """Force the owned-prefill class flag on (install-order precondition) and - snapshot the three methods mtp_thinking wraps.""" - from gmlx.spec.engine import _FULL_PREFILL_FLAG - gen = importlib.import_module("mlx_vlm.server.generation") - ar = importlib.import_module("mlx_vlm.generate.ar") - cls = gen.ResponseGenerator - had_flag = getattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG, False) - setattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG, True) - saved = (cls.generate, cls._make_thinking_budget_criteria, - ar.PromptProcessingBatch.generate) - yield gen, ar - (cls.generate, cls._make_thinking_budget_criteria, - ar.PromptProcessingBatch.generate) = saved - if not had_flag: - delattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG) - - -def _mtp_self(): - return types.SimpleNamespace( - draft_model=object(), draft_kind="mtp", - tokenizer=_FakeThinkTok(), - _thinking_token_ids=lambda args: (99, 100)) - - -def _budget_args(**kw): - base = dict(thinking_budget=6, enable_thinking=False, seed=None, - temperature=1.0, thinking_start_token=None, - thinking_end_token=None) - base.update(kw) - return types.SimpleNamespace(**base) - - -def test_mtp_thinking_install_refuses_without_owned_prefill(): - from gmlx.spec.engine import _FULL_PREFILL_FLAG - gen = importlib.import_module("mlx_vlm.server.generation") - ar = importlib.import_module("mlx_vlm.generate.ar") - had_flag = getattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG, False) - if had_flag: - delattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG) - before = gen.ResponseGenerator.generate - try: - sp_mtp.install_mtp_thinking_budget() - assert gen.ResponseGenerator.generate is before # refused, unbound - finally: - if had_flag: - setattr(ar.PromptProcessingBatch, _FULL_PREFILL_FLAG, True) - - -def test_mtp_thinking_defers_budget_only_for_mtp(_mtp_seams): - gen, _ar = _mtp_seams - cls = gen.ResponseGenerator - seen = [] - - def stub(self, prompt, images=None, audio=None, args=None, videos=None): - seen.append(args.thinking_budget if args is not None else None) - return "gen" - - cls.generate = stub - sp_mtp.install_mtp_thinking_budget() - args = _budget_args() - assert cls.generate(_mtp_self(), "p", args=args) == "gen" - assert seen[-1] is None # moved aside - assert getattr(args, sp_mtp._DEFERRED_ATTR) == 6 - # Non-MTP drafter: untouched, so the upstream raise still fires there. - eagle = types.SimpleNamespace(draft_model=object(), draft_kind="eagle") - args2 = _budget_args() - cls.generate(eagle, "p", args=args2) - assert seen[-1] == 6 and not hasattr(args2, sp_mtp._DEFERRED_ATTR) - # Plain model and args=None: untouched. - plain = types.SimpleNamespace(draft_model=None, draft_kind=None) - args3 = _budget_args() - cls.generate(plain, "p", args=args3) - assert seen[-1] == 6 - cls.generate(_mtp_self(), "p") # args=None tolerated - - -def test_mtp_thinking_criteria_restores_even_on_early_out(_mtp_seams): - gen, _ar = _mtp_seams - cls = gen.ResponseGenerator - cls._make_thinking_budget_criteria = lambda self, args, input_ids: None - sp_mtp.install_mtp_thinking_budget() - make = cls._make_thinking_budget_criteria - args = _budget_args(thinking_budget=None) - setattr(args, sp_mtp._DEFERRED_ATTR, 6) - crit = make(_mtp_self(), args, [1, 2, 3]) - assert args.thinking_budget == 6 # restored - assert not hasattr(args, sp_mtp._DEFERRED_ATTR) - # Delegate returned None: the hook rides a duck-shaped carrier that the - # plain batch loop can call without raising. - hook = crit._kq_mtp_hook - assert hook is not None and hook.budget == 6 - assert crit(5) is None and crit.pop_forced_token_id() is None - # Prompt ending inside an open think block seeds the hook in-thinking. - args_open = _budget_args(thinking_budget=None) - setattr(args_open, sp_mtp._DEFERRED_ATTR, 6) - assert make(_mtp_self(), args_open, [1, 2, 99])._kq_mtp_hook.in_thinking - - -def test_mtp_thinking_criteria_restores_on_raise(_mtp_seams): - gen, _ar = _mtp_seams - cls = gen.ResponseGenerator - - def boom(self, args, input_ids): - raise RuntimeError("delegate failed") - - cls._make_thinking_budget_criteria = boom - sp_mtp.install_mtp_thinking_budget() - args = _budget_args(thinking_budget=None) - setattr(args, sp_mtp._DEFERRED_ATTR, 6) - with pytest.raises(RuntimeError): - cls._make_thinking_budget_criteria(_mtp_self(), args, [1]) - assert args.thinking_budget == 6 # not stranded - - -def test_mtp_thinking_full_chain_with_seed_and_tbfix(_mtp_seams): - # Runtime chain mtp -> seed -> tbfix: one call restores the deferred - # budget, stashes the seed, builds the armed criteria, and attaches the - # rounds hook to it. - import gmlx.serve.seed_rows as sr - gen, ar = _mtp_seams - cls = gen.ResponseGenerator - saved_insert = (ar.BatchGenerator.insert, ar.GenerationBatch._step, - ar.SpeculativeGenerationBatch.next) - sr._PENDING.clear() - try: - sp.install_thinking_budget_fix() - sr.install_per_request_seed() - sp_mtp.install_mtp_thinking_budget() - args = _budget_args(thinking_budget=None, seed=11) - setattr(args, sp_mtp._DEFERRED_ATTR, 6) - crit = cls._make_thinking_budget_criteria(_mtp_self(), args, [1, 2]) - assert args.thinking_budget == 6 - assert sr._PENDING == [11] - assert crit is not None and crit.in_thinking is False # tbfix armed - assert crit._kq_mtp_hook is not None and crit._kq_mtp_hook.budget == 6 - finally: - (ar.BatchGenerator.insert, ar.GenerationBatch._step, - ar.SpeculativeGenerationBatch.next) = saved_insert - sr._PENDING.clear() - - -def test_mtp_thinking_transport_stash_and_batch_drop(_mtp_seams): - _gen, ar = _mtp_seams - ar.PromptProcessingBatch.generate = \ - lambda self, sampler, *a, **k: self._out - sp_mtp.install_mtp_thinking_budget() - wrapper = ar.PromptProcessingBatch.generate - hook = object() - crit = types.SimpleNamespace(_kq_mtp_hook=hook) - cache_entry = types.SimpleNamespace() - batch = types.SimpleNamespace(prompt_cache=[cache_entry], uids=["u"]) - me = types.SimpleNamespace( - draft_model=object(), draft_kind="mtp", - thinking_budget_criteria=[crit], _out=batch) - assert wrapper(me, None) is batch - assert cache_entry._kq_mtp_thinking_hook is hook # B==1 stash - # B>1: dropped, nothing stashed. - c2 = types.SimpleNamespace() - batch2 = types.SimpleNamespace(prompt_cache=[c2], uids=["u", "v"]) - me2 = types.SimpleNamespace( - draft_model=object(), draft_kind="mtp", - thinking_budget_criteria=[crit, crit], _out=batch2) - wrapper(me2, None) - assert not hasattr(c2, "_kq_mtp_thinking_hook") - # Criteria/rows mismatch: dropped, not indexed blindly. - c3 = types.SimpleNamespace() - batch3 = types.SimpleNamespace(prompt_cache=[c3], uids=["u"]) - me3 = types.SimpleNamespace( - draft_model=object(), draft_kind="mtp", - thinking_budget_criteria=[crit, crit], _out=batch3) - wrapper(me3, None) - assert not hasattr(c3, "_kq_mtp_thinking_hook") - # Non-MTP batch: untouched. - c4 = types.SimpleNamespace() - batch4 = types.SimpleNamespace(prompt_cache=[c4], uids=["u"]) - me4 = types.SimpleNamespace( - draft_model=None, draft_kind=None, - thinking_budget_criteria=[crit], _out=batch4) - wrapper(me4, None) - assert not hasattr(c4, "_kq_mtp_thinking_hook") - - -def test_mtp_thinking_flags_carry_through_preflight(_mtp_seams): - # The defer wrap carries earlier flags forward and stamps its own, so a - # later mem_preflight re-install must see its flag and not double-wrap. - from gmlx.serve import mem_preflight as mp - gen, _ar = _mtp_seams - cls = gen.ResponseGenerator - - def stub(self, prompt, images=None, audio=None, args=None, videos=None): - return "gen" - - stub.__dict__[mp._INSTALLED_FLAG] = True # preflight installed - cls.generate = stub - sp_mtp.install_mtp_thinking_budget() - wrapped = cls.generate - assert wrapped is not stub - assert getattr(wrapped, mp._INSTALLED_FLAG, False) # carried forward - mp.install_memory_preflight() - assert cls.generate is wrapped # no double wrap - sp_mtp._install_defer(cls) - assert cls.generate is wrapped # own re-install no-ops # 2. gen-args wrapper reference swap @@ -738,209 +127,6 @@ def _register(doc): serving.register_resolved_models(build_config(doc)) -def test_models_payload_lists_configured_ids_not_hf(): - _register({"models": { - "qwen": {"path": "/abs/qwen.gguf"}, - "gemma-vlm": {"path": "/abs/g.gguf", "mmproj": "/abs/mm.gguf"}, - }}) - payload = sp_routes._models_payload() - ids = {m["id"] for m in payload["data"]} - assert ids == {"qwen", "gemma-vlm"} - vlm = next(m for m in payload["data"] if m["id"] == "gemma-vlm") - assert vlm["vlm"] is True - assert all(m["resident"] is False for m in payload["data"]) # no pool - - -def test_models_payload_marks_resident_from_pool(): - _register({"models": {"qwen": {"path": "/abs/qwen.gguf", "pin": True}}}) - - class _FakePool: - def stats(self): - return {"resident": [{"model_path": "/abs/qwen.gguf", "pinned": True, - "footprint_bytes": 10, "idle_s": 3.0, - "ttl_s": 900}]} - - _PKG._kq_residency_pool = _FakePool() - m = sp_routes._models_payload()["data"][0] - assert m["resident"] is True and m["pinned"] is True - - -def test_models_payload_lists_aliases_as_pickable_entries(): - _register({ - "profiles": {"coder": {"sampling": {"temperature": 0.2}}}, - "models": {"qwen": {"path": "/abs/qwen.gguf", "speculative": False}}, - "aliases": {"big": "qwen", "coder-preset": "qwen@coder"}, - }) - payload = sp_routes._models_payload() - by_id = {m["id"]: m for m in payload["data"]} - assert set(by_id) == {"qwen", "big", "coder-preset"} # aliases listed - assert by_id["big"]["alias_of"] == "qwen" - assert by_id["coder-preset"]["alias_of"] == "qwen" - assert by_id["coder-preset"]["profile"] == "coder" # baked profile shown - assert "alias_of" not in by_id["qwen"] # real model unmarked - - -def test_models_payload_marks_default(): - _register({ - "server": {"defaults": {"model": "qwen"}}, - "models": {"qwen": {"path": "/abs/qwen.gguf"}, - "gemma": {"path": "/abs/g.gguf"}}, - }) - by_id = {m["id"]: m for m in sp_routes._models_payload()["data"]} - assert by_id["qwen"]["default"] is True - assert by_id["gemma"]["default"] is False - - -def test_models_override_registers_single_route(): - sp.install_models_endpoint_override() - paths = [getattr(r, "path", None) for r in _APP.app.router.routes] - assert paths.count("/v1/models") == 1 - sp.install_models_endpoint_override() # idempotent-ish - paths = [getattr(r, "path", None) for r in _APP.app.router.routes] - assert paths.count("/v1/models") == 1 - - -# 4. HF gate -def test_gate_allows_local_and_gguf(tmp_path): - calls = [] - orig = lambda p, *a, **k: calls.append(p) or "OK" - local = tmp_path / "f" - local.write_text("x") - assert sp_routes._gate_model_path(str(local), False, orig) == "OK" - assert sp_routes._gate_model_path("/x/model.gguf", False, orig) == "OK" - assert len(calls) == 2 - - -def test_gate_blocks_hf_id_when_disabled(): - orig = lambda p, *a, **k: "OK" - with pytest.raises(sp.HFAccessDisabled): - sp_routes._gate_model_path("org/model", False, orig) - - -def test_gate_allows_hf_id_when_cache_on(): - calls = [] - orig = lambda p, *a, **k: calls.append(p) or "OK" - assert sp_routes._gate_model_path("org/model", True, orig) == "OK" - assert calls == ["org/model"] - - -def test_install_hf_gate_sets_offline_env(monkeypatch): - monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) - sp.install_hf_download_gate(hf_cache=True) - import os - assert os.environ.get("HF_HUB_OFFLINE") == "1" - - -# 5. runtime-snapshot enrichment -def test_snapshot_enrichment_adds_resident_models(): - _APP._server_runtime_snapshot = lambda: {"loaded_model": "x"} - - class _FakePool: - def stats(self): - return {"resident": [{"model_path": "/abs/qwen.gguf", "pinned": False, - "busy": 3, "footprint_bytes": 99, - "idle_s": 1.234, "ttl_s": 900}]} - - _PKG._kq_residency_pool = _FakePool() - serving._PATH_TO_IDS["/abs/qwen.gguf"] = ["qwen"] - try: - sp.install_runtime_snapshot_enrichment() - snap = _APP._server_runtime_snapshot() - finally: - serving._PATH_TO_IDS.pop("/abs/qwen.gguf", None) - assert snap["loaded_model"] == "x" # base preserved - assert snap["resident_models"][0]["ids"] == ["qwen"] - assert snap["resident_models"][0]["idle_s"] == 1.2 # rounded - assert snap["resident_models"][0]["busy"] == 3 # in-flight count - - -# 6. error handlers + reload + unload route -def test_error_content_dialect_shapes(): - # One condition, two envelopes: OpenAI-style everywhere, Anthropic's - # {"type": "error", ...} with its fixed taxonomy on /v1/messages. - openai = sp_common._error_content( - "/v1/chat/completions", 404, "model_not_found", "no such model", - available_models=["a"]) - assert openai == {"error": {"type": "model_not_found", - "message": "no such model", - "available_models": ["a"]}} - anthropic = sp_common._error_content( - "/v1/messages", 404, "model_not_found", "no such model") - assert anthropic["type"] == "error" - assert anthropic["error"] == {"type": "not_found_error", - "message": "no such model"} - assert sp_common._error_content("/v1/messages", 500, "server_error", - "x")["error"]["type"] == "api_error" - - -def test_http_exception_envelope_unwrapped(): - # The residency resolver path raises HTTPException carrying the unified - # {"error": {...}} detail; the app-level handler must serve that body - # directly (no {"detail": ...} wrapper) and wrap plain-string details. - from fastapi import HTTPException - from fastapi.testclient import TestClient - - app = _APP.app - if not any(getattr(r, "path", None) == "/test/raise-envelope" - for r in app.router.routes): - @app.get("/test/raise-envelope") - async def _raise_envelope(): - raise HTTPException(status_code=404, detail={"error": { - "type": "model_not_found", "message": "no such model", - "available_models": ["a"]}}) - - @app.get("/test/raise-string") - async def _raise_string(): - raise HTTPException(status_code=500, detail="it broke") - - sp.install_resolver_error_handlers() - client = TestClient(app) - r = client.get("/test/raise-envelope") - assert r.status_code == 404 - assert r.json() == {"error": {"type": "model_not_found", - "message": "no such model", - "available_models": ["a"]}} - r2 = client.get("/test/raise-string") - assert r2.status_code == 500 - assert r2.json() == {"error": {"type": "server_error", - "message": "it broke"}} - - -def test_resolver_error_handlers_registered(): - sp.install_resolver_error_handlers() - handlers = _APP.app.exception_handlers - assert serving.ModelNotFound in handlers - assert serving.ModelFileMissing in handlers - assert serving.UnknownProfile in handlers - assert sp.HFAccessDisabled in handlers - - -def test_unload_and_reload_routes_register(): - sp.install_pool_aware_unload() - sp.install_reload_route(lambda: {"reloaded": 1}) - paths = [getattr(r, "path", None) for r in _APP.app.router.routes] - assert paths.count("/unload") == 1 - assert paths.count("/v1/reload") == 1 - - -def test_unload_accepts_body_and_empty_post(): - """Regression: the ``request: Request`` annotation must resolve at module level. - Under ``from __future__ import annotations`` a locally-imported ``Request`` left - FastAPI treating ``request`` as a required query param, so every POST 422'd before - the body was read (caught only at e2e). A route-count check can't see this - POST - it for real and assert the body is actually consumed.""" - from fastapi.testclient import TestClient - - sp.install_pool_aware_unload() - client = TestClient(_APP.app) - # no pool registered -> handler runs (no 422) and reports the absence - # cleanly, as a 503: the unload cannot be honored without a pool - r_body = client.post("/unload", json={"model": "m"}) - assert r_body.status_code == 503, r_body.text - assert r_body.json() == {"status": "error", "message": "no residency pool"} - r_empty = client.post("/unload") - assert r_empty.status_code == 200, r_empty.text - assert r_empty.json()["status"] == "no_model_loaded" # 6b. /v1/keep - the keep tier (TTL-exempt, LRU-eligible) @@ -952,97 +138,6 @@ def set_keep(self, path, keep): self.kept.append((path, keep)) -def test_json_content_type_tolerance(): - # `curl -d '{...}'` (every doc example) sends form-encoded; the middleware - # must rewrite it to application/json so pydantic parses the body instead - # of 422ing. Multipart (audio uploads) must pass through untouched. - from fastapi.testclient import TestClient - - app = _APP.app - if not any(getattr(r, "path", None) == "/test/echo-ct" - for r in app.router.routes): - # `dict` (a builtin) survives this module's stringized annotations; a - # test-local pydantic class would resolve as a query param instead. - @app.post("/test/echo-ct") - async def _echo_ct(body: dict): - return {"model": body.get("model")} - - sp.install_json_content_type_tolerance() - client = TestClient(app) - r = client.post("/test/echo-ct", content=b'{"model": "m1"}', - headers={"Content-Type": "application/x-www-form-urlencoded"}) - assert r.status_code == 200 and r.json() == {"model": "m1"} - r = client.post("/test/echo-ct", content=b'{"model": "m2"}', - headers={"Content-Type": "text/plain"}) - assert r.status_code == 200 and r.json() == {"model": "m2"} - r = client.post("/test/echo-ct", json={"model": "m3"}) # normal path intact - assert r.status_code == 200 - r = client.post("/test/echo-ct", files={"file": ("a.txt", b"x")}) - assert r.status_code == 422 # multipart not rewritten - - -def test_keep_route_registers(): - sp.install_keep_route() - paths = [getattr(r, "path", None) for r in _APP.app.router.routes] - assert paths.count("/v1/keep") == 1 - - -def test_keep_no_pool_reports_error(): - from fastapi.testclient import TestClient - - sp.install_keep_route() - client = TestClient(_APP.app) - r = client.post("/v1/keep", json={"model": "m"}) - assert r.status_code == 503, r.text - assert r.json() == {"status": "error", "message": "no residency pool"} - - -def test_keep_marks_resolved_model(monkeypatch): - from fastapi.testclient import TestClient - - monkeypatch.setattr(sp_routes, "_spawn_keep_warm", lambda model_id: None) - _register({"models": {"qwen": {"path": "/abs/qwen.gguf"}}}) - pool = _FakeKeepPool() - _PKG._kq_residency_pool = pool - sp.install_keep_route() - client = TestClient(_APP.app) - r = client.post("/v1/keep", json={"model": "qwen", "warm": False}) - assert r.status_code == 200, r.text - assert r.json() == {"status": "kept", "model": "qwen", "warming": False} - assert pool.kept == [("/abs/qwen.gguf", True)] - - -def test_keep_warm_default_spawns_warm(monkeypatch): - from fastapi.testclient import TestClient - - warmed = [] - monkeypatch.setattr(sp_routes, "_spawn_keep_warm", lambda model_id: warmed.append(model_id)) - _register({"models": {"qwen": {"path": "/abs/qwen.gguf"}}}) - _PKG._kq_residency_pool = _FakeKeepPool() - sp.install_keep_route() - client = TestClient(_APP.app) - r = client.post("/v1/keep", json={"model": "qwen"}) # warm omitted -> default True - assert r.json() == {"status": "kept", "model": "qwen", "warming": True} - assert warmed == ["qwen"] - - -def test_keep_false_releases_without_evicting(monkeypatch): - # A voice session ending releases its hold; the model stays resident - # under normal LRU/TTL rather than being dumped. - from fastapi.testclient import TestClient - - warmed = [] - monkeypatch.setattr(sp_routes, "_spawn_keep_warm", lambda model_id: warmed.append(model_id)) - _register({"models": {"qwen": {"path": "/abs/qwen.gguf"}}}) - pool = _FakeKeepPool() - _PKG._kq_residency_pool = pool - sp.install_keep_route() - client = TestClient(_APP.app) - r = client.post("/v1/keep", json={"model": "qwen", "keep": False}) - assert r.status_code == 200, r.text - assert r.json() == {"status": "released", "model": "qwen"} - assert pool.kept == [("/abs/qwen.gguf", False)] - assert warmed == [] # release never warms # 7. vanilla streaming chunks (exclude_none) - the Open WebUI blank-render fix @@ -1177,29 +272,6 @@ def test_transcriptions_route_emits_req_line(monkeypatch, capsys): assert "file=a.wav" in out and "in_bytes=4" in out -def test_keep_unknown_model_graceful(): - from fastapi.testclient import TestClient - - _register({"models": {"qwen": {"path": "/abs/qwen.gguf"}}}) - _PKG._kq_residency_pool = _FakeKeepPool() - sp.install_keep_route() - client = TestClient(_APP.app) - r = client.post("/v1/keep", json={"model": "nope"}) - # 404, not 200: a typo'd keep must not read as success (launch checks the - # status code); the body still names the id for older/other clients. - assert r.status_code == 404, r.text - assert r.json() == {"status": "unknown_model", "model": "nope"} - - -def test_keep_missing_model_field(): - from fastapi.testclient import TestClient - - _PKG._kq_residency_pool = _FakeKeepPool() - sp.install_keep_route() - client = TestClient(_APP.app) - r = client.post("/v1/keep", json={}) - assert r.status_code == 400, r.text - assert r.json() == {"status": "error", "message": "missing 'model'"} # 6c. /v1/audio/speech - TTS route (mlx-audio stubbed) @@ -1497,474 +569,10 @@ def test_model_not_found_str_is_plain(): assert not msg.startswith('"') -# 7. XTC sampling injection -def test_attach_xtc_noop_without_request_or_profile(): - args = types.SimpleNamespace(logits_processors=None) - request = types.SimpleNamespace(model_fields_set=set()) - sp_sampling._attach_xtc(args, request, None) - assert args.logits_processors is None - - -def test_attach_xtc_appends_processor_from_request_extras(): - args = types.SimpleNamespace(logits_processors=None) - request = types.SimpleNamespace(xtc_probability=1.0, xtc_threshold=0.2) - sp_sampling._attach_xtc(args, request, None) - assert args.logits_processors is not None and len(args.logits_processors) == 1 - # functional: prob=1.0 always triggers; threshold 0.2 with probs ~[.6,.3,.1] - # masks the top token, so argmax moves to the runner-up. - import math - - import mlx.core as mx - logits = mx.log(mx.array([[0.6, 0.3, 0.1]])) - out = args.logits_processors[0](mx.array([0]), logits) - assert int(mx.argmax(out, axis=-1).item()) == 1 - assert math.isinf(float(out[0, 0].item())) - - -def test_attach_xtc_profile_fallback_and_request_precedence(): - spec = _spec(xtc_probability=1.0, xtc_threshold=0.3) - token = serving.set_active_spec(spec) - try: - args = types.SimpleNamespace(logits_processors=None) - sp_sampling._attach_xtc(args, types.SimpleNamespace(), None) - assert args.logits_processors and len(args.logits_processors) == 1 - # an explicit client 0.0 wins over the profile and disables XTC - args2 = types.SimpleNamespace(logits_processors=None) - sp_sampling._attach_xtc(args2, types.SimpleNamespace(xtc_probability=0.0), None) - assert args2.logits_processors is None - finally: - serving.reset_active_spec(token) - - -def test_attach_xtc_string_zero_disables(): - # extra="allow" preserves raw JSON types: a client's "0" (string) is truthy, - # but must still disable XTC after coercion - the live bug this pins down. - for raw in ("0", "0.0", 0, 0.0): - args = types.SimpleNamespace(logits_processors=None) - sp_sampling._attach_xtc(args, types.SimpleNamespace(xtc_probability=raw), None) - assert args.logits_processors is None, f"xtc_probability={raw!r}" - - -def test_attach_xtc_string_prob_attaches(): - args = types.SimpleNamespace(logits_processors=None) - request = types.SimpleNamespace(xtc_probability="0.5", xtc_threshold="0.2") - sp_sampling._attach_xtc(args, request, None) - assert args.logits_processors is not None and len(args.logits_processors) == 1 - - -def test_attach_xtc_garbage_prob_rejects_400(): - # matches the neighboring coercion behavior (_sampling_float): typed 400, - # never a 500 out of the handler, and args stay untouched. - from fastapi import HTTPException - args = types.SimpleNamespace(logits_processors=None) - request = types.SimpleNamespace(xtc_probability="lots") - with pytest.raises(HTTPException) as ei: - sp_sampling._attach_xtc(args, request, None) - assert ei.value.status_code == 400 - assert args.logits_processors is None - - -def test_xtc_special_tokens_dedup_and_defensive(): - class _Tok: - eos_token_id = 7 - - def encode(self, s, add_special_tokens=True): - return [7] - - assert sp_sampling._xtc_special_tokens(types.SimpleNamespace(tokenizer=_Tok())) == [7] - assert sp_sampling._xtc_special_tokens(None) == [] - class _IntEosTok(_Tok): - # regression (live server): TokenizersBackend exposes eos_token_ids as - # a bare int - iterating it raised "'int' object is not iterable" - eos_token_ids = 9 - assert sp_sampling._xtc_special_tokens( - types.SimpleNamespace(tokenizer=_IntEosTok())) == [7, 9] -def test_install_xtc_wraps_and_stacks_with_profile_injection(): - sp.install_gen_args_profile_injection() - sp.install_xtc_sampling() - fn = _APP._build_gen_args - assert getattr(fn, sp_common._PATCH_FLAG, False) # carried forward - assert getattr(fn, sp_sampling._XTC_FLAG, False) - sp.install_xtc_sampling() # idempotent - assert _APP._build_gen_args is fn - - -# 7a2. top_k / min_p aware batch sampler (the historical dropped-top_k bug class) -def _kept_ids(sampler, probs): - """Vocab ids surviving the sampler's filter for one row of probs, plus the - masked [1, k] logits (sorted desc by prob).""" - import mlx.core as mx - logits = mx.log(mx.array([probs])) - masked, part, order = sampler._filtered(logits) - kept = [] - for j in range(masked.shape[-1]): - if float(masked[0, j].item()) != float("-inf"): - kept.append(int(part[0, int(order[0, j].item())].item())) - return kept, masked - - -def test_fast_sampler_hierarchical_topk_matches_flat(): - # Large vocabs route _filtered's top-k through the hierarchical id - # selector; the surviving id SET must equal the flat argpartition's - # (order within the set is re-sorted downstream either way). - import mlx.core as mx - for v, seed in ((201088, 0), (200005, 1), (131072, 2)): - lp = mx.random.normal((1, v), key=mx.random.key(seed)) - lp = lp.astype(mx.float32) - mx.eval(lp) - hier = set(sp_sampling._topk_ids(lp, 20)[0].tolist()) - flat = set(mx.argpartition(-lp, kth=19, axis=-1)[:, :20][0].tolist()) - assert hier == flat - - -def test_fast_sampler_masking(): - S = sp_sampling._FastPositionedSampler - probs = [0.4, 0.3, 0.2, 0.1] - # top_k=2: exactly the two most probable survive - kept, _ = _kept_ids(S(temperature=1.0, top_k=2), probs) - assert kept == [0, 1] - # top_p=0.5: nucleus keeps ids 0,1 (mass-before 0.0 and 0.4 < 0.5) - kept, _ = _kept_ids(S(temperature=1.0, top_p=0.5), probs) - assert kept == [0, 1] - # min_p=0.6: threshold 0.4*0.6=0.24 -> 0.3 stays, 0.2 pruned - kept, _ = _kept_ids(S(temperature=1.0, min_p=0.6), probs) - assert kept == [0, 1] - # llama.cpp order: top_k FIRST, top_p over the k renormalized survivors. - # [0.36, 0.34, 0.30] @ top_k=2 renorms to [0.514, 0.486]; top_p=0.45 then - # drops the runner-up (mass-before 0.514 > 0.45). Vocab-order top_p would - # have kept it (0.36 < 0.45). - kept, _ = _kept_ids(S(temperature=1.0, top_k=2, top_p=0.45), - [0.36, 0.34, 0.30]) - assert kept == [0] - # the argmax can never be filtered away (_MIN_KEEP) - kept, _ = _kept_ids(S(temperature=1.0, top_p=1e-9), probs) - assert kept == [0] - # temperature scales the surviving logits (applied last) - import math - _, masked = _kept_ids(S(temperature=0.5, top_k=2), probs) - assert math.isclose(float(masked[0, 0].item()), math.log(0.4) / 0.5, - rel_tol=1e-5) - - -def test_fast_sampler_call_shapes_and_determinism(): - import mlx.core as mx - s = sp_sampling._FastPositionedSampler(temperature=0.7, top_k=1) - logits = mx.log(mx.array([[0.1, 0.2, 0.6, 0.1]])) - # top_k=1 leaves a single candidate -> always the argmax id - assert int(s(logits).item()) == 2 - # a drafter's [B, 1, V] block keeps its leading shape - assert s(logits[:, None, :]).shape == (1, 1) - - -def test_fast_sampler_install_lands(): - """Identity check on the REAL upstream class: an mlx-vlm rename of - ResponseGenerator._make_sampler must fail here, not silently no-op.""" - gen = importlib.import_module("mlx_vlm.server.generation") - cls = gen.ResponseGenerator - original = cls._make_sampler - try: - sp.install_fast_sampler() - patched = cls._make_sampler - assert patched is not original # actually swapped - assert getattr(patched, sp_sampling._FAST_SAMPLER_FLAG, False) - sp.install_fast_sampler() # idempotent - assert cls._make_sampler is patched - me = types.SimpleNamespace() - # greedy keeps the batch engine's argmax fast path - assert patched(me, types.SimpleNamespace(temperature=0)) is None - s = patched(me, types.SimpleNamespace(temperature=0.6, top_p=0.9, - top_k=40, min_p=0.05, seed=3)) - assert isinstance(s, sp_sampling._FastPositionedSampler) - assert (s.top_k, s.min_p, s.seed) == (40, 0.05, 3) - finally: - cls._make_sampler = original - - -# 7b. chat_template_kwargs passthrough -def _spec_ctkw(**ctkw): - return ResolvedModel(id="m", path="/p", sampling={}, load={}, cache={}, - system=None, speculative=False, mmproj=None, - draft_gguf=None, pin=False, ttl_s=None, - chat_template_kwargs=ctkw) - - -def test_merged_template_kwargs_request_wins_over_profile(): - spec = _spec_ctkw(preserve_thinking=True, foo="profile") - request = types.SimpleNamespace(chat_template_kwargs={"foo": "request"}) - merged = sp_chat._merged_template_kwargs(request, spec) - assert merged == {"preserve_thinking": True, "foo": "request"} - - -def test_merged_template_kwargs_each_side_alone_and_empty(): - # request only (single-model mode: no active spec) - req = types.SimpleNamespace(chat_template_kwargs={"preserve_thinking": True}) - assert sp_chat._merged_template_kwargs(req, None) == {"preserve_thinking": True} - # profile only (request carries nothing) - spec = _spec_ctkw(preserve_thinking=False) - assert sp_chat._merged_template_kwargs(types.SimpleNamespace(), spec) == { - "preserve_thinking": False} - # neither => {} - assert sp_chat._merged_template_kwargs(types.SimpleNamespace(), None) == {} - - -def test_merged_template_kwargs_spec_thinking_controls_mapped(): - """Profile-level thinking/reasoning_effort are dedicated controls: mapped - onto whatever switch the serving model's template reads.""" - spec = _spec_ctkw() - spec.thinking = "off" - req = types.SimpleNamespace() - assert sp_chat._merged_template_kwargs( - req, spec, "{% if enable_thinking %}...{% endif %}") == \ - {"enable_thinking": False} - assert sp_chat._merged_template_kwargs( - req, spec, "reasoning_effort in ['low','high','no_think']") == \ - {"reasoning_effort": "no_think"} - spec.thinking = "adaptive" - assert sp_chat._merged_template_kwargs( - req, spec, 'thinking_mode == "adaptive"') == \ - {"thinking_mode": "adaptive"} - spec.thinking = None - spec.reasoning_effort = "high" - assert sp_chat._merged_template_kwargs( - req, spec, 'set reasoning_effort = "medium"') == \ - {"reasoning_effort": "high"} - - -def test_merged_template_kwargs_request_kwargs_beat_spec_controls(): - """A request's explicit chat_template_kwargs pass through verbatim and win - over the profile's mapped controls.""" - spec = _spec_ctkw() - spec.thinking = "off" - spec.reasoning_effort = "low" - req = types.SimpleNamespace( - chat_template_kwargs={"enable_thinking": True}) - merged = sp_chat._merged_template_kwargs( - req, spec, "{% if enable_thinking %}{% endif %} reasoning_effort") - assert merged["enable_thinking"] is True - assert merged["reasoning_effort"] == "low" - - -def test_install_chat_template_kwargs_forwards_into_to_template_kwargs(): - """End-to-end seam: the gen-args wrapper stashes the merged dict and the - patched to_template_kwargs folds it into what mlx-vlm hands the template.""" - gen = importlib.import_module("mlx_vlm.server.generation") - - def stub(request, processor=None, tenant_id=None): - return gen.GenerationArguments() - - _APP._build_gen_args = stub - sp.install_gen_args_profile_injection() - sp.install_chat_template_kwargs() - fn = _APP._build_gen_args - assert getattr(fn, sp_chat._CTKW_FLAG, False) # stash carried on the chain - - spec = _spec_ctkw(preserve_thinking=True) - tok = serving.set_active_spec(spec) - try: - req = types.SimpleNamespace(model_fields_set=set(), - chat_template_kwargs={"foo": "bar"}) - args = _APP._build_gen_args(req) - finally: - serving.reset_active_spec(tok) - kw = args.to_template_kwargs() - assert kw["preserve_thinking"] is True # from the profile - assert kw["foo"] == "bar" # from the request - # enable_thinking was not explicit (request/spec/env) -> dropped from the - # template kwargs so the chat template's own default governs (b90aa60), - # while the args flag stays True for the generation path. - assert "enable_thinking" not in kw - assert args.enable_thinking is True - - # explicitly set on the request -> preserved verbatim - spec = _spec_ctkw() - tok = serving.set_active_spec(spec) - try: - req = types.SimpleNamespace(model_fields_set={"enable_thinking"}, - chat_template_kwargs=None, - enable_thinking=False) - args = _APP._build_gen_args(req) - finally: - serving.reset_active_spec(tok) - assert "enable_thinking" in args.to_template_kwargs() - - -def test_install_chat_template_kwargs_idempotent_and_noop_default(): - sp.install_chat_template_kwargs() - gen = importlib.import_module("mlx_vlm.server.generation") - first = gen.GenerationArguments.to_template_kwargs - sp.install_chat_template_kwargs() - assert gen.GenerationArguments.to_template_kwargs is first - # a request/spec with no kwargs leaves to_template_kwargs untouched (stock keys) - assert gen.GenerationArguments().to_template_kwargs() == { - "enable_thinking": gen.GenerationArguments().enable_thinking} - - -_KIMI_TAIL = ( - "{%- if thinking is defined and thinking is false -%}" - "{%- else -%}{%- endif -%}") - - -def test_request_thinking_off_maps_onto_kimi_bare_switch(): - """A plain --thinking value forwarded by the chat client (or any client's - `thinking: "off"`) must reach the template as the model's own switch - spelling - Kimi K2.x reads a bare `thinking` variable.""" - gen = importlib.import_module("mlx_vlm.server.generation") - proc = types.SimpleNamespace(chat_template=_KIMI_TAIL) - req = types.SimpleNamespace(model_fields_set=set(), - chat_template_kwargs=None, thinking="off") - out = sp_chat._stash_template_kwargs(gen.GenerationArguments(), req, proc) - assert out._kq_template_kwargs == {"thinking": False} - assert out.enable_thinking is False - assert out._kq_thinking_explicit is True - - req = types.SimpleNamespace(model_fields_set=set(), - chat_template_kwargs=None, thinking="on") - out = sp_chat._stash_template_kwargs(gen.GenerationArguments(), req, proc) - assert out._kq_template_kwargs == {"thinking": True} - assert out.enable_thinking is True - - -def test_request_reasoning_effort_field_maps_onto_template(): - gen = importlib.import_module("mlx_vlm.server.generation") - proc = types.SimpleNamespace(chat_template="reads reasoning_effort") - req = types.SimpleNamespace(model_fields_set=set(), - chat_template_kwargs=None, - reasoning_effort="low") - out = sp_chat._stash_template_kwargs(gen.GenerationArguments(), req, proc) - assert out._kq_template_kwargs == {"reasoning_effort": "low"} - assert out.enable_thinking is True # effort alone: not a switch - assert out._kq_thinking_explicit is False - - -class _RecordingTok: - """A tokenizer stand-in whose **kwargs signature makes mlx-vlm's - enable_thinking capability probe say yes (the 0.6.15 injection path).""" - chat_template = "{{ messages }}" - - def __init__(self): - self.kwargs = {} - - def apply_chat_template(self, messages, **kwargs): - self.kwargs = kwargs - return "rendered" - - -def test_template_default_thinking_blocks_0615_false_injection(): - """Regression: mlx-vlm >= 0.6.15 get_chat_template injects - enable_thinking=False when the kwarg is absent, so served models with - default-on reasoning templates rendered the dead think prefill and - stopped thinking. The seeded Jinja Undefined must reach the tokenizer - instead of False, and an explicit value must pass through verbatim.""" - import jinja2 - - import gmlx.tui.reasoning as reasoning - - sp.install_chat_template_kwargs() # installs the render guard too - pu = importlib.import_module("mlx_vlm.prompt_utils") - assert getattr(pu.get_chat_template, reasoning._TEMPLATE_DEFAULT_FLAG, False) - - tok = _RecordingTok() - msgs = [{"role": "user", "content": "hi"}] - assert pu.get_chat_template(tok, msgs, True) == "rendered" - assert isinstance(tok.kwargs["enable_thinking"], jinja2.Undefined) - - pu.get_chat_template(tok, msgs, True, enable_thinking=False) - assert tok.kwargs["enable_thinking"] is False - pu.get_chat_template(tok, msgs, True, enable_thinking=True) - assert tok.kwargs["enable_thinking"] is True - - # idempotent: a second install keeps the same wrapper - fn = pu.get_chat_template - reasoning.install_template_default_thinking() - assert pu.get_chat_template is fn - - -# The old-style role dispatch from issue #66: no developer alias, else-raise. -_ROLE_RAISE_TMPL = ("{% if m.role == 'system' %}{% elif m.role == 'user' %}" - "{% else %}{{ raise_exception('Unexpected message role.') }}" - "{% endif %}") - - -def test_normalize_developer_roles(): - msgs = [{"role": "developer", "content": "terse"}, - {"role": "user", "content": "hi"}, "not-a-dict"] - out = sp_chat._normalize_developer_roles(msgs, _ROLE_RAISE_TMPL) - assert out[0] == {"role": "system", "content": "terse"} - assert out[1]["role"] == "user" and out[2] == "not-a-dict" - assert msgs[0]["role"] == "developer" # input not mutated - # a template that handles developer gets the messages verbatim - assert sp_chat._normalize_developer_roles( - msgs, "role == 'developer'") is msgs - # nothing to rewrite -> same object - plain = [{"role": "user", "content": "hi"}] - assert sp_chat._normalize_developer_roles(plain, _ROLE_RAISE_TMPL) is plain - assert sp_chat._normalize_developer_roles("prompt", _ROLE_RAISE_TMPL) == \ - "prompt" - - -def test_install_role_normalization_rewrites_before_render(): - """Issue #66: a developer-role request against a template without the - alias must render as system instead of raising in the template.""" - class _Recorder: - chat_template = _ROLE_RAISE_TMPL - - def apply_chat_template(self, messages, **kwargs): - self.messages = messages - return "rendered" - - sp.install_role_normalization() - openai = importlib.import_module("mlx_vlm.server.openai") - assert getattr(openai.apply_chat_template, - sp_chat._ROLE_NORM_FLAG, False) - - tok = _Recorder() - out = openai.apply_chat_template( - tok, {"model_type": "gguf-llama"}, - [{"role": "developer", "content": "terse"}, - {"role": "user", "content": "hi"}]) - assert out == "rendered" - assert [m["role"] for m in tok.messages - if isinstance(m, dict) and "role" in m][:2] == ["system", "user"] - - # idempotent - fn = openai.apply_chat_template - sp.install_role_normalization() - assert openai.apply_chat_template is fn - - -def test_template_error_becomes_clean_400(): - """A raise_exception from the chat template answers 400 with the - template's message; template bugs (subclasses) stay 500, clean body.""" - import jinja2 - from fastapi.testclient import TestClient - - app = _APP.app - if not any(getattr(r, "path", None) == "/test/raise-template" - for r in app.router.routes): - @app.get("/test/raise-template") - async def _raise_template(): - raise jinja2.exceptions.TemplateError("Unexpected message role.") - - @app.get("/test/raise-template-bug") - async def _raise_template_bug(): - raise jinja2.exceptions.TemplateSyntaxError("bad", 1) - - sp.install_resolver_error_handlers() - client = TestClient(app, raise_server_exceptions=False) - r = client.get("/test/raise-template") - assert r.status_code == 400 - assert r.json()["error"] == { - "type": "invalid_request_error", - "message": "chat template rejected the conversation: " - "Unexpected message role."} - r2 = client.get("/test/raise-template-bug") - assert r2.status_code == 500 - assert r2.json()["error"]["type"] == "server_error" - assert "chat template failed to render" in r2.json()["error"]["message"] # 8. OpenAI stop sequences @@ -3053,108 +1661,6 @@ def test_sampling_float_rejects_non_numeric(): assert ei.value.status_code == 400 -# harmony (gpt-oss) serve-side split -_HARMONY_PROMPT = ("<|start|>system<|message|>You are helpful.<|end|>" - "<|start|>user<|message|>hi<|end|><|start|>assistant") -_HARMONY_REPLY = ('<|channel|>analysis<|message|>User greets; keep it short.' - "<|end|><|start|>assistant<|channel|>final<|message|>" - "Hello! How can I help?") - - -def test_nonstream_split_harmony_reply(): - app_mod = importlib.import_module("mlx_vlm.server.app") - sp.install_stream_thinking_seed() - split = app_mod._split_thinking_text - tok = sp_chat._LAST_RENDERED_PROMPT.set(None) - try: - r, c = split(_HARMONY_REPLY) - assert r == "User greets; keep it short." - assert c == "Hello! How can I help?" - assert "<|" not in c and "<|" not in r - # Length-capped inside analysis: all reasoning, empty content - # (the truncated-thinking convention). - r, c = split("<|channel|>analysis<|message|>Entry 55 reads 4") - assert r == "Entry 55 reads 4" - assert c == "" - # Gemma's lopsided spelling must not take the harmony branch. - r, c = split("<|channel>thought\nplan\nHi there.") - assert c and "<|channel|>" not in c - finally: - sp_chat._LAST_RENDERED_PROMPT.reset(tok) - - -def test_stream_harmony_filter_routes_channels(): - rs = importlib.import_module("mlx_vlm.server.responses_state") - cls = rs.ThinkingStreamState - original_init = cls.__init__ - original_feed = cls.feed - try: - sp.install_stream_thinking_seed() - tok = sp_chat._LAST_RENDERED_PROMPT.set(_HARMONY_PROMPT) - try: - st = cls(True) - assert getattr(st, "_kq_harmony", None) is not None - reasoning, content, closes = [], [], 0 - for i in range(0, len(_HARMONY_REPLY), 7): - d = st.feed(_HARMONY_REPLY[i:i + 7]) - if d.reasoning: - reasoning.append(d.reasoning) - if d.content: - content.append(d.content) - closes += bool(d.thinking_closed) - assert "".join(reasoning) == "User greets; keep it short." - assert "".join(content) == "Hello! How can I help?" - assert closes == 1 - # Non-harmony prompt: the stock state machine still drives. - sp_chat._LAST_RENDERED_PROMPT.set("<|im_start|>assistant\n\n") - st = cls(True) - assert getattr(st, "_kq_harmony", None) is None - d = st.feed("plananswer") - assert d.reasoning == "plan" and d.content == "answer" - finally: - sp_chat._LAST_RENDERED_PROMPT.reset(tok) - finally: - cls.__init__ = original_init - cls.feed = original_feed - - -def test_faithful_history_aliases_gpt_oss_thinking(): - from gmlx.serve.patches import _common as sp_common - from gmlx.serve.patches import render as sp_render - - def fake(processor, config, prompt, add_generation_prompt=True, - return_messages=False, num_images=0, num_audios=0, **kwargs): - return [dict(m) for m in prompt] - - # Exercise through the installer against a stub target module. - target = types.SimpleNamespace(apply_chat_template=fake) - orig_targets = sp_common._render_target_modules - sp_common._render_target_modules = lambda: [target] - try: - sp_render.install_faithful_history() - finally: - sp_common._render_target_modules = orig_targets - wrapped = target.apply_chat_template - assert wrapped is not fake - msgs = wrapped( - "processor", {"model_type": "gpt_oss"}, - [{"role": "assistant", "content": "Hi.", - "reasoning_content": "Short greeting."}], - return_messages=True) - assert msgs[0]["thinking"] == "Short greeting." - # Explicit thinking key wins; non-gpt-oss untouched. - msgs = wrapped( - "processor", {"model_type": "gpt_oss"}, - [{"role": "assistant", "content": "Hi.", "thinking": "keep", - "reasoning_content": "drop"}], - return_messages=True) - assert msgs[0]["thinking"] == "keep" - msgs = wrapped( - "processor", {"model_type": "qwen3"}, - [{"role": "assistant", "content": "Hi.", - "reasoning_content": "r"}], - return_messages=True) - assert "thinking" not in msgs[0] def test_unload_409_keeps_the_preload_hold(monkeypatch): diff --git a/tests/spec/test_full_prompt_prefill.py b/tests/spec/test_full_prompt_prefill.py index 632c5eca..03cc9144 100644 --- a/tests/spec/test_full_prompt_prefill.py +++ b/tests/spec/test_full_prompt_prefill.py @@ -20,7 +20,7 @@ import mlx.core as mx -pytestmark = [pytest.mark.integration, pytest.mark.slow] +pytestmark = pytest.mark.integration GREEDY = lambda x: mx.argmax(x, axis=-1) diff --git a/tests/test_eval_guard.py b/tests/test_eval_guard.py index db525e0c..e04041c5 100644 --- a/tests/test_eval_guard.py +++ b/tests/test_eval_guard.py @@ -11,12 +11,12 @@ from __future__ import annotations import ast -import os -import re from pathlib import Path import pytest +from helpers import _real_apple_gpu + import mlx.core as mx import gmlx.eval_guard as eg @@ -30,16 +30,6 @@ GMLX_DIR = Path(eg.__file__).resolve().parent -def _real_apple_gpu() -> bool: - if os.environ.get("KQUANT_FORCE_CPU"): - return False - try: - name = str(mx.device_info().get("device_name", "")) - except Exception: - return False - return bool(re.search(r"Apple M\d", name)) - - def _trip(): """Oversized allocation that throws at one malloc with nothing committed (same construction as tests/invariant_probes/_common.py: diff --git a/tests/test_pinned_invariants.py b/tests/test_pinned_invariants.py index 2c4fe37c..26d956c8 100644 --- a/tests/test_pinned_invariants.py +++ b/tests/test_pinned_invariants.py @@ -23,7 +23,6 @@ under KQUANT_FORCE_CPU (CI) and on paravirtual Metal devices. """ -import os import re import subprocess import sys @@ -31,6 +30,8 @@ import pytest +from helpers import _real_apple_gpu + PROBES = Path(__file__).parent / "invariant_probes" @@ -44,18 +45,6 @@ def _mlx_version(): return tuple(parts) -def _real_apple_gpu() -> bool: - if os.environ.get("KQUANT_FORCE_CPU"): - return False - try: - import mlx.core as mx - - name = str(mx.device_info().get("device_name", "")) - except Exception: - return False - return "Apple" in name and "Paravirtual" not in name - - pytestmark = pytest.mark.skipif( not _real_apple_gpu(), reason="pinned-invariant probes need a real Apple GPU",