ci: make model prefetch fail loudly instead of stalling

[solderzzc]

Fixes the integration failures currently red on #114, #115 and #116. None of them are caused by those PRs — all three touch disjoint files and reproduce the identical four failures, so the problem is the pipeline.

What was failing

modality	model it loads	status
server	Qwen2.5-0.5B-Instruct-4bit	pass
vision	Qwen2-VL-2B-Instruct-4bit	pass
graph	none	pass
audio / omni / opencode	gemma-4-e4b-it-4bit (~4.9 GB)	fail

The three failing modalities are exactly the three that need the large model. Their errors:

Error: Mismatched parameter language_model.model.per_layer_model_projection.weight … Actual [10752, 320], expected [10752, 2560]
Error: Key language_model.model.layers.24.self_attn.k_proj.weight not found
Error: Server failed to start or respond on port 15413 within 180 seconds.

The missing key changes between retries inside a single job — k_proj, then v_proj, then k_norm. That is a file growing underneath the loader, not a model-definition bug.

dflash fails differently but from the same cause; its log shows the download frozen:

Download: [=>    ]  10% (2919.3 MB / 29193.2 MB) | Speed: 0.0 MB/s
Download: [===>  ]  20% (2919.3 MB / 14596.6 MB) | Speed: 0.0 MB/s

Byte count pinned, the total recomputed against a moving fraction, 0.0 MB/s throughout.

Root causes

1. integration_matrix had no prefetch step. The server downloaded the model itself inside the test's 180 s startup window. 0.5 B and 2 B make it; 4.9 GB does not.

2. The cache key could never be refreshed. It was the static mlx-model-qwen2.5-0.5b-4bit — named for a model only one modality uses. actions/cache does not overwrite an existing key, so once that entry was created it was frozen: the gemma-4 jobs kept restoring a cache that had never held their model.

3. hf download <repo> || true swallowed every failure, and HF_HUB_DOWNLOAD_TIMEOUT was set on the test step but not on the download. An unbounded request can hang forever — I reproduced exactly this locally while fetching a model for #108: sockets stuck in CLOSE_WAIT, zero bytes for more than ten minutes, no timeout to break it, process alive at 0% CPU.

Changes

• scripts/ci-download-models.sh — retries with backoff, bounds each request via HF_HUB_DOWNLOAD_TIMEOUT (resume continues from existing .incomplete blobs), treats leftover .incomplete files as failure, and exits non-zero so a partial model never reaches the tests or gets saved into the cache.
• integration_matrix — per-modality prefetch driven by a matrix include, so each job fetches exactly what its test script loads.
• Cache keys — per-modality and derived from the test script via hashFiles, so changing a test's model rotates the key automatically; restore-keys still allows reuse. The four static keys in the other jobs are bumped to -v2 to drop any entry poisoned by a partial download.
• Every prefetch step now calls the script rather than || true.

Verification

The script was exercised locally on all four paths:

case	expected	result
already-cached model	exit 0, fast	✅
unknown repo	exit 1, ::error::	✅
planted .incomplete file	exit 1, rejected as partial	✅
multi-model invocation	each repo reported	✅

ci.yml parses, and the matrix resolves to the six modalities with the right model each.

What this does not fix

If the gemma-4 weight errors turn out to be a genuine model-definition mismatch rather than a truncated download, this will surface it as a clean, reproducible failure instead of one that varies per retry. That is the point — the current setup cannot tell the two apart.

I cannot run this pipeline locally, so the proof is in the next CI run on this branch.

🤖 Generated with Claude Code

[solderzzc]

CI results on this branch: the download hypothesis was wrong, and this PR does not turn the jobs green. It did do what it was built for — the failures are now clean and specific instead of ambiguous.

Prefetch works

--- mlx-community/gemma-4-e4b-it-4bit (attempt 1/3, per-request timeout 60s)
--- mlx-community/gemma-4-e4b-it-4bit OK

Complete, verified, no .incomplete files, no stall. The model is fully present before the test starts — so a partial download is now ruled out as the cause.

The real failures

audio / omni / opencode still fail, identically:

Error: Mismatched parameter language_model.model.per_layer_model_projection.weight … Actual [10752, 320], expected [10752, 2560]
Error: Key language_model.model.layers.24.self_attn.v_proj.weight not found

2560 / 320 = 8, which is exactly 4-bit packing density (eight values per uint32). So the checkpoint's weights are quantized while the module expecting them is a plain Linear — the quantization predicate did not fire for that path. That is a weight-loading bug in gemma-4-e4b-it-4bit, not a transport problem.

Correction to this PR's description. I attributed the varying missing key across retries to "a file growing underneath the loader". With the model now verified complete before the run, that explanation is dead. Swift Dictionary iteration order varies between processes, so each run simply reports whichever mismatch it reaches first. Same defect every time, different messenger.

dflash also fails cleanly now, and it is a different bug entirely:

[SwiftLM] ⚠️  Failed to load DFlash draft model: keyNotFound(CodingKeys(stringValue: "rope_theta"))

The download reaches 100%; z-lab/Qwen3.5-4B-DFlash's config simply has no rope_theta key for the decoder to find.

What this PR is still worth

It does not fix CI. It fixes three real fragilities and removes the download as a variable:

• the cache key that could never be refreshed once created
• hf download … || true swallowing every failure
• no timeout on the prefetch, which is what hung my local download for 5h43m with sockets in CLOSE_WAIT

Merging it means the next person reading these logs sees the actual bug instead of a download that may or may not have finished. Both remaining failures now reproduce deterministically and deserve their own issues.

Happy to close this instead if you would rather fix the model-loading bugs first — the CI hardening stands on its own either way.