1122 lines
45 KiB
Python
1122 lines
45 KiB
Python
"""Full node startup sequence — self-configuring, non-interactive."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import socket
|
||
import sys
|
||
import threading
|
||
import time
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from .downloader import compute_shard_checksum, download_shard
|
||
from .hardware import detect_hardware, benchmark_throughput_checked
|
||
from .model_catalog import model_metadata_for
|
||
from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
|
||
from .server import StubNodeServer
|
||
from .torch_server import TorchNodeServer
|
||
from .wallet import load_or_create_wallet
|
||
|
||
|
||
_DEFAULT_BYTES_PER_LAYER = 30 * 1024 * 1024
|
||
|
||
|
||
def _downloaded_model_inventory(
|
||
model: str,
|
||
shard_start: int,
|
||
shard_end: int,
|
||
shard_path: Path,
|
||
hf_repo: str | None = None,
|
||
model_sources: list[dict] | None = None,
|
||
) -> list[dict]:
|
||
"""Return a cheap local inventory record without reading model file contents."""
|
||
file_count = 0
|
||
total_bytes = 0
|
||
existing_rel_files: set[str] = set()
|
||
if shard_path.exists():
|
||
for path in shard_path.rglob("*"):
|
||
if not path.is_file():
|
||
continue
|
||
file_count += 1
|
||
try:
|
||
existing_rel_files.add(path.relative_to(shard_path).as_posix())
|
||
except ValueError:
|
||
pass
|
||
try:
|
||
total_bytes += path.stat().st_size
|
||
except OSError:
|
||
pass
|
||
record = {
|
||
"model": model,
|
||
"shard_start": shard_start,
|
||
"shard_end": shard_end,
|
||
"path": str(shard_path),
|
||
"file_count": file_count,
|
||
"total_bytes": total_bytes,
|
||
}
|
||
if hf_repo is not None:
|
||
record["hf_repo"] = hf_repo
|
||
expected_files: set[str] = set()
|
||
file_sizes: dict[str, int] = {}
|
||
for source in model_sources or []:
|
||
for rel in source.get("full_files") or source.get("files") or []:
|
||
if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts:
|
||
expected_files.add(rel)
|
||
sizes = source.get("file_sizes")
|
||
if isinstance(sizes, dict):
|
||
for rel, size in sizes.items():
|
||
if isinstance(rel, str) and isinstance(size, int):
|
||
file_sizes[rel] = size
|
||
if expected_files:
|
||
expected_bytes = sum(file_sizes.get(rel, 0) for rel in expected_files)
|
||
local_expected_files = existing_rel_files & expected_files
|
||
local_expected_bytes = sum(file_sizes.get(rel, 0) for rel in local_expected_files)
|
||
record["expected_file_count"] = len(expected_files)
|
||
record["local_expected_file_count"] = len(local_expected_files)
|
||
record["expected_bytes"] = expected_bytes
|
||
record["local_expected_bytes"] = local_expected_bytes
|
||
record["local_model_percentage"] = (
|
||
round((local_expected_bytes / expected_bytes) * 100, 4)
|
||
if expected_bytes > 0
|
||
else 0.0
|
||
)
|
||
return [record]
|
||
|
||
|
||
def _registration_shard_checksum(model: str, shard_path: Path) -> str | None:
|
||
"""Only checksum tiny stub shards; real model folders are too large to hash at startup."""
|
||
if model != "stub-model":
|
||
return None
|
||
return compute_shard_checksum(shard_path)
|
||
|
||
|
||
def _model_cache_path(model_id: str, cache_dir: Path | None) -> Path | None:
|
||
if cache_dir is None:
|
||
return None
|
||
if (cache_dir / "config.json").exists():
|
||
return cache_dir
|
||
candidate = cache_dir / model_id.split("/")[-1]
|
||
if candidate.exists():
|
||
return candidate
|
||
return None
|
||
|
||
|
||
def _memory_budget(device: str, vram_mb: int, ram_mb: int, shared_vram_mb: int = 0) -> tuple[int, str]:
|
||
"""Return the capacity budget in MB and whether it came from VRAM or RAM."""
|
||
if device == "cuda" and vram_mb > 0:
|
||
if shared_vram_mb > 0:
|
||
return vram_mb + shared_vram_mb, "VRAM + shared RAM"
|
||
return vram_mb, "VRAM"
|
||
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,
|
||
# full_files (when advertised) enables robust per-file download
|
||
# of the whole snapshot; empty list falls back to the tar stream.
|
||
"files": source.get("full_files") or [],
|
||
"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"
|
||
if gpu_name:
|
||
return "CPU (CUDA inactive)"
|
||
return "CPU"
|
||
|
||
|
||
def _positive_int(value: int | str | None, name: str) -> int | None:
|
||
if value is None or value == "":
|
||
return None
|
||
try:
|
||
parsed = int(value)
|
||
except (TypeError, ValueError) as exc:
|
||
raise ValueError(f"{name} must be a positive integer") from exc
|
||
if parsed < 1:
|
||
raise ValueError(f"{name} must be a positive integer")
|
||
return parsed
|
||
|
||
|
||
def _configure_torch_threads(
|
||
torch_threads: int | None = None,
|
||
torch_interop_threads: int | None = None,
|
||
) -> dict[str, int]:
|
||
"""Apply PyTorch CPU thread settings before model load/benchmark."""
|
||
intra_threads = _positive_int(
|
||
torch_threads if torch_threads is not None else os.environ.get("MESHNET_TORCH_THREADS"),
|
||
"--torch-threads",
|
||
)
|
||
interop_threads = _positive_int(
|
||
torch_interop_threads
|
||
if torch_interop_threads is not None
|
||
else os.environ.get("MESHNET_TORCH_INTEROP_THREADS"),
|
||
"--torch-interop-threads",
|
||
)
|
||
|
||
if intra_threads is not None:
|
||
os.environ.setdefault("OMP_NUM_THREADS", str(intra_threads))
|
||
os.environ.setdefault("MKL_NUM_THREADS", str(intra_threads))
|
||
try:
|
||
import torch
|
||
except ModuleNotFoundError:
|
||
return {}
|
||
|
||
if intra_threads is not None:
|
||
torch.set_num_threads(intra_threads)
|
||
if interop_threads is not None:
|
||
torch.set_num_interop_threads(interop_threads)
|
||
|
||
active: dict[str, int] = {}
|
||
try:
|
||
active["torch_threads"] = int(torch.get_num_threads())
|
||
except Exception:
|
||
pass
|
||
try:
|
||
active["torch_interop_threads"] = int(torch.get_num_interop_threads())
|
||
except Exception:
|
||
pass
|
||
return active
|
||
|
||
|
||
def _max_assignable_layers(
|
||
memory_mb: int,
|
||
total_layers: int | None,
|
||
bytes_per_layer: int | None = None,
|
||
) -> int:
|
||
if total_layers is None or total_layers <= 0 or memory_mb <= 0:
|
||
return 0
|
||
budget_bytes = memory_mb * 1024 * 1024
|
||
layer_bytes = bytes_per_layer or _DEFAULT_BYTES_PER_LAYER
|
||
return min(total_layers, int((budget_bytes * 0.8) // layer_bytes))
|
||
|
||
|
||
def _shard_budget_line(
|
||
memory_mb: int,
|
||
memory_source: str,
|
||
total_layers: int | None,
|
||
quantization: str,
|
||
bytes_per_layer: int | None = None,
|
||
) -> str:
|
||
memory_gb = memory_mb / 1024
|
||
gb_str = f"{memory_gb:.1f} GB"
|
||
budget_quantization = "bfloat16" if quantization == "auto" else quantization
|
||
if total_layers is None or total_layers <= 0:
|
||
return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count"
|
||
max_layers = _max_assignable_layers(memory_mb, total_layers, bytes_per_layer=bytes_per_layer)
|
||
# Remaining capacity after one full model load (rough estimate)
|
||
shard_bytes = max_layers * (bytes_per_layer or _DEFAULT_BYTES_PER_LAYER)
|
||
remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3)
|
||
remaining_str = f"; {remaining_gb:.1f} GB remaining after full load" if remaining_gb > 1 else ""
|
||
return (
|
||
f"Memory budget: {gb_str} {memory_source}; "
|
||
f"Shard budget: up to {max_layers}/{total_layers} layers at {budget_quantization}"
|
||
f"{remaining_str}"
|
||
)
|
||
|
||
|
||
def _assignment_bytes_per_layer(assignment: dict, quantization: str) -> int | None:
|
||
bytes_per_layer = assignment.get("bytes_per_layer")
|
||
if isinstance(bytes_per_layer, int) and bytes_per_layer > 0:
|
||
return bytes_per_layer
|
||
if not isinstance(bytes_per_layer, dict):
|
||
return None
|
||
keys = [quantization, "bfloat16", "bf16", "int8", "nf4"]
|
||
for key in keys:
|
||
value = bytes_per_layer.get(key)
|
||
if isinstance(value, int) and value > 0:
|
||
return value
|
||
for value in bytes_per_layer.values():
|
||
if isinstance(value, int) and value > 0:
|
||
return value
|
||
return None
|
||
|
||
|
||
def _post_json(url: str, payload: dict, timeout: float = 10.0) -> dict:
|
||
data = json.dumps(payload).encode()
|
||
req = urllib.request.Request(
|
||
url, data=data, headers={"Content-Type": "application/json"}, method="POST"
|
||
)
|
||
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||
return json.loads(r.read())
|
||
|
||
|
||
def _get_json(url: str, timeout: float = 10.0) -> dict:
|
||
with urllib.request.urlopen(url, timeout=timeout) as r:
|
||
return json.loads(r.read())
|
||
|
||
|
||
def _infer_relay_url_from_tracker(tracker_url: str) -> str | None:
|
||
"""Infer relay WebSocket URL from a public HTTPS tracker origin.
|
||
|
||
Public deployments colocate relay at /ws on the same host as the tracker API
|
||
(see QUICKSTART nginx layout). Local LAN trackers use a separate relay port
|
||
and must advertise relay_url explicitly via /v1/network/map.
|
||
"""
|
||
parsed = urllib.parse.urlparse(tracker_url)
|
||
if parsed.scheme != "https":
|
||
return None
|
||
host = parsed.hostname
|
||
if not host or host in ("127.0.0.1", "localhost"):
|
||
return None
|
||
return f"wss://{parsed.netloc}/ws"
|
||
|
||
|
||
def _discover_relay_url(tracker_url: str) -> str | None:
|
||
relay_url: str | None = None
|
||
try:
|
||
network_map = _get_json(f"{tracker_url}/v1/network/map", timeout=5.0)
|
||
raw = network_map.get("relay_url")
|
||
if isinstance(raw, str) and raw:
|
||
relay_url = raw
|
||
except Exception:
|
||
pass
|
||
return relay_url or _infer_relay_url_from_tracker(tracker_url)
|
||
|
||
|
||
def _start_relay_bridge_if_available(
|
||
tracker_url: str,
|
||
wallet_address: str,
|
||
local_base_url: str,
|
||
advertised_endpoint: str,
|
||
relay_url: str | None = None,
|
||
) -> tuple[RelayHttpBridge | None, dict]:
|
||
relay_url = relay_url or _discover_relay_url(tracker_url)
|
||
if not relay_url:
|
||
return None, {}
|
||
peer_id = peer_id_from_wallet(wallet_address)
|
||
bridge = RelayHttpBridge(
|
||
relay_url=relay_url,
|
||
peer_id=peer_id,
|
||
local_base_url=local_base_url,
|
||
advertised_addr=advertised_endpoint,
|
||
)
|
||
info = bridge.start()
|
||
if bridge.wait_connected(timeout=5.0):
|
||
print(f" Relay connected — {info.relay_addr}", flush=True)
|
||
else:
|
||
print(f" Relay configured but not connected yet — {info.relay_addr}", flush=True)
|
||
return bridge, {
|
||
"relay_addr": info.relay_addr,
|
||
"peer_id": info.peer_id,
|
||
}
|
||
|
||
|
||
def _attach_relay_bridge(node: StubNodeServer | TorchNodeServer, bridge: RelayHttpBridge | None) -> None:
|
||
setattr(node, "relay_bridge", bridge)
|
||
if bridge is None:
|
||
return
|
||
original_stop = node.stop
|
||
|
||
def _stop_with_bridge() -> None:
|
||
try:
|
||
bridge.stop()
|
||
finally:
|
||
original_stop()
|
||
|
||
node.stop = _stop_with_bridge # type: ignore[method-assign]
|
||
|
||
|
||
def _start_heartbeat(
|
||
tracker_url: str,
|
||
node_id: str,
|
||
register_payload: dict,
|
||
interval: float = 20.0,
|
||
node_ref: Any | None = None,
|
||
start_time: float | None = None,
|
||
) -> threading.Thread:
|
||
"""Daemon thread: sends heartbeats and re-registers automatically after tracker restarts.
|
||
|
||
Heartbeat body carries cumulative stats (total_requests, failed_requests,
|
||
queue_depth, uptime_seconds, status). Stats are buffered locally during
|
||
outage and flushed on next successful heartbeat.
|
||
|
||
Heartbeat response may include new_assignment: {model, shard_start, shard_end}
|
||
which is logged for now (hot-reload implemented in US-026).
|
||
"""
|
||
_start_time = start_time or time.monotonic()
|
||
|
||
def _get_stats() -> dict:
|
||
uptime = time.monotonic() - _start_time
|
||
stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
|
||
if node_ref is not None:
|
||
stats["total_requests"] = getattr(
|
||
node_ref,
|
||
"total_requests",
|
||
getattr(node_ref, "chat_completion_count", 0),
|
||
)
|
||
stats["failed_requests"] = getattr(node_ref, "failed_requests", 0)
|
||
stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
|
||
return stats
|
||
|
||
def _reregister() -> bool:
|
||
nonlocal node_id
|
||
try:
|
||
resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload)
|
||
node_id = resp.get("node_id", node_id)
|
||
if node_ref is not None:
|
||
setattr(node_ref, "tracker_node_id", node_id)
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
def _register_additional_assignment(applied: dict) -> None:
|
||
model_id = str(applied.get("model") or register_payload.get("hf_repo") or register_payload.get("model"))
|
||
extra_payload = {
|
||
**register_payload,
|
||
"model": model_id.split("/")[-1],
|
||
"hf_repo": model_id if "/" in model_id else register_payload.get("hf_repo"),
|
||
"shard_start": applied["shard_start"],
|
||
"shard_end": applied["shard_end"],
|
||
"quantization": applied.get("quantization", register_payload.get("quantization")),
|
||
"tracker_mode": bool(applied.get("tracker_mode", False)),
|
||
"managed_assignment": True,
|
||
}
|
||
try:
|
||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", extra_payload)
|
||
print(
|
||
f" [node] registered additional model — node ID: {reg_resp.get('node_id')}",
|
||
flush=True,
|
||
)
|
||
except Exception as exc:
|
||
print(f" [node] WARNING: additional model registration failed: {exc}", flush=True)
|
||
|
||
def _apply_directives(directives: list[dict]) -> None:
|
||
if not directives:
|
||
return
|
||
if node_ref is None or not hasattr(node_ref, "apply_tracker_directives"):
|
||
print(f" [node] tracker directives received: {directives}", flush=True)
|
||
return
|
||
try:
|
||
applied = node_ref.apply_tracker_directives(directives)
|
||
except Exception as exc:
|
||
print(f" [node] WARNING: failed to apply tracker directives: {exc}", flush=True)
|
||
return
|
||
if applied:
|
||
if applied.get("action") == "ADD_SHARD":
|
||
_register_additional_assignment(applied)
|
||
return
|
||
model_id = applied.get("model", register_payload.get("hf_repo") or register_payload.get("model"))
|
||
register_payload["model"] = str(model_id).split("/")[-1]
|
||
register_payload["hf_repo"] = model_id
|
||
register_payload["shard_start"] = applied["shard_start"]
|
||
register_payload["shard_end"] = applied["shard_end"]
|
||
register_payload["quantization"] = applied.get("quantization", register_payload.get("quantization"))
|
||
register_payload["tracker_mode"] = bool(applied.get("tracker_mode", False))
|
||
|
||
def _loop() -> None:
|
||
nonlocal node_id
|
||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||
outage_streak = 0 # consecutive intervals where tracker was unreachable
|
||
|
||
while True:
|
||
time.sleep(interval)
|
||
|
||
if outage_streak > 0:
|
||
# Tracker was down — attempt re-registration first (it may have restarted
|
||
# with a clean slate and won't know this node).
|
||
if _reregister():
|
||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||
print(f" [node] re-registered after outage — node ID: {node_id}", flush=True)
|
||
outage_streak = 0
|
||
else:
|
||
outage_streak += 1
|
||
if outage_streak <= 3 or outage_streak % 10 == 0:
|
||
print(
|
||
f" [node] WARNING: tracker still unreachable "
|
||
f"({outage_streak * interval:.0f}s)",
|
||
flush=True,
|
||
)
|
||
continue
|
||
|
||
try:
|
||
resp = _post_json(hb_url, _get_stats())
|
||
_apply_directives(resp.get("directives", []))
|
||
new_asgn = resp.get("new_assignment")
|
||
if new_asgn:
|
||
print(
|
||
f" [node] tracker assignment received: "
|
||
f"action={new_asgn.get('action')!r} model={new_asgn.get('model')!r} "
|
||
f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
|
||
flush=True,
|
||
)
|
||
_apply_directives([new_asgn])
|
||
except urllib.error.HTTPError as exc:
|
||
if exc.code == 404:
|
||
# Node was purged (e.g. long gap before restart noticed) — re-register now.
|
||
print(" [node] tracker lost registration — re-registering...", flush=True)
|
||
if _reregister():
|
||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||
print(f" [node] re-registered — node ID: {node_id}", flush=True)
|
||
else:
|
||
print(" [node] WARNING: re-registration failed", flush=True)
|
||
outage_streak = 1
|
||
else:
|
||
print(f" [node] WARNING: heartbeat failed ({exc.code}): {exc}", flush=True)
|
||
except Exception as exc:
|
||
outage_streak = 1
|
||
print(f" [node] WARNING: tracker unreachable: {exc}", flush=True)
|
||
|
||
t = threading.Thread(target=_loop, daemon=True, name="heartbeat")
|
||
t.start()
|
||
return t
|
||
|
||
|
||
_PENDING_NODE_ID = "pending"
|
||
|
||
|
||
def _register_with_tracker(
|
||
tracker_url: str,
|
||
reg_payload: dict,
|
||
node: Any,
|
||
start_time: float,
|
||
) -> str | None:
|
||
"""Register with the tracker, or start background retries when it is unreachable."""
|
||
try:
|
||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
|
||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||
setattr(node, "tracker_node_id", tracker_node_id)
|
||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||
_start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=start_time)
|
||
return tracker_node_id
|
||
except Exception as exc:
|
||
setattr(node, "tracker_node_id", None)
|
||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||
print(" [node] will retry registration in the background", flush=True)
|
||
_start_heartbeat(
|
||
tracker_url,
|
||
_PENDING_NODE_ID,
|
||
reg_payload,
|
||
node_ref=node,
|
||
start_time=start_time,
|
||
)
|
||
return None
|
||
|
||
|
||
def _warn_virtual_network_ip(ip: str | None) -> None:
|
||
"""Print a warning when the auto-detected advertise IP is in a known virtual-network range.
|
||
|
||
172.16.0.0/12 is used by Docker, WSL2, and most hypervisors. Nodes behind these
|
||
adapters are NOT directly reachable from other physical machines on the LAN, so
|
||
cross-host pipeline hops will time out. The user must pass --advertise-host with
|
||
their actual LAN IP (e.g. 192.168.x.x) to fix this.
|
||
"""
|
||
if ip is None:
|
||
return
|
||
try:
|
||
parts = [int(p) for p in ip.split(".")]
|
||
if len(parts) != 4:
|
||
return
|
||
a, b = parts[0], parts[1]
|
||
# 172.16.0.0/12 → 172.16–31.x.x
|
||
if a == 172 and 16 <= b <= 31:
|
||
print(
|
||
f"\n WARNING: auto-detected endpoint IP {ip} is in 172.16.0.0/12.\n"
|
||
f" This range is used by Docker, WSL2, and virtual machines and is\n"
|
||
f" NOT reachable from other physical machines on your LAN.\n"
|
||
f" Cross-host pipeline hops WILL time out.\n"
|
||
f" Fix: use a public tracker with relay (wss://…/ws), or pass\n"
|
||
f" --advertise-host <your-LAN-ip> (e.g. 192.168.x.x).\n",
|
||
flush=True,
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def run_startup(
|
||
tracker_url: str,
|
||
port: int = 0,
|
||
model: str | None = None,
|
||
model_id: str | None = None,
|
||
shard_start: int | None = None,
|
||
shard_end: int | None = None,
|
||
quantization: str = "auto",
|
||
wallet_path: Path | None = None,
|
||
cache_dir: Path | None = None,
|
||
host: str = "127.0.0.1",
|
||
advertise_host: str | None = None,
|
||
contracts: Any | None = None,
|
||
route_timeout: float = 30.0,
|
||
vram_mb_override: int | None = None,
|
||
max_loaded_shards: int = 1,
|
||
debug: bool = False,
|
||
tracker_source_disabled: bool = False,
|
||
torch_threads: int | None = None,
|
||
torch_interop_threads: int | None = None,
|
||
) -> StubNodeServer | TorchNodeServer:
|
||
"""Execute the full startup sequence and return a running node server.
|
||
|
||
Steps (all non-interactive):
|
||
1. Detect GPU / hardware profile
|
||
2. Load or generate Solana wallet keypair
|
||
3. Query tracker for optimal shard assignment
|
||
4. Download (or stub) the assigned shard from peers, then HuggingFace
|
||
5. Start local HTTP server
|
||
6. Register with tracker
|
||
|
||
Prints a compact status summary on completion.
|
||
"""
|
||
|
||
tracker_url = tracker_url.rstrip("/")
|
||
relay_url = _discover_relay_url(tracker_url)
|
||
if max_loaded_shards < 1:
|
||
raise ValueError("--max-shards must be at least 1")
|
||
|
||
# 1. Hardware detection
|
||
if advertise_host is None and host == "0.0.0.0":
|
||
# socket.getfqdn() returns an mDNS name (.local / .localdomain) that remote
|
||
# machines on a different OS or subnet often can't resolve. Instead, probe the
|
||
# outbound IP by opening a UDP socket toward the tracker — no data is sent.
|
||
try:
|
||
_tracker_host = urllib.parse.urlparse(tracker_url).hostname or "8.8.8.8"
|
||
_s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||
_s.connect((_tracker_host, 80))
|
||
advertise_host = _s.getsockname()[0]
|
||
_s.close()
|
||
except Exception:
|
||
advertise_host = socket.getfqdn()
|
||
|
||
if relay_url:
|
||
print(f"Relay advertised by tracker — using outbound tunnel {relay_url}", flush=True)
|
||
else:
|
||
_warn_virtual_network_ip(advertise_host)
|
||
|
||
print("Detecting hardware...", flush=True)
|
||
hw = detect_hardware()
|
||
torch_thread_config = _configure_torch_threads(torch_threads, torch_interop_threads)
|
||
if torch_thread_config:
|
||
hw.update(torch_thread_config)
|
||
intra = torch_thread_config.get("torch_threads", "?")
|
||
interop = torch_thread_config.get("torch_interop_threads", "?")
|
||
print(f" PyTorch threads: intra-op={intra}, inter-op={interop}", flush=True)
|
||
device: str = hw["device"]
|
||
gpu_name: str | None = hw.get("gpu_name")
|
||
vram_mb: int = hw.get("vram_mb", 0)
|
||
shared_vram_mb: int = hw.get("shared_vram_mb", 0)
|
||
ram_mb: int = hw.get("ram_mb", 16 * 1024)
|
||
|
||
if vram_mb_override is not None:
|
||
vram_mb = vram_mb_override
|
||
shared_vram_mb = 0
|
||
print(f" Memory budget overridden to {vram_mb / 1024:.1f} GB via --memory", flush=True)
|
||
elif device == "cpu":
|
||
gpu_suffix = ""
|
||
if gpu_name and vram_mb > 0:
|
||
gpu_suffix = (
|
||
f"; CUDA inactive; detected {gpu_name} "
|
||
f"({vram_mb / 1024:.1f} GB dedicated VRAM, {shared_vram_mb / 1024:.1f} GB shared)"
|
||
)
|
||
print(f" WARNING: No CUDA GPU detected — running in CPU mode ({ram_mb / 1024:.1f} GB RAM{gpu_suffix})", flush=True)
|
||
else:
|
||
shared_suffix = f", {shared_vram_mb / 1024:.1f} GB shared" if shared_vram_mb > 0 else ""
|
||
print(f" GPU: {gpu_name} ({vram_mb / 1024:.1f} GB dedicated VRAM{shared_suffix}, {ram_mb / 1024:.1f} GB RAM)", flush=True)
|
||
|
||
if vram_mb_override is not None:
|
||
memory_budget_mb = vram_mb
|
||
memory_budget_source = "memory override"
|
||
else:
|
||
memory_budget_mb, memory_budget_source = _memory_budget(device, vram_mb, ram_mb, shared_vram_mb)
|
||
assignment_vram_mb = memory_budget_mb if device == "cuda" or vram_mb_override is not None else 0
|
||
print(f" Memory budget: {memory_budget_mb / 1024:.1f} GB {memory_budget_source}", flush=True)
|
||
|
||
print("Benchmarking compute...", flush=True)
|
||
if device != "cuda" and gpu_name:
|
||
_cuda_score, cuda_ok, cuda_error = benchmark_throughput_checked("cuda")
|
||
hw["cuda_benchmark_ok"] = cuda_ok
|
||
if cuda_error:
|
||
hw["cuda_benchmark_error"] = cuda_error
|
||
if not cuda_ok:
|
||
print(f" CUDA benchmark unavailable: {cuda_error}; using CPU benchmark", flush=True)
|
||
bench_tps, bench_ok, bench_error = benchmark_throughput_checked(device)
|
||
hw["benchmark_device"] = device
|
||
hw["benchmark_ok"] = bench_ok
|
||
if bench_error:
|
||
hw["benchmark_error"] = bench_error
|
||
device_label = "GPU" if device == "cuda" else "CPU"
|
||
print(f" {device_label} throughput index: {bench_tps:,.0f}", flush=True)
|
||
|
||
registration_capabilities = {
|
||
"vram_bytes": max(0, int(assignment_vram_mb)) * 1024 * 1024,
|
||
"ram_bytes": max(0, int(ram_mb)) * 1024 * 1024,
|
||
"max_loaded_shards": max_loaded_shards,
|
||
"benchmark_tokens_per_sec": bench_tps,
|
||
}
|
||
# 2. Wallet
|
||
print("Loading wallet...", flush=True)
|
||
wallet_kwargs: dict = {}
|
||
if wallet_path is not None:
|
||
wallet_kwargs["path"] = wallet_path
|
||
_, _, address = load_or_create_wallet(**wallet_kwargs)
|
||
print(f" Wallet: {address}", flush=True)
|
||
probationary_line = _probationary_status_line(contracts, address)
|
||
if probationary_line is not None:
|
||
print(f" {probationary_line}", flush=True)
|
||
|
||
pinned_shard_start = shard_start
|
||
pinned_shard_end = shard_end
|
||
user_pinned_shard = pinned_shard_start is not None or pinned_shard_end is not None
|
||
|
||
if model_id: # treat "" the same as None — no explicit model given
|
||
full_sources: list[dict] = []
|
||
# Auto-detect shard range from model config if not explicitly provided
|
||
if shard_start is None or shard_end is None:
|
||
try:
|
||
detected = _detect_num_layers(model_id, cache_dir=cache_dir)
|
||
except TypeError:
|
||
detected = _detect_num_layers(model_id)
|
||
if detected is None:
|
||
raise ValueError(
|
||
f"Could not read num_hidden_layers from {model_id} config. "
|
||
"Pass --shard-start and --shard-end explicitly."
|
||
)
|
||
# When no explicit shard range given, ask the tracker if there's a gap for this model.
|
||
if shard_start is None and shard_end is None:
|
||
try:
|
||
qs = urllib.parse.urlencode({
|
||
"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb, "hf_repo": model_id,
|
||
})
|
||
net_asgn = _get_json(f"{tracker_url}/v1/network/assign?{qs}", timeout=5.0)
|
||
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 = (
|
||
[] if tracker_source_disabled
|
||
else _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})",
|
||
flush=True,
|
||
)
|
||
except Exception:
|
||
pass # No other nodes registered yet — default to full model below
|
||
shard_start = shard_start if shard_start is not None else 0
|
||
shard_end = shard_end if shard_end is not None else detected - 1
|
||
print(f" Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True)
|
||
|
||
print("Loading real PyTorch model shard...", flush=True)
|
||
node = TorchNodeServer(
|
||
host=host,
|
||
port=port,
|
||
model_id=model_id,
|
||
shard_start=shard_start,
|
||
shard_end=shard_end,
|
||
quantization=quantization,
|
||
tracker_url=tracker_url,
|
||
route_timeout=route_timeout,
|
||
cache_dir=cache_dir,
|
||
debug=debug,
|
||
max_loaded_shards=max_loaded_shards,
|
||
)
|
||
_node_start_time = time.monotonic()
|
||
actual_port = node.start()
|
||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
|
||
if isinstance(total_layers, int) and total_layers > 0:
|
||
layer_count = shard_end - shard_start + 1
|
||
shard_label = f"layers {shard_start}–{shard_end}; {layer_count} of {total_layers}"
|
||
else:
|
||
shard_label = f"layers {shard_start}–{shard_end}"
|
||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||
endpoint = f"http://{public_host}:{actual_port}"
|
||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||
tracker_url,
|
||
address,
|
||
local_base_url,
|
||
endpoint,
|
||
relay_url=relay_url,
|
||
)
|
||
_attach_relay_bridge(node, relay_bridge)
|
||
# Register with tracker so other nodes can auto-join this model.
|
||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
|
||
model_cache_path = _model_cache_path(model_id, cache_dir)
|
||
reg_payload = {
|
||
"endpoint": endpoint,
|
||
"model": model_id.split("/")[-1],
|
||
"hf_repo": model_id,
|
||
"num_layers": total_layers,
|
||
"shard_start": shard_start,
|
||
"shard_end": shard_end,
|
||
"hardware_profile": hw,
|
||
"wallet_address": address,
|
||
"quantization": quantization,
|
||
"score": 1.0,
|
||
"tracker_mode": (shard_start == 0),
|
||
"managed_assignment": not user_pinned_shard,
|
||
"model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
|
||
"downloaded_models": (
|
||
_downloaded_model_inventory(
|
||
model_id.split("/")[-1],
|
||
shard_start,
|
||
shard_end,
|
||
model_cache_path,
|
||
hf_repo=model_id,
|
||
model_sources=full_sources,
|
||
)
|
||
if model_cache_path is not None
|
||
else []
|
||
),
|
||
**registration_capabilities,
|
||
**relay_fields,
|
||
}
|
||
tracker_node_id: str | None = None
|
||
try:
|
||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
|
||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||
setattr(node, "tracker_node_id", tracker_node_id)
|
||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||
_start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=_node_start_time)
|
||
except Exception as exc:
|
||
setattr(node, "tracker_node_id", None)
|
||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||
|
||
print(
|
||
f"\n{'=' * 32}\n"
|
||
f"meshnet-node ready\n"
|
||
f" Wallet: {address}\n"
|
||
f" Model ID: {model_id}\n"
|
||
f" Shard: {shard_label}\n"
|
||
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization)}\n"
|
||
f" Quantization: {quantization}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Node ID: {tracker_node_id or 'unregistered'}\n"
|
||
f" Hardware: {_hardware_label(device, gpu_name)}\n"
|
||
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
|
||
f"{'=' * 32}",
|
||
flush=True,
|
||
)
|
||
return node
|
||
if user_pinned_shard and not model:
|
||
raise ValueError("--shard-start / --shard-end require --model")
|
||
|
||
# 3a. Auto-join: query tracker for network-wide HF model assignment.
|
||
# Skipped when the user explicitly requested a model — the shard-assignment
|
||
# query below (/v1/nodes/assign?model=…) is authoritative there, and a fresh
|
||
# tracker would otherwise print a scary 503 for the model-less auto-join.
|
||
net_assignment: dict = {}
|
||
if model_id or (model and model != "stub-model"):
|
||
if model:
|
||
print(f"Model {model!r} requested explicitly — skipping network auto-join.", flush=True)
|
||
else:
|
||
print("Querying tracker for network assignment...", flush=True)
|
||
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb})
|
||
try:
|
||
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
|
||
except Exception as exc:
|
||
print(f" (auto-join unavailable: {exc})", flush=True)
|
||
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
|
||
_gap_found: bool = bool(net_assignment.get("gap_found", False))
|
||
|
||
if assigned_hf_repo:
|
||
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", [])
|
||
if _gap_found:
|
||
print(
|
||
f" Assigned gap: {assigned_hf_repo} "
|
||
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||
f"(of {assigned_num_layers})",
|
||
flush=True,
|
||
)
|
||
else:
|
||
print(
|
||
f" Assigned redundant copy: {assigned_hf_repo} "
|
||
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||
f"(of {assigned_num_layers})",
|
||
flush=True,
|
||
)
|
||
full_sources = [] if tracker_source_disabled else _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,
|
||
port=port,
|
||
model_id=assigned_hf_repo,
|
||
shard_start=assigned_shard_start,
|
||
shard_end=assigned_shard_end,
|
||
quantization=quantization,
|
||
tracker_url=tracker_url,
|
||
route_timeout=route_timeout,
|
||
cache_dir=cache_dir,
|
||
debug=debug,
|
||
max_loaded_shards=max_loaded_shards,
|
||
)
|
||
_node_start_time = time.monotonic()
|
||
actual_port = node.start()
|
||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||
endpoint = f"http://{public_host}:{actual_port}"
|
||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||
tracker_url,
|
||
address,
|
||
local_base_url,
|
||
endpoint,
|
||
relay_url=relay_url,
|
||
)
|
||
_attach_relay_bridge(node, relay_bridge)
|
||
model_cache_path = _model_cache_path(assigned_hf_repo, cache_dir)
|
||
auto_reg_payload = {
|
||
"endpoint": endpoint,
|
||
"model": assigned_hf_repo.split("/")[-1],
|
||
"hf_repo": assigned_hf_repo,
|
||
"num_layers": assigned_num_layers,
|
||
"shard_start": assigned_shard_start,
|
||
"shard_end": assigned_shard_end,
|
||
"hardware_profile": hw,
|
||
"wallet_address": address,
|
||
"quantization": quantization,
|
||
"score": 1.0,
|
||
"tracker_mode": (assigned_shard_start == 0),
|
||
"managed_assignment": True,
|
||
"model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
|
||
"downloaded_models": (
|
||
_downloaded_model_inventory(
|
||
assigned_hf_repo.split("/")[-1],
|
||
assigned_shard_start,
|
||
assigned_shard_end,
|
||
model_cache_path,
|
||
hf_repo=assigned_hf_repo,
|
||
model_sources=full_sources,
|
||
)
|
||
if model_cache_path is not None
|
||
else []
|
||
),
|
||
**registration_capabilities,
|
||
**relay_fields,
|
||
}
|
||
tracker_node_id = None
|
||
try:
|
||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload)
|
||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||
setattr(node, "tracker_node_id", tracker_node_id)
|
||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||
_start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload, node_ref=node, start_time=_node_start_time)
|
||
except Exception as exc:
|
||
setattr(node, "tracker_node_id", None)
|
||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||
shard_count = assigned_shard_end - assigned_shard_start + 1
|
||
print(
|
||
f"\n{'=' * 32}\n"
|
||
f"meshnet-node ready (auto-joined)\n"
|
||
f" Wallet: {address}\n"
|
||
f" Model ID: {assigned_hf_repo}\n"
|
||
f" Shard: layers {assigned_shard_start}–{assigned_shard_end} "
|
||
f"({shard_count} of {assigned_num_layers})\n"
|
||
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n"
|
||
f" Quantization: {quantization}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Node ID: {tracker_node_id or 'unregistered'}\n"
|
||
f" Hardware: {_hardware_label(device, gpu_name)}\n"
|
||
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
|
||
f"{'=' * 32}",
|
||
flush=True,
|
||
)
|
||
return node
|
||
|
||
if not assigned_hf_repo and model is None:
|
||
raise RuntimeError(
|
||
"Tracker did not assign a model. Join a network that already serves one, "
|
||
"or start with --model <HF_REPO>."
|
||
)
|
||
|
||
# 3b. Stub preset path (tests / explicit stub-model) or named preset models.
|
||
print("Querying tracker for shard assignment...", flush=True)
|
||
assign_qs = urllib.parse.urlencode({
|
||
"model": model or "stub-model",
|
||
"device": device,
|
||
# CPU-mode nodes must be sized by RAM: a detected-but-unusable GPU's
|
||
# VRAM would otherwise cap the shard (e.g. 8 GB VRAM → 3 layers on a
|
||
# 79 GB box whose Torch has no CUDA).
|
||
"vram_mb": assignment_vram_mb,
|
||
"ram_mb": ram_mb,
|
||
})
|
||
try:
|
||
assignment = _get_json(f"{tracker_url}/v1/nodes/assign?{assign_qs}")
|
||
except urllib.error.URLError as exc:
|
||
print(f" ERROR: Cannot reach tracker at {tracker_url}: {exc}", file=sys.stderr, flush=True)
|
||
raise
|
||
|
||
shard_start = assignment["shard_start"]
|
||
shard_end = assignment["shard_end"]
|
||
if user_pinned_shard:
|
||
if pinned_shard_start is not None:
|
||
shard_start = pinned_shard_start
|
||
if pinned_shard_end is not None:
|
||
shard_end = pinned_shard_end
|
||
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] = [] if tracker_source_disabled else assignment.get("model_sources", [])
|
||
assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization)
|
||
if user_pinned_shard:
|
||
print(
|
||
f" Shard: layers {shard_start}-{shard_end} of {assigned_model} (pinned)",
|
||
flush=True,
|
||
)
|
||
else:
|
||
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
|
||
|
||
# 4. Download shard
|
||
print("Downloading shard...", flush=True)
|
||
dl_kwargs: dict = {}
|
||
if cache_dir is not None:
|
||
dl_kwargs["cache_dir"] = cache_dir
|
||
if hf_repo is not None:
|
||
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 = _registration_shard_checksum(assigned_model, shard_path)
|
||
downloaded_models = _downloaded_model_inventory(
|
||
assigned_model,
|
||
shard_start,
|
||
shard_end,
|
||
shard_path,
|
||
hf_repo=hf_repo,
|
||
model_sources=model_sources,
|
||
)
|
||
print(f" Cached at: {shard_path}", flush=True)
|
||
|
||
# 5. Start HTTP server
|
||
is_last = shard_end >= assignment.get("model_layers_end", shard_end)
|
||
node = StubNodeServer(
|
||
host=host,
|
||
port=port,
|
||
shard_start=shard_start,
|
||
shard_end=shard_end,
|
||
is_last_shard=is_last,
|
||
model=assigned_model,
|
||
shard_path=shard_path,
|
||
)
|
||
_node_start_time = time.monotonic()
|
||
actual_port = node.start()
|
||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||
endpoint = f"http://{public_host}:{actual_port}"
|
||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||
tracker_url,
|
||
address,
|
||
local_base_url,
|
||
endpoint,
|
||
relay_url=relay_url,
|
||
)
|
||
_attach_relay_bridge(node, relay_bridge)
|
||
|
||
# 6. Register with tracker
|
||
print("Registering with tracker...", flush=True)
|
||
reg_payload = {
|
||
"endpoint": endpoint,
|
||
"model": assigned_model,
|
||
"shard_start": shard_start,
|
||
"shard_end": shard_end,
|
||
"shard_checksum": shard_checksum,
|
||
"downloaded_models": downloaded_models,
|
||
"hardware_profile": hw,
|
||
"wallet_address": address,
|
||
"score": 1.0,
|
||
"managed_assignment": not user_pinned_shard,
|
||
**registration_capabilities,
|
||
**relay_fields,
|
||
}
|
||
try:
|
||
reg_resp = _post_json(
|
||
f"{tracker_url}/v1/nodes/register",
|
||
reg_payload,
|
||
)
|
||
node_id = str(reg_resp["node_id"])
|
||
setattr(node, "tracker_node_id", node_id)
|
||
_start_heartbeat(tracker_url, node_id, reg_payload, node_ref=node, start_time=_node_start_time)
|
||
except Exception:
|
||
node.stop()
|
||
raise
|
||
|
||
# Status summary
|
||
hw_str = device.upper()
|
||
if gpu_name:
|
||
hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
|
||
print(
|
||
f"\n{'=' * 32}\n"
|
||
f"meshnet-node ready\n"
|
||
f" Wallet: {address}\n"
|
||
f" Shard: layers {shard_start}-{shard_end} ({assigned_model})\n"
|
||
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n"
|
||
f" Endpoint: {endpoint}\n"
|
||
f" Node ID: {node_id}\n"
|
||
f" Hardware: {hw_str}\n"
|
||
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
|
||
f"{'=' * 32}",
|
||
flush=True,
|
||
)
|
||
|
||
return node
|
||
|
||
|
||
def _detect_num_layers(model_id: str, cache_dir: Path | None = None) -> int | None:
|
||
"""Fetch num_hidden_layers from HuggingFace model config (downloads ~1 KB config.json only)."""
|
||
try:
|
||
from transformers import AutoConfig # type: ignore[import]
|
||
|
||
from .model_catalog import layers_from_config
|
||
|
||
local_model = _model_cache_path(model_id, cache_dir)
|
||
load_source = str(local_model) if local_model is not None else model_id
|
||
cfg = AutoConfig.from_pretrained(
|
||
load_source,
|
||
cache_dir=str(cache_dir) if cache_dir is not None and local_model is None else None,
|
||
)
|
||
layers = layers_from_config(cfg)
|
||
if layers is None:
|
||
print(
|
||
f" Warning: no layer count in {model_id} config "
|
||
"(checked top level and text_config)", flush=True,
|
||
)
|
||
return layers
|
||
except Exception as exc:
|
||
print(f" Warning: could not read model config from HF: {exc}", flush=True)
|
||
return None
|
||
|
||
|
||
def _probationary_status_line(contracts: Any | None, wallet_address: str) -> str | None:
|
||
if contracts is None:
|
||
return None
|
||
remaining = contracts.registry.probationary_jobs_remaining(wallet_address)
|
||
if remaining <= 0:
|
||
return "Probationary period complete: earning enabled"
|
||
suffix = "job" if remaining == 1 else "jobs"
|
||
return f"Probationary period: {remaining} {suffix} remaining before earning"
|