From 2e696be80f856fa04d46073b16701b28e35b85cc Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Mon, 6 Jul 2026 17:31:11 +0300 Subject: [PATCH] dual billing; tracker to node model sharing --- .gitignore | 2 + billing.sqlite | Bin 12288 -> 12288 bytes ...4-tracker-shard-source-partial-download.md | 32 ++- docs/issues/45-dual-rate-billing.md | 60 ++++++ packages/node/meshnet_node/downloader.py | 133 +++++++++++- packages/node/meshnet_node/model_backend.py | 13 +- .../meshnet_node/safetensors_selection.py | 175 ++++++++++++++++ packages/node/meshnet_node/startup.py | 40 ++++ packages/tracker/meshnet_tracker/billing.py | 56 ++++- packages/tracker/meshnet_tracker/cli.py | 23 ++- .../tracker/meshnet_tracker/model_files.py | 172 ++++++++++++++++ packages/tracker/meshnet_tracker/server.py | 193 +++++++++++++++++- tests/test_node_startup.py | 148 ++++++++++++++ tests/test_safetensors_selection.py | 86 ++++++++ 14 files changed, 1092 insertions(+), 41 deletions(-) create mode 100644 docs/issues/45-dual-rate-billing.md create mode 100644 packages/node/meshnet_node/safetensors_selection.py create mode 100644 packages/tracker/meshnet_tracker/model_files.py create mode 100644 tests/test_safetensors_selection.py diff --git a/.gitignore b/.gitignore index d0dd023..940dcd4 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,5 @@ dist/ .env.* !.env.example !.env.testnet +.rocm-local/* +billing.sqlite diff --git a/billing.sqlite b/billing.sqlite index 543290aea6d3366cdaf5336e33822df5835f4f15..3e98f578a7913beba0c3cdf9b0b25a7dc446be69 100644 GIT binary patch delta 271 zcmZojXh@hK&B!xR#+i|4W5N=CX&y!f27X=M^*lyg2RZk16tc6h9%gA~c4A^;oLKlU zuCbnnHM&`uJ-H|~C9_1=#5ma~HPOT{InmO>#5lz=$-vCg%pfh*Bqh&jmYX%YMV@`KqMrI@L!ryO zGCaH=8Tf1YSMn|7ljSSrJInWz&yzop_apyRe$|bI7kPEGnWGtrw3%nKtlncr0DcEc Az5oCK delta 102 zcmV-s0Ga=QV1Qtd8vzTE976#Mv0$VR8Vdpf01qt (input, output)` (new); `price_for(model)` returns + the blended average for back-compat (estimators/logs). + - `charge_request(...)` gains keyword `input_tokens`/`output_tokens`; when + provided, `cost = in_rate·in/1k + out_rate·out/1k` and the event records + the split. Without them, legacy behavior (blended × total) — old events + and gossip replicas replay unchanged (`cost` stays the applied field). +2. **Token counting** (`server.py`): + - Non-stream: prefer `usage.prompt_tokens`/`completion_tokens`; fall back + to content estimates (`_estimate_prompt_tokens`, observed completion), + capped by `max_tokens` bounds as today. + - Stream (direct + relay): output = observed deltas as today; input = + `usage.prompt_tokens` when a usage chunk appears, else the prompt + estimate from the request body. `_stream_line_tokens` returns the parsed + usage triple instead of just the total. +3. **Presets**: `input_price_per_1k_tokens` / `output_price_per_1k_tokens` + (dual keys win; `price_per_1k_tokens` alone still means "both rates"). + Qwen3.6-35B-A3B: input 0.00012, output 0.00076 (80% of deepinfra). +4. **HF refresher**: applies 80% of each side separately via `set_prices` + (all alias keys); change log keeps recording the blended pair for history + continuity. +5. **Spend cap** (`--max-charge-per-request`): estimate = + `in_rate·prompt_estimate + out_rate·completion_limit`. + +## Acceptance criteria + +- Streamed and non-streamed requests for the same exchange bill the same + split (input charged in both) +- A model with asymmetric provider rates bills input and output differently; + `usage_for` / billing events expose the split +- Old persisted billing events replay byte-identically (balances unchanged) +- HF refresh sets both rates from the marketplace row, not the average +- Spend cap uses the dual rates +- `python -m pytest` passes from repo root diff --git a/packages/node/meshnet_node/downloader.py b/packages/node/meshnet_node/downloader.py index 3c8c0f4..b7758a7 100644 --- a/packages/node/meshnet_node/downloader.py +++ b/packages/node/meshnet_node/downloader.py @@ -13,6 +13,7 @@ import tarfile import tempfile import urllib.parse import urllib.request +from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from typing import Any @@ -105,6 +106,113 @@ def _download_shard_from_peer( return False +def _download_model_source( + source: dict, + shard_dir: Path, + timeout: float, +) -> Path | None: + url = source.get("url") + if not isinstance(url, str) or not url: + endpoint = source.get("endpoint") + if not isinstance(endpoint, str): + return None + url = f"{endpoint.rstrip('/')}/v1/model-files/download" + shard_dir.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="meshnet-model-source-", dir=shard_dir.parent) as tmp: + tmp_root = Path(tmp) + archive_path = tmp_root / "model-files.tar" + extract_dir = tmp_root / "extract" + extract_dir.mkdir() + try: + with urllib.request.urlopen(url, timeout=timeout) as resp, archive_path.open("wb") as out: + while True: + chunk = resp.read(1024 * 1024) + if not chunk: + break + out.write(chunk) + _safe_extract_shard(archive_path, extract_dir) + shutil.move(str(extract_dir), str(shard_dir)) + return shard_dir + except Exception: + return None + + +def _download_huggingface_subset( + hf_repo: str, + cache_dir: Path, + shard_dir: Path, + allow_patterns: list[str] | None, +) -> Path: + from huggingface_hub import snapshot_download # type: ignore[import] + + kwargs = { + "repo_id": hf_repo, + "cache_dir": str(cache_dir), + "local_dir": str(shard_dir), + } + if allow_patterns: + kwargs["allow_patterns"] = allow_patterns + try: + return Path(snapshot_download(**kwargs)) + except TypeError: + kwargs.pop("allow_patterns", None) + return Path(snapshot_download(**kwargs)) + + +def _download_from_fastest_source( + *, + model_sources: list[dict], + hf_repo: str, + cache_dir: Path, + shard_dir: Path, + progress: bool, + timeout: float, +) -> tuple[str, Path] | None: + shard_dir.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="meshnet-race-", dir=shard_dir.parent) as tmp: + tmp_root = Path(tmp) + jobs: dict[Any, tuple[str, Path]] = {} + pool = ThreadPoolExecutor(max_workers=min(4, len(model_sources) + 1)) + try: + for index, source in enumerate(model_sources): + label = str(source.get("type") or "model-source") + candidate = tmp_root / f"source-{index}" + jobs[pool.submit(_download_model_source, source, candidate, timeout)] = (label, candidate) + allow_patterns = _allow_patterns_from_sources(model_sources) + hf_candidate = tmp_root / "huggingface" + jobs[pool.submit(_download_huggingface_subset, hf_repo, cache_dir, hf_candidate, allow_patterns)] = ( + "HuggingFace", + hf_candidate, + ) + for future in as_completed(jobs): + label, candidate = jobs[future] + try: + result = future.result() + except Exception: + continue + if result is None: + continue + if shard_dir.exists(): + shutil.rmtree(shard_dir) + shutil.move(str(candidate), str(shard_dir)) + if progress: + print(f" download source: {label}", flush=True) + pool.shutdown(wait=False, cancel_futures=True) + return label, shard_dir + finally: + pool.shutdown(wait=False, cancel_futures=True) + return None + + +def _allow_patterns_from_sources(model_sources: list[dict]) -> list[str] | None: + patterns: set[str] = set() + for source in model_sources: + for rel in source.get("files") or []: + if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts: + patterns.add(rel) + return sorted(patterns) if patterns else None + + def download_shard( model: str, shard_start: int, @@ -113,6 +221,7 @@ def download_shard( hf_repo: str | None = None, progress: bool = True, peers: list[dict] | None = None, + model_sources: list[dict] | None = None, peer_timeout: float = _PEER_TIMEOUT_SECONDS, ) -> Path: """Ensure the shard is present in *cache_dir* and return its local path. @@ -157,18 +266,26 @@ def download_shard( print(f" [stub] shard already cached at {shard_dir}", flush=True) return shard_dir - from huggingface_hub import snapshot_download # type: ignore[import] - if progress: print( f" Downloading layers {shard_start}-{shard_end} from {hf_repo} ...", flush=True, ) + if model_sources: + if progress: + print(" Racing tracker model source against HuggingFace ...", flush=True) + raced = _download_from_fastest_source( + model_sources=model_sources, + hf_repo=hf_repo, + cache_dir=cache_dir, + shard_dir=shard_dir, + progress=progress, + timeout=peer_timeout, + ) + if raced is not None: + return raced[1] + + if progress: print(" download source: HuggingFace", flush=True) - local_dir = snapshot_download( - repo_id=hf_repo, - cache_dir=str(cache_dir), - local_dir=str(shard_dir), - ) - return Path(local_dir) + return _download_huggingface_subset(hf_repo, cache_dir, shard_dir, None) diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index 35b98d1..ece10d8 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -85,24 +85,25 @@ class TorchModelShard: self.torch = torch self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id quant_config, dtype, uses_quantized_weights = _model_load_plan( AutoConfig, - model_id, + load_source, quantization, torch, - cache_dir, + 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 else None, + "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( - model_id, + load_source, **load_kwargs, ) if not uses_quantized_weights: @@ -117,8 +118,8 @@ class TorchModelShard: self.model.eval() self.tokenizer = AutoTokenizer.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 load_source == model_id else None, ) self.layers = _model_layers(self.model) self.total_layers = len(self.layers) diff --git a/packages/node/meshnet_node/safetensors_selection.py b/packages/node/meshnet_node/safetensors_selection.py new file mode 100644 index 0000000..532b6eb --- /dev/null +++ b/packages/node/meshnet_node/safetensors_selection.py @@ -0,0 +1,175 @@ +"""Layer-aware SafeTensors snapshot file selection.""" + +from __future__ import annotations + +import json +import re +from pathlib import Path +from typing import Any + +INDEX_FILENAME = "model.safetensors.index.json" + +_LAYER_RE = re.compile( + r"(?:^|\.)" + r"(?:model\.layers|layers|h|blocks|decoder\.layers|encoder\.layers)" + r"\.(\d+)(?:\.|$)" +) + +_METADATA_FILENAMES = { + INDEX_FILENAME, + "config.json", + "generation_config.json", + "preprocessor_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer.model", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "added_tokens.json", +} + +_METADATA_PREFIXES = ("config.", "tokenizer.", "tokenizer_", "vocab.") + +_HEAD_MARKERS = ( + "embed", + "embedding", + "embed_tokens", + "wte", + "wpe", +) + +_TAIL_EXACT = { + "lm_head.weight", + "lm_head.bias", + "model.norm.weight", + "model.norm.bias", + "transformer.ln_f.weight", + "transformer.ln_f.bias", + "decoder.final_layer_norm.weight", + "decoder.final_layer_norm.bias", +} + +_TAIL_MARKERS = ( + ".lm_head.", + ".norm.", + ".ln_f.", + ".final_layer_norm.", +) + + +def select_safetensors_files_for_layers( + model_dir: str | Path, + start_layer: int, + end_layer: int, + *, + total_layers: int | None = None, +) -> list[str]: + """Return relative snapshot files needed for an inclusive layer range. + + The returned list always includes root-level config/tokenizer metadata and + the SafeTensors index. Weight shard files are included only when at least one + tensor in the index belongs to the assigned layer range, or when the tensor + is needed by the head/tail shard. + """ + if start_layer < 0: + raise ValueError("start_layer must be non-negative") + if end_layer < start_layer: + raise ValueError("end_layer must be greater than or equal to start_layer") + + root = Path(model_dir) + index_path = root / INDEX_FILENAME + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise FileNotFoundError(f"missing SafeTensors index: {index_path}") from exc + + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"{INDEX_FILENAME} must contain a weight_map object") + + inferred_total_layers = total_layers if total_layers is not None else _read_total_layers(root) + selected = _metadata_files(root) + + 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, inferred_total_layers): + selected.add(_normalise_relative_file(rel_file)) + + return sorted(selected) + + +def _tensor_belongs_to_range( + tensor_name: str, + start_layer: int, + end_layer: int, + total_layers: int | None, +) -> bool: + layer = _layer_index(tensor_name) + if layer is not None: + return start_layer <= layer <= end_layer + + if start_layer == 0 and _is_head_tensor(tensor_name): + return True + + if total_layers is not None and end_layer >= total_layers - 1 and _is_tail_tensor(tensor_name): + return True + + return False + + +def _layer_index(tensor_name: str) -> int | None: + match = _LAYER_RE.search(tensor_name) + if match is None: + return None + return int(match.group(1)) + + +def _is_head_tensor(tensor_name: str) -> bool: + lowered = tensor_name.lower() + return any(marker in lowered for marker in _HEAD_MARKERS) + + +def _is_tail_tensor(tensor_name: str) -> bool: + lowered = tensor_name.lower() + return lowered in _TAIL_EXACT or any(marker in lowered for marker in _TAIL_MARKERS) + + +def _metadata_files(root: Path) -> set[str]: + files = {INDEX_FILENAME} + for path in root.iterdir(): + if not path.is_file(): + continue + name = path.name + if name in _METADATA_FILENAMES or name.startswith(_METADATA_PREFIXES): + files.add(name) + return files + + +def _read_total_layers(root: Path) -> int | None: + config_path = root / "config.json" + if not config_path.exists(): + return None + config = json.loads(config_path.read_text(encoding="utf-8")) + return _layers_from_config(config) + + +def _layers_from_config(config: dict[str, Any]) -> int | None: + for key in ("num_hidden_layers", "num_layers", "n_layer", "n_layers"): + value = config.get(key) + if isinstance(value, int) and value > 0: + return value + + text_config = config.get("text_config") + if isinstance(text_config, dict): + return _layers_from_config(text_config) + + return None + + +def _normalise_relative_file(rel_file: str) -> str: + path = Path(rel_file) + if path.is_absolute() or ".." in path.parts: + raise ValueError(f"unsafe relative file in {INDEX_FILENAME}: {rel_file}") + return path.as_posix() diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 1b0f45e..5700cf3 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -34,6 +34,21 @@ def _memory_budget(device: str, vram_mb: int, ram_mb: int, shared_vram_mb: int = return max(0, ram_mb), "RAM" +def _full_model_sources(model_sources: list[dict]) -> list[dict]: + """Use tracker full-snapshot URLs for real HF model loading.""" + full_sources: list[dict] = [] + for source in model_sources: + full_url = source.get("full_url") + if isinstance(full_url, str) and full_url: + full_sources.append({ + **source, + "url": full_url, + "files": [], + "type": f"{source.get('type') or 'model-source'}-full", + }) + return full_sources + + def _hardware_label(device: str, gpu_name: str | None = None) -> str: if device == "cuda": return "CUDA" @@ -443,6 +458,16 @@ def run_startup( if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"): shard_start = net_asgn["shard_start"] shard_end = net_asgn["shard_end"] + full_sources = _full_model_sources(net_asgn.get("model_sources", [])) + if full_sources: + cache_dir = download_shard( + model_id.split("/")[-1], + shard_start, + shard_end, + cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards", + hf_repo=model_id, + model_sources=full_sources, + ) print( f" Tracker found uncovered shard: " f"layers {shard_start}–{shard_end} (of {detected})", @@ -550,12 +575,24 @@ def run_startup( assigned_shard_start: int = net_assignment["shard_start"] assigned_shard_end: int = net_assignment["shard_end"] assigned_num_layers: int = net_assignment["num_layers"] + assigned_model_sources: list[dict] = net_assignment.get("model_sources", []) print( f" Assigned: {assigned_hf_repo} " f"layers {assigned_shard_start}–{assigned_shard_end} " f"(of {assigned_num_layers})", flush=True, ) + full_sources = _full_model_sources(assigned_model_sources) + if full_sources: + print("Downloading assigned model snapshot...", flush=True) + cache_dir = download_shard( + assigned_hf_repo.split("/")[-1], + assigned_shard_start, + assigned_shard_end, + cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards", + hf_repo=assigned_hf_repo, + model_sources=full_sources, + ) print("Loading real PyTorch model shard...", flush=True) node = TorchNodeServer( host=host, @@ -647,6 +684,7 @@ def run_startup( assigned_model: str = assignment.get("model", model) hf_repo: str | None = assignment.get("hf_repo") peers: list[dict] = assignment.get("peers", []) + model_sources: list[dict] = assignment.get("model_sources", []) print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True) # 4. Download shard @@ -658,6 +696,8 @@ def run_startup( dl_kwargs["hf_repo"] = hf_repo if peers: dl_kwargs["peers"] = peers + 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) print(f" Cached at: {shard_path}", flush=True) diff --git a/packages/tracker/meshnet_tracker/billing.py b/packages/tracker/meshnet_tracker/billing.py index 0d99d48..ddef7d3 100644 --- a/packages/tracker/meshnet_tracker/billing.py +++ b/packages/tracker/meshnet_tracker/billing.py @@ -24,6 +24,18 @@ DEFAULT_BILLING_DB_PATH = "billing.sqlite" NODE_REVENUE_SHARE = 0.90 # nodes 90% / protocol cut 10% (ADR-0015) +def _normalize_rates(value: "float | tuple[float, float] | dict") -> tuple[float, float]: + """Coerce a price spec into an (input_per_1k, output_per_1k) pair.""" + if isinstance(value, dict): + base = value.get("price") + inp = value.get("input", base) + out = value.get("output", base) + return (float(inp), float(out)) + if isinstance(value, (tuple, list)): + return (float(value[0]), float(value[1])) + return (float(value), float(value)) + + class BillingLedger: """Thread-safe USDT ledger with SQLite persistence and event replication.""" @@ -33,13 +45,17 @@ class BillingLedger: self, db_path: str | None = None, *, - prices: dict[str, float] | None = None, + prices: dict[str, float | tuple[float, float]] | None = None, default_price_per_1k: float = DEFAULT_PRICE_PER_1K_TOKENS, starting_credit: float = DEFAULT_STARTING_CREDIT, node_share: float = NODE_REVENUE_SHARE, ) -> None: self._db_path = db_path - self._prices = dict(prices) if prices else {} + # US-045: per-model (input_per_1k, output_per_1k). A bare float in + # ``prices`` sets both rates (legacy single-rate models). + self._prices: dict[str, tuple[float, float]] = { + model: _normalize_rates(value) for model, value in (prices or {}).items() + } self._default_price_per_1k = default_price_per_1k self._starting_credit = starting_credit self._node_share = node_share @@ -62,11 +78,23 @@ class BillingLedger: # ---- pricing ---- def price_for(self, model: str) -> float: - return self._prices.get(model, self._default_price_per_1k) + """Blended (average) per-1k rate — kept for estimators and history logs.""" + rates = self._prices.get(model) + if rates is None: + return self._default_price_per_1k + return (rates[0] + rates[1]) / 2.0 + + def prices_for(self, model: str) -> tuple[float, float]: + """(input_per_1k, output_per_1k) for a model (US-045).""" + return self._prices.get(model, (self._default_price_per_1k, self._default_price_per_1k)) def set_price(self, model: str, price_per_1k: float) -> None: + """Legacy single-rate setter — applies the same rate to input and output.""" + self.set_prices(model, price_per_1k, price_per_1k) + + def set_prices(self, model: str, input_per_1k: float, output_per_1k: float) -> None: with self._lock: - self._prices[model] = price_per_1k + self._prices[model] = (float(input_per_1k), float(output_per_1k)) # ---- local operations (create + apply + log an event) ---- @@ -130,9 +158,17 @@ class BillingLedger: model: str, total_tokens: int, node_work: list[tuple[str | None, int]], + *, + input_tokens: int | None = None, + output_tokens: int | None = None, ) -> dict: """Debit the client and split the fee 90/10. + With ``input_tokens``/``output_tokens`` (US-045) the cost is + ``input·in_rate + output·out_rate``; without them, legacy behavior — + ``total_tokens`` at the blended rate. Replayed/gossiped events apply + their recorded ``cost``, so old events are unaffected either way. + ``node_work`` is ``[(wallet_address | None, work_units), ...]`` for the nodes that served the request. Work units of nodes without a wallet accrue to the protocol cut — there is nowhere to pay them out. @@ -140,7 +176,14 @@ class BillingLedger: request is then rejected by ``has_funds`` (post-pay drift, standard metered-billing behavior). """ - cost = self.price_for(model) * max(0, total_tokens) / 1000.0 + if input_tokens is not None or output_tokens is not None: + in_rate, out_rate = self.prices_for(model) + in_tokens = max(0, input_tokens or 0) + out_tokens = max(0, output_tokens or 0) + cost = (in_tokens * in_rate + out_tokens * out_rate) / 1000.0 + total_tokens = in_tokens + out_tokens + else: + cost = self.price_for(model) * max(0, total_tokens) / 1000.0 total_work = sum(max(0, w) for _, w in node_work) node_pool = cost * self._node_share shares: dict[str, float] = {} @@ -165,6 +208,9 @@ class BillingLedger: "api_key": api_key, "model": model, "total_tokens": total_tokens, + **({"input_tokens": max(0, input_tokens or 0), + "output_tokens": max(0, output_tokens or 0)} + if (input_tokens is not None or output_tokens is not None) else {}), "cost": cost, "shares": shares, "protocol_amount": protocol_amount, diff --git a/packages/tracker/meshnet_tracker/cli.py b/packages/tracker/meshnet_tracker/cli.py index 88f8289..e14ba30 100644 --- a/packages/tracker/meshnet_tracker/cli.py +++ b/packages/tracker/meshnet_tracker/cli.py @@ -214,12 +214,18 @@ def main() -> None: "enables GET /v1/pricing/hf/history)" ), ) - common.add_argument( - "--hf-pricing-refresh-interval", - type=float, - default=86400.0, - help="Seconds between dynamic pricing refresh passes (default: daily)", - ) + common.add_argument( + "--hf-pricing-refresh-interval", + type=float, + default=86400.0, + help="Seconds between dynamic pricing refresh passes (default: daily)", + ) + common.add_argument( + "--models-dir", + default=None, + metavar="PATH", + help="Local HuggingFace snapshot root advertised as tracker model-file source (default: MESHNET_MODELS_DIR)", + ) parser = argparse.ArgumentParser( prog="meshnet-tracker", @@ -276,8 +282,9 @@ def main() -> None: args.hf_pricing_log_db or (DEFAULT_HF_PRICING_LOG_DB_PATH if args.enable_hf_pricing else None) ), - hf_pricing_refresh_interval=args.hf_pricing_refresh_interval, - ) + hf_pricing_refresh_interval=args.hf_pricing_refresh_interval, + models_dir=args.models_dir, + ) port = server.start() print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True) try: diff --git a/packages/tracker/meshnet_tracker/model_files.py b/packages/tracker/meshnet_tracker/model_files.py new file mode 100644 index 0000000..4e97a49 --- /dev/null +++ b/packages/tracker/meshnet_tracker/model_files.py @@ -0,0 +1,172 @@ +"""Helpers for serving layer-scoped model files from tracker-local snapshots.""" + +from __future__ import annotations + +import json +import re +from pathlib import Path +from typing import Any + +INDEX_FILENAME = "model.safetensors.index.json" + +_LAYER_RE = re.compile( + r"(?:^|\.)" + r"(?:model\.layers|layers|h|blocks|decoder\.layers|encoder\.layers)" + r"\.(\d+)(?:\.|$)" +) + +_METADATA_FILENAMES = { + INDEX_FILENAME, + "config.json", + "generation_config.json", + "preprocessor_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer.model", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "added_tokens.json", +} + +_METADATA_PREFIXES = ("config.", "tokenizer.", "tokenizer_", "vocab.") + +_HEAD_MARKERS = ("embed", "embedding", "embed_tokens", "wte", "wpe") + +_TAIL_EXACT = { + "lm_head.weight", + "lm_head.bias", + "model.norm.weight", + "model.norm.bias", + "transformer.ln_f.weight", + "transformer.ln_f.bias", + "decoder.final_layer_norm.weight", + "decoder.final_layer_norm.bias", +} + +_TAIL_MARKERS = (".lm_head.", ".norm.", ".ln_f.", ".final_layer_norm.") + + +def snapshot_dir_for_repo(models_dir: Path, repo_id: str) -> Path | None: + """Return the most likely local HF snapshot directory for *repo_id*.""" + candidates = [ + models_dir / repo_id, + models_dir / repo_id.replace("/", "--"), + models_dir / f"models--{repo_id.replace('/', '--')}", + ] + for candidate in candidates: + if (candidate / "snapshots").is_dir(): + snapshots = sorted(p for p in (candidate / "snapshots").iterdir() if p.is_dir()) + if snapshots: + return snapshots[-1] + if candidate.is_dir(): + return candidate + return None + + +def files_for_layer_range(snapshot_dir: Path, shard_start: int, shard_end: int) -> list[str]: + """Select files needed to load a conservative safetensors shard subset.""" + return select_safetensors_files_for_layers(snapshot_dir, shard_start, shard_end) + + +def select_safetensors_files_for_layers( + model_dir: str | Path, + start_layer: int, + end_layer: int, + *, + total_layers: int | None = None, +) -> list[str]: + if start_layer < 0: + raise ValueError("start_layer must be non-negative") + if end_layer < start_layer: + raise ValueError("end_layer must be greater than or equal to start_layer") + + root = Path(model_dir) + index_path = root / INDEX_FILENAME + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except FileNotFoundError: + return sorted(p.name for p in root.glob("*.safetensors") if p.is_file()) + + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"{INDEX_FILENAME} must contain a weight_map object") + + inferred_total_layers = total_layers if total_layers is not None else _read_total_layers(root) + selected = _metadata_files(root) + + 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, inferred_total_layers): + selected.add(_normalise_relative_file(rel_file)) + + return sorted(rel for rel in selected if (root / rel).is_file()) + + +def _tensor_belongs_to_range( + tensor_name: str, + start_layer: int, + end_layer: int, + total_layers: int | None, +) -> bool: + layer = _layer_index(tensor_name) + if layer is not None: + return start_layer <= layer <= end_layer + if start_layer == 0 and _is_head_tensor(tensor_name): + return True + if total_layers is not None and end_layer >= total_layers - 1 and _is_tail_tensor(tensor_name): + return True + return False + + +def _layer_index(tensor_name: str) -> int | None: + match = _LAYER_RE.search(tensor_name) + return int(match.group(1)) if match else None + + +def _is_head_tensor(tensor_name: str) -> bool: + lowered = tensor_name.lower() + return any(marker in lowered for marker in _HEAD_MARKERS) + + +def _is_tail_tensor(tensor_name: str) -> bool: + lowered = tensor_name.lower() + return lowered in _TAIL_EXACT or any(marker in lowered for marker in _TAIL_MARKERS) + + +def _metadata_files(root: Path) -> set[str]: + files = {INDEX_FILENAME} + for path in root.iterdir(): + if not path.is_file(): + continue + name = path.name + if name in _METADATA_FILENAMES or name.startswith(_METADATA_PREFIXES): + files.add(name) + return files + + +def _read_total_layers(root: Path) -> int | None: + config_path = root / "config.json" + if not config_path.exists(): + return None + config = json.loads(config_path.read_text(encoding="utf-8")) + return _layers_from_config(config) + + +def _layers_from_config(config: dict[str, Any]) -> int | None: + for key in ("num_hidden_layers", "num_layers", "n_layer", "n_layers"): + value = config.get(key) + if isinstance(value, int) and value > 0: + return value + text_config = config.get("text_config") + if isinstance(text_config, dict): + return _layers_from_config(text_config) + return None + + +def _normalise_relative_file(rel_file: str) -> str: + path = Path(rel_file) + if path.is_absolute() or ".." in path.parts: + raise ValueError(f"unsafe relative file in {INDEX_FILENAME}: {rel_file}") + return path.as_posix() diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 633d2b0..32779f7 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -28,12 +28,14 @@ import json import os import socketserver import sqlite3 +import tarfile import threading import time import urllib.parse import urllib.request import uuid from importlib.resources import files +from pathlib import Path from typing import Any from .accounts import DEFAULT_ACCOUNTS_DB_PATH, AccountStore @@ -43,6 +45,7 @@ from .billing import DEFAULT_BILLING_DB_PATH, BillingLedger from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH, HfPricingLog, refresh_preset_price from .gossip import NodeGossip +from .model_files import files_for_layer_range, snapshot_dir_for_repo from .raft import RaftNode @@ -118,6 +121,14 @@ def _clone_model_presets(presets: dict[str, dict]) -> dict[str, dict]: DEFAULT_VRAM_BYTES = 8 * 1024 * 1024 * 1024 DEFAULT_RAM_BYTES = 16 * 1024 * 1024 * 1024 + + +def _snapshot_regular_files(snapshot_dir: Path) -> list[str]: + return sorted( + path.relative_to(snapshot_dir).as_posix() + for path in snapshot_dir.rglob("*") + if path.is_file() + ) DEFAULT_QUANTIZATIONS = ["bfloat16"] DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0 # US-039/US-040 — single source of truth for the credit defaults (referenced by @@ -1008,8 +1019,25 @@ def _relay_http_request( } -def _stream_line_tokens(line: bytes) -> tuple[int, int | None]: - """Token accounting for one SSE line: (observed delta, reported total or None).""" +def _usage_split(payload: dict) -> dict | None: + """Parse a usage block into {"prompt", "completion", "total"} (ints or None).""" + usage = payload.get("usage") + if not isinstance(usage, dict): + return None + + def _num(key: str) -> int | None: + value = usage.get(key) + return int(value) if isinstance(value, (int, float)) else None + + return { + "prompt": _num("prompt_tokens"), + "completion": _num("completion_tokens"), + "total": _usage_total_tokens(payload), + } + + +def _stream_line_tokens(line: bytes) -> tuple[int, dict | None]: + """Token accounting for one SSE line: (observed output delta, usage split or None).""" if not line.startswith(b"data:"): return 0, None payload = line[5:].strip() @@ -1019,7 +1047,49 @@ def _stream_line_tokens(line: bytes) -> tuple[int, int | None]: chunk_payload = json.loads(payload) except json.JSONDecodeError: return 1, None - return _observed_stream_tokens(chunk_payload), _usage_total_tokens(chunk_payload) + return _observed_stream_tokens(chunk_payload), _usage_split(chunk_payload) + + +def _stream_billable_split( + observed_output: int, usage: dict | None, request_body: dict +) -> tuple[int, int]: + """(input_tokens, output_tokens) for a streamed response (US-045). + + Output: observed deltas, capped by reported completion count when the + stream carried a usage chunk. Input: reported prompt count, else the + prompt estimate from the request body. + """ + prompt = (usage or {}).get("prompt") + completion = (usage or {}).get("completion") + total = (usage or {}).get("total") + if prompt is None: + prompt = _estimate_prompt_tokens(request_body) or 0 + if completion is None and total is not None: + completion = max(0, total - prompt) + return max(0, prompt), _billable_stream_tokens(observed_output, completion) + + +def _billable_non_stream_split(payload: dict, request_body: dict) -> tuple[int, int]: + """(input_tokens, output_tokens) for a buffered response (US-045). + + Prefers the response usage block; falls back to content estimates. + Completion stays capped by the request's max-tokens bound, as before. + """ + usage = _usage_split(payload) + prompt = (usage or {}).get("prompt") + completion = (usage or {}).get("completion") + if prompt is None: + prompt = _estimate_prompt_tokens(request_body) or 0 + if completion is None: + total = (usage or {}).get("total") + if total is not None: + completion = max(0, total - prompt) + else: + completion = _observed_non_stream_completion_tokens(payload) + limit = _requested_completion_token_limit(request_body) + if limit is not None: + completion = min(completion, limit) + return max(0, prompt), max(0, completion) def _find_pinned_route( @@ -1598,6 +1668,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): toploc_calibration_gate_min_hardware_profiles: int = 1, toploc_backend: Any | None = None, hf_pricing_log: "HfPricingLog | None" = None, + models_dir: Path | None = None, ) -> None: super().__init__(addr, handler) self.registry = registry @@ -1628,6 +1699,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): self.toploc_calibration_gate_min_hardware_profiles = toploc_calibration_gate_min_hardware_profiles self.toploc_backend = toploc_backend self.hf_pricing_log: HfPricingLog | None = hf_pricing_log + self.models_dir = models_dir class _TrackerHandler(http.server.BaseHTTPRequestHandler): @@ -1800,6 +1872,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._handle_network_map() elif parsed.path == "/v1/models": self._handle_models() + elif parsed.path == "/v1/model-files/download": + self._handle_model_files_download(parsed) elif parsed.path.startswith("/v1/coverage/"): model = urllib.parse.unquote(parsed.path.removeprefix("/v1/coverage/")) self._handle_coverage(model) @@ -2482,8 +2556,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): model: str, total_tokens: int, node_work: list[tuple[str | None, int]], + *, + input_tokens: int | None = None, + output_tokens: int | None = None, ) -> None: - """Charge a completed request against the billing ledger (ADR-0015).""" + """Charge a completed request against the billing ledger (ADR-0015). + + With ``input_tokens``/``output_tokens`` the ledger bills each side at + its own rate (US-045); ``total_tokens`` alone falls back to the + blended rate. + """ server: _TrackerHTTPServer = self.server # type: ignore[assignment] if server.billing is None or api_key is None: return @@ -2501,11 +2583,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): adjusted.append((wallet, work)) node_work = adjusted try: - event = server.billing.charge_request(api_key, model, total_tokens, node_work) + event = server.billing.charge_request( + api_key, model, total_tokens, node_work, + input_tokens=input_tokens, output_tokens=output_tokens, + ) print( f"[tracker] billed api_key=…{api_key[-6:]}: model={model!r} " - f"tokens={total_tokens} cost={event['cost']:.6f} USDT " - f"shares={event['shares']}", + f"tokens={event['total_tokens']} " + f"(in={event.get('input_tokens', '?')} out={event.get('output_tokens', '?')}) " + f"cost={event['cost']:.6f} USDT shares={event['shares']}", flush=True, ) except Exception as exc: @@ -3699,9 +3785,83 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "model": resolved_name, "model_layers_end": required_end, "peers": peers, + "model_sources": self._model_sources( + resolved_name, + preset, + shard_start, + shard_end, + ), **({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}), }) + def _model_sources(self, model: str, preset: dict, shard_start: int, shard_end: int) -> list[dict]: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + hf_repo = preset.get("hf_repo") + if not server.models_dir or not isinstance(hf_repo, str) or not hf_repo: + return [] + snapshot_dir = snapshot_dir_for_repo(server.models_dir, hf_repo) + if snapshot_dir is None: + return [] + files = files_for_layer_range(snapshot_dir, shard_start, shard_end) + if not files: + return [] + host = self.headers.get("Host") or f"{self.server.server_address[0]}:{self.server.server_address[1]}" + base_url = f"http://{host}" + query = urllib.parse.urlencode({ + "model": model, + "shard_start": shard_start, + "shard_end": shard_end, + }) + full_query = urllib.parse.urlencode({"model": model, "full": 1}) + return [{ + "type": "tracker", + "endpoint": base_url, + "url": f"{base_url}/v1/model-files/download?{query}", + "full_url": f"{base_url}/v1/model-files/download?{full_query}", + "files": files, + }] + + def _handle_model_files_download(self, parsed: urllib.parse.ParseResult) -> None: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.models_dir is None: + self._send_json(404, {"error": "tracker model-file source is not enabled"}) + return + params = urllib.parse.parse_qs(parsed.query) + model = params.get("model", [""])[0] + resolved_name, preset = _resolve_model_preset(server.model_presets, model) + if preset is None: + self._send_json(404, {"error": f"unknown model preset: {model!r}"}) + return + hf_repo = preset.get("hf_repo") + if not isinstance(hf_repo, str) or not hf_repo: + self._send_json(404, {"error": f"model preset has no hf_repo: {resolved_name!r}"}) + return + snapshot_dir = snapshot_dir_for_repo(server.models_dir, hf_repo) + if snapshot_dir is None: + self._send_json(404, {"error": f"local snapshot not found for {hf_repo}"}) + return + full_download = params.get("full", ["0"])[0] in {"1", "true", "yes"} + if full_download: + rel_files = _snapshot_regular_files(snapshot_dir) + else: + try: + shard_start = int(params.get("shard_start", [""])[0]) + shard_end = int(params.get("shard_end", [""])[0]) + except ValueError: + self._send_json(400, {"error": "shard_start and shard_end must be integers"}) + return + rel_files = files_for_layer_range(snapshot_dir, shard_start, shard_end) + if not rel_files: + self._send_json(404, {"error": "no local files matched the assigned shard"}) + return + self.send_response(200) + self.send_header("Content-Type", "application/x-tar") + self.send_header("X-Meshnet-Model-Source", "tracker") + self.end_headers() + with tarfile.open(fileobj=self.wfile, mode="w|") as archive: + for rel in rel_files: + archive.add(snapshot_dir / rel, arcname=rel) + def _handle_network_assign(self, parsed: urllib.parse.ParseResult): """Assign a new node to fill the biggest uncovered shard gap across HF-model nodes. @@ -3771,6 +3931,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "gap_found": True, "price_per_token": 0.0, "deployment": _deployment_summary(all_nodes, preset), + "model_sources": self._model_sources( + str(resolved_name), + preset, + shard_start, + shard_end, + ), }) return @@ -3842,7 +4008,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): # Capacity: use the same 80%-of-memory rule as registered node planning. total_l = best_num_layers memory_mb = vram_mb if vram_mb > 0 else ram_mb - _resolved_name, best_preset = _resolve_model_preset(server.model_presets, str(best_repo)) + resolved_name, best_preset = _resolve_model_preset(server.model_presets, str(best_repo)) if memory_mb > 0: max_layers = _max_layers_for_memory(memory_mb, total_l, best_preset) elif device == "cuda" and vram_mb >= 8192: @@ -3855,11 +4021,18 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(200, { "hf_repo": best_repo, + "model": resolved_name, "shard_start": shard_start, "shard_end": shard_end, "num_layers": total_l, "gap_found": gap_found, "price_per_token": 0.0, + "model_sources": self._model_sources( + str(resolved_name or best_repo), + best_preset or {"hf_repo": best_repo}, + shard_start, + shard_end, + ), }) def _handle_route(self, parsed: urllib.parse.ParseResult): @@ -4053,6 +4226,7 @@ class TrackerServer: hf_pricing_log_db: str | None = None, hf_pricing_refresh_interval: float = 86400.0, hf_pricing_fetch_html: Any | None = None, + models_dir: str | Path | None = None, ) -> None: self._host = host self._requested_port = port @@ -4140,6 +4314,8 @@ class TrackerServer: self._enable_hf_pricing = enable_hf_pricing self._hf_pricing_refresh_interval = hf_pricing_refresh_interval self._hf_pricing_fetch_html = hf_pricing_fetch_html + raw_models_dir = models_dir if models_dir is not None else os.environ.get("MESHNET_MODELS_DIR") + self._models_dir = Path(raw_models_dir).expanduser() if raw_models_dir else None self._hf_pricing_stop = threading.Event() self._hf_pricing_thread: threading.Thread | None = None self.port: int | None = None @@ -4178,6 +4354,7 @@ class TrackerServer: toploc_calibration_gate_min_hardware_profiles=self._toploc_calibration_gate_min_hardware_profiles, toploc_backend=self._toploc_backend, hf_pricing_log=self._hf_pricing_log, + models_dir=self._models_dir, ) self.port = self._server.server_address[1] diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 804507b..b6980be 100644 --- a/tests/test_node_startup.py +++ b/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 diff --git a/tests/test_safetensors_selection.py b/tests/test_safetensors_selection.py new file mode 100644 index 0000000..335a7b4 --- /dev/null +++ b/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)