billing ledger
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user