diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index d8d5e3b..4b8b30d 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -89,6 +89,17 @@ jobs: # renaming a model_type - the weekly cron catches those) must not leave # the committed doc lying about what loads. run: python scripts/check-coverage.py --check --strict + - name: Seam drift check (pyright) + # Pyright infers upstream signatures from the installed mlx-vlm / + # mlx-lm source, so it needs this job's environment, not the + # deps-free lint job. One interpreter is enough: the checked + # surface does not vary by Python version. On the weekly cron a + # new upstream release can fail this step with no repo change; + # that is the drift signal, handled per docs/internals/upstream-upgrades.md. + if: matrix.python-version == '3.12' + run: | + pip install "pyright[nodejs]==1.1.414" + pyright - name: Run CPU logic tests env: # Hosted runners have no dependable Metal device; the conftest flips diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 36c7f35..1735680 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -48,8 +48,20 @@ tests, in ruff check . python scripts/check-docs.py # docs style and link check, also a CI step pre-commit install # optional, runs ruff on each commit +pip install "pyright[nodejs]==1.1.414" && pyright # seam drift check, needs the deps installed ``` +pyright is a guard against upstream symbol and signature drift, not a type +checker for the tree. It covers the files in `[tool.pyright].include`, all +of them at zero errors; the other files that import mlx_vlm or mlx_lm +statically are listed in `[tool.gmlx.pyright].backlog` and join the gate +one file at a time as each reaches zero. `tests/test_pyright_scope.py` +holds both lists to the tree, so a new static import lands in one list or +the other. A `# pyright: ignore[rule]` is allowed only where the false +positive comes from upstream typing (a stub that omits the attribute, a +scalar union, a field added to an upstream dataclass), with a comment +saying which; keep the count under twenty repo-wide. + ## Things to know before you patch - The serving stack is stock mlx-vlm with late-bound patches over its diff --git a/docs/internals/upstream-upgrades.md b/docs/internals/upstream-upgrades.md index 469181e..d24f25f 100644 --- a/docs/internals/upstream-upgrades.md +++ b/docs/internals/upstream-upgrades.md @@ -23,6 +23,7 @@ environment inside those bounds: | declared versions | `pyproject.toml` | the exact mlx-vlm pin and the mlx-lm and mlx-kquant floors described above | | seam contract | `tests/upstream/test_upstream_seams.py` | every patched symbol is pinned to a source fingerprint. Drift fails CI naming the seam | | runtime gate | `check_upstream_versions`, at CLI entry | mlx, mlx-lm or mlx-vlm below its floor refuses to run with an upgrade message, and newer than the qualified set warns once. `gmlx doctor` is exempt | +| static-import check | `pyright`, in the macOS CI job | the files in `[tool.pyright].include` are checked against the installed upstream source; a moved symbol or changed signature there fails CI | ## Watching upstream releases diff --git a/gmlx/cache/apc_manager.py b/gmlx/cache/apc_manager.py index 8cd4162..3779a66 100644 --- a/gmlx/cache/apc_manager.py +++ b/gmlx/cache/apc_manager.py @@ -81,7 +81,7 @@ class GmlxAPCManager(_apc.APCManager): override defers to the stock store instead. """ - def autosize(self, model, budget_fraction: float = None) -> None: + def autosize(self, model, budget_fraction: float | None = None) -> None: """Size the caches to the box post-load. Pool: raise the block cap to a working-budget share when APC_NUM_BLOCKS is unset (blocks allocate lazily, so the cap costs nothing until @@ -166,7 +166,8 @@ def entry_bytes(e): total -= sizes.pop(k, 0) # Mirror of pool_bytes for the exact tier: lets harnesses # separate budgeted, evictable retention from real residue. - self.stats.exact_bytes = int(total) + # gmlx side counter on the stock stats dataclass. + self.stats.exact_bytes = int(total) # pyright: ignore[reportAttributeAccessIssue] def stats_snapshot(self) -> dict: """Stock snapshot plus the gmlx ckpt-tier side counters (pure @@ -524,7 +525,7 @@ def _flush_pending(force=False): "APC disk save scheduling failed: %s", e) self.stats.pool_used = sum( 1 for x in self.pool if x.block_hash is not None) - self.stats.pool_bytes = int( + self.stats.pool_bytes = int( # pyright: ignore[reportAttributeAccessIssue] self.stats.pool_used * self.block_size * getattr(self, "_pool_per_token_bytes", 0)) return new_blocks diff --git a/gmlx/load/modules.py b/gmlx/load/modules.py index 79dcb2b..d31bcb6 100644 --- a/gmlx/load/modules.py +++ b/gmlx/load/modules.py @@ -1594,6 +1594,9 @@ class LoRAKQuantLinear(nn.Module): with nothing published the static ``scale`` applies to every row. """ + _kq_tables: dict + _kq_extra: list + def __init__(self, base: nn.Module, a: mx.array, b: mx.array, scale: float, slot: int = 0): super().__init__() diff --git a/gmlx/serve/bridge_vlm.py b/gmlx/serve/bridge_vlm.py index db5aa4e..d591b36 100644 --- a/gmlx/serve/bridge_vlm.py +++ b/gmlx/serve/bridge_vlm.py @@ -48,7 +48,9 @@ import logging import os import sys +from collections.abc import Callable from contextvars import ContextVar +from typing import Any from mlx_vlm import tokenizer_utils as _mlxvlm_tok from gmlx.models.vlm_text_only import Model as TextOnlyModel @@ -233,11 +235,11 @@ def __getattr__(self, name): return getattr(self._wrapper, name) -def _as_dict(config) -> dict: +def _as_dict(config: Any) -> dict: if isinstance(config, dict): return config for attr in ("to_dict", "__dict__"): - value = getattr(config, attr, None) + value: Callable[..., Any] | dict | None = getattr(config, attr, None) if callable(value): return dict(value()) if isinstance(value, dict): @@ -372,9 +374,10 @@ def _load_serveable_vlm( """ from gmlx.load.vlm import load_vlm_model - model, _config_dict, processor = load_vlm_model( + loaded = load_vlm_model( gguf_path, mmproj_path, hf_source=hf_source, verbose=False ) + model, processor = loaded[0], loaded[2] # Return the model's own dataclass config (what stock load_model_resources # returns as the 3rd element), not the synthesized dict. return model, processor, model.config @@ -387,7 +390,7 @@ def _make_text_processor(tokenizer) -> "_GgufServerProcessor": detokenizer. (The VLM path gets an engine-ready processor from the loader and does not use this.)""" backend = getattr(tokenizer, "_tokenizer", tokenizer) - eos = getattr(tokenizer, "eos_token_ids", None) or getattr( + eos: Any = getattr(tokenizer, "eos_token_ids", None) or getattr( tokenizer, "eos_token_id", None ) # StoppingCriteria.add_eos_token_ids() mutates this list in place, so it must @@ -662,7 +665,7 @@ def _reject_unwired(base_kind: str, *, streamable: bool = False, ) moe_experts = moe_expert_mass = None moe_miss_shed = moe_layer_shed = moe_prestage = None - _levers = dict( + _levers: dict[str, Any] = dict( moe_experts=moe_experts, moe_expert_mass=moe_expert_mass, moe_miss_shed=moe_miss_shed, moe_prestage=moe_prestage, moe_layer_shed=moe_layer_shed) @@ -761,7 +764,7 @@ def load_drafter(path_or_repo, kind=None, **kwargs): _apply_draft_block_size_override(result) return result - drafters.load_drafter = load_drafter + setattr(drafters, "load_drafter", load_drafter) setattr(drafters, _DRAFTER_PATCH_FLAG, True) @@ -825,7 +828,7 @@ class _DrafterSourceFilter(logging.Filter): def filter(self, record): if str(record.msg).startswith("Loading speculative drafter"): args = record.args or () - path = args[-1] if args else None + path = args[-1] if isinstance(args, tuple) and args else None if isinstance(path, str) and os.path.abspath(path) in _MTP_DRAFTER_STASH: return False return True @@ -951,7 +954,7 @@ def _bridge_load(model_path, adapter_path=None): # `moe_layer_shed:`/`moe_prestage:` / the paired serve flags) ride # along; None keeps the loader default / trained fan-out. stream = getattr(spec, "stream", None) - feeders = dict( + feeders: dict[str, Any] = dict( moe_experts=getattr(spec, "moe_experts", None), moe_expert_mass=getattr(spec, "moe_expert_mass", None), moe_miss_shed=getattr(spec, "moe_miss_shed", None), @@ -1053,7 +1056,7 @@ def load_model_resources(model_path, adapter_path=None): _raise_if_first_party_import(e) raise - generation.load_model_resources = load_model_resources + setattr(generation, "load_model_resources", load_model_resources) setattr(generation, _BRIDGE_FLAG, True) # generation.py logs on the parent "mlx_vlm.server" logger. engine_log = logging.getLogger("mlx_vlm.server") @@ -1077,7 +1080,7 @@ def load_model_resources(model_path, adapter_path=None): # finds it, and exposes the resolved spec for *this* request through a ContextVar # (mirroring residency's ``_active_entry`` discipline). -_RESOLVED_MODELS: dict[str, "object"] = {} # id -> ResolvedModel +_RESOLVED_MODELS: dict[str, Any] = {} # id -> ResolvedModel _PATH_TO_IDS: dict[str, list[str]] = {} # abspath -> [id, ...] _SERVER_CFG = None # the live ServerCfg (for re-resolve) # The ResolvedModel for the request in flight - set at the residency seam, read at diff --git a/gmlx/spec/dflash_drafter.py b/gmlx/spec/dflash_drafter.py index 0e4ce14..a7a6b95 100644 --- a/gmlx/spec/dflash_drafter.py +++ b/gmlx/spec/dflash_drafter.py @@ -47,7 +47,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, List, Optional +from typing import TYPE_CHECKING, Any, List, Optional import mlx.core as mx import mlx.nn as nn @@ -501,6 +501,9 @@ def make_cache(self, left_padding: Optional[List[int]] = None) -> list: caches = [] for layer_type in self.config.layer_types: if layer_type == "sliding_attention": + if self.config.sliding_window is None: + raise ValueError( + "sliding_attention layer without a sliding_window") # Temporal and slack-backed: the draft path reads cache.state # directly and indexes its rows as time. caches.append(BufferedRotatingKVCache( @@ -692,7 +695,15 @@ def draft_block( # --- target side -------------------------------------------------------------- -class DFlashCaptureHooks: +if TYPE_CHECKING: + # The hooks are mixed in ahead of this class; giving pyright the real + # base resolves the super() calls against upstream's method set. + from mlx_vlm.models.qwen3_5.language import LanguageModel as _CaptureBase +else: + _CaptureBase = object + + +class DFlashCaptureHooks(_CaptureBase): """Packed-hidden capture for owned qwen3.5 LanguageModels. While armed, every hidden the engine sees is ``[trunk | cap ...]``: the diff --git a/gmlx/spec/helpers.py b/gmlx/spec/helpers.py index 7399c0f..c23a621 100644 --- a/gmlx/spec/helpers.py +++ b/gmlx/spec/helpers.py @@ -48,13 +48,19 @@ def _generation_stream(): # --- draft/target sampler RNG coupling ------------------------------------- +def _rng_state() -> list[mx.array]: + # The mlx.core.random stub does not type ``state``. + return mx.random.state # pyright: ignore[reportReturnType] + + def _copy_rng_state() -> list[mx.array]: - return [mx.array(state) for state in mx.random.state] + return [mx.array(state) for state in _rng_state()] def _restore_rng_state(state: list[mx.array]) -> None: + live = _rng_state() for i, value in enumerate(state): - mx.random.state[i] = value + live[i] = value def _append_arrays(value: Any, arrays: list[mx.array]) -> None: @@ -91,8 +97,10 @@ class _SpeculativeSamplerRNG: def __init__(self, draft_model: nn.Module, *, enabled: bool): self.draft_model = draft_model self.enabled = bool(enabled) - self._target_rng_state = _copy_rng_state() if self.enabled else None - self._draft_rng_state = _copy_rng_state() if self.enabled else None + self._target_rng_state: list[mx.array] = ( + _copy_rng_state() if self.enabled else []) + self._draft_rng_state: list[mx.array] = ( + _copy_rng_state() if self.enabled else []) def draft_call(self, fn: Callable, *args, **kwargs): if not self.enabled: @@ -109,7 +117,7 @@ def draft_call(self, fn: Callable, *args, **kwargs): result = fn(*args, **kwargs) arrays = _draft_sampler_state_arrays(self.draft_model) - arrays.extend(mx.random.state) + arrays.extend(_rng_state()) if arrays: mx.async_eval(*arrays) @@ -133,7 +141,7 @@ def draft_tokens(self, fn: Callable, *args, **kwargs): arrays = [] _append_arrays(result, arrays) arrays.extend(_draft_sampler_state_arrays(self.draft_model)) - arrays.extend(mx.random.state) + arrays.extend(_rng_state()) if arrays: mx.async_eval(*arrays) @@ -239,10 +247,11 @@ def _mtp_next_block_size( return min(budget, native) if getattr(draft_model, "prefer_requested_block_size", False): return budget + accept_lens: Any = draft_model.accept_lens return _effective_mtp_block_size( requested_block_total, configured_block_total, - draft_model.accept_lens, + accept_lens, remaining_budget, ) @@ -291,7 +300,7 @@ def _mtp_cache_offset(prompt_cache: list[Any]) -> Any: def _mtp_cache_offset_max(prompt_cache: list[Any]) -> int: offset = _mtp_cache_offset(prompt_cache) - return int(offset.max().item()) if isinstance(offset, mx.array) else int(offset) + return int(offset.max()) if isinstance(offset, mx.array) else int(offset) def _mtp_draft_position(kv_valid_len: Any) -> Any: @@ -350,7 +359,7 @@ class _MTPVerifyResult: def _mtp_draft_hidden(lm: nn.Module, hidden: mx.array) -> mx.array: - prepare = getattr(lm, "speculative_draft_hidden", None) + prepare: Callable[..., Any] | None = getattr(lm, "speculative_draft_hidden", None) return prepare(hidden) if callable(prepare) else hidden @@ -399,7 +408,8 @@ def _mtp_verify_without_logits( verify_input: mx.array, prompt_cache: list[Any], ) -> _MTPVerifyResult | None: - verify_hidden = getattr(lm, "speculative_verify_hidden", None) + verify_hidden: Callable[..., Any] | None = getattr( + lm, "speculative_verify_hidden", None) if callable(verify_hidden): _note_verify_branch("hook:speculative_verify_hidden", lm) result = verify_hidden(verify_input, prompt_cache) @@ -481,7 +491,8 @@ def _mtp_verify_target( sample_target_tokens: bool = True, ) -> _MTPVerifyResult: if sample_target_tokens: - argmax_from_hidden = getattr(lm, "speculative_argmax_from_hidden", None) + argmax_from_hidden: Callable[..., Any] | None = getattr( + lm, "speculative_argmax_from_hidden", None) if callable(argmax_from_hidden): result = _mtp_verify_without_logits(lm, verify_input, prompt_cache) if result is not None: diff --git a/gmlx/spec/mtp_drafter.py b/gmlx/spec/mtp_drafter.py index 776ad72..b12ad9f 100644 --- a/gmlx/spec/mtp_drafter.py +++ b/gmlx/spec/mtp_drafter.py @@ -83,7 +83,7 @@ def _cache_offset(caches) -> int: return 0 off = getattr(caches[0], "offset", 0) if isinstance(off, mx.array): - return int(off.max().item()) if off.size else 0 + return int(off.max()) if off.size else 0 return int(off) @@ -459,7 +459,7 @@ def accept_verified_tokens_batch( for cache in self._cache: cache.trim(self._round_appended) - draft_rows = draft_tokens.tolist() + draft_rows: Any = draft_tokens.tolist() row_tokens: list[list[int]] = [] row_hiddens: list[list[mx.array]] = [] for row, accepted_i in enumerate(accepted): diff --git a/pyproject.toml b/pyproject.toml index 5f8c4f5..4025c82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -159,6 +159,113 @@ markers = [ "needs_kvarn_row_ends: dispatches the decode kernels with per-row ends; skipped unless the installed mlx-kquant is 0.4.9 or later", ] +[tool.pyright] +# Structural drift guard at the static upstream import sites: pyright +# infers signatures from the installed mlx-vlm / mlx-lm source, so a +# moved symbol, changed signature or class-identity move fails here +# before it fails at runtime. Scope is exactly the files that import +# mlx_vlm / mlx_lm statically, minus gmlx/models (mirrors, annotation +# exempt); tests/test_pyright_scope.py keeps the list in step with the +# tree. Seams reached through importlib.import_module are opaque to +# pyright and stay covered by gmlx/upstream/seams.py. Runs in the macOS +# CI job, not the deps-free lint job: without the libraries installed +# there is nothing to infer from. +include = [ + "gmlx/cache/apc_manager.py", + "gmlx/load/modules.py", + "gmlx/serve/bridge_lm.py", + "gmlx/serve/bridge_vlm.py", + "gmlx/spec/dflash_drafter.py", + "gmlx/spec/helpers.py", + "gmlx/spec/mtp_drafter.py", +] +typeCheckingMode = "standard" +pythonVersion = "3.11" +pythonPlatform = "Darwin" +reportMissingTypeStubs = false +reportMissingModuleSource = false +# Deep-importing upstream privates is the design; the contract is the +# seam registry, not the upstream __all__. +reportPrivateImportUsage = false +# mlx ships py.typed with an unannotated nn.Module.__getattr__ whose else +# branch never returns, so every undeclared attribute of any Module infers +# as "Any | None". These rules then fire on ordinary submodule access and +# carry no drift signal (a moved symbol is an attribute-access, import or +# call error, never an Optional one). +reportOptionalMemberAccess = "none" +reportOptionalSubscript = "none" +reportOptionalCall = "none" +reportOptionalIterable = "none" +reportOptionalOperand = "none" + +[tool.gmlx.pyright] +# Static mlx_vlm / mlx_lm importers not yet in the pyright gate. A file +# moves from here to [tool.pyright].include once it checks clean; the +# list only shrinks. tests/test_pyright_scope.py holds both lists to +# the tree. +backlog = [ + "gmlx/cache/apc_pooling.py", + "gmlx/cache/fresh_gate.py", + "gmlx/cache/kv_policy.py", + "gmlx/cache/kvarn_cache.py", + "gmlx/cache/kvarn_serve.py", + "gmlx/cache/retire_key.py", + "gmlx/cache/snapshot.py", + "gmlx/commands/cli.py", + "gmlx/commands/train.py", + "gmlx/eval_guard.py", + "gmlx/gen/benchmarks.py", + "gmlx/gen/diffusion.py", + "gmlx/gen/generation.py", + "gmlx/gen/media_spans.py", + "gmlx/gen/prefill_decay.py", + "gmlx/gen/thinking_budget.py", + "gmlx/load/loader.py", + "gmlx/load/mtp_target.py", + "gmlx/load/vlm.py", + "gmlx/lora_rows.py", + "gmlx/serve/admit_gate.py", + "gmlx/serve/batch_sched.py", + "gmlx/serve/capacity.py", + "gmlx/serve/cb_phase.py", + "gmlx/serve/decode_batch.py", + "gmlx/serve/estimate.py", + "gmlx/serve/governor.py", + "gmlx/serve/kv_policy.py", + "gmlx/serve/live_requests.py", + "gmlx/serve/mem_preflight.py", + "gmlx/serve/memtrace.py", + "gmlx/serve/patches/chat_behavior.py", + "gmlx/serve/patches/mtp_thinking.py", + "gmlx/serve/patches/row_failed.py", + "gmlx/serve/patches/sampling.py", + "gmlx/serve/queue_cap.py", + "gmlx/serve/residency.py", + "gmlx/serve/seed_rows.py", + "gmlx/serve/step_timing.py", + "gmlx/serve/tick_guard.py", + "gmlx/spec/admission.py", + "gmlx/spec/ckpt.py", + "gmlx/spec/engine.py", + "gmlx/spec/kv_quant.py", + "gmlx/spec/mtp_load.py", + "gmlx/spec/mtp_prefill.py", + "gmlx/spec/ragged_decode.py", + "gmlx/tui/chat.py", + "gmlx/upstream/cascade_sdpa.py", + "gmlx/upstream/dsv32_patches.py", + "gmlx/upstream/gdn_patches.py", + "gmlx/upstream/occupancy_fuse.py", + "gmlx/upstream/qkv_fuse.py", + "gmlx/upstream/quantized_cache_pack_fix.py", + "gmlx/upstream/quantized_sdpa_fix.py", + "gmlx/upstream/rotating_cache_fix.py", + "gmlx/upstream/seams.py", + "gmlx/upstream/softcap_f32.py", + "scripts/check-coverage.py", + "scripts/kld_harness.py", +] + [tool.ruff] # Vendored third-party code is linted upstream, not here. extend-exclude = ["gmlx/_vendor"] diff --git a/tests/spec/test_full_prompt_prefill.py b/tests/spec/test_full_prompt_prefill.py index 03cc914..30d67ab 100644 --- a/tests/spec/test_full_prompt_prefill.py +++ b/tests/spec/test_full_prompt_prefill.py @@ -1449,7 +1449,7 @@ def test_apc_hit_on_injected_request(mtp_model): apc_hit_seen = False # Capture APC log to verify the hit actually fired. - apc_log = logging.getLogger("gmlx.spec.ckpt") + apc_log = logging.getLogger("gmlx.spec.mtp_prefill") log_messages = [] handler = logging.Handler() handler.emit = lambda record: log_messages.append(record.getMessage()) diff --git a/tests/test_pyright_scope.py b/tests/test_pyright_scope.py new file mode 100644 index 0000000..73d8ae5 --- /dev/null +++ b/tests/test_pyright_scope.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""The pyright gate ([tool.pyright].include) and its backlog +([tool.gmlx.pyright].backlog) together name every file that imports +mlx_vlm or mlx_lm statically, minus gmlx/models (upstream mirrors, +annotation exempt). A new static import anywhere lands in one list or the +other; a file leaves the tree, it leaves both. Seams reached through +importlib.import_module are opaque to pyright and are not the gate's +business. +""" +from __future__ import annotations + +import re +import subprocess +import tomllib +from pathlib import Path + +import pytest + +_ROOT = Path(__file__).resolve().parents[1] +_STATIC = re.compile( + r"^\s*(from (mlx_vlm|mlx_lm)(\.[\w.]+)? import\b" + r"|import (mlx_vlm|mlx_lm)(\.[\w.]+)?( as \w+)?\s*$)", + re.MULTILINE, +) + + +def _static_importers() -> set[str]: + tracked = subprocess.run( + ["git", "ls-files", "gmlx/*.py", "scripts/*.py"], + cwd=_ROOT, capture_output=True, text=True, check=True, + ).stdout.split() + out = set() + for rel in tracked: + if rel.startswith(("gmlx/_vendor/", "gmlx/models/")): + continue + if _STATIC.search((_ROOT / rel).read_text()): + out.add(rel) + return out + + +@pytest.fixture(scope="module") +def lists() -> tuple[set[str], set[str], set[str]]: + with open(_ROOT / "pyproject.toml", "rb") as f: + tool = tomllib.load(f)["tool"] + gated = set(tool["pyright"]["include"]) + backlog = set(tool["gmlx"]["pyright"]["backlog"]) + return gated, backlog, _static_importers() + + +def test_gate_and_backlog_are_disjoint(lists): + gated, backlog, _ = lists + assert not gated & backlog, sorted(gated & backlog) + + +def test_ceiling_every_listed_file_is_a_static_importer(lists): + gated, backlog, importers = lists + stale = (gated | backlog) - importers + assert not stale, f"no longer a static upstream importer: {sorted(stale)}" + + +def test_floor_every_static_importer_is_listed(lists): + gated, backlog, importers = lists + missing = importers - gated - backlog + assert not missing, ( + f"static upstream importer in neither [tool.pyright].include nor " + f"[tool.gmlx.pyright].backlog: {sorted(missing)}")