from __future__ import annotations import math import os import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from statistics import mean from typing import Any, Protocol from openai import OpenAI from llmserve_env.models import MetricsSnapshot, QuantizationTier, ServeAction, WorkloadSnapshot from server.trace_simulator import TraceSimulator class ServingBackend(Protocol): mode: str def reset(self, seed: int | None = None) -> None: ... def run_step(self, task_id: str, action: ServeAction, workload: WorkloadSnapshot) -> MetricsSnapshot: ... def describe(self) -> dict[str, Any]: ... class SimulatedServingBackend: mode = "sim" def __init__(self, seed: int = 42) -> None: try: self.simulator = TraceSimulator(seed=seed) except Exception as e: raise RuntimeError(f"Failed to initialize TraceSimulator in SimulatedServingBackend: {e}") from e def reset(self, seed: int | None = None) -> None: self.simulator.reset(seed=seed) def run_step(self, task_id: str, action: ServeAction, workload: WorkloadSnapshot) -> MetricsSnapshot: return self.simulator.simulate_step(task_id, action, workload) def describe(self) -> dict[str, Any]: return {"mode": self.mode, "provider": "simulator"} @dataclass class _RequestResult: latency_s: float ttft_ms: float itl_ms: float prompt_tokens: int completion_tokens: int total_tokens: int cost_usd: float truncated: bool class RealOpenAIBackend: mode = "real" def __init__( self, seed: int = 42, api_key: str | None = None, model: str | None = None, base_url: str | None = None, max_requests_per_step: int | None = None, max_prompt_tokens: int | None = None, max_completion_tokens: int | None = None, client: OpenAI | None = None, ) -> None: resolved_key = api_key or os.getenv("OPENAI_API_KEY") if not resolved_key and client is None: raise RuntimeError("OPENAI_API_KEY is required when LLMSERVE_MODE=real.") env_base_url = os.getenv("OPENAI_BASE_URL") resolved_base_url = (base_url or env_base_url or "").strip() or None self.seed = seed self.model = (model or os.getenv("LLMSERVE_REAL_MODEL", "gpt-4.1-mini")).strip() self.base_url = resolved_base_url self.max_requests_per_step = max_requests_per_step or int(os.getenv("LLMSERVE_REAL_MAX_REQUESTS_PER_STEP", "4")) self.max_prompt_tokens = max_prompt_tokens or int(os.getenv("LLMSERVE_REAL_MAX_PROMPT_TOKENS", "512")) self.max_completion_tokens = max_completion_tokens or int(os.getenv("LLMSERVE_REAL_MAX_COMPLETION_TOKENS", "64")) self.client = client or OpenAI( api_key=resolved_key, base_url=self.base_url, timeout=60.0, max_retries=2, ) self.pricing = { "gpt-4.1-mini": {"input_per_million": 0.40, "output_per_million": 1.60}, "gpt-4.1": {"input_per_million": 2.00, "output_per_million": 8.00}, "gpt-4o-mini": {"input_per_million": 0.15, "output_per_million": 0.60}, "gpt-4o": {"input_per_million": 2.50, "output_per_million": 10.00}, } def reset(self, seed: int | None = None) -> None: if seed is not None: self.seed = seed def run_step(self, task_id: str, action: ServeAction, workload: WorkloadSnapshot) -> MetricsSnapshot: request_count = max( 1, min( self.max_requests_per_step, action.batch_cap, max(1, workload.queue_depth + int(math.ceil(workload.arrival_rate))), ), ) prompts = [ self._build_request_payload(task_id, workload, action, request_index=index, request_count=request_count) for index in range(request_count) ] batch_start = time.perf_counter() with ThreadPoolExecutor(max_workers=request_count) as executor: results = list(executor.map(self._execute_request, prompts)) batch_latency_s = max(time.perf_counter() - batch_start, 1e-6) total_completion_tokens = sum(result.completion_tokens for result in results) total_prompt_tokens = sum(result.prompt_tokens for result in results) total_cost = sum(result.cost_usd for result in results) mean_total_tokens = sum(result.total_tokens for result in results) / len(results) throughput_tps = total_completion_tokens / batch_latency_s mean_prompt = workload.mean_prompt_length memory_factor = { "gpt-4.1-mini": 0.010, "gpt-4.1": 0.018, "gpt-4o-mini": 0.009, "gpt-4o": 0.016, }.get(self.model, 0.012) quant_factor = { QuantizationTier.FP16.value: 1.00, QuantizationTier.INT8.value: 0.84, QuantizationTier.INT4.value: 0.72, }[action.quantization_tier] gpu_memory_used_gb = max( 2.0, (total_prompt_tokens * memory_factor * quant_factor * max(action.kv_budget_fraction, 0.1)) / 10.0 + request_count * 0.35, ) cost_per_1k = max(0.0001, (total_cost / max(total_prompt_tokens + total_completion_tokens, 1)) * 1000.0) evictions = sum(1 for result in results if result.truncated) return MetricsSnapshot( p50_ttft_ms=_percentile([result.ttft_ms for result in results], 0.50), p99_ttft_ms=_percentile([result.ttft_ms for result in results], 0.99), p50_itl_ms=_percentile([result.itl_ms for result in results], 0.50), throughput_tps=max(1.0, throughput_tps), gpu_memory_used_gb=gpu_memory_used_gb, estimated_cost_per_1k=cost_per_1k, spec_acceptance_rate=min(0.6, action.speculation_depth / 8.0 * 0.35), eviction_events=evictions, slo_violations=0, requests_served=request_count, ) def describe(self) -> dict[str, Any]: return { "mode": self.mode, "provider": "openai", "model": self.model, "max_requests_per_step": self.max_requests_per_step, "max_prompt_tokens": self.max_prompt_tokens, "max_completion_tokens": self.max_completion_tokens, } def _build_request_payload( self, task_id: str, workload: WorkloadSnapshot, action: ServeAction, request_index: int, request_count: int, ) -> dict[str, Any]: priority_cutoff = max(1, int(round(request_count * workload.priority_fraction))) is_priority = request_index < priority_cutoff if action.priority_routing else False spread = (request_index - (request_count / 2.0)) / max(request_count, 1) target_prompt_tokens = max(32, int(workload.mean_prompt_length * (1.0 + spread * 0.35))) effective_prompt_tokens = max(16, int(target_prompt_tokens * action.kv_budget_fraction)) truncated = effective_prompt_tokens < target_prompt_tokens prompt = self._build_prompt(task_id, workload.phase, effective_prompt_tokens, is_priority=is_priority) return { "prompt": prompt, "target_prompt_tokens": target_prompt_tokens, "effective_prompt_tokens": effective_prompt_tokens, "truncated": truncated, "priority": is_priority, } def _execute_request(self, payload: dict[str, Any]) -> _RequestResult: start = time.perf_counter() response = self.client.chat.completions.create( model=self.model, temperature=0, max_completion_tokens=self.max_completion_tokens, messages=[ { "role": "system", "content": "You are a concise assistant. Answer the request directly in plain text.", }, {"role": "user", "content": payload["prompt"]}, ], ) latency_s = max(time.perf_counter() - start, 1e-6) usage = response.usage prompt_tokens = int(getattr(usage, "prompt_tokens", payload["effective_prompt_tokens"])) completion_tokens = int(getattr(usage, "completion_tokens", self.max_completion_tokens // 2)) total_tokens = int(getattr(usage, "total_tokens", prompt_tokens + completion_tokens)) ttft_ms = latency_s * 1000.0 * 0.35 itl_ms = max(1.0, ((latency_s * 1000.0) - ttft_ms) / max(completion_tokens, 1)) pricing = self.pricing.get(self.model, {"input_per_million": 0.40, "output_per_million": 1.60}) cost_usd = ( (prompt_tokens / 1_000_000.0) * pricing["input_per_million"] + (completion_tokens / 1_000_000.0) * pricing["output_per_million"] ) return _RequestResult( latency_s=latency_s, ttft_ms=ttft_ms, itl_ms=itl_ms, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens, cost_usd=cost_usd, truncated=bool(payload["truncated"]), ) def _build_prompt(self, task_id: str, phase: str, target_tokens: int, is_priority: bool) -> str: header = ( f"Task: {task_id}\n" f"Phase: {phase}\n" f"Priority: {is_priority}\n" "Summarize the impact of serving-policy changes on latency, throughput, and user experience.\n" ) filler_unit = "latency throughput queue kv cache scheduling token generation " filler = (filler_unit * ((target_tokens // 8) + 8)).strip() words = filler.split() return header + " ".join(words[:target_tokens]) def create_serving_backend(mode: str | None = None, seed: int = 42) -> ServingBackend: resolved_mode = (mode or os.getenv("LLMSERVE_MODE", "sim")).strip().lower() if resolved_mode == "sim": return SimulatedServingBackend(seed=seed) if resolved_mode == "real": provider = os.getenv("LLMSERVE_REAL_PROVIDER", "openai").strip().lower() if provider != "openai": raise RuntimeError(f"Unsupported LLMSERVE_REAL_PROVIDER: {provider}") return RealOpenAIBackend(seed=seed) raise RuntimeError(f"Unsupported LLMSERVE_MODE: {resolved_mode}") def _percentile(values: list[float], pct: float) -> float: ordered = sorted(values) if not ordered: return 0.0 index = min(len(ordered) - 1, max(0, int(round((len(ordered) - 1) * pct)))) return ordered[index]