diff --git a/packages/tracker/meshnet_tracker/model_files.py b/packages/tracker/meshnet_tracker/model_files.py index 4e97a49..70cc503 100644 --- a/packages/tracker/meshnet_tracker/model_files.py +++ b/packages/tracker/meshnet_tracker/model_files.py @@ -99,9 +99,18 @@ def select_safetensors_files_for_layers( if not isinstance(tensor_name, str) or not isinstance(rel_file, str): continue if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, inferred_total_layers): - selected.add(_normalise_relative_file(rel_file)) + rel = _normalise_relative_file(rel_file) + if not (root / rel).is_file(): + # A required weight file is missing from this local snapshot: the + # snapshot is incomplete for this layer range, so it must not be + # advertised as a source at all. Silently dropping just this file + # (as before) made a partial snapshot look "fully cached" to + # downstream clients, which then skipped the real download and + # only discovered the missing weights at model-load time. + return [] + selected.add(rel) - return sorted(rel for rel in selected if (root / rel).is_file()) + return sorted(selected) def _tensor_belongs_to_range( diff --git a/tests/test_tracker_model_files.py b/tests/test_tracker_model_files.py new file mode 100644 index 0000000..bc7cfe9 --- /dev/null +++ b/tests/test_tracker_model_files.py @@ -0,0 +1,71 @@ +"""Tests for tracker-side layer-aware SafeTensors file selection. + +The tracker advertises its local snapshot as a downloadable "model source" for +whatever layer range a node needs. If a required weight file for that range is +missing from the tracker's own disk, the tracker must refuse to advertise +itself as a source for that range — not quietly report the subset of files it +does have as if that were the complete, correct set. A partial-but-"complete" +answer makes the requesting node believe the download is already satisfied, +so it never fetches the real weights and only discovers the gap much later +at model-load time. +""" + +import json + +from meshnet_tracker.model_files import select_safetensors_files_for_layers + + +def _write_index(tmp_path, *, config=None): + (tmp_path / "config.json").write_text( + json.dumps(config or {"num_hidden_layers": 5}), + encoding="utf-8", + ) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({ + "weight_map": { + "model.embed_tokens.weight": "model-00001-of-00004.safetensors", + "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors", + "model.layers.4.self_attn.q_proj.weight": "model-00004-of-00004.safetensors", + "model.norm.weight": "model-00004-of-00004.safetensors", + "lm_head.weight": "model-00004-of-00004.safetensors", + }, + }), + encoding="utf-8", + ) + + +def _touch(path, size=1): + path.write_bytes(b"0" * size) + + +def test_selects_files_when_snapshot_is_complete(tmp_path): + "Selects files when snapshot is complete\n\nTags: general" + _write_index(tmp_path) + _touch(tmp_path / "model-00001-of-00004.safetensors") + + files = select_safetensors_files_for_layers(tmp_path, 0, 0) + + assert files == ["config.json", "model-00001-of-00004.safetensors", "model.safetensors.index.json"] + + +def test_returns_empty_when_a_required_weight_file_is_missing_on_disk(tmp_path): + "Returns empty when a required weight file is missing on disk\n\nTags: general" + _write_index(tmp_path) + # model-00001-of-00004.safetensors is required for layer 0 but was never + # downloaded onto this tracker host — the snapshot is incomplete. + + files = select_safetensors_files_for_layers(tmp_path, 0, 0) + + assert files == [] + + +def test_tail_range_returns_empty_when_only_head_shard_present(tmp_path): + "Tail range returns empty when only head shard present\n\nTags: general" + _write_index(tmp_path) + _touch(tmp_path / "model-00001-of-00004.safetensors") + # model-00004-of-00004.safetensors (norm/lm_head) is required for the tail + # range but missing — must not be silently dropped from the result. + + files = select_safetensors_files_for_layers(tmp_path, 4, 4) + + assert files == []