popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

tracker download fix

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-06 18:36:42 +03:00
parent 2e696be80f
commit d151dd5484
6 changed files with 159 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
docs/issues/44-tracker-shard-source-partial-download.md
View File

@@ -38,6 +38,17 @@ What exists already (build on it, don't duplicate):
subset (new) → HF `snapshot_download` with `allow_patterns` for the same
subset (new — stop downloading the whole repo even from HF) → full snapshot
(last resort).
- The `allow_patterns` subset must not depend on the tracker having a local
snapshot: when the tracker has no `--models-dir` match for a repo (or
hasn't cached it yet — the common case for a fresh public tracker),
`model_sources` comes back empty and `download_shard` falls straight to
`_download_huggingface_subset(..., allow_patterns=None)`, i.e. the full
repo. Reported 2026-07-06: a CPU node assigned layers 0–2 of
`unsloth/Qwen3.6-35B-A3B` (42 safetensor shards) sat downloading the
entire model unauthenticated because of this. Fix: fetch
`model.safetensors.index.json` + `config.json` directly from HF (a few
KB) and compute the same layer-scoped file subset client-side, so the
HF-fallback path is filtered even with an empty `model_sources`.
4. **Partial LOAD (the hard half).** Downloading a subset is wasted unless the
node stops instantiating the full model: build the model skeleton on the
`meta` device, materialize only assigned layers (+embeddings/norm/head as
@@ -64,6 +75,11 @@ What exists already (build on it, don't duplicate):
- [x] Node downloader keeps exact-shard peers first, then races tracker model
sources against a HuggingFace `snapshot_download(..., allow_patterns=...)`
subset download, using the first successful source.
- [ ] When no tracker model source is available at all, the HuggingFace
fallback still computes `allow_patterns` from the repo's own
`model.safetensors.index.json` (fetched directly, not via the tracker) —
it never silently downloads the full model just because the tracker has
nothing cached.
- [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.

52
packages/node/meshnet_node/downloader.py
View File

@@ -213,6 +213,47 @@ def _allow_patterns_from_sources(model_sources: list[dict]) -> list[str] | None:
return sorted(patterns) if patterns else None
def _allow_patterns_from_remote_index(
hf_repo: str,
cache_dir: Path,
shard_start: int,
shard_end: int,
) -> list[str] | None:
"""Fetch just the SafeTensors index + config (a few KB) from HF and compute
which weight files the assigned layer range needs, so a HuggingFace fallback
download stays layer-scoped even when the tracker has no model_sources
(e.g. it has no local snapshot for this repo cached yet)."""
try:
from huggingface_hub import hf_hub_download # type: ignore[import]
from .safetensors_selection import (
INDEX_FILENAME,
METADATA_FILENAMES,
layers_from_config_dict,
select_files_for_layers_from_index,
)
index_path = hf_hub_download(repo_id=hf_repo, filename=INDEX_FILENAME, cache_dir=str(cache_dir))
weight_map = json.loads(Path(index_path).read_text(encoding="utf-8")).get("weight_map")
except Exception:
return None
if not isinstance(weight_map, dict):
return None
total_layers: int | None = None
try:
config_path = hf_hub_download(repo_id=hf_repo, filename="config.json", cache_dir=str(cache_dir))
config = json.loads(Path(config_path).read_text(encoding="utf-8"))
total_layers = layers_from_config_dict(config)
except Exception:
pass
selected = select_files_for_layers_from_index(
weight_map, shard_start, shard_end, total_layers=total_layers
)
return sorted(selected | METADATA_FILENAMES)
def download_shard(
model: str,
shard_start: int,
@@ -285,7 +326,14 @@ def download_shard(
if raced is not None:
return raced[1]
allow_patterns = _allow_patterns_from_remote_index(hf_repo, cache_dir, shard_start, shard_end)
if progress:
print(" download source: HuggingFace", flush=True)
if allow_patterns:
print(" download source: HuggingFace (layer-filtered)", flush=True)
else:
print(
" download source: HuggingFace (full snapshot — no SafeTensors index found)",
flush=True,
)
return _download_huggingface_subset(hf_repo, cache_dir, shard_dir, None)
return _download_huggingface_subset(hf_repo, cache_dir, shard_dir, allow_patterns)

34
packages/node/meshnet_node/safetensors_selection.py
View File

@@ -15,7 +15,7 @@ _LAYER_RE = re.compile(
r"\.(\d+)(?:\.|$)"
)
_METADATA_FILENAMES = {
METADATA_FILENAMES = {
INDEX_FILENAME,
"config.json",
"generation_config.json",
@@ -90,14 +90,32 @@ def select_safetensors_files_for_layers(
inferred_total_layers = total_layers if total_layers is not None else _read_total_layers(root)
selected = _metadata_files(root)
selected |= select_files_for_layers_from_index(
weight_map, start_layer, end_layer, total_layers=inferred_total_layers
)
return sorted(selected)
def select_files_for_layers_from_index(
weight_map: dict[str, str],
start_layer: int,
end_layer: int,
*,
total_layers: int | None = None,
) -> set[str]:
"""Pure variant of the weight-file selection: takes an already-parsed
``weight_map`` (no local snapshot directory needed), so callers that only
have the index fetched over the network — not a full local snapshot — can
still compute which shard files they need. Combine the result with
``METADATA_FILENAMES`` for a complete download pattern set.
"""
selected: set[str] = set()
for tensor_name, rel_file in weight_map.items():
if not isinstance(tensor_name, str) or not isinstance(rel_file, str):
continue
if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, inferred_total_layers):
if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, total_layers):
selected.add(_normalise_relative_file(rel_file))
return sorted(selected)
return selected
def _tensor_belongs_to_range(
@@ -142,7 +160,7 @@ def _metadata_files(root: Path) -> set[str]:
if not path.is_file():
continue
name = path.name
if name in _METADATA_FILENAMES or name.startswith(_METADATA_PREFIXES):
if name in METADATA_FILENAMES or name.startswith(_METADATA_PREFIXES):
files.add(name)
return files
@@ -152,10 +170,10 @@ def _read_total_layers(root: Path) -> int | None:
if not config_path.exists():
return None
config = json.loads(config_path.read_text(encoding="utf-8"))
return _layers_from_config(config)
return layers_from_config_dict(config)
def _layers_from_config(config: dict[str, Any]) -> int | None:
def layers_from_config_dict(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:
@@ -163,7 +181,7 @@ def _layers_from_config(config: dict[str, Any]) -> int | None:
text_config = config.get("text_config")
if isinstance(text_config, dict):
return _layers_from_config(text_config)
return layers_from_config_dict(text_config)
return None

4
packages/tracker/meshnet_tracker/hf_pricing.py
View File

@@ -295,6 +295,10 @@ def refresh_preset_price(
"model": model_name,
"old_price_per_1k": current_price,
"new_price_per_1k": new_price,
# US-045: per-side rates (per 1k tokens) so the ledger bills input
# and output at the provider's actual asymmetry, not the average.
"new_input_price_per_1k": round(quote.input_per_1m * price_fraction / 1000.0, 6),
"new_output_price_per_1k": round(quote.output_per_1m * price_fraction / 1000.0, 6),
"source_repo_id": quote.repo_id,
"source_provider": quote.provider,
}

6
packages/tracker/meshnet_tracker/model_presets.json
View File

@@ -55,7 +55,7 @@
"hf_aliases": [
"qwen/qwen3.6-35b-a3b"
],
"hf_verified_match_note": "Verified 2026-07-06: unsloth/Qwen3.6-35B-A3B is a bf16 mirror of Qwen/Qwen3.6-35B-A3B; deepinfra and featherless-ai serve the official weights on the HF inference marketplace, so their rates are a fair comparable. Static price 0.00044 = 80% of deepinfra's blended $0.55/1M ($0.15 in / $0.95 out); the nightly refresher keeps it tracking.",
"hf_verified_match_note": "Verified 2026-07-06: unsloth/Qwen3.6-35B-A3B is a bf16 mirror of Qwen/Qwen3.6-35B-A3B; deepinfra and featherless-ai serve the official weights on the HF inference marketplace, so their rates are a fair comparable. Rates are 80% of deepinfra: input 0.00012/1k ($0.15/1M), output 0.00076/1k ($0.95/1M); price_per_1k_tokens keeps the blended 0.00044 for display/back-compat. The nightly refresher tracks both sides.",
"required_model_bytes": 71903776776,
"download_size_bytes": 71903776776,
"native_quantization": "bfloat16",
@@ -72,7 +72,9 @@
"context_length": 262144,
"native_quantization": "bfloat16",
"download_size_gb": 72
}
},
"input_price_per_1k_tokens": 0.00012,
"output_price_per_1k_tokens": 0.00076
}
}
}

90
packages/tracker/meshnet_tracker/server.py
View File

@@ -2201,8 +2201,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"code": "missing_token_limit",
}})
return
total_token_bound = _request_total_token_upper_bound(body) or token_limit
estimated_charge = server.billing.price_for(model) * total_token_bound / 1000.0
in_rate, out_rate = server.billing.prices_for(model)
prompt_estimate = _estimate_prompt_tokens(body) or 0
estimated_charge = (in_rate * prompt_estimate + out_rate * token_limit) / 1000.0
if estimated_charge > server.max_charge_per_request:
self._send_json(402, {"error": {
"message": (
@@ -2368,6 +2369,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self._stream_relayed_frames(
first, frames, started,
model, route_model, route_nodes, api_key, node_work,
request_body=body,
)
return
if first is not None:
@@ -2376,11 +2378,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if int(first.get("status", 503)) < 400:
body_text = first.get("body") or ""
try:
tokens = _billable_non_stream_tokens(json.loads(body_text), body)
in_tokens, out_tokens = _billable_non_stream_split(json.loads(body_text), body)
except (json.JSONDecodeError, TypeError):
tokens = 0
in_tokens, out_tokens = 0, 0
tokens = in_tokens + out_tokens
self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes)
self._bill_completed(api_key, model, tokens, node_work)
self._bill_completed(
api_key, model, tokens, node_work,
input_tokens=in_tokens, output_tokens=out_tokens,
)
return
print(
f"[tracker] relay proxy failed {request_id}: {node.relay_addr}; "
@@ -2428,7 +2434,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.send_header("Cache-Control", "no-cache")
self.end_headers()
reported_stream_tokens: int | None = None
stream_usage: dict | None = None
observed_stream_tokens = 0
try:
while True:
@@ -2437,22 +2443,25 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
break
self.wfile.write(line)
self.wfile.flush()
observed, reported = _stream_line_tokens(line)
observed, usage = _stream_line_tokens(line)
observed_stream_tokens += observed
if reported is not None:
reported_stream_tokens = reported
if usage is not None:
stream_usage = usage
except BrokenPipeError:
pass
elapsed = time.monotonic() - started
# Bill even on client disconnect — the nodes did the work.
# Observed stream chunks are authoritative for the upper bound;
# upstream usage may only lower that count.
observed_tokens = _billable_stream_tokens(observed_stream_tokens, reported_stream_tokens)
self._record_observed_throughput(model, route_model, observed_tokens, elapsed, route_nodes)
in_tokens, out_tokens = _stream_billable_split(
observed_stream_tokens, stream_usage, body
)
self._record_observed_throughput(
model, route_model, in_tokens + out_tokens, elapsed, route_nodes
)
self._bill_completed(
api_key, model,
observed_tokens,
node_work,
api_key, model, in_tokens + out_tokens, node_work,
input_tokens=in_tokens, output_tokens=out_tokens,
)
else:
# Non-streaming: buffer and relay
@@ -2465,13 +2474,17 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
observed_output = ""
try:
response_payload = json.loads(resp_body)
tokens = _billable_non_stream_tokens(response_payload, body)
in_tokens, out_tokens = _billable_non_stream_split(response_payload, body)
observed_output = _observed_output_from_non_stream_payload(response_payload)
except json.JSONDecodeError:
tokens = 0
in_tokens, out_tokens = 0, 0
tokens = in_tokens + out_tokens
self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes)
self._record_validation_event(request_id, model, body, observed_output, route_nodes)
self._bill_completed(api_key, model, tokens, node_work)
self._bill_completed(
api_key, model, tokens, node_work,
input_tokens=in_tokens, output_tokens=out_tokens,
)
self.send_response(200)
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(resp_body)))
@@ -2607,6 +2620,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
route_nodes: list,
api_key: str | None,
node_work: list,
request_body: dict,
) -> None:
"""Forward a streamed relay response (US-036) to the client as SSE,
billing with the same accounting as the direct stream path."""
@@ -2615,7 +2629,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self.send_header("Content-Type", headers.get("Content-Type", "text/event-stream; charset=utf-8"))
self.send_header("Cache-Control", "no-cache")
self.end_headers()
reported_stream_tokens: int | None = None
stream_usage: dict | None = None
observed_stream_tokens = 0
client_gone = False
for frame in itertools.chain([first], frames):
@@ -2631,14 +2645,21 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
# Keep draining frames — the nodes did the work; bill it.
client_gone = True
for line in data.splitlines():
observed, reported = _stream_line_tokens(line)
observed, usage = _stream_line_tokens(line)
observed_stream_tokens += observed
if reported is not None:
reported_stream_tokens = reported
if usage is not None:
stream_usage = usage
elapsed = time.monotonic() - started
observed_tokens = _billable_stream_tokens(observed_stream_tokens, reported_stream_tokens)
self._record_observed_throughput(model, route_model, observed_tokens, elapsed, route_nodes)
self._bill_completed(api_key, model, observed_tokens, node_work)
in_tokens, out_tokens = _stream_billable_split(
observed_stream_tokens, stream_usage, request_body
)
self._record_observed_throughput(
model, route_model, in_tokens + out_tokens, elapsed, route_nodes
)
self._bill_completed(
api_key, model, in_tokens + out_tokens, node_work,
input_tokens=in_tokens, output_tokens=out_tokens,
)
def _send_relayed_response(self, response: dict) -> None:
status = int(response.get("status", 503))
@@ -4256,12 +4277,17 @@ class TrackerServer:
if db_path is None and enable_billing:
db_path = DEFAULT_BILLING_DB_PATH
if db_path:
preset_prices = {
key: float(preset["price_per_1k_tokens"])
for name, preset in self._model_presets.items()
if isinstance(preset, dict) and "price_per_1k_tokens" in preset
for key in _preset_price_keys(name, preset)
}
preset_prices: dict[str, tuple[float, float]] = {}
for name, preset in self._model_presets.items():
if not isinstance(preset, dict):
continue
base = preset.get("price_per_1k_tokens")
input_rate = preset.get("input_price_per_1k_tokens", base)
output_rate = preset.get("output_price_per_1k_tokens", base)
if input_rate is None or output_rate is None:
continue
for key in _preset_price_keys(name, preset):
preset_prices[key] = (float(input_rate), float(output_rate))
billing = BillingLedger(db_path=db_path, prices=preset_prices)
self._billing: BillingLedger | None = billing
self._billing_gossip_cursor = 0
@@ -4528,8 +4554,10 @@ class TrackerServer:
continue
if result is None:
continue
new_input = result.get("new_input_price_per_1k", result["new_price_per_1k"])
new_output = result.get("new_output_price_per_1k", result["new_price_per_1k"])
for key in _preset_price_keys(name, preset):
billing.set_price(key, result["new_price_per_1k"])
billing.set_prices(key, new_input, new_output)
preset["hf_last_price_per_1k"] = result["new_price_per_1k"]
preset["hf_last_updated"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
if self._hf_pricing_log is not None: