billing ledger

This commit is contained in:
Dobromir Popov
2026-07-02 21:27:23 +02:00
parent 416ceba0f6
commit 57ec7c1e4b
6 changed files with 758 additions and 5 deletions

View File

@@ -34,6 +34,7 @@ import uuid
from importlib.resources import files
from typing import Any
from .billing import BillingLedger
from .gossip import NodeGossip
from .raft import RaftNode
@@ -975,6 +976,29 @@ def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None:
_rebalance_hf_model_locked(server, hf_repo)
def _api_key_from_headers(headers) -> str | None:
auth = headers.get("Authorization")
if not auth:
return None
if auth.lower().startswith("bearer "):
return auth.split(" ", 1)[1].strip() or None
return auth.strip() or None
def _usage_total_tokens(payload: dict) -> int | None:
usage = payload.get("usage")
if not isinstance(usage, dict):
return None
total = usage.get("total_tokens")
if isinstance(total, (int, float)):
return int(total)
prompt = usage.get("prompt_tokens")
completion = usage.get("completion_tokens")
if isinstance(prompt, (int, float)) or isinstance(completion, (int, float)):
return int(prompt or 0) + int(completion or 0)
return None
def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None:
if contracts is None or not wallet_address:
return None
@@ -1021,6 +1045,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
raft: "RaftNode | None" = None,
gossip: "NodeGossip | None" = None,
stats: "_StatsCollector | None" = None,
billing: "BillingLedger | None" = None,
) -> None:
super().__init__(addr, handler)
self.registry = registry
@@ -1033,6 +1058,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
self.raft = raft
self.gossip = gossip
self.stats: _StatsCollector | None = stats
self.billing: BillingLedger | None = billing
class _TrackerHandler(http.server.BaseHTTPRequestHandler):
@@ -1085,6 +1111,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if self.path == "/v1/stats/gossip":
self._handle_stats_gossip()
return
if self.path == "/v1/billing/gossip":
self._handle_billing_gossip()
return
parts = self.path.split("/")
# /v1/nodes/<node_id>/heartbeat -> ['', 'v1', 'nodes', '<id>', 'heartbeat']
if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat":
@@ -1117,6 +1146,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self._handle_raft_status()
elif parsed.path == "/v1/stats":
self._handle_stats()
elif parsed.path == "/v1/billing/summary":
self._handle_billing_summary()
elif parsed.path == "/v1/health":
self._send_json(200, {"status": "ok"})
else:
@@ -1345,6 +1376,24 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if model and server.stats is not None:
server.stats.record_request(model)
# Billing gate (ADR-0015): reject before any routing — no free work.
api_key = _api_key_from_headers(self.headers)
if server.billing is not None:
if api_key is None:
self._send_json(401, {"error": {
"message": "missing API key: send Authorization: Bearer <key>",
"type": "invalid_request_error",
"code": "missing_api_key",
}})
return
if not server.billing.has_funds(api_key):
self._send_json(402, {"error": {
"message": "insufficient balance: deposit USDT to continue",
"type": "insufficient_quota",
"code": "insufficient_balance",
}})
return
# Find a live tracker-mode node for this model
with server.lock:
self._purge_expired_nodes()
@@ -1395,12 +1444,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
# This allows overlapping shard registrations without double-computation.
covered_up_to = rs - 1
route_hops: list[dict] = []
node_work: list[tuple[str | None, int]] = []
for rn in route_nodes:
hop: dict = {"endpoint": rn.endpoint, "start_layer": covered_up_to + 1}
if rn.relay_addr:
hop["relay_addr"] = rn.relay_addr
route_hops.append(hop)
covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to
effective_end = rn.shard_end if rn.shard_end is not None else covered_up_to
node_work.append((rn.wallet_address, max(0, effective_end - covered_up_to)))
covered_up_to = effective_end
# Strip the first-shard node we're about to proxy to — it's already handling the request.
downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")]
downstream_urls = json.dumps(downstream_hops)
@@ -1449,6 +1501,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
)
if relayed is not None:
self._send_relayed_response(relayed)
if int(relayed.get("status", 503)) < 400:
body_text = relayed.get("body") or ""
try:
tokens = _usage_total_tokens(json.loads(body_text)) or 0
except (json.JSONDecodeError, TypeError):
tokens = 0
self._bill_completed(api_key, model, tokens, node_work)
return
print(
f"[tracker] relay proxy failed {request_id}: {node.relay_addr}; "
@@ -1495,6 +1554,8 @@ 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()
stream_tokens: int | None = None
chunk_count = 0
try:
while True:
line = upstream.readline()
@@ -1502,8 +1563,27 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
break
self.wfile.write(line)
self.wfile.flush()
if line.startswith(b"data:"):
payload = line[5:].strip()
if payload and payload != b"[DONE]":
chunk_count += 1
if b'"usage"' in payload:
try:
found = _usage_total_tokens(json.loads(payload))
if found:
stream_tokens = found
except json.JSONDecodeError:
pass
except BrokenPipeError:
pass
# Bill even on client disconnect — the nodes did the work.
# Chunk count approximates generated tokens when the stream
# carries no usage record.
self._bill_completed(
api_key, model,
stream_tokens if stream_tokens is not None else chunk_count,
node_work,
)
else:
# Non-streaming: buffer and relay
resp_body = upstream.read()
@@ -1519,6 +1599,33 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self.wfile.write(resp_body)
except BrokenPipeError:
pass
try:
tokens = _usage_total_tokens(json.loads(resp_body)) or 0
except json.JSONDecodeError:
tokens = 0
self._bill_completed(api_key, model, tokens, node_work)
def _bill_completed(
self,
api_key: str | None,
model: str,
total_tokens: int,
node_work: list[tuple[str | None, int]],
) -> None:
"""Charge a completed request against the billing ledger (ADR-0015)."""
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if server.billing is None or api_key is None:
return
try:
event = server.billing.charge_request(api_key, model, total_tokens, node_work)
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']}",
flush=True,
)
except Exception as exc:
print(f"[tracker] billing failed for model={model!r}: {exc}", flush=True)
def _send_relayed_response(self, response: dict) -> None:
status = int(response.get("status", 503))
@@ -1870,6 +1977,28 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
server.stats.merge_peer_rpms(tracker_url, rpms)
self._send_json(200, {})
def _handle_billing_summary(self):
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if server.billing is None:
self._send_json(404, {"error": "billing is not enabled on this tracker"})
return
self._send_json(200, server.billing.snapshot())
def _handle_billing_gossip(self):
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
body = self._read_json_body()
if body is None:
return
if server.billing is None:
self._send_json(200, {"applied": 0})
return
events = body.get("events")
if not isinstance(events, list):
self._send_json(400, {"error": "events must be a list"})
return
applied = server.billing.apply_events([e for e in events if isinstance(e, dict)])
self._send_json(200, {"applied": applied})
def _handle_assign(self, parsed: urllib.parse.ParseResult):
"""Return an optimal shard assignment for a node given its hardware profile.
@@ -2299,6 +2428,8 @@ class TrackerServer:
cluster_self_url: str | None = None,
stats_db: str | None = None,
relay_url: str | None = None,
billing: BillingLedger | None = None,
billing_db: str | None = None,
) -> None:
self._host = host
self._requested_port = port
@@ -2323,6 +2454,15 @@ class TrackerServer:
self._stats: _StatsCollector | None = _StatsCollector(db_path=stats_db) if stats_db else _StatsCollector()
self._stats_stop = threading.Event()
self._stats_thread: threading.Thread | None = None
if billing is None and billing_db:
preset_prices = {
name: 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
}
billing = BillingLedger(db_path=billing_db, prices=preset_prices)
self._billing: BillingLedger | None = billing
self._billing_gossip_cursor = 0
self.port: int | None = None
def start(self) -> int:
@@ -2346,6 +2486,7 @@ class TrackerServer:
self._minimum_stake,
relay_url=effective_relay_url,
stats=self._stats,
billing=self._billing,
)
self.port = self._server.server_address[1]
@@ -2416,10 +2557,12 @@ class TrackerServer:
_rebalance_all_locked(server)
def _stats_loop(self) -> None:
"""Periodically save stats to DB and push local slice to cluster peers."""
"""Periodically save stats/billing to DB and push local slices to cluster peers."""
while not self._stats_stop.wait(_StatsCollector.SAVE_INTERVAL):
if self._stats is not None:
self._stats.save_to_db()
if self._billing is not None:
self._billing.save_to_db()
if self._stats is not None and self._cluster_peers:
self_url = self._cluster_self_url or f"http://{self._host}:{self.port}"
local_rpms = self._stats.get_local_rpms()
@@ -2436,6 +2579,28 @@ class TrackerServer:
r.read()
except Exception:
pass
if self._billing is not None and self._cluster_peers:
events, cursor = self._billing.events_since(self._billing_gossip_cursor)
if events:
delivered_all = True
body = json.dumps({"events": events}).encode()
for peer in self._cluster_peers:
try:
req = urllib.request.Request(
f"{peer}/v1/billing/gossip",
data=body,
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req, timeout=2.0) as r:
r.read()
except Exception:
delivered_all = False
# Only advance past events every peer has seen; unreachable
# peers get the full backlog on the next tick (idempotent
# via event-id dedupe on the receiving side).
if delivered_all:
self._billing_gossip_cursor = cursor
def stop(self) -> None:
if self._raft is not None:
@@ -2448,6 +2613,8 @@ class TrackerServer:
self._stats_stop.set()
if self._stats is not None:
self._stats.save_to_db()
if self._billing is not None:
self._billing.save_to_db()
self._server.shutdown()
self._server.server_close()
if self._thread is not None: