src/inferscale/simulator.py

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()