From d151dd5484131ce74e974c2cba3bc9165cee765b Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Mon, 6 Jul 2026 18:36:42 +0300 Subject: [PATCH] tracker download fix --- ...4-tracker-shard-source-partial-download.md | 16 ++++ packages/node/meshnet_node/downloader.py | 52 ++++++++++- .../meshnet_node/safetensors_selection.py | 34 +++++-- .../tracker/meshnet_tracker/hf_pricing.py | 4 + .../meshnet_tracker/model_presets.json | 6 +- packages/tracker/meshnet_tracker/server.py | 90 ++++++++++++------- 6 files changed, 159 insertions(+), 43 deletions(-) diff --git a/docs/issues/44-tracker-shard-source-partial-download.md b/docs/issues/44-tracker-shard-source-partial-download.md index 576d054..a9cca40 100644 --- a/docs/issues/44-tracker-shard-source-partial-download.md +++ b/docs/issues/44-tracker-shard-source-partial-download.md @@ -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. diff --git a/packages/node/meshnet_node/downloader.py b/packages/node/meshnet_node/downloader.py index b7758a7..b68d3bf 100644 --- a/packages/node/meshnet_node/downloader.py +++ b/packages/node/meshnet_node/downloader.py @@ -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) diff --git a/packages/node/meshnet_node/safetensors_selection.py b/packages/node/meshnet_node/safetensors_selection.py index 532b6eb..f9e87d3 100644 --- a/packages/node/meshnet_node/safetensors_selection.py +++ b/packages/node/meshnet_node/safetensors_selection.py @@ -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 diff --git a/packages/tracker/meshnet_tracker/hf_pricing.py b/packages/tracker/meshnet_tracker/hf_pricing.py index 94fc062..075ab85 100644 --- a/packages/tracker/meshnet_tracker/hf_pricing.py +++ b/packages/tracker/meshnet_tracker/hf_pricing.py @@ -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, } diff --git a/packages/tracker/meshnet_tracker/model_presets.json b/packages/tracker/meshnet_tracker/model_presets.json index 05ada0a..40ee0ed 100644 --- a/packages/tracker/meshnet_tracker/model_presets.json +++ b/packages/tracker/meshnet_tracker/model_presets.json @@ -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 } } } \ No newline at end of file diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 32779f7..acd07e4 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -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: