dual billing; tracker to node model sharing
This commit is contained in:
2
.gitignore
vendored
2
.gitignore
vendored
@@ -17,3 +17,5 @@ dist/
|
||||
.env.*
|
||||
!.env.example
|
||||
!.env.testnet
|
||||
.rocm-local/*
|
||||
billing.sqlite
|
||||
|
||||
BIN
billing.sqlite
BIN
billing.sqlite
Binary file not shown.
@@ -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 0–k; 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.
|
||||
|
||||
60
docs/issues/45-dual-rate-billing.md
Normal file
60
docs/issues/45-dual-rate-billing.md
Normal 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
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
175
packages/node/meshnet_node/safetensors_selection.py
Normal file
175
packages/node/meshnet_node/safetensors_selection.py
Normal 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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -220,6 +220,12 @@ def main() -> None:
|
||||
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",
|
||||
@@ -277,6 +283,7 @@ def main() -> None:
|
||||
or (DEFAULT_HF_PRICING_LOG_DB_PATH if args.enable_hf_pricing else None)
|
||||
),
|
||||
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)
|
||||
|
||||
172
packages/tracker/meshnet_tracker/model_files.py
Normal file
172
packages/tracker/meshnet_tracker/model_files.py
Normal 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()
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
86
tests/test_safetensors_selection.py
Normal file
86
tests/test_safetensors_selection.py
Normal 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)
|
||||
Reference in New Issue
Block a user