dual billing; tracker to node model sharing

This commit is contained in:
Dobromir Popov
2026-07-06 17:31:11 +03:00
parent ccb69c41e3
commit 2e696be80f
14 changed files with 1092 additions and 41 deletions

2
.gitignore vendored
View File

@@ -17,3 +17,5 @@ dist/
.env.*
!.env.example
!.env.testnet
.rocm-local/*
billing.sqlite

Binary file not shown.

View File

@@ -1,6 +1,6 @@
# US-044 — Tracker as model-file source; nodes download only their shard
Status: planned
Status: in progress
Priority: High (blocks multi-machine big-model serving; pairs with US-042)
Stage: Designed (grill remaining decisions before build)
@@ -56,10 +56,30 @@ What exists already (build on it, don't duplicate):
## Acceptance criteria
- Two-machine test: machine A (tracker + node, holds full snapshot) serves
- [x] Tracker can be started with `--models-dir PATH` / `MESHNET_MODELS_DIR`
and advertises a local model-file source in assignment responses when it has
a matching HF snapshot.
- [x] Tracker serves a tar stream containing only the safetensors files selected
for the assigned layer range plus config/tokenizer/index metadata.
- [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.
- [x] Real PyTorch model startup can use tracker `full_url` sources to fetch
the full local snapshot over LAN before `from_pretrained`, so local-network
testing no longer has to pull from HuggingFace first.
- [ ] Two-machine test: machine A (tracker + node, holds full snapshot) serves
layers 0k; machine B joins with no model and receives **only** the files
for its assigned range from A — nothing fetched from HF
- Machine B's resident memory scales with its shard size, not model size
- Checksums verified end-to-end; corrupted transfer falls back cleanly
- Single-node/full-model flows unchanged
- `python -m pytest` passes from repo root
- [ ] Machine B's resident memory scales with its shard size, not model size
- [ ] Checksums verified end-to-end; corrupted transfer falls back cleanly
- [x] Single-node/full-model flows unchanged
- [ ] `python -m pytest` passes from repo root
## Implementation notes
- 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.

View File

@@ -0,0 +1,60 @@
# US-045 — Dual-rate billing: separate input and output token prices
Status: in progress
Priority: High (billing correctness before friends test; providers all price this way)
Stage: Designed
## Context
Today the ledger has one `price_per_1k_tokens` per model, and the two proxy
paths don't even agree on what they count:
- **Non-streaming** bills `usage.total_tokens` (prompt + completion) at the
blended rate (`_billable_non_stream_tokens`).
- **Streaming** bills `min(observed output deltas, reported total)` — output
only in practice (`_billable_stream_tokens`).
- The HF pricing refresher (issue 23) averages a provider's input/output
rates 50/50 (`blended_price_per_1k_tokens`), which misprices asymmetric
models — e.g. Qwen3.6-35B-A3B on deepinfra is $0.15/1M in, $0.95/1M out.
Decision (user, 2026-07-06): charge **both** input and output tokens, at
**two separate rates**, same as other providers.
## Design
1. **`BillingLedger`** stores `{model: (input_per_1k, output_per_1k)}`.
- `set_prices(model, input_per_1k, output_per_1k)` (new);
`set_price(model, p)` keeps working and sets both.
- `prices_for(model) -> (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

View File

@@ -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)

View File

@@ -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)

View File

@@ -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()

View File

@@ -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)

View File

@@ -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,

View File

@@ -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:

View File

@@ -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()

View File

@@ -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]

View File

@@ -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

View File

@@ -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)