monitor_server.py

"""Monitor server for vLLM Prefix Cache metrics."""

import time
import json
import threading
from dataclasses import dataclass, field
from http.server import HTTPServer, BaseHTTPRequestHandler


@dataclass
class CacheMetrics:
    """Thread-safe cache metrics collector."""
    total_queries: int = 0
    cache_hits: int = 0
    cache_misses: int = 0
    total_tokens_hit: int = 0
    total_tokens_missed: int = 0
    block_usage: int = 0
    total_blocks: int = 0
    lock: threading.Lock = field(default_factory=threading.Lock)

    def record_hit(self, num_tokens: int):
        with self.lock:
            self.total_queries += 1
            self.cache_hits += 1
            self.total_tokens_hit += num_tokens

    def record_miss(self, num_tokens: int):
        with self.lock:
            self.total_queries += 1
            self.cache_misses += 1
            self.total_tokens_missed += num_tokens

    def update_block_usage(self, used: int, total: int):
        with self.lock:
            self.block_usage = used
            self.total_blocks = total

    def snapshot(self) -> dict:
        with self.lock:
            total = max(self.cache_hits + self.cache_misses, 1)
            token_hit_rate = 0.0
            total_tokens = self.total_tokens_hit + self.total_tokens_missed
            if total_tokens > 0:
                token_hit_rate = self.total_tokens_hit / total_tokens

            return {
                "timestamp": time.time(),
                "query_hit_rate": self.cache_hits / total,
                "token_hit_rate": round(token_hit_rate, 4),
                "cache_hits": self.cache_hits,
                "cache_misses": self.cache_misses,
                "tokens_hit": self.total_tokens_hit,
                "tokens_missed": self.total_tokens_missed,
                "block_usage": self.block_usage,
                "total_blocks": self.total_blocks,
                "block_usage_pct": round(self.block_usage / max(self.total_blocks, 1) * 100, 2),
            }


# Singleton metrics collector
METRICS = CacheMetrics()


def get_metrics():
    return METRICS


class MetricsHandler(BaseHTTPRequestHandler):
    """HTTP handler for /cache_stats endpoint."""

    def do_GET(self):
        if self.path == "/cache_stats":
            metrics = METRICS.snapshot()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps(metrics, indent=2).encode())
        elif self.path == "/health":
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.end_headers()
            self.wfile.write(b"ok")
        else:
            self.send_response(404)
            self.end_headers()

    def log_message(self, format, *args):
        pass  # Suppress default logging


def start_monitor_server(host: str = "0.0.0.0", port: int = 9090):
    """Start the metrics HTTP server in a daemon thread."""
    server = HTTPServer((host, port), MetricsHandler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    print(f"Cache monitor running at http://{host}:{port}/cache_stats")
    return server