Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

packages/node/meshnet_node/downloader.py

@@ -1,6 +1,6 @@
 """Shard downloader — fetches model shards from peers or HuggingFace Hub.
 
-Cache layout: ~/.cache/meshnet/shards/<model>/layers_<start>-<end>/
+Cache layout: ~/.cache/meshnet/shards/<model>/
 
 For "stub-model" (no HF repo), a placeholder JSON file is written so the
 test suite never touches the network.
@@ -106,9 +106,7 @@ def _download_shard_from_peer(
             _safe_extract_shard(archive_path, extract_dir)
             if compute_shard_checksum(extract_dir) != checksum:
                 return False
-            if shard_dir.exists():
-                shutil.rmtree(shard_dir)
-            shutil.move(str(extract_dir), str(shard_dir))
+            _merge_tree(extract_dir, shard_dir)
             return True
         except Exception:
             return False
@@ -217,6 +215,38 @@ def _reuse_local_file(expected: int, dest: Path, final: Path) -> bool:
     return False
 
 
+def _valid_source_rel_files(source: dict) -> list[str]:
+    return [
+        rel for rel in (source.get("files") or [])
+        if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts
+    ]
+
+
+def _source_files_cached(source: dict, shard_dir: Path) -> bool:
+    rel_files = _valid_source_rel_files(source)
+    if not rel_files:
+        return False
+    sizes = source.get("file_sizes")
+    if not isinstance(sizes, dict):
+        return False
+    for rel in rel_files:
+        expected = sizes.get(rel)
+        path = shard_dir / rel
+        if not isinstance(expected, int) or not path.exists() or path.stat().st_size != expected:
+            return False
+    return True
+
+
+def _merge_tree(src: Path, dest: Path) -> None:
+    dest.mkdir(parents=True, exist_ok=True)
+    for path in sorted(p for p in src.rglob("*") if p.is_file()):
+        target = dest / path.relative_to(src)
+        target.parent.mkdir(parents=True, exist_ok=True)
+        if target.exists():
+            target.unlink()
+        shutil.move(str(path), str(target))
+
+
 def _download_source_files(
     source: dict,
     shard_dir: Path,
@@ -231,10 +261,7 @@ def _download_source_files(
     the download is retried or the node is restarted.
     """
     url = source.get("url")
-    rel_files = [
-        rel for rel in (source.get("files") or [])
-        if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts
-    ]
+    rel_files = _valid_source_rel_files(source)
     if not isinstance(url, str) or not url or not rel_files:
         return None
     sizes = source.get("file_sizes")
