dual billing; tracker to node model sharing

packages/tracker/meshnet_tracker/server.py

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