dual billing; tracker to node model sharing

tests/test_safetensors_selection.py

@@ -0,0 +1,86 @@
+"""Tests for layer-aware SafeTensors snapshot file selection."""
+
+import json
+
+import pytest
+
+from meshnet_node.safetensors_selection import select_safetensors_files_for_layers
+
+
+def _write_snapshot(tmp_path, *, config=None):
+    (tmp_path / "config.json").write_text(
+        json.dumps(config or {"num_hidden_layers": 5}),
+        encoding="utf-8",
+    )
+    (tmp_path / "tokenizer.json").write_text("{}", encoding="utf-8")
+    (tmp_path / "tokenizer_config.json").write_text("{}", encoding="utf-8")
+    (tmp_path / "README.md").write_text("not part of runtime snapshot", encoding="utf-8")
+    (tmp_path / "model.safetensors.index.json").write_text(
+        json.dumps({
+            "metadata": {"total_size": 123},
+            "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.1.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+                "model.layers.2.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+                "model.layers.3.mlp.down_proj.weight": "nested/model-00003-of-00004.safetensors",
+                "model.layers.4.self_attn.q_proj.weight": "nested/model-00003-of-00004.safetensors",
+                "model.norm.weight": "model-00004-of-00004.safetensors",
+                "lm_head.weight": "model-00004-of-00004.safetensors",
+            },
+        }),
+        encoding="utf-8",
+    )
+
+
+def test_selects_only_weight_shards_for_middle_layer_range(tmp_path):
+    _write_snapshot(tmp_path)
+
+    files = select_safetensors_files_for_layers(tmp_path, 2, 3)
+
+    assert files == [
+        "config.json",
+        "model-00002-of-00004.safetensors",
+        "model.safetensors.index.json",
+        "nested/model-00003-of-00004.safetensors",
+        "tokenizer.json",
+        "tokenizer_config.json",
+    ]
+
+
+def test_head_range_includes_embeddings(tmp_path):
+    _write_snapshot(tmp_path)
+
+    files = select_safetensors_files_for_layers(tmp_path, 0, 0)
+
+    assert "model-00001-of-00004.safetensors" in files
+    assert "model-00004-of-00004.safetensors" not in files
+
+
+def test_tail_range_includes_norm_and_lm_head_from_inferred_layer_count(tmp_path):
+    _write_snapshot(tmp_path, config={"text_config": {"num_hidden_layers": 5}})
+
+    files = select_safetensors_files_for_layers(tmp_path, 4, 4)
+
+    assert "nested/model-00003-of-00004.safetensors" in files
+    assert "model-00004-of-00004.safetensors" in files
+    assert "model-00001-of-00004.safetensors" not in files
+
+
+def test_tail_files_are_not_selected_without_total_layer_count(tmp_path):
+    _write_snapshot(tmp_path, config={"architectures": ["UnknownForTest"]})
+
+    files = select_safetensors_files_for_layers(tmp_path, 4, 4)
+
+    assert "nested/model-00003-of-00004.safetensors" in files
+    assert "model-00004-of-00004.safetensors" not in files
+
+
+def test_rejects_unsafe_weight_map_paths(tmp_path):
+    (tmp_path / "model.safetensors.index.json").write_text(
+        json.dumps({"weight_map": {"model.layers.0.weight": "../escape.safetensors"}}),
+        encoding="utf-8",
+    )
+
+    with pytest.raises(ValueError, match="unsafe relative file"):
+        select_safetensors_files_for_layers(tmp_path, 0, 0)