@@ -283,9 +310,11 @@ def _download_source_files(
                     time.sleep(1.0 * attempt)
     finally:
         tracker_bar.close()
-    if shard_dir.exists():
-        shutil.rmtree(shard_dir)
-    shutil.move(str(partial_dir), str(shard_dir))
+    _merge_tree(partial_dir, shard_dir)
+    try:
+        partial_dir.rmdir()
+    except OSError:
+        pass
     return shard_dir
 
 
@@ -331,9 +360,7 @@ def _download_model_source(
             if progress:
                 print(f"  {label}: transfer complete ({received / 1e9:.2f} GB), extracting ...", flush=True)
             _safe_extract_shard(archive_path, extract_dir)
-            if shard_dir.exists():
-                shutil.rmtree(shard_dir)
-            shutil.move(str(extract_dir), str(shard_dir))
+            _merge_tree(extract_dir, shard_dir)
             return shard_dir
         except Exception as exc:
             if progress:
@@ -431,10 +458,17 @@ def download_shard(
     the test suite hermetic while the real download path is exercised by
     passing a non-stub *hf_repo*.
     """
-    shard_dir = cache_dir / model / f"layers_{shard_start}-{shard_end}"
+    shard_dir = cache_dir / model
     if progress:
         print(f"  Target location: {shard_dir}", flush=True)
 
+    for source in model_sources or []:
+        label = str(source.get("type") or "model-source")
+        if _source_files_cached(source, shard_dir):
+            if progress:
+                print(f"  [{label}] requested files already cached at {shard_dir}", flush=True)
+            return shard_dir
+
     for peer in peers or []:
         if progress:
             print(f"  Trying peer shard download from {peer.get('endpoint')} ...", flush=True)

packages/node/meshnet_node/model_catalog.py

@@ -64,6 +64,17 @@ def _load_model_metadata() -> dict[str, dict]:
 _MODEL_METADATA = _load_model_metadata()
 
 
+def _local_model_path(hf_repo: str, cache_dir: Path | None) -> Path | None:
+    if cache_dir is None:
+        return None
+    if (cache_dir / "config.json").exists():
+        return cache_dir
+    candidate = cache_dir / hf_repo.split("/")[-1]
+    if (candidate / "config.json").exists():
+        return candidate
+    return None
+
+
 CURATED_MODELS: list[ModelPreset] = [
     ModelPreset(
         name="Qwen2.5-0.5B-Instruct",
@@ -215,9 +226,11 @@ def model_metadata_for(
     try:
         from transformers import AutoConfig  # type: ignore[import]
 
+        local_model = _local_model_path(hf_repo, cache_dir)
+        load_source = str(local_model) if local_model is not None else hf_repo
         cfg = AutoConfig.from_pretrained(
-            hf_repo,
-            cache_dir=str(cache_dir) if cache_dir is not None else None,
+            load_source,
+            cache_dir=str(cache_dir) if cache_dir is not None and local_model is None else None,
         )
         # Composite configs (VLM/MoE) nest decoder fields in text_config.
         text_cfg = getattr(cfg, "text_config", None) or cfg

packages/node/meshnet_node/startup.py

@@ -26,6 +26,86 @@ from .wallet import load_or_create_wallet
 _DEFAULT_BYTES_PER_LAYER = 30 * 1024 * 1024
 
 
+def _downloaded_model_inventory(
+    model: str,
+    shard_start: int,
+    shard_end: int,
+    shard_path: Path,
+    hf_repo: str | None = None,
+    model_sources: list[dict] | None = None,
+) -> list[dict]:
+    """Return a cheap local inventory record without reading model file contents."""
+    file_count = 0
+    total_bytes = 0
+    existing_rel_files: set[str] = set()
+    if shard_path.exists():
+        for path in shard_path.rglob("*"):
+            if not path.is_file():
+                continue
+            file_count += 1
+            try:
+                existing_rel_files.add(path.relative_to(shard_path).as_posix())
+            except ValueError:
+                pass
+            try:
+                total_bytes += path.stat().st_size
+            except OSError:
+                pass
+    record = {
+        "model": model,
+        "shard_start": shard_start,
+        "shard_end": shard_end,
+        "path": str(shard_path),
+        "file_count": file_count,
+        "total_bytes": total_bytes,
+    }
+    if hf_repo is not None:
+        record["hf_repo"] = hf_repo
+    expected_files: set[str] = set()
+    file_sizes: dict[str, int] = {}
+    for source in model_sources or []:
+        for rel in source.get("full_files") or source.get("files") or []:
+            if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts:
+                expected_files.add(rel)
+        sizes = source.get("file_sizes")
+        if isinstance(sizes, dict):
+            for rel, size in sizes.items():
+                if isinstance(rel, str) and isinstance(size, int):
+                    file_sizes[rel] = size
+    if expected_files:
+        expected_bytes = sum(file_sizes.get(rel, 0) for rel in expected_files)
+        local_expected_files = existing_rel_files & expected_files
+        local_expected_bytes = sum(file_sizes.get(rel, 0) for rel in local_expected_files)
+        record["expected_file_count"] = len(expected_files)
+        record["local_expected_file_count"] = len(local_expected_files)
+        record["expected_bytes"] = expected_bytes
+        record["local_expected_bytes"] = local_expected_bytes
+        record["local_model_percentage"] = (
+            round((local_expected_bytes / expected_bytes) * 100, 4)
+            if expected_bytes > 0
+            else 0.0
+        )
+    return [record]
+
+
+def _registration_shard_checksum(model: str, shard_path: Path) -> str | None:
+    """Only checksum tiny stub shards; real model folders are too large to hash at startup."""
+    if model != "stub-model":
+        return None
+    return compute_shard_checksum(shard_path)
+
+
+def _model_cache_path(model_id: str, cache_dir: Path | None) -> Path | None:
+    if cache_dir is None:
+        return None
+    if (cache_dir / "config.json").exists():
+        return cache_dir
+    candidate = cache_dir / model_id.split("/")[-1]
+    if candidate.exists():
+        return candidate
+    return None
+
+
 def _memory_budget(device: str, vram_mb: int, ram_mb: int, shared_vram_mb: int = 0) -> tuple[int, str]:
     """Return the capacity budget in MB and whether it came from VRAM or RAM."""
     if device == "cuda" and vram_mb > 0:
@@ -502,6 +582,7 @@ def run_startup(
 
     if model_id:  # treat "" the same as None — no explicit model given
         user_pinned_shard = shard_start is not None or shard_end is not None
+        full_sources: list[dict] = []
         # Auto-detect shard range from model config if not explicitly provided
         if shard_start is None or shard_end is None:
             try:
@@ -581,6 +662,7 @@ def run_startup(
         _attach_relay_bridge(node, relay_bridge)
         # Register with tracker so other nodes can auto-join this model.
         total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
+        model_cache_path = _model_cache_path(model_id, cache_dir)
         reg_payload = {
             "endpoint": endpoint,
             "model": model_id.split("/")[-1],
@@ -595,6 +677,18 @@ def run_startup(
             "tracker_mode": (shard_start == 0),
             "managed_assignment": not user_pinned_shard,
             "model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
+            "downloaded_models": (
+                _downloaded_model_inventory(
+                    model_id.split("/")[-1],
+                    shard_start,
+                    shard_end,
+                    model_cache_path,
+                    hf_repo=model_id,
+                    model_sources=full_sources,
+                )
+                if model_cache_path is not None
+                else []
+            ),
             **registration_capabilities,
             **relay_fields,
         }
@@ -693,6 +787,7 @@ def run_startup(
             relay_url=relay_url,
         )
         _attach_relay_bridge(node, relay_bridge)
+        model_cache_path = _model_cache_path(assigned_hf_repo, cache_dir)
         auto_reg_payload = {
             "endpoint": endpoint,
             "model": assigned_hf_repo.split("/")[-1],
@@ -707,6 +802,18 @@ def run_startup(
             "tracker_mode": (assigned_shard_start == 0),
             "managed_assignment": True,
             "model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
+            "downloaded_models": (
+                _downloaded_model_inventory(
+                    assigned_hf_repo.split("/")[-1],
+                    assigned_shard_start,
+                    assigned_shard_end,
+                    model_cache_path,
+                    hf_repo=assigned_hf_repo,
+                    model_sources=full_sources,
+                )
+                if model_cache_path is not None
+                else []
+            ),
             **registration_capabilities,
             **relay_fields,
         }
@@ -776,7 +883,15 @@ def run_startup(
     if model_sources:
         dl_kwargs["model_sources"] = model_sources
     shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
-    shard_checksum = compute_shard_checksum(shard_path)
+    shard_checksum = _registration_shard_checksum(assigned_model, shard_path)
+    downloaded_models = _downloaded_model_inventory(
+        assigned_model,
+        shard_start,
+        shard_end,
+        shard_path,
+        hf_repo=hf_repo,
+        model_sources=model_sources,
+    )
     print(f"  Cached at: {shard_path}", flush=True)
 
     # 5. Start HTTP server
@@ -790,6 +905,7 @@ def run_startup(
         model=assigned_model,
         shard_path=shard_path,
     )
+    _node_start_time = time.monotonic()
     actual_port = node.start()
     public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
     endpoint = f"http://{public_host}:{actual_port}"
@@ -805,24 +921,27 @@ def run_startup(
 
     # 6. Register with tracker
     print("Registering with tracker...", flush=True)
+    reg_payload = {
+        "endpoint": endpoint,
+        "model": assigned_model,
+        "shard_start": shard_start,
+        "shard_end": shard_end,
+        "shard_checksum": shard_checksum,
+        "downloaded_models": downloaded_models,
+        "hardware_profile": hw,
+        "wallet_address": address,
+        "score": 1.0,
+        **registration_capabilities,
+        **relay_fields,
+    }
     try:
         reg_resp = _post_json(
             f"{tracker_url}/v1/nodes/register",
-            {
-                "endpoint": endpoint,
-                "model": assigned_model,
-                "shard_start": shard_start,
-                "shard_end": shard_end,
-                "shard_checksum": shard_checksum,
-                "hardware_profile": hw,
-                "wallet_address": address,
-                "score": 1.0,
-                **registration_capabilities,
-                **relay_fields,
-            },
+            reg_payload,
         )
         node_id = str(reg_resp["node_id"])
         setattr(node, "tracker_node_id", node_id)
+        _start_heartbeat(tracker_url, node_id, reg_payload, node_ref=node, start_time=_node_start_time)
     except Exception:
         node.stop()
         raise
@@ -855,9 +974,11 @@ def _detect_num_layers(model_id: str, cache_dir: Path | None = None) -> int | No
 
         from .model_catalog import layers_from_config
 
+        local_model = _model_cache_path(model_id, cache_dir)
+        load_source = str(local_model) if local_model is not None else model_id
         cfg = AutoConfig.from_pretrained(
-            model_id,
-            cache_dir=str(cache_dir) if cache_dir is not None else None,
+            load_source,
+            cache_dir=str(cache_dir) if cache_dir is not None and local_model is None else None,
         )
         layers = layers_from_config(cfg)
         if layers is None: