from __future__ import annotations from dataclasses import asdict from .diagnostics import diagnose_run from .kv_cache import KVCacheModel from .latency import AnalyticalLatencyModel from .metrics import summarize from .models import Request, SimulationConfig, SimulationResult, TimelinePoint from .profiles import get_accelerator, get_model from .workloads import generate_workload SCHEDULERS = { "static_fcfs", "continuous_fcfs", "continuous_sjf", "continuous_slo", "chunked_slo", } class Simulator: def __init__(self, cfg: SimulationConfig): if cfg.scheduler not in SCHEDULERS: raise ValueError(f"Unsupported scheduler: {cfg.scheduler}") self.cfg = cfg self.model = get_model(cfg.model) self.accelerator = get_accelerator(cfg.accelerator) self.latency = AnalyticalLatencyModel(self.model, self.accelerator, cfg.quantization) self.kv = KVCacheModel(self.latency, cfg) self.requests = generate_workload(cfg) self.pending_idx = 0 self.waiting: list[Request] = [] self.prefill_pending: list[Request] = [] self.active: list[Request] = [] self.completed: list[Request] = [] self.now = 0.0 self.busy_time = 0.0 self.peak_kv_gb = 0.0 self.timeline: list[TimelinePoint] = [] self.warnings: list[str] = [] def _admit_arrivals(self) -> None: while self.pending_idx < len(self.requests) and self.requests[self.pending_idx].arrival_time <= self.now + 1e-12: self.waiting.append(self.requests[self.pending_idx]) self.pending_idx += 1 def _next_arrival(self) -> float | None: if self.pending_idx >= len(self.requests): return None return self.requests[self.pending_idx].arrival_time def _waiting_sorted(self) -> list[Request]: if self.cfg.scheduler == "continuous_sjf": return sorted(self.waiting, key=lambda r: (r.prompt_tokens + r.output_tokens, r.arrival_time)) if self.cfg.scheduler in {"continuous_slo", "chunked_slo"}: # Least-slack-first proxy: deadline minus an analytical estimate of # remaining standalone service. Unlike plain EDF, this distinguishes # requests with the same relative SLO but heterogeneous token lengths. def slack(req: Request) -> tuple[float, float]: prefill = self.latency.prefill_seconds([max(req.remaining_prefill, 1)]) midpoint_context = req.prompt_tokens + max(req.output_tokens // 2, 1) decode = req.output_tokens * self.latency.decode_step_seconds([midpoint_context]) return (req.deadline_time - self.now - prefill - decode, req.arrival_time) return sorted(self.waiting, key=slack) return sorted(self.waiting, key=lambda r: r.arrival_time) def _record_timeline(self, force: bool = False) -> None: # Keep result payload bounded. This is display telemetry, not the event log. total_target = max(self.cfg.timeline_points, 20) if not force and len(self.timeline) >= total_target: stride = max(2, len(self.timeline) // total_target + 1) self.timeline = self.timeline[::stride] kv_used = self.kv.used_gb(self.active, self.prefill_pending) self.peak_kv_gb = max(self.peak_kv_gb, kv_used) point = TimelinePoint( time_s=self.now, waiting=len(self.waiting), prefill_pending=len(self.prefill_pending), decoding=len(self.active), completed=len(self.completed), kv_used_gb=kv_used, kv_capacity_gb=self.kv.capacity_gb, ) if not self.timeline or force or self.now - self.timeline[-1].time_s >= max(self.cfg.duration_s / total_target, 0.05): self.timeline.append(point) def _advance(self, delta: float) -> None: delta = max(delta, 0.0) self.busy_time += delta self.now += delta self._admit_arrivals() self._record_timeline() def _idle_to_next_arrival(self) -> bool: nxt = self._next_arrival() if nxt is None: return False self.now = max(self.now, nxt) self._admit_arrivals() self._record_timeline() return True def _mark_complete(self) -> None: done = [r for r in self.active if r.complete] for r in done: r.completion_time = self.now self.completed.append(r) if done: done_ids = {r.request_id for r in done} self.active = [r for r in self.active if r.request_id not in done_ids] def _prefill_full_requests(self) -> bool: slots = self.cfg.max_batch_size - len(self.active) if slots <= 0 or not self.waiting: return False selected: list[Request] = [] total_tokens = 0 for req in self._waiting_sorted(): if len(selected) >= slots: break if selected and total_tokens + req.remaining_prefill > self.cfg.max_batch_tokens: continue if not self.kv.can_admit(req, self.active, selected): continue selected.append(req) total_tokens += req.remaining_prefill if not selected: return False selected_ids = {r.request_id for r in selected} self.waiting = [r for r in self.waiting if r.request_id not in selected_ids] for req in selected: if req.first_prefill_time is None: req.first_prefill_time = self.now self._advance(self.latency.prefill_seconds([r.remaining_prefill for r in selected])) for req in selected: req.remaining_prefill = 0 self.active.append(req) return True def _prefill_chunked(self) -> bool: slots = self.cfg.max_batch_size - len(self.active) - len(self.prefill_pending) if slots > 0 and self.waiting: for req in self._waiting_sorted(): if slots <= 0: break if not self.kv.can_admit(req, self.active, self.prefill_pending): continue self.waiting.remove(req) if req.first_prefill_time is None: req.first_prefill_time = self.now self.prefill_pending.append(req) slots -= 1 if not self.prefill_pending: return False chunks: list[int] = [] selected: list[Request] = [] token_budget = self.cfg.max_batch_tokens for req in list(self.prefill_pending): if token_budget <= 0: break chunk = min(req.remaining_prefill, self.cfg.chunk_size, token_budget) if chunk <= 0: continue selected.append(req) chunks.append(chunk) token_budget -= chunk if not selected: return False # One prefill chunk. Decode is serviced on the next loop iteration, # producing the intended prefill/decode interleaving. self._advance(self.latency.prefill_seconds(chunks)) for req, chunk in zip(selected, chunks, strict=True): req.remaining_prefill -= chunk if req.remaining_prefill <= 0: self.prefill_pending.remove(req) self.active.append(req) return True def _decode_step(self) -> bool: if not self.active: return False contexts = [r.context_tokens for r in self.active] self._advance(self.latency.decode_step_seconds(contexts)) for req in self.active: req.generated_tokens += 1 if req.first_token_time is None: req.first_token_time = self.now self._mark_complete() return True def _run_static(self) -> None: # Static batching deliberately refuses new admission while a batch is # decoding. New arrivals queue until every member of the current batch # completes, giving a clean baseline against continuous batching. while len(self.completed) < len(self.requests): self._admit_arrivals() if not self.active: if not self.waiting and not self._idle_to_next_arrival(): break selected = self._waiting_sorted()[: self.cfg.max_batch_size] admitted: list[Request] = [] for req in selected: if self.kv.can_admit(req, admitted, None): admitted.append(req) if not admitted: self.warnings.append("No static batch could fit in the configured KV budget.") break ids = {r.request_id for r in admitted} self.waiting = [r for r in self.waiting if r.request_id not in ids] for req in admitted: req.first_prefill_time = self.now self._advance(self.latency.prefill_seconds([r.remaining_prefill for r in admitted])) for req in admitted: req.remaining_prefill = 0 self.active.append(req) # Finish this batch without admitting queued work into free slots. while self.active: contexts = [r.context_tokens for r in self.active] delta = self.latency.decode_step_seconds(contexts) self.busy_time += delta self.now += delta # Arrivals are queued but never admitted until the batch drains. self._admit_arrivals() for req in self.active: req.generated_tokens += 1 if req.first_token_time is None: req.first_token_time = self.now self._mark_complete() self._record_timeline() def _run_continuous(self) -> None: while len(self.completed) < len(self.requests): self._admit_arrivals() progressed = False if self.cfg.scheduler == "chunked_slo": # Decode first if work is active, then execute one prefill chunk. # This prevents long prompts from monopolizing the device. if self.active: progressed = self._decode_step() or progressed progressed = self._prefill_chunked() or progressed else: progressed = self._prefill_full_requests() or progressed progressed = self._decode_step() or progressed if not progressed: if self.waiting or self.prefill_pending: self.warnings.append( "Simulation stalled: queued requests could not fit within the configured KV budget." ) break if not self._idle_to_next_arrival(): break def run(self) -> SimulationResult: if not self.requests: self.warnings.append("The workload generator produced zero requests; increase duration or request rate.") self._record_timeline(force=True) if self.cfg.scheduler == "static_fcfs": self._run_static() else: self._run_continuous() self._record_timeline(force=True) makespan = max(self.now, self.cfg.duration_s if self.requests else 0.0) summary, latency = summarize(self.completed, self.cfg, makespan, self.busy_time) summary["requests_generated"] = len(self.requests) summary["requests_unfinished"] = len(self.requests) - len(self.completed) resource = { "model_weight_gb": self.latency.model_weight_gb, "kv_capacity_gb": self.kv.capacity_gb, "peak_kv_gb": self.peak_kv_gb, "peak_kv_utilization": self.peak_kv_gb / self.kv.capacity_gb if self.kv.capacity_gb > 0 else 0.0, "accelerator_vram_gb": self.accelerator.vram_gb, "topology": "colocated", "accelerator_instances": 1, "prefix_cache_gb": self.kv.shared_prefix_gb, "prefix_cache_hits": sum(1 for r in self.requests if r.prefix_cache_hit), "prefix_cache_hit_rate": (sum(1 for r in self.requests if r.prefix_cache_hit) / len(self.requests)) if self.requests else 0.0, "prefill_tokens_saved": sum(r.cached_prefix_tokens for r in self.requests), } request_rows = [] # Preserve a bounded sample for scatterplots/export. Aggregate metrics # still cover every completed request. for req in self.completed[:2000]: request_rows.append({ "request_id": req.request_id, "arrival_time": req.arrival_time, "prompt_tokens": req.prompt_tokens, "output_tokens": req.output_tokens, "cached_prefix_tokens": req.cached_prefix_tokens, "ttft_ms": (req.first_token_time - req.arrival_time) * 1000.0 if req.first_token_time is not None else None, "e2e_ms": (req.completion_time - req.arrival_time) * 1000.0 if req.completion_time is not None else None, }) diagnostics = diagnose_run(summary, latency, resource, self.cfg) provenance = { "simulator": "InferScale-Sim", "version": "0.3.0", "latency_profile_type": "analytical-reference", "profile_warning": "Reference profiles are analytical proxies, not measured hardware benchmarks.", "model_profile_source": self.model.source, "accelerator_profile_source": self.accelerator.source, "topology": "colocated", } return SimulationResult( config=self.cfg.to_dict(), provenance=provenance, summary=summary, latency=latency, resource=resource, diagnostics=diagnostics, requests=request_rows, timeline=[asdict(p) for p in self.timeline], warnings=self.warnings, ) def run_simulation(config: dict) -> dict: cfg = SimulationConfig.from_dict(config) if cfg.topology == "disaggregated_pd": from .disaggregated import DisaggregatedSimulator return DisaggregatedSimulator(cfg).run().to_dict() if cfg.topology != "colocated": raise ValueError(f"Unsupported topology: {cfg.topology}") return Simulator(cfg).run().to_dict()