diff --git a/billing.sqlite b/billing.sqlite index 3e98f57..85ceb51 100644 Binary files a/billing.sqlite and b/billing.sqlite differ diff --git a/docs/issues/44-tracker-shard-source-partial-download.md b/docs/issues/44-tracker-shard-source-partial-download.md index a9cca40..2f19111 100644 --- a/docs/issues/44-tracker-shard-source-partial-download.md +++ b/docs/issues/44-tracker-shard-source-partial-download.md @@ -75,7 +75,7 @@ What exists already (build on it, don't duplicate): - [x] Node downloader keeps exact-shard peers first, then races tracker model sources against a HuggingFace `snapshot_download(..., allow_patterns=...)` subset download, using the first successful source. -- [ ] When no tracker model source is available at all, the HuggingFace +- [x] When no tracker model source is available at all, the HuggingFace fallback still computes `allow_patterns` from the repo's own `model.safetensors.index.json` (fetched directly, not via the tracker) — it never silently downloads the full model just because the tracker has @@ -95,7 +95,9 @@ What exists already (build on it, don't duplicate): - 2026-07-06: Added the tracker/node download path. For immediate Qwen3.6-35B LAN testing, real PyTorch nodes fetch the full snapshot from the tracker via - `full_url` and race HuggingFace as fallback. Remaining hard half is true - partial model materialization: the backend can prefer a downloaded local - model directory, but Transformers still needs a `meta`-device load path that - materializes only assigned layers. + `full_url`; HuggingFace remains fallback-only, and when it is used the node + computes `allow_patterns` from the repo's remote SafeTensors index so it + stays layer-filtered even without tracker-cached files. Remaining hard half + is true partial model materialization: the backend can prefer a downloaded + local model directory, but Transformers still needs a `meta`-device load + path that materializes only assigned layers. diff --git a/meshnet_registry.sqlite3 b/meshnet_registry.sqlite3 index ee3c390..7a6d3a4 100644 Binary files a/meshnet_registry.sqlite3 and b/meshnet_registry.sqlite3 differ diff --git a/packages/node/meshnet_node/downloader.py b/packages/node/meshnet_node/downloader.py index 0e62086..1d4c4ba 100644 --- a/packages/node/meshnet_node/downloader.py +++ b/packages/node/meshnet_node/downloader.py @@ -1,4 +1,4 @@ -"""Shard downloader — fetches model shards from peers or HuggingFace Hub. +"""Shard downloader — fetches model files from peers, tracker sources, or HuggingFace. Cache layout: ~/.cache/meshnet/shards// diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index ece10d8..008cf32 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -4,6 +4,7 @@ from __future__ import annotations import base64 from dataclasses import dataclass +import json from pathlib import Path from typing import Any, Literal @@ -22,6 +23,10 @@ class InsufficientVRAMError(ModelBackendError): """Raised when a requested shard cannot fit in available CUDA memory.""" +class PartialModelLoadUnsupported(ModelBackendError): + """Raised when a shard cannot be materialized from a local snapshot subset.""" + + @dataclass(frozen=True) class TensorPayload: body: bytes @@ -94,20 +99,39 @@ class TorchModelShard: None if load_source != model_id else cache_dir, ) try: - load_kwargs = { - "device_map": "auto" if uses_quantized_weights else None, - "dtype": dtype, - "low_cpu_mem_usage": True, - "cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None, - } - if quant_config is not None: - load_kwargs["quantization_config"] = quant_config - self.model = AutoModelForCausalLM.from_pretrained( + total_layers_hint = _total_layers_for_local_snapshot(AutoConfig, load_source) + if _should_partial_materialize_shard( load_source, - **load_kwargs, - ) - if not uses_quantized_weights: - self.model.to(self.device) + shard_start, + shard_end, + total_layers_hint=total_layers_hint, + uses_quantized_weights=uses_quantized_weights, + ): + self.model = _load_partial_model_from_snapshot( + AutoConfig, + AutoModelForCausalLM, + torch, + load_source, + shard_start, + shard_end, + dtype, + self.device, + ) + else: + load_kwargs = { + "device_map": "auto" if uses_quantized_weights else None, + "dtype": dtype, + "low_cpu_mem_usage": True, + "cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None, + } + if quant_config is not None: + load_kwargs["quantization_config"] = quant_config + self.model = AutoModelForCausalLM.from_pretrained( + load_source, + **load_kwargs, + ) + if not uses_quantized_weights: + self.model.to(self.device) except Exception as exc: if _looks_like_oom(exc): raise InsufficientVRAMError( @@ -357,6 +381,135 @@ def load_torch_shard( return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir) +def _total_layers_for_local_snapshot(auto_config: Any, load_source: str) -> int | None: + snapshot_dir = Path(load_source) + if not (snapshot_dir / "config.json").exists(): + return None + from .model_catalog import layers_from_config + + try: + cfg = auto_config.from_pretrained(str(snapshot_dir)) + except Exception: + return None + return layers_from_config(cfg) + + +def _should_partial_materialize_shard( + load_source: str, + shard_start: int, + shard_end: int, + *, + total_layers_hint: int | None, + uses_quantized_weights: bool, +) -> bool: + if uses_quantized_weights: + return False + snapshot_dir = Path(load_source) + if not snapshot_dir.exists() or not (snapshot_dir / "config.json").exists(): + return False + if not (snapshot_dir / "model.safetensors.index.json").exists(): + return False + if total_layers_hint is None: + return False + return not (shard_start == 0 and shard_end >= total_layers_hint - 1) + + +def _load_partial_model_from_snapshot( + auto_config: Any, + auto_model_for_causal_lm: Any, + torch: Any, + load_source: str, + shard_start: int, + shard_end: int, + dtype: Any, + device: Any, + *, + init_empty_weights_fn: Any | None = None, + set_tensor_fn: Any | None = None, + safe_open_fn: Any | None = None, +) -> Any: + from .model_catalog import layers_from_config + from .safetensors_selection import ( + INDEX_FILENAME, + select_tensor_names_for_layers_from_index, + ) + + if init_empty_weights_fn is None: + from accelerate import init_empty_weights as init_empty_weights_fn + if set_tensor_fn is None: + from accelerate.utils import set_module_tensor_to_device as set_tensor_fn + if safe_open_fn is None: + from safetensors import safe_open as safe_open_fn + + snapshot_dir = Path(load_source) + cfg = auto_config.from_pretrained(str(snapshot_dir)) + total_layers = layers_from_config(cfg) + if total_layers is None: + raise PartialModelLoadUnsupported( + f"could not determine num_hidden_layers for local snapshot {snapshot_dir}" + ) + if shard_end >= total_layers: + raise ValueError( + f"shard_end {shard_end} exceeds last layer index {total_layers - 1}" + ) + + index_path = snapshot_dir / INDEX_FILENAME + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise PartialModelLoadUnsupported( + f"missing SafeTensors index for partial load: {index_path}" + ) from exc + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise PartialModelLoadUnsupported(f"{INDEX_FILENAME} must contain a weight_map object") + + tensor_names = select_tensor_names_for_layers_from_index( + weight_map, + shard_start, + shard_end, + total_layers=total_layers, + ) + if not tensor_names: + raise PartialModelLoadUnsupported( + f"no checkpoint tensors matched layers {shard_start}-{shard_end} in {snapshot_dir}" + ) + + with init_empty_weights_fn(): + model = auto_model_for_causal_lm.from_config(cfg, torch_dtype=dtype) + tie_weights = getattr(model, "tie_weights", None) + if callable(tie_weights): + tie_weights() + + tensors_by_file: dict[str, list[str]] = {} + for tensor_name in sorted(tensor_names): + rel_file = weight_map.get(tensor_name) + if not isinstance(rel_file, str): + continue + tensors_by_file.setdefault(rel_file, []).append(tensor_name) + + for rel_file, names in tensors_by_file.items(): + checkpoint_file = snapshot_dir / rel_file + if not checkpoint_file.exists(): + raise PartialModelLoadUnsupported( + f"checkpoint file advertised in {INDEX_FILENAME} is missing: {checkpoint_file}" + ) + with safe_open_fn(str(checkpoint_file), framework="pt", device="cpu") as handle: + for tensor_name in names: + set_tensor_fn( + model, + tensor_name, + device, + value=handle.get_tensor(tensor_name), + dtype=dtype, + ) + + for module in _active_modules_for_shard(model, shard_start, shard_end): + if hasattr(module, "to"): + module.to(device) + return model + + def _model_load_plan( auto_config: Any, model_id: str, @@ -442,6 +595,37 @@ def _position_embeddings(model: Any) -> Any | None: return None +def _rotary_embedding_module(model: Any) -> Any | None: + if hasattr(model, "model") and hasattr(model.model, "rotary_emb"): + return model.model.rotary_emb + if hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"): + return model.transformer.rotary_emb + return None + + +def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]: + active: list[Any] = [] + + def add(module: Any | None) -> None: + if module is None: + return + if any(existing is module for existing in active): + return + active.append(module) + + if shard_start == 0: + add(_embed_tokens(model)) + add(_position_embeddings(model)) + add(_rotary_embedding_module(model)) + for layer in _model_layers(model)[shard_start:shard_end + 1]: + add(layer) + total_layers = len(_model_layers(model)) + if shard_end >= total_layers - 1: + add(_final_norm(model)) + add(getattr(model, "lm_head", None)) + return active + + def _final_norm(model: Any) -> Any | None: if hasattr(model, "model") and hasattr(model.model, "norm"): return model.model.norm @@ -485,11 +669,7 @@ def _rotary_position_embeddings(model: Any, hidden_states: Any, position_ids: An """Return model-level rotary embeddings required by newer HF decoder layers.""" if position_ids is None: return None - rotary = None - if hasattr(model, "model") and hasattr(model.model, "rotary_emb"): - rotary = model.model.rotary_emb - elif hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"): - rotary = model.transformer.rotary_emb + rotary = _rotary_embedding_module(model) if rotary is None: return None return rotary(hidden_states, position_ids) diff --git a/packages/node/meshnet_node/safetensors_selection.py b/packages/node/meshnet_node/safetensors_selection.py index f9e87d3..bfe80f1 100644 --- a/packages/node/meshnet_node/safetensors_selection.py +++ b/packages/node/meshnet_node/safetensors_selection.py @@ -118,6 +118,23 @@ def select_files_for_layers_from_index( return selected +def select_tensor_names_for_layers_from_index( + weight_map: dict[str, str], + start_layer: int, + end_layer: int, + *, + total_layers: int | None = None, +) -> set[str]: + """Pure variant that returns checkpoint tensor names instead of file paths.""" + selected: set[str] = set() + for tensor_name, rel_file in weight_map.items(): + if not isinstance(tensor_name, str) or not isinstance(rel_file, str): + continue + if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, total_layers): + selected.add(tensor_name) + return selected + + def _tensor_belongs_to_range( tensor_name: str, start_layer: int, diff --git a/packages/tracker/meshnet_tracker/billing.py b/packages/tracker/meshnet_tracker/billing.py index ddef7d3..0a87865 100644 --- a/packages/tracker/meshnet_tracker/billing.py +++ b/packages/tracker/meshnet_tracker/billing.py @@ -453,13 +453,13 @@ class BillingLedger: with self._lock: return self._node_pending.get(wallet, 0.0) - def usage_for(self, api_keys: list[str], *, recent_limit: int = 20) -> dict: + def usage_for(self, api_keys: list[str], *, recent_limit: int | None = None) -> dict: """Aggregate charge history for a set of API keys (dashboard view).""" keys = set(api_keys) requests = 0 total_tokens = 0 total_cost = 0.0 - recent: list[dict] = [] + records: list[dict] = [] with self._lock: for event in self._event_log: if event.get("type") != "charge" or event.get("api_key") not in keys: @@ -467,18 +467,20 @@ class BillingLedger: requests += 1 total_tokens += int(event.get("total_tokens", 0)) total_cost += float(event.get("cost", 0.0)) - recent.append({ + records.append({ "api_key": event["api_key"], "model": event.get("model"), "total_tokens": event.get("total_tokens", 0), "cost": event.get("cost", 0.0), "ts": event.get("ts", 0.0), }) + recent = records[-recent_limit:] if recent_limit is not None else records return { "requests": requests, "total_tokens": total_tokens, "total_cost": total_cost, - "recent": recent[-recent_limit:], + "records": records, + "recent": recent, } def snapshot(self) -> dict: diff --git a/packages/tracker/meshnet_tracker/dashboard.html b/packages/tracker/meshnet_tracker/dashboard.html index e50dbd0..35c126a 100644 --- a/packages/tracker/meshnet_tracker/dashboard.html +++ b/packages/tracker/meshnet_tracker/dashboard.html @@ -39,12 +39,41 @@ .form-row { display:flex; gap:8px; } .form-row button { white-space:nowrap; } .error-msg { color:var(--bad); font-size:12px; min-height:16px; } - .keybox { word-break:break-all; background:var(--bg); border:1px solid var(--border); + .keybox { display:flex; flex-wrap:wrap; align-items:center; gap:6px; + position:relative; + word-break:break-all; background:var(--bg); border:1px solid var(--border); border-radius:6px; padding:4px 8px; margin:4px 0; font-size:11px; } + .key-text { cursor:text; flex:1 1 auto; min-width:12rem; } + .copy-tooltip { + position:absolute; right:8px; top:-26px; + background:var(--panel); border:1px solid var(--border); color:var(--ok); + padding:2px 8px; border-radius:4px; font-size:11px; + pointer-events:none; z-index:1; white-space:nowrap; + } .tabs { display:flex; gap:10px; margin-bottom:8px; } .tabs a { color:var(--dim); cursor:pointer; } .tabs a.active { color:var(--accent); border-bottom:1px solid var(--accent); } + .dashboard-tabs { display:flex; gap:10px; padding:10px 20px 0; border-bottom:1px solid var(--border); } + .dashboard-tabs button { border:0; border-bottom:1px solid transparent; border-radius:0; + background:transparent; color:var(--dim); padding:5px 0 8px; } + .dashboard-tabs button.active { color:var(--accent); border-bottom-color:var(--accent); } .wide { grid-column:1 / -1; } + section[hidden] { display:none !important; } + .chat-shell { display:grid; grid-template-columns:minmax(0, 1.35fr) minmax(320px, 0.65fr); gap:12px; } + .chat-pane { display:flex; flex-direction:column; gap:10px; min-width:0; } + .chat-panel { background:var(--bg); border:1px solid var(--border); border-radius:6px; padding:10px; } + .chat-controls { display:flex; gap:10px; align-items:end; flex-wrap:wrap; } + .chat-controls label { display:flex; flex-direction:column; gap:4px; color:var(--dim); } + .chat-controls select { min-width:220px; } + .chat-history { display:flex; flex-direction:column; gap:8px; min-height:220px; max-height:420px; overflow:auto; } + .chat-message { border:1px solid #21262d; border-radius:6px; padding:8px 10px; background:#10151d; } + .chat-role { color:var(--dim); font-size:11px; text-transform:uppercase; letter-spacing:.06em; margin-bottom:4px; } + .chat-role-user { color:var(--accent); } + .chat-role-assistant { color:var(--ok); } + .chat-role-error { color:var(--bad); } + .chat-compose { display:flex; flex-direction:column; gap:8px; } + .chat-compose textarea { min-height:112px; resize:vertical; width:100%; } + .chat-status { color:var(--dim); font-size:12px; } .console { background:var(--bg); border:1px solid var(--border); border-radius:6px; min-height:160px; max-height:280px; overflow:auto; padding:7px 9px; @@ -55,6 +84,10 @@ .console-level-info { color:var(--accent); } .console-level-warn { color:var(--warn); } .console-level-error { color:var(--bad); } + .status-pending { color:var(--warn); } + .status-processing { color:var(--accent); } + .status-failed { color:var(--bad); } + .status-complete { color:var(--ok); } @@ -63,19 +96,50 @@ +
-

Account

loading…
- -

Tracker hive

loading…
-

Nodes & coverage

loading…
-

Client balances

loading…
-

Node pending payouts

loading…
-

Settlement history

loading…
-

Strikes / bans / forfeitures

loading…
-

Model usage (RPM)

loading…
-

Node throughput

loading…
-

Console output

loading…
-

Inference history

loading...
+

Account

loading…
+

Tracker hive

loading…
+

Nodes & coverage

loading…
+

Strikes / bans / forfeitures

loading…
+

Model usage (RPM)

loading…
+

Call wall

loading...
+
+

Chat / inference

+
+
+
+ + +
+
+ +
+ +
+
+
+
+
+
select a model to start
+
no messages yet
+
+
+
+
+

Request history

login required
+

Node pending payouts

admin login required
+

Settlement history

admin login required
+

All accounts (admin)

+

Client balances

admin login required
+

Console output

admin login required
diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 5d42ba9..18b4d71 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -142,18 +142,18 @@ DEFAULT_CALLER_CREDIT_USDT = 1.0 DEFAULT_DEVNET_TOPUP_USDT = 1.0 -def _model_aliases(model: str | None) -> set[str]: - """Return stable lookup aliases for a model repo or display name.""" - if not model: - return set() - normalized = model.strip() - if not normalized: - return set() - aliases = {normalized} - short = normalized.rsplit("/", 1)[-1] - aliases.add(short) - lowered = short.lower() - aliases.add(lowered) +def _model_aliases(model: str | None) -> set[str]: + """Return stable lookup aliases for a model repo or display name.""" + if not model: + return set() + normalized = model.strip() + if not normalized: + return set() + aliases = {normalized} + short = normalized.rsplit("/", 1)[-1] + aliases.add(short) + lowered = short.lower() + aliases.add(lowered) if lowered.endswith("-instruct"): aliases.add(lowered.removesuffix("-instruct")) return aliases @@ -542,8 +542,8 @@ class _NodeEntry: "model_tokens_per_sec", "pending_directives", "last_heartbeat", "tracker_mode", "relay_addr", "cert_fingerprint", "peer_id", - # heartbeat stats (reported by node, cumulative) - "total_requests", "failed_requests", "queue_depth", "proxy_inflight", "uptime_seconds", + # heartbeat stats (reported by node, cumulative) + "total_requests", "failed_requests", "queue_depth", "proxy_inflight", "uptime_seconds", "status", # "ready" | "loading" "heartbeats_expected", "heartbeats_received", # dynamic reassignment queued by the tracker @@ -604,18 +604,18 @@ class _NodeEntry: self.peer_id = peer_id self.pending_directives: list[dict] = [] self.last_heartbeat: float = time.monotonic() - self.total_requests: int = 0 - self.failed_requests: int = 0 - self.queue_depth: int = 0 - self.proxy_inflight: int = 0 - self.uptime_seconds: float = 0.0 + self.total_requests: int = 0 + self.failed_requests: int = 0 + self.queue_depth: int = 0 + self.proxy_inflight: int = 0 + self.uptime_seconds: float = 0.0 self.status: str = "ready" self.heartbeats_expected: int = 0 self.heartbeats_received: int = 0 self.pending_new_assignment: dict | None = None -def _effective_throughput(node: "_NodeEntry", model: str | None = None) -> float: +def _effective_throughput(node: "_NodeEntry", model: str | None = None) -> float: """Effective tokens/s accounting for current queue depth.""" observed = None if model: @@ -626,22 +626,22 @@ def _effective_throughput(node: "_NodeEntry", model: str | None = None) -> float if observed is not None: break base = observed if observed is not None and observed > 0 else node.benchmark_tokens_per_sec - return base / (_effective_queue_depth(node) + 1) - - -def _effective_queue_depth(node: "_NodeEntry") -> int: - """Best current load estimate: heartbeat queue or tracker-routed in-flight.""" - return max(node.queue_depth, node.proxy_inflight) - - -def _record_proxy_inflight( - server: "_TrackerHTTPServer", - nodes: list["_NodeEntry"], - delta: int, -) -> None: - with server.lock: - for node in nodes: - node.proxy_inflight = max(0, node.proxy_inflight + delta) + return base / (_effective_queue_depth(node) + 1) + + +def _effective_queue_depth(node: "_NodeEntry") -> int: + """Best current load estimate: heartbeat queue or tracker-routed in-flight.""" + return max(node.queue_depth, node.proxy_inflight) + + +def _record_proxy_inflight( + server: "_TrackerHTTPServer", + nodes: list["_NodeEntry"], + delta: int, +) -> None: + with server.lock: + for node in nodes: + node.proxy_inflight = max(0, node.proxy_inflight + delta) def _reputation_multiplier(node: "_NodeEntry", contracts: Any | None) -> float: @@ -707,12 +707,12 @@ def _node_route_summary(nodes: list["_NodeEntry"]) -> list[dict]: "model": node.model, "hf_repo": node.hf_repo, "endpoint": node.endpoint, - "shard": f"{node.shard_start}-{node.shard_end}", - "num_layers": node.num_layers, - "queue_depth": _effective_queue_depth(node), - "heartbeat_queue_depth": node.queue_depth, - "proxy_inflight": node.proxy_inflight, - } + "shard": f"{node.shard_start}-{node.shard_end}", + "num_layers": node.num_layers, + "queue_depth": _effective_queue_depth(node), + "heartbeat_queue_depth": node.queue_depth, + "proxy_inflight": node.proxy_inflight, + } for node in nodes ] @@ -749,11 +749,11 @@ def _coverage_percentage( return round((covered / required_layers) * 100, 2) -def _served_model_copies( - nodes: list[_NodeEntry], - required_start: int, - required_end: int, -) -> float: +def _served_model_copies( + nodes: list[_NodeEntry], + required_start: int, + required_end: int, +) -> float: required_layers = required_end - required_start + 1 if required_layers <= 0: return 0.0 @@ -772,60 +772,60 @@ def _served_model_copies( return 0.0 complete_copies = min(layer_counts) - residual_layers = sum(1 for count in layer_counts if count > complete_copies) - return round(complete_copies + (residual_layers / required_layers), 2) - - -def _model_health_summary( - server: "_TrackerHTTPServer", - model: str | None, - hf_repo: str | None = None, -) -> dict: - lookup = hf_repo or model - if not lookup: - return {} - - resolved_name, preset = _resolve_model_preset(server.model_presets, lookup) - if preset is not None: - required_start, required_end = _preset_layer_bounds(preset) - model_nodes = [ - node for node in server.registry.values() - if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] - ] - model_id = resolved_name or lookup - else: - model_nodes = [ - node for node in server.registry.values() - if _node_matches_model(node, lookup) - and node.shard_start is not None - and node.shard_end is not None - and node.num_layers is not None - ] - if not model_nodes: - return { - "model": lookup, - "served_model_copies": 0.0, - "coverage_percentage": 0.0, - "node_count": 0, - } - required_start = 0 - required_end = max(node.num_layers for node in model_nodes if node.num_layers is not None) - 1 - model_id = hf_repo or model or lookup - - return { - "model": model_id, - "required_start": required_start, - "required_end": required_end, - "served_model_copies": _served_model_copies(model_nodes, required_start, required_end), - "coverage_percentage": _coverage_percentage(model_nodes, required_start, required_end), - "node_count": len(model_nodes), - } - - -def _preset_layer_bounds(preset: dict) -> tuple[int, int]: - start = int(preset.get("layers_start", 0)) - if "layers_end" in preset: - return start, int(preset["layers_end"]) + residual_layers = sum(1 for count in layer_counts if count > complete_copies) + return round(complete_copies + (residual_layers / required_layers), 2) + + +def _model_health_summary( + server: "_TrackerHTTPServer", + model: str | None, + hf_repo: str | None = None, +) -> dict: + lookup = hf_repo or model + if not lookup: + return {} + + resolved_name, preset = _resolve_model_preset(server.model_presets, lookup) + if preset is not None: + required_start, required_end = _preset_layer_bounds(preset) + model_nodes = [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] + ] + model_id = resolved_name or lookup + else: + model_nodes = [ + node for node in server.registry.values() + if _node_matches_model(node, lookup) + and node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + ] + if not model_nodes: + return { + "model": lookup, + "served_model_copies": 0.0, + "coverage_percentage": 0.0, + "node_count": 0, + } + required_start = 0 + required_end = max(node.num_layers for node in model_nodes if node.num_layers is not None) - 1 + model_id = hf_repo or model or lookup + + return { + "model": model_id, + "required_start": required_start, + "required_end": required_end, + "served_model_copies": _served_model_copies(model_nodes, required_start, required_end), + "coverage_percentage": _coverage_percentage(model_nodes, required_start, required_end), + "node_count": len(model_nodes), + } + + +def _preset_layer_bounds(preset: dict) -> tuple[int, int]: + start = int(preset.get("layers_start", 0)) + if "layers_end" in preset: + return start, int(preset["layers_end"]) return start, start + int(preset["total_layers"]) - 1 @@ -986,13 +986,13 @@ def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict: return { "node_id": node.node_id, "endpoint": node.endpoint, - "alive": alive, - "last_seen_seconds_ago": round(age, 1), - "status": node.status, - "queue_depth": _effective_queue_depth(node), - "heartbeat_queue_depth": node.queue_depth, - "proxy_inflight": node.proxy_inflight, - "total_requests": node.total_requests, + "alive": alive, + "last_seen_seconds_ago": round(age, 1), + "status": node.status, + "queue_depth": _effective_queue_depth(node), + "heartbeat_queue_depth": node.queue_depth, + "proxy_inflight": node.proxy_inflight, + "total_requests": node.total_requests, "heartbeat_success_rate": hb_rate, "inference_success_rate": inf_rate, "capacity": _node_capacity_summary(node), @@ -1357,32 +1357,32 @@ def _drop_directive(node: _NodeEntry, model: str, start: int, end: int, quantiza } -def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]: - now = time.monotonic() - expired_ids = [ - node_id for node_id, entry in server.registry.items() - if (now - entry.last_heartbeat) > server.heartbeat_timeout - ] - expired_entries: list[tuple[str, _NodeEntry]] = [] - for node_id in expired_ids: - entry = server.registry.pop(node_id) - expired_entries.append((node_id, entry)) - if expired_ids: - _rebalance_all_locked(server) - for node_id, entry in expired_entries: - _tracker_log( - server, - "warn", - "node expired", - node_id=node_id, - endpoint=entry.endpoint, - model=entry.model, - hf_repo=entry.hf_repo, - shard=f"{entry.shard_start}-{entry.shard_end}", - heartbeat_timeout_seconds=server.heartbeat_timeout, - model_health=_model_health_summary(server, entry.model, entry.hf_repo), - ) - return expired_ids +def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]: + now = time.monotonic() + expired_ids = [ + node_id for node_id, entry in server.registry.items() + if (now - entry.last_heartbeat) > server.heartbeat_timeout + ] + expired_entries: list[tuple[str, _NodeEntry]] = [] + for node_id in expired_ids: + entry = server.registry.pop(node_id) + expired_entries.append((node_id, entry)) + if expired_ids: + _rebalance_all_locked(server) + for node_id, entry in expired_entries: + _tracker_log( + server, + "warn", + "node expired", + node_id=node_id, + endpoint=entry.endpoint, + model=entry.model, + hf_repo=entry.hf_repo, + shard=f"{entry.shard_start}-{entry.shard_end}", + heartbeat_timeout_seconds=server.heartbeat_timeout, + model_health=_model_health_summary(server, entry.model, entry.hf_repo), + ) + return expired_ids def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None: @@ -1745,7 +1745,7 @@ def _registration_ban_error(contracts: Any | None, wallet_address: str | None) - return None -def _tracker_log(server: "_TrackerHTTPServer", level: str, message: str, **fields: Any) -> None: +def _tracker_log(server: "_TrackerHTTPServer", level: str, message: str, **fields: Any) -> None: event = { "ts": time.time(), "level": level, @@ -1759,40 +1759,40 @@ def _tracker_log(server: "_TrackerHTTPServer", level: str, message: str, **field with server.console_lock: server.console_events.append(event) extras = " ".join(f"{key}={value}" for key, value in event["fields"].items()) - suffix = f" {extras}" if extras else "" - print(f"[tracker] {level}: {message}{suffix}", flush=True) - - -def _tracker_log_proxy_progress( - server: "_TrackerHTTPServer", - *, - request_id: str, - model: str, - route_model: str, - tokens: int, - started: float, - route_nodes: list["_NodeEntry"], - stream: bool = True, - relay: bool = False, -) -> None: - elapsed = time.monotonic() - started - _tracker_log( - server, - "info", - "proxy progress", - request_id=request_id, - model=model, - route_model=route_model, - stream=stream, - relay=relay or None, - tokens=tokens, - elapsed_seconds=round(elapsed, 4), - tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, - route=_node_route_summary(route_nodes), - ) - - -def _node_id_for_registration( + suffix = f" {extras}" if extras else "" + print(f"[tracker] {level}: {message}{suffix}", flush=True) + + +def _tracker_log_proxy_progress( + server: "_TrackerHTTPServer", + *, + request_id: str, + model: str, + route_model: str, + tokens: int, + started: float, + route_nodes: list["_NodeEntry"], + stream: bool = True, + relay: bool = False, +) -> None: + elapsed = time.monotonic() - started + _tracker_log( + server, + "info", + "proxy progress", + request_id=request_id, + model=model, + route_model=route_model, + stream=stream, + relay=relay or None, + tokens=tokens, + elapsed_seconds=round(elapsed, 4), + tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, + route=_node_route_summary(route_nodes), + ) + + +def _node_id_for_registration( endpoint: str, model: str, wallet_address: str | None, @@ -2273,8 +2273,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): result[model] = server.stats.get_node_model_stats(node.node_id, model) return result - def model_supply_for(node: _NodeEntry) -> dict: - return _model_health_summary(server, node.model, node.hf_repo) + def model_supply_for(node: _NodeEntry) -> dict: + return _model_health_summary(server, node.model, node.hf_repo) self._send_json(200, { "relay_url": server.relay_url, @@ -2537,23 +2537,23 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): # Strip the first-shard node we're about to proxy to — it's already handling the request. downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")] downstream_urls = json.dumps(downstream_hops) - route_debug = " -> ".join( - f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]" - for n in route_nodes - ) - inflight_nodes = route_nodes or [node] - inflight_recorded = True - _record_proxy_inflight(server, inflight_nodes, 1) - - def finish_proxy_inflight() -> None: - nonlocal inflight_recorded - if inflight_recorded: - _record_proxy_inflight(server, inflight_nodes, -1) - inflight_recorded = False - - _tracker_log( - server, - "info", + route_debug = " -> ".join( + f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]" + for n in route_nodes + ) + inflight_nodes = route_nodes or [node] + inflight_recorded = True + _record_proxy_inflight(server, inflight_nodes, 1) + + def finish_proxy_inflight() -> None: + nonlocal inflight_recorded + if inflight_recorded: + _record_proxy_inflight(server, inflight_nodes, -1) + inflight_recorded = False + + _tracker_log( + server, + "info", "proxy route selected", request_id=request_id, model=model, @@ -2604,15 +2604,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if first is not None and first.get("stream"): # Streamed response (US-036): forward SSE chunks as they arrive # and run the same token accounting as the direct stream path. - self._stream_relayed_frames( - first, frames, started, - model, route_model, route_nodes, api_key, node_work, - request_body=body, - ) - finish_proxy_inflight() - return - if first is not None: - elapsed = time.monotonic() - started + self._stream_relayed_frames( + first, frames, started, + model, route_model, route_nodes, api_key, node_work, + request_body=body, + ) + finish_proxy_inflight() + return + if first is not None: + elapsed = time.monotonic() - started self._send_relayed_response(first) if int(first.get("status", 503)) < 400: body_text = first.get("body") or "" @@ -2635,12 +2635,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, route=_node_route_summary(route_nodes), ) - self._bill_completed( - api_key, model, tokens, node_work, - input_tokens=in_tokens, output_tokens=out_tokens, - ) - finish_proxy_inflight() - return + self._bill_completed( + api_key, model, tokens, node_work, + input_tokens=in_tokens, output_tokens=out_tokens, + ) + finish_proxy_inflight() + return _tracker_log( server, "warn", @@ -2661,12 +2661,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(err_body))) self.end_headers() - try: - self.wfile.write(err_body) - except BrokenPipeError: - pass - finish_proxy_inflight() - return + try: + self.wfile.write(err_body) + except BrokenPipeError: + pass + finish_proxy_inflight() + return except Exception as exc: if node.relay_addr: _tracker_log( @@ -2687,13 +2687,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): target_url=target_url, error=repr(exc), ) - self._send_json(503, {"error": { - "message": f"upstream node unreachable: {exc}", - "type": "service_unavailable", - "code": "upstream_error", - }}) - finish_proxy_inflight() - return + self._send_json(503, {"error": { + "message": f"upstream node unreachable: {exc}", + "type": "service_unavailable", + "code": "upstream_error", + }}) + finish_proxy_inflight() + return with upstream: content_type = upstream.headers.get("Content-Type", "application/json") @@ -2784,13 +2784,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.send_header("Content-Type", content_type) self.send_header("Content-Length", str(len(resp_body))) self.end_headers() - try: - self.wfile.write(resp_body) - except BrokenPipeError: - pass - finish_proxy_inflight() - - def _record_observed_throughput( + try: + self.wfile.write(resp_body) + except BrokenPipeError: + pass + finish_proxy_inflight() + + def _record_observed_throughput( self, requested_model: str, route_model: str, @@ -3175,50 +3175,50 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): with server.lock: self._purge_expired_nodes() # Dedup: if this endpoint is already registered, remove the old entry first. - stale_ids = [ - eid for eid, e in server.registry.items() - if e.endpoint == entry.endpoint.rstrip("/") - ] - stale_entries: list[tuple[str, _NodeEntry]] = [] - for eid in stale_ids: - old = server.registry.pop(eid) - stale_entries.append((eid, old)) - server.registry[node_id] = entry - if entry.managed_assignment: - if entry.hf_repo: - _rebalance_hf_model_locked(server, entry.hf_repo) - else: + stale_ids = [ + eid for eid, e in server.registry.items() + if e.endpoint == entry.endpoint.rstrip("/") + ] + stale_entries: list[tuple[str, _NodeEntry]] = [] + for eid in stale_ids: + old = server.registry.pop(eid) + stale_entries.append((eid, old)) + server.registry[node_id] = entry + if entry.managed_assignment: + if entry.hf_repo: + _rebalance_hf_model_locked(server, entry.hf_repo) + else: _rebalance_model_locked(server, model) assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None - if assignment_directive is not None: - entry.pending_directives.clear() - - model_health = _model_health_summary(server, entry.model, entry.hf_repo) - for eid, old in stale_entries: - _tracker_log( - server, - "info", - "node re-registered", - old_node_id=eid, - node_id=node_id, - endpoint=old.endpoint, - old_model=old.model, - old_hf_repo=old.hf_repo, - old_model_health=_model_health_summary(server, old.model, old.hf_repo), - model_health=model_health, - ) - _tracker_log( - server, - "info", - "node registered", - node_id=node_id, + if assignment_directive is not None: + entry.pending_directives.clear() + + model_health = _model_health_summary(server, entry.model, entry.hf_repo) + for eid, old in stale_entries: + _tracker_log( + server, + "info", + "node re-registered", + old_node_id=eid, + node_id=node_id, + endpoint=old.endpoint, + old_model=old.model, + old_hf_repo=old.hf_repo, + old_model_health=_model_health_summary(server, old.model, old.hf_repo), + model_health=model_health, + ) + _tracker_log( + server, + "info", + "node registered", + node_id=node_id, endpoint=entry.endpoint, model=entry.model, - hf_repo=entry.hf_repo, - shard=f"{entry.shard_start}-{entry.shard_end}", - tracker_mode=entry.tracker_mode, - model_health=model_health, - ) + hf_repo=entry.hf_repo, + shard=f"{entry.shard_start}-{entry.shard_end}", + tracker_mode=entry.tracker_mode, + model_health=model_health, + ) shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded" repo_info = f" [{hf_repo}]" if hf_repo else "" @@ -3560,7 +3560,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): return keys = server.accounts.keys_for(account["account_id"]) # type: ignore[union-attr] balances = {} - usage: dict = {"requests": 0, "total_tokens": 0, "total_cost": 0.0, "recent": []} + usage: dict = {"requests": 0, "total_tokens": 0, "total_cost": 0.0, "recent": [], "records": []} if server.billing is not None: balances = {key: server.billing.get_client_balance(key) for key in keys} usage = server.billing.usage_for(keys) diff --git a/tests/test_dashboard.py b/tests/test_dashboard.py index 52fec68..4a6cabd 100644 --- a/tests/test_dashboard.py +++ b/tests/test_dashboard.py @@ -12,7 +12,9 @@ from meshnet_tracker.server import TrackerServer PANELS = [ "Tracker hive", "Nodes & coverage", "Client balances", "Node pending payouts", "Settlement history", - "Strikes / bans / forfeitures", "Model usage", "Node throughput", + "Strikes / bans / forfeitures", "Model usage", "Call wall", + "Request history", "node throughput", + "Chat / inference", "Console output", ] diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index a1ff4f8..59a6e6a 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -566,6 +566,78 @@ def test_download_shard_prefers_tracker_model_source_over_huggingface( assert hf_calls == [] +def test_download_shard_prefers_tracker_full_model_source_over_huggingface( + tmp_path, + monkeypatch, +): + """A tracker-advertised full snapshot is sufficient on its own — HF is never contacted.""" + contents = { + "config.json": b"{}", + "weights-a.safetensors": b"tracker-a", + "weights-b.safetensors": b"tracker-b", + } + + class FakeFileResponse: + def __init__(self, payload: bytes): + self._payload = io.BytesIO(payload) + self._length = len(payload) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def getheader(self, name: str): + if name == "Content-Length": + return str(self._length) + if name == "Content-Type": + return "application/octet-stream" + return None + + def read(self, size: int = -1) -> bytes: + return self._payload.read(size) + + def fake_urlopen(url, *args, **kwargs): + query = urllib.parse.parse_qs(urllib.parse.urlparse(url).query) + rel = query.get("file", [None])[0] + assert rel in contents, f"unexpected per-file request: {url}" + return FakeFileResponse(contents[rel]) + + monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen) + hf_calls = [] + + def fake_snapshot_download(**kwargs): + hf_calls.append(kwargs) + raise AssertionError("HuggingFace should not be contacted when tracker full_files are available") + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=fake_snapshot_download), + ) + + shard_dir = download_shard( + "tiny-llama", + 0, + 3, + cache_dir=tmp_path / "cache", + hf_repo="org/tiny-llama-shards", + model_sources=[{ + "type": "tracker-full", + "url": "http://tracker/v1/model-files/download?model=tiny-llama&full=1", + "files": ["config.json", "weights-a.safetensors", "weights-b.safetensors"], + "full_files": ["config.json", "weights-a.safetensors", "weights-b.safetensors"], + }], + progress=False, + ) + + assert (shard_dir / "config.json").read_text() == "{}" + assert (shard_dir / "weights-a.safetensors").read_text() == "tracker-a" + assert (shard_dir / "weights-b.safetensors").read_text() == "tracker-b" + assert hf_calls == [] + + def test_download_shard_falls_back_to_huggingface_when_tracker_source_fails( tmp_path, monkeypatch, diff --git a/tests/test_real_model_backend.py b/tests/test_real_model_backend.py index de3b740..c03050e 100644 --- a/tests/test_real_model_backend.py +++ b/tests/test_real_model_backend.py @@ -11,8 +11,12 @@ import pytest from meshnet_node.model_backend import ( InsufficientVRAMError, + PartialModelLoadUnsupported, TensorPayload, + TorchModelShard, _call_layer, + _load_partial_model_from_snapshot, + _should_partial_materialize_shard, _decoder_attention_mask, _int_tensor_header, build_quantization_config, @@ -334,6 +338,295 @@ def test_call_layer_passes_rotary_position_embeddings(): ) == "hidden" +def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path): + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text("{}") + (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}') + + assert _should_partial_materialize_shard( + str(snapshot_dir), + 4, + 7, + total_layers_hint=40, + uses_quantized_weights=False, + ) is True + assert _should_partial_materialize_shard( + str(snapshot_dir), + 0, + 39, + total_layers_hint=40, + uses_quantized_weights=False, + ) is False + assert _should_partial_materialize_shard( + str(snapshot_dir), + 4, + 7, + total_layers_hint=40, + uses_quantized_weights=True, + ) is False + assert _should_partial_materialize_shard( + "repo/model", + 4, + 7, + total_layers_hint=40, + uses_quantized_weights=False, + ) is False + + +def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path): + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text("{}") + (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ + "weight_map": { + "model.embed_tokens.weight": "shard-1.safetensors", + "model.layers.0.self_attn.q_proj.weight": "shard-1.safetensors", + "model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors", + "model.layers.2.self_attn.q_proj.weight": "shard-3.safetensors", + "model.norm.weight": "shard-3.safetensors", + "lm_head.weight": "shard-3.safetensors", + } + })) + for rel in ("shard-1.safetensors", "shard-2.safetensors", "shard-3.safetensors"): + (snapshot_dir / rel).write_bytes(b"stub") + + class FakeModule: + def __init__(self, name): + self.name = name + self.to_calls = [] + + def to(self, device): + self.to_calls.append(device) + return self + + class FakeModel: + def __init__(self): + self.model = types.SimpleNamespace( + embed_tokens=FakeModule("embed"), + layers=[FakeModule("layer0"), FakeModule("layer1"), FakeModule("layer2")], + rotary_emb=FakeModule("rotary"), + norm=FakeModule("norm"), + ) + self.lm_head = FakeModule("lm_head") + self.tie_weights_called = 0 + + def tie_weights(self): + self.tie_weights_called += 1 + + class AutoConfigStub: + @staticmethod + def from_pretrained(model_id): + assert model_id == str(snapshot_dir) + return types.SimpleNamespace(num_hidden_layers=3) + + class AutoModelStub: + @staticmethod + def from_config(cfg, torch_dtype=None): + assert cfg.num_hidden_layers == 3 + assert torch_dtype == "bf16" + return FakeModel() + + class EmptyWeights: + def __init__(self): + self.entered = 0 + self.exited = 0 + + def __call__(self): + return self + + def __enter__(self): + self.entered += 1 + return None + + def __exit__(self, exc_type, exc, tb): + self.exited += 1 + return False + + init_empty_weights = EmptyWeights() + set_calls = [] + + def fake_set_tensor(module, tensor_name, device, value=None, dtype=None): + set_calls.append((tensor_name, device, value, dtype)) + + tensors = { + "shard-1.safetensors": { + "model.embed_tokens.weight": "embed", + "model.layers.0.self_attn.q_proj.weight": "layer0", + }, + "shard-2.safetensors": { + "model.layers.1.self_attn.q_proj.weight": "layer1", + }, + "shard-3.safetensors": { + "model.layers.2.self_attn.q_proj.weight": "layer2", + "model.norm.weight": "norm", + "lm_head.weight": "lm_head", + }, + } + + class FakeSafeOpen: + def __init__(self, filename, framework, device): + assert framework == "pt" + assert device == "cpu" + self.filename = Path(filename).name + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get_tensor(self, tensor_name): + return tensors[self.filename][tensor_name] + + model = _load_partial_model_from_snapshot( + AutoConfigStub, + AutoModelStub, + types.SimpleNamespace(), + str(snapshot_dir), + 1, + 1, + "bf16", + "cpu:0", + init_empty_weights_fn=init_empty_weights, + set_tensor_fn=fake_set_tensor, + safe_open_fn=FakeSafeOpen, + ) + + assert init_empty_weights.entered == 1 + assert init_empty_weights.exited == 1 + assert model.tie_weights_called == 1 + assert [call[0] for call in set_calls] == ["model.layers.1.self_attn.q_proj.weight"] + assert model.model.layers[1].to_calls == ["cpu:0"] + assert model.model.layers[0].to_calls == [] + assert model.model.layers[2].to_calls == [] + assert model.model.embed_tokens.to_calls == [] + assert model.model.norm.to_calls == [] + assert model.lm_head.to_calls == [] + assert model.model.rotary_emb.to_calls == ["cpu:0"] + + +def test_partial_snapshot_loader_requires_known_layer_count(tmp_path): + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text("{}") + (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ + "weight_map": {"model.layers.0.self_attn.q_proj.weight": "shard.safetensors"} + })) + (snapshot_dir / "shard.safetensors").write_bytes(b"stub") + + class AutoConfigStub: + @staticmethod + def from_pretrained(model_id): + return types.SimpleNamespace() + + class AutoModelStub: + @staticmethod + def from_config(cfg, torch_dtype=None): + raise AssertionError("from_config should not run without a known layer count") + + class UnusedContext: + def __enter__(self): + return None + + def __exit__(self, exc_type, exc, tb): + return False + + with pytest.raises(PartialModelLoadUnsupported, match="num_hidden_layers"): + _load_partial_model_from_snapshot( + AutoConfigStub, + AutoModelStub, + types.SimpleNamespace(), + str(snapshot_dir), + 0, + 0, + "bf16", + "cpu:0", + init_empty_weights_fn=lambda: UnusedContext(), + set_tensor_fn=lambda *args, **kwargs: None, + safe_open_fn=lambda *args, **kwargs: None, + ) + + +def test_torch_model_shard_prefers_partial_loader_for_local_snapshot(tmp_path, monkeypatch): + import meshnet_node.model_backend as backend + + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text("{}") + (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}') + + class FakeModel: + def __init__(self): + self.model = types.SimpleNamespace( + layers=[object(), object(), object()], + embed_tokens=object(), + ) + self.config = types.SimpleNamespace(hidden_size=8) + self.eval_called = 0 + + def eval(self): + self.eval_called += 1 + + fake_model = FakeModel() + partial_calls = [] + + class AutoConfigStub: + @staticmethod + def from_pretrained(model_id, cache_dir=None): + return types.SimpleNamespace(num_hidden_layers=3, text_config=types.SimpleNamespace(dtype="torch.bfloat16")) + + class AutoModelStub: + @staticmethod + def from_pretrained(*args, **kwargs): + raise AssertionError("full model load should not run for partial local shards") + + class AutoTokenizerStub: + @staticmethod + def from_pretrained(model_id, cache_dir=None): + assert model_id == str(snapshot_dir) + return types.SimpleNamespace() + + monkeypatch.setitem( + sys.modules, + "torch", + types.SimpleNamespace( + cuda=types.SimpleNamespace(is_available=lambda: False), + device=lambda value: value, + bfloat16="bf16", + ), + ) + monkeypatch.setitem( + sys.modules, + "transformers", + types.SimpleNamespace( + AutoConfig=AutoConfigStub, + AutoModelForCausalLM=AutoModelStub, + AutoTokenizer=AutoTokenizerStub, + ), + ) + monkeypatch.setattr( + backend, + "_load_partial_model_from_snapshot", + lambda *args, **kwargs: partial_calls.append((args, kwargs)) or fake_model, + ) + + shard = TorchModelShard( + "repo/model", + 1, + 1, + quantization="auto", + cache_dir=snapshot_dir, + ) + + assert len(partial_calls) == 1 + assert shard.model is fake_model + assert fake_model.eval_called == 1 + assert shard.total_layers == 3 + assert shard.is_head is False + assert shard.is_tail is False + + @pytest.mark.integration def test_two_node_gpt2_completion_is_deterministic(): if os.environ.get("CI"):