From 5240166e0eec7c19a3ce2727e3af590750b53f30 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:06:35 -0700 Subject: [PATCH 1/5] fix(serve): drafter GGUF reload no longer double-counts headroom weights --- CHANGELOG.md | 3 +++ gmlx/loader.py | 5 ++++- gmlx/prefill_decay.py | 27 ++++++++++++++++++++------- tests/test_prefill_decay.py | 24 ++++++++++++++++++++++-- 4 files changed, 49 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 01cbbb1..8b6a685 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -63,6 +63,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Fixed +- serve --speculative no longer serializes concurrent requests: the + drafter's reload of the target GGUF counted the shared weights twice, + holding admission headroom negative for the server's lifetime. - Chat's bottom toolbar no longer clips the live tok/s readout on narrow terminals; sampling knobs are dropped first instead. - The serve free-headroom estimate went negative on models whose load diff --git a/gmlx/loader.py b/gmlx/loader.py index 6646add..d6cd70c 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -2622,7 +2622,10 @@ def _warm_mmap_residency( tracked = max(0.0, mx.get_active_memory() - active_before) except Exception: tracked = 0.0 - note_untracked_weights(max(0.0, total - min(tracked, total))) + # Keyed by shard paths so a drafter reloading the target's GGUF + # replaces its earlier registration instead of double-counting. + key = tuple(os.path.abspath(p) for p in paths) if paths else None + note_untracked_weights(max(0.0, total - min(tracked, total)), key=key) def _warm_touch_pass( diff --git a/gmlx/prefill_decay.py b/gmlx/prefill_decay.py index e697a27..56f1646 100644 --- a/gmlx/prefill_decay.py +++ b/gmlx/prefill_decay.py @@ -338,16 +338,29 @@ def _tick_step(base: int) -> int | None: return step -_UNTRACKED_WEIGHTS = 0.0 +# Untracked weight bytes keyed by source (GGUF shard paths). Reloading the +# same file maps the same pages, so same-key registrations replace (max) +# rather than accumulate; a drafter reload used to double the count and +# hold headroom negative, serializing admissions. +_UNTRACKED_WEIGHTS: dict[object, float] = {} +_UNTRACKED_ANON = "\0anon" _HEADROOM_FRACTION = 0.5 -def note_untracked_weights(nbytes: float) -> None: +def note_untracked_weights(nbytes: float, key: object = None) -> None: """Loader hook: bytes wired at inference but invisible to - mx.get_active_memory (zero-copy mmap weights). Accumulates across loads - (target model + drafter).""" - global _UNTRACKED_WEIGHTS - _UNTRACKED_WEIGHTS += float(nbytes) + mx.get_active_memory (zero-copy mmap weights). Same-key registrations + keep the max, distinct keys sum; key=None accumulates.""" + if key is None: + _UNTRACKED_WEIGHTS[_UNTRACKED_ANON] = ( + _UNTRACKED_WEIGHTS.get(_UNTRACKED_ANON, 0.0) + float(nbytes)) + else: + _UNTRACKED_WEIGHTS[key] = max( + _UNTRACKED_WEIGHTS.get(key, 0.0), float(nbytes)) + + +def untracked_weight_bytes() -> float: + return sum(_UNTRACKED_WEIGHTS.values()) def headroom_bytes() -> float | None: @@ -361,7 +374,7 @@ def headroom_bytes() -> float | None: active = float(mx.get_active_memory()) except Exception: return None - return ws - _UNTRACKED_WEIGHTS - active + return ws - untracked_weight_bytes() - active _headroom_bytes = headroom_bytes diff --git a/tests/test_prefill_decay.py b/tests/test_prefill_decay.py index 5a357a2..8ba360b 100644 --- a/tests/test_prefill_decay.py +++ b/tests/test_prefill_decay.py @@ -168,10 +168,30 @@ def test_seed_step_honors_kill_switch(monkeypatch): def test_note_untracked_weights_accumulates(monkeypatch): - monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", 0.0) + monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) pd.note_untracked_weights(10 * GB) pd.note_untracked_weights(5 * GB) - assert pd._UNTRACKED_WEIGHTS == 15 * GB + assert pd.untracked_weight_bytes() == 15 * GB + + +def test_note_untracked_weights_same_key_replaces(monkeypatch): + # a drafter reloading the target's GGUF maps the same pages: the + # re-registration must not double the count (it serialized admissions) + monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) + key = ("/models/a.gguf",) + pd.note_untracked_weights(87 * GB, key=key) + pd.note_untracked_weights(87 * GB, key=key) + assert pd.untracked_weight_bytes() == 87 * GB + pd.note_untracked_weights(2 * GB, key=key) # subset re-walk keeps max + assert pd.untracked_weight_bytes() == 87 * GB + + +def test_note_untracked_weights_distinct_keys_sum(monkeypatch): + monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) + pd.note_untracked_weights(10 * GB, key=("/models/a.gguf",)) + pd.note_untracked_weights(4 * GB, key=("/models/b.gguf",)) + pd.note_untracked_weights(1 * GB) + assert pd.untracked_weight_bytes() == 15 * GB # --- arch score profiles ------------------------------------------------- From d9de89b3f6e831b24899d905ff340e57d5153424 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:32:44 -0700 Subject: [PATCH 2/5] fix(serve): key MTP and VLM install routes by weight source --- gmlx/loader.py | 15 ++++++++++++--- gmlx/mtp_load.py | 11 ++++++++++- gmlx/vlm.py | 4 +++- 3 files changed, 25 insertions(+), 5 deletions(-) diff --git a/gmlx/loader.py b/gmlx/loader.py index d6cd70c..94d0b3a 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -2576,10 +2576,17 @@ def _active_now() -> float | None: return None +def weights_source_key(*paths: str) -> tuple | None: + """Identity of a load's weight bytes (absolute file paths) for the + untracked-headroom registry: reloads of the same file replace their + earlier registration instead of double-counting.""" + return tuple(os.path.abspath(p) for p in paths) or None + + def _warm_mmap_residency( model, *, log=print, paths: list[str] | None = None, batch_bytes: int = 4 << 30, threshold_bytes: int | None = None, - active_before: float | None = None, + active_before: float | None = None, source_key: tuple | None = None, ) -> None: """Pre-wire GPU residency of mmap-backed weights in small batches. @@ -2624,7 +2631,7 @@ def _warm_mmap_residency( tracked = 0.0 # Keyed by shard paths so a drafter reloading the target's GGUF # replaces its earlier registration instead of double-counting. - key = tuple(os.path.abspath(p) for p in paths) if paths else None + key = source_key or (weights_source_key(*paths) if paths else None) note_untracked_weights(max(0.0, total - min(tracked, total)), key=key) @@ -2782,6 +2789,7 @@ def _install_and_load( sanitize: bool = True, no_alias: set[str] | None = None, fp32_keep: tuple[str, ...] = (), + source_key: tuple | None = None, ) -> None: """Sanitize -> de-interleave native-fp -> swap kquant leaves -> cast -> load. @@ -2914,7 +2922,8 @@ def _install_and_load( model.load_weights(list(loadable.items()), strict=False) log(f"[load_weights] loaded {len(loadable)} / {len(model_params)} model parameters") - _warm_mmap_residency(model, log=log, active_before=active_before) + _warm_mmap_residency(model, log=log, active_before=active_before, + source_key=source_key) missing = sorted(model_params - set(loadable.keys())) if missing: diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index 3ee531f..f93e848 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -37,6 +37,7 @@ remap_arrays, remap_gemma4_assistant_arrays, remap_mtp_arrays, + weights_source_key, ) from .native_fp import _strip_weight from .populate import maybe_populate_for_load @@ -136,6 +137,7 @@ def _load_mtp_drafter( n_head: int | None = None, n_head_kv: int | None = None, log=loadlog.verbose_print, + source_key: tuple | None = None, ): """Build + load + bind the native-head MTP drafter (seam 4). @@ -213,6 +215,7 @@ def _load_mtp_drafter( log=log, sanitize=False, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE.get(model_type, ()), + source_key=source_key, ) drafter.bind(target) from .drafter_protocol import validate_drafter @@ -325,7 +328,8 @@ def _load_gemma4_assistant_drafter( d_weights, d_meta, d_stats = remap_gemma4_assistant_arrays(arrays, kquant_meta) log(f"[mtp] drafter remap: {d_stats}") - _install_and_load(drafter, d_weights, d_meta, log=log, sanitize=False) + _install_and_load(drafter, d_weights, d_meta, log=log, sanitize=False, + source_key=weights_source_key(draft_gguf_path)) # Ordered-embeddings drafters (E2B/E4B) route the LM head through a # MaskedEmbedder that reads embed_tokens.weight as a [vocab, hidden] float # matrix (gathers candidate rows then a dense matmul). A kquant wire-byte @@ -531,6 +535,7 @@ def _load_deepseek4_mtp_drafter( log=log, sanitize=False, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE["deepseek_v4"], + source_key=weights_source_key(draft_gguf_path), ) drafter.bind(target) @@ -846,6 +851,7 @@ def _load_deepseek4_dspark_drafter( sanitize=False, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE["deepseek_v4"] + ("confidence_proj.",), + source_key=weights_source_key(draft_gguf_path), ) finally: if force_wire: @@ -996,6 +1002,7 @@ def load_mtp_model( sanitize=(_mt in ("deepseek_v4", "hy_v3")), no_alias=owned_names, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE.get(_mt, ()), + source_key=weights_source_key(*pf.shards), ) # 2b. fused gated-delta verify kernel. The multi-position verify forward is the @@ -1056,6 +1063,7 @@ def load_mtp_model( model, n_head=n_head, n_head_kv=n_head_kv, + source_key=weights_source_key(*pf.shards), log=_log, ) @@ -1233,6 +1241,7 @@ def load_vlm_mtp_model( model, n_head=n_head, n_head_kv=n_head_kv, + source_key=weights_source_key(*pf.shards), log=_log, ) diff --git a/gmlx/vlm.py b/gmlx/vlm.py index c9f5fe8..06a448d 100644 --- a/gmlx/vlm.py +++ b/gmlx/vlm.py @@ -38,6 +38,7 @@ load_gguf_wire_bytes, materialize_module_arrays, remap_arrays, + weights_source_key, ) from .preflight import preflight from .transforms import coalesce_split_experts @@ -1735,7 +1736,8 @@ def load_vlm_model( # codec and stay native. The remap already produced final mlx-vlm names # (text under [thinker.]language_model.model.*, vision/audio under their # towers), so model.sanitize must not run - it would re-prefix text keys. - _install_and_load(model, hf_weights, hf_kquant_meta, log=_log, sanitize=False) + _install_and_load(model, hf_weights, hf_kquant_meta, log=_log, sanitize=False, + source_key=weights_source_key(*pf.shards, mmproj_path)) materialize_module_arrays(model) # 5. processor (image preprocessing + tokenizer + chat template). Synthesized From 01ccf07355f92202887a68cbd8892664796cb114 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:48:14 -0700 Subject: [PATCH 3/5] fix(serve): bracket untracked-weights baseline around the wire read --- gmlx/loader.py | 11 +++++++++-- gmlx/mtp_load.py | 12 +++++++++++- gmlx/prefill_decay.py | 12 ++++++------ gmlx/vlm.py | 5 ++++- tests/test_prefill_decay.py | 15 +++++++++++++-- 5 files changed, 43 insertions(+), 12 deletions(-) diff --git a/gmlx/loader.py b/gmlx/loader.py index 94d0b3a..bec74d9 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -2630,7 +2630,7 @@ def _warm_mmap_residency( except Exception: tracked = 0.0 # Keyed by shard paths so a drafter reloading the target's GGUF - # replaces its earlier registration instead of double-counting. + # cannot register the same pages twice (first registration wins). key = source_key or (weights_source_key(*paths) if paths else None) note_untracked_weights(max(0.0, total - min(tracked, total)), key=key) @@ -2790,6 +2790,7 @@ def _install_and_load( no_alias: set[str] | None = None, fp32_keep: tuple[str, ...] = (), source_key: tuple | None = None, + active_before: float | None = None, ) -> None: """Sanitize -> de-interleave native-fp -> swap kquant leaves -> cast -> load. @@ -2809,9 +2810,15 @@ def _install_and_load( ``fp32_keep``: target-name substrings pinned to float32 through the bf16 cast (see ``_FP32_KEEP_BY_MODEL_TYPE``). + + ``active_before``: active-memory baseline for the untracked-weights split. + Callers that read wire bytes before installing must pass the pre-read + value; wire reads can grow active memory, and a post-read baseline makes + those tracked bytes register as untracked on top of it. """ loadlog.stage("loading weights") - active_before = _active_now() + if active_before is None: + active_before = _active_now() # 5. sanitize first - model.sanitize may rename keys; rebuild meta. if sanitize and hasattr(model, "sanitize"): hf_weights = model.sanitize(hf_weights) diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index f93e848..4e4ece9 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -27,6 +27,7 @@ from .gguf_meta import first_nonzero_int, read_int from .loader import ( _FP32_KEEP_BY_MODEL_TYPE, + _active_now, _install_and_load, _resolve_chat_template, build_model, @@ -298,6 +299,7 @@ def _load_gemma4_assistant_drafter( """ import importlib + active_before = _active_now() arrays, kquant_meta, d_arch, meta, tensor_shapes = load_gguf_wire_bytes( draft_gguf_path, zero_copy=zero_copy ) @@ -329,7 +331,8 @@ def _load_gemma4_assistant_drafter( log(f"[mtp] drafter remap: {d_stats}") _install_and_load(drafter, d_weights, d_meta, log=log, sanitize=False, - source_key=weights_source_key(draft_gguf_path)) + source_key=weights_source_key(draft_gguf_path), + active_before=active_before) # Ordered-embeddings drafters (E2B/E4B) route the LM head through a # MaskedEmbedder that reads embed_tokens.weight as a [vocab, hidden] float # matrix (gathers candidate rows then a dense matmul). A kquant wire-byte @@ -478,6 +481,7 @@ def _load_deepseek4_mtp_drafter( loader shape; the block config is the target's config with ``compress_ratios`` post-init extended by the MTP layer's ratio 0 (``ModelArgs.__post_init__`` truncates to num_hidden_layers).""" + active_before = _active_now() arrays, kquant_meta, d_arch, _meta, _shapes = load_gguf_wire_bytes( draft_gguf_path, zero_copy=zero_copy ) @@ -494,6 +498,7 @@ def _load_deepseek4_mtp_drafter( arrays=arrays, kquant_meta=kquant_meta, meta=_meta, + active_before=active_before, log=log, ) if d_arch != "deepseek4_mtp_support": @@ -536,6 +541,7 @@ def _load_deepseek4_mtp_drafter( sanitize=False, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE["deepseek_v4"], source_key=weights_source_key(draft_gguf_path), + active_before=active_before, ) drafter.bind(target) @@ -761,6 +767,7 @@ def _load_deepseek4_dspark_drafter( arrays: dict, kquant_meta: dict, meta: dict, + active_before: float | None = None, log=loadlog.verbose_print, ): """Build + load + bind the DSpark drafter from its companion GGUF (arch @@ -852,6 +859,7 @@ def _load_deepseek4_dspark_drafter( fp32_keep=_FP32_KEEP_BY_MODEL_TYPE["deepseek_v4"] + ("confidence_proj.",), source_key=weights_source_key(draft_gguf_path), + active_before=active_before, ) finally: if force_wire: @@ -920,6 +928,7 @@ def load_mtp_model( maybe_populate_for_load(pf.shards, log=_log) loadlog.stage("reading tensors") + active_before = _active_now() arrays, kquant_meta, _arch_meta, meta, tensor_shapes = load_gguf_wire_bytes( gguf_path, zero_copy=zero_copy, shards=pf.shards ) @@ -1003,6 +1012,7 @@ def load_mtp_model( no_alias=owned_names, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE.get(_mt, ()), source_key=weights_source_key(*pf.shards), + active_before=active_before, ) # 2b. fused gated-delta verify kernel. The multi-position verify forward is the diff --git a/gmlx/prefill_decay.py b/gmlx/prefill_decay.py index 56f1646..5767b30 100644 --- a/gmlx/prefill_decay.py +++ b/gmlx/prefill_decay.py @@ -339,9 +339,9 @@ def _tick_step(base: int) -> int | None: # Untracked weight bytes keyed by source (GGUF shard paths). Reloading the -# same file maps the same pages, so same-key registrations replace (max) -# rather than accumulate; a drafter reload used to double the count and -# hold headroom negative, serializing admissions. +# same file maps the same pages, so only the first registration per key +# counts; a drafter reload used to double the count and hold headroom +# negative, serializing admissions. _UNTRACKED_WEIGHTS: dict[object, float] = {} _UNTRACKED_ANON = "\0anon" _HEADROOM_FRACTION = 0.5 @@ -350,13 +350,13 @@ def _tick_step(base: int) -> int | None: def note_untracked_weights(nbytes: float, key: object = None) -> None: """Loader hook: bytes wired at inference but invisible to mx.get_active_memory (zero-copy mmap weights). Same-key registrations - keep the max, distinct keys sum; key=None accumulates.""" + keep the first (the walk that read the file; re-walks see pages already + counted), distinct keys sum; key=None accumulates.""" if key is None: _UNTRACKED_WEIGHTS[_UNTRACKED_ANON] = ( _UNTRACKED_WEIGHTS.get(_UNTRACKED_ANON, 0.0) + float(nbytes)) else: - _UNTRACKED_WEIGHTS[key] = max( - _UNTRACKED_WEIGHTS.get(key, 0.0), float(nbytes)) + _UNTRACKED_WEIGHTS.setdefault(key, float(nbytes)) def untracked_weight_bytes() -> float: diff --git a/gmlx/vlm.py b/gmlx/vlm.py index 06a448d..41338f1 100644 --- a/gmlx/vlm.py +++ b/gmlx/vlm.py @@ -34,6 +34,7 @@ ) from .gguf_meta import first_nonzero_int, read_int from .loader import ( + _active_now, _install_and_load, load_gguf_wire_bytes, materialize_module_arrays, @@ -1655,6 +1656,7 @@ def load_vlm_model( loadlog.fact_file_size(pf.shards) _log(f"[vlm] llm arch={llm_arch}") loadlog.stage("reading tensors") + active_before = _active_now() arrays, kquant_meta, _arch, llm_meta, llm_shapes = load_gguf_wire_bytes( gguf_path, zero_copy=zero_copy, shards=pf.shards) arrays, kquant_meta, n_coalesced = coalesce_split_experts(arrays, kquant_meta) @@ -1737,7 +1739,8 @@ def load_vlm_model( # (text under [thinker.]language_model.model.*, vision/audio under their # towers), so model.sanitize must not run - it would re-prefix text keys. _install_and_load(model, hf_weights, hf_kquant_meta, log=_log, sanitize=False, - source_key=weights_source_key(*pf.shards, mmproj_path)) + source_key=weights_source_key(*pf.shards, mmproj_path), + active_before=active_before) materialize_module_arrays(model) # 5. processor (image preprocessing + tokenizer + chat template). Synthesized diff --git a/tests/test_prefill_decay.py b/tests/test_prefill_decay.py index 8ba360b..251bf9d 100644 --- a/tests/test_prefill_decay.py +++ b/tests/test_prefill_decay.py @@ -174,7 +174,7 @@ def test_note_untracked_weights_accumulates(monkeypatch): assert pd.untracked_weight_bytes() == 15 * GB -def test_note_untracked_weights_same_key_replaces(monkeypatch): +def test_note_untracked_weights_same_key_first_wins(monkeypatch): # a drafter reloading the target's GGUF maps the same pages: the # re-registration must not double the count (it serialized admissions) monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) @@ -182,10 +182,21 @@ def test_note_untracked_weights_same_key_replaces(monkeypatch): pd.note_untracked_weights(87 * GB, key=key) pd.note_untracked_weights(87 * GB, key=key) assert pd.untracked_weight_bytes() == 87 * GB - pd.note_untracked_weights(2 * GB, key=key) # subset re-walk keeps max + pd.note_untracked_weights(2 * GB, key=key) # subset re-walk ignored assert pd.untracked_weight_bytes() == 87 * GB +def test_note_untracked_weights_first_wins_over_phantom(monkeypatch): + # when active memory tracks the wire bytes, the bracketed first walk + # registers ~0; a re-walk with no read of its own sees no delta and + # would register the full bytes again - it must not override the first + monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) + key = ("/models/a.gguf",) + pd.note_untracked_weights(0, key=key) + pd.note_untracked_weights(28 * GB, key=key) + assert pd.untracked_weight_bytes() == 0 + + def test_note_untracked_weights_distinct_keys_sum(monkeypatch): monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) pd.note_untracked_weights(10 * GB, key=("/models/a.gguf",)) From 65d3db02dd36d86403040d8973389861370f4533 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:48:51 -0700 Subject: [PATCH 4/5] docs: changelog wording for the headroom double-count fix --- CHANGELOG.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b6a685..5497bfe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -63,9 +63,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Fixed -- serve --speculative no longer serializes concurrent requests: the - drafter's reload of the target GGUF counted the shared weights twice, - holding admission headroom negative for the server's lifetime. +- serve --speculative no longer serializes concurrent requests: MTP + loads counted the model weights twice in the admission headroom + estimate, holding it negative for the server's lifetime. - Chat's bottom toolbar no longer clips the live tok/s readout on narrow terminals; sampling knobs are dropped first instead. - The serve free-headroom estimate went negative on models whose load From c70be66468d345f7bd8bb289f7bc3262d7db2ad5 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Mon, 10 Aug 2026 07:23:24 -0700 Subject: [PATCH 5/5] test(apc): drain the disk writer before the skeleton-recovery lookup --- tests/test_ckpt_tier.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/test_ckpt_tier.py b/tests/test_ckpt_tier.py index f2911f0..c545876 100644 --- a/tests/test_ckpt_tier.py +++ b/tests/test_ckpt_tier.py @@ -10,6 +10,7 @@ import os import subprocess import sys +import time as _time from types import SimpleNamespace import mlx.core as mx @@ -76,6 +77,25 @@ def assert_warm_matches(warm, orig, p): assert mx.array_equal(a, b).item() +def drain_disk(disk, timeout=30.0): + """Block until the APC disk writer has published every queued shard. + + ``save_exact_cache`` only enqueues; a background writer thread indexes + the shard later, and the lookup side (``find_exact_prefix``) reads the + index with no in-flight wait. Any test that stores and then expects to + read the result back must drain first, or it races the writer -- fast + locally, lost on a loaded CI runner. + """ + disk._q.join() + deadline = _time.monotonic() + timeout + while _time.monotonic() < deadline: + with disk._in_flight_lock: + if not disk._in_flight: + return + _time.sleep(0.005) + raise AssertionError("APC disk writer did not drain within %.1fs" % timeout) + + def test_ckpt_supported_shapes(): assert ckpt_supported(make_hybrid_cache(4)) assert not ckpt_supported([KVCache(), KVCache()]) @@ -823,6 +843,9 @@ def test_stripped_boundary_recovers_via_skeleton(tmp_path): cache = make_swa_cache(p, seed=12) ids = list(range(300, 300 + p)) assert ckpt_store(man, ids, cache, extra_hash=0) + # Recovery reads the skeleton off disk, so the async writer has to + # have published it before the lookup runs. + drain_disk(disk) idx = _ckpt_records(man) (key, rec), = list(idx.items()) _release_record(man, idx.pop(key))