dual billing; tracker to node model sharing

tests/test_node_startup.py

@@ -4,6 +4,8 @@ import json
 import io
 import os
 import sys
+import tarfile
+import time
 import types
 import urllib.request
 from pathlib import Path
@@ -405,6 +407,75 @@ def test_download_shard_uses_huggingface_when_repo_is_assigned(tmp_path, monkeyp
     }]
 
 
+def test_download_shard_races_tracker_model_source_against_huggingface(
+    tmp_path,
+    monkeypatch,
+):
+    """Tracker-hosted model files can win while HF receives the same allow_patterns."""
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+    (source_dir / "config.json").write_text("{}")
+    (source_dir / "model-00002-of-00004.safetensors").write_text("tracker")
+    archive = io.BytesIO()
+    with tarfile.open(fileobj=archive, mode="w") as tf:
+        tf.add(source_dir / "config.json", arcname="config.json")
+        tf.add(
+            source_dir / "model-00002-of-00004.safetensors",
+            arcname="model-00002-of-00004.safetensors",
+        )
+
+    class FakeTrackerResponse:
+        def __init__(self, payload: bytes):
+            self._payload = io.BytesIO(payload)
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+        def read(self, size: int = -1) -> bytes:
+            return self._payload.read(size)
+
+    monkeypatch.setattr(
+        urllib.request,
+        "urlopen",
+        lambda *args, **kwargs: FakeTrackerResponse(archive.getvalue()),
+    )
+    hf_calls = []
+
+    def fake_snapshot_download(**kwargs):
+        hf_calls.append(kwargs)
+        time.sleep(0.05)
+        local_dir = Path(kwargs["local_dir"])
+        local_dir.mkdir(parents=True, exist_ok=True)
+        (local_dir / "model-00002-of-00004.safetensors").write_text("hf")
+        return str(local_dir)
+
+    monkeypatch.setitem(
+        sys.modules,
+        "huggingface_hub",
+        types.SimpleNamespace(snapshot_download=fake_snapshot_download),
+    )
+
+    shard_dir = download_shard(
+        "tiny-llama",
+        2,
+        3,
+        cache_dir=tmp_path / "cache",
+        hf_repo="org/tiny-llama-shards",
+        model_sources=[{
+            "type": "tracker",
+            "url": "http://tracker/v1/model-files/download?model=tiny-llama",
+            "files": ["config.json", "model-00002-of-00004.safetensors"],
+        }],
+        progress=False,
+    )
+
+    assert (shard_dir / "model-00002-of-00004.safetensors").read_text() == "tracker"
+    assert hf_calls[0]["allow_patterns"] == ["config.json", "model-00002-of-00004.safetensors"]
+
+
 def test_download_shard_logs_huggingface_source(tmp_path, monkeypatch, capsys):
     """Shard download status tells the node operator when HuggingFace was used."""
 
@@ -585,6 +656,83 @@ def test_tracker_assign_returns_huggingface_repo_when_configured():
         tracker.stop()
 
 
+def test_tracker_assign_advertises_local_model_source_and_serves_subset(tmp_path):
+    """Tracker with models_dir advertises and serves only files needed for the shard."""
+    snapshot = tmp_path / "models" / "models--org--tiny-llama-shards" / "snapshots" / "abc"
+    nested = snapshot / "nested"
+    nested.mkdir(parents=True)
+    (snapshot / "config.json").write_text(json.dumps({"num_hidden_layers": 4}))
+    (snapshot / "tokenizer.json").write_text("{}")
+    (snapshot / "model.safetensors.index.json").write_text(json.dumps({
+        "weight_map": {
+            "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+            "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+            "model.layers.1.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+            "model.layers.2.self_attn.q_proj.weight": "nested/model-00002-of-00003.safetensors",
+            "model.layers.3.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+            "lm_head.weight": "model-00003-of-00003.safetensors",
+        },
+    }))
+    for rel in [
+        "model-00001-of-00003.safetensors",
+        "model-00002-of-00003.safetensors",
+        "nested/model-00002-of-00003.safetensors",
+        "model-00003-of-00003.safetensors",
+    ]:
+        (snapshot / rel).write_text(rel)
+
+    tracker = TrackerServer(
+        model_presets={
+            "tiny-llama": {
+                "layers_start": 0,
+                "layers_end": 3,
+                "hf_repo": "org/tiny-llama-shards",
+                "bytes_per_layer": {"bfloat16": 1024 * 1024},
+            },
+        },
+        models_dir=tmp_path / "models",
+    )
+    port = tracker.start()
+    try:
+        data = json.dumps({
+            "endpoint": "http://127.0.0.1:9100",
+            "model": "tiny-llama",
+            "shard_start": 0,
+            "shard_end": 0,
+            "hardware_profile": {},
+            "score": 1.0,
+        }).encode()
+        req = urllib.request.Request(
+            f"http://127.0.0.1:{port}/v1/nodes/register",
+            data=data,
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        with urllib.request.urlopen(req) as r:
+            r.read()
+        resp = _get_json(
+            f"http://127.0.0.1:{port}/v1/nodes/assign?model=tiny-llama&device=cpu&ram_mb=3"
+        )
+        assert resp["shard_start"] == 1
+        assert resp["shard_end"] == 2
+        assert resp["model_sources"]
+        source = resp["model_sources"][0]
+        assert source["files"] == [
+            "config.json",
+            "model-00002-of-00003.safetensors",
+            "model.safetensors.index.json",
+            "nested/model-00002-of-00003.safetensors",
+            "tokenizer.json",
+        ]
+        with urllib.request.urlopen(source["url"], timeout=5) as response:
+            payload = io.BytesIO(response.read())
+        with tarfile.open(fileobj=payload, mode="r") as tf:
+            names = sorted(tf.getnames())
+        assert names == source["files"]
+    finally:
+        tracker.stop()
+
+
 def test_tracker_assign_lists_peers_for_same_model_shard():
     """A registered node with a completed shard is returned as a same-shard peer."""
     import json as _json