File size: 5,567 Bytes
0cd599e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """Bounded, best-effort usage telemetry; never part of inference success."""
import logging
import sqlite3
import threading
import time
from collections import deque
from dataclasses import dataclass
from .database import Database
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class UsageEvent:
created_at: str
source: str
key_id: int | None
model: str
input_tokens: int
output_tokens: int
duration_ms: int
class UsageWriter:
def __init__(self, db: Database, *, capacity: int = 1000, batch_size: int = 100,
flush_interval: float = 1.0, shutdown_timeout: float = 5.0):
if min(capacity, batch_size, flush_interval, shutdown_timeout) <= 0:
raise ValueError("Usage writer limits must be positive")
self.db = db
self.capacity = capacity
self.batch_size = batch_size
self.flush_interval = flush_interval
self.shutdown_timeout = shutdown_timeout
self._condition = threading.Condition()
self._pending: deque[UsageEvent] = deque()
self._thread: threading.Thread | None = None
self._accepting = False
self._deadline: float | None = None
self.dropped = 0
def start(self) -> None:
with self._condition:
if self._thread is not None and self._thread.is_alive():
raise RuntimeError("Previous usage writer is still running")
self._deadline = None
self._accepting = True
self._thread = threading.Thread(target=self._run, name="laya-usage-writer", daemon=True)
self._thread.start()
def submit(self, event: UsageEvent) -> bool:
with self._condition:
if not self._accepting or len(self._pending) >= self.capacity:
self.dropped += 1
# Avoid flooding logs when storage cannot keep up.
if self.dropped == 1 or self.dropped & (self.dropped - 1) == 0:
logger.warning("Usage queue unavailable/full; dropped events=%d", self.dropped)
return False
self._pending.append(event)
self._condition.notify()
return True
def stop(self) -> None:
with self._condition:
self._accepting = False
self._deadline = time.monotonic() + self.shutdown_timeout
self._condition.notify_all()
if self._thread is not None:
self._thread.join(self.shutdown_timeout)
if self._thread.is_alive():
logger.warning("Usage shutdown timed out; unflushed statistics may be lost")
def _next_batch(self) -> list[UsageEvent]:
with self._condition:
while not self._pending and self._accepting:
self._condition.wait()
flush_at = time.monotonic() + self.flush_interval
while self._accepting and len(self._pending) < self.batch_size:
remaining = flush_at - time.monotonic()
if remaining <= 0:
break
self._condition.wait(remaining)
if self._deadline is not None and time.monotonic() >= self._deadline:
self.dropped += len(self._pending)
self._pending.clear()
return [self._pending.popleft() for _ in range(min(self.batch_size, len(self._pending)))]
@staticmethod
def _write_batch(connection: sqlite3.Connection, batch: list[UsageEvent]) -> None:
latest: dict[int, str] = {}
for event in batch:
if event.key_id is not None:
latest[event.key_id] = max(latest.get(event.key_id, ""), event.created_at)
with connection:
connection.executemany(
"INSERT INTO usage_events (created_at,source,key_id,model,input_tokens,output_tokens,duration_ms) "
"VALUES (?,?,?,?,?,?,?)",
[(e.created_at, e.source, e.key_id, e.model, e.input_tokens, e.output_tokens, e.duration_ms)
for e in batch],
)
connection.executemany(
"UPDATE api_keys SET last_used_at=MAX(COALESCE(last_used_at,''),?) WHERE id=?",
[(created_at, key_id) for key_id, created_at in latest.items()],
)
def _run(self) -> None:
connection = None
try:
while batch := self._next_batch():
try:
if connection is None:
connection = self.db.open_connection()
self._write_batch(connection, batch)
except sqlite3.Error:
with self._condition:
self.dropped += len(batch)
logger.exception("Failed to persist %d usage events; dropping batch", len(batch))
# Do not retry an ambiguously committed batch and double-count it.
if connection is not None:
connection.close()
connection = None
finally:
with self._condition:
self._accepting = False
if self._pending:
self.dropped += len(self._pending)
logger.warning("Usage writer stopped; discarding %d pending events", len(self._pending))
self._pending.clear()
if connection is not None:
connection.close()
|