src/inferscale/models.py

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Any


@dataclass
class ModelProfile:
    name: str
    params_b: float
    layers: int
    hidden_size: int
    attention_heads: int
    kv_heads: int
    default_dtype_bytes: float = 2.0
    source: str = "analytical-reference"

    @property
    def head_dim(self) -> int:
        return self.hidden_size // self.attention_heads


@dataclass
class AcceleratorProfile:
    name: str
    vram_gb: float
    peak_tflops_fp16: float
    bandwidth_gbps: float
    compute_efficiency: float
    bandwidth_efficiency: float
    source: str = "vendor-spec-reference"


@dataclass
class SimulationConfig:
    model: str = "Llama-3.1-8B"
    accelerator: str = "L4"
    scheduler: str = "continuous_fcfs"
    topology: str = "colocated"
    arrival_process: str = "poisson"
    request_rate_rps: float = 4.0
    duration_s: float = 60.0
    prompt_tokens_mean: int = 512
    prompt_tokens_cv: float = 0.35
    output_tokens_mean: int = 128
    output_tokens_cv: float = 0.35
    max_batch_size: int = 16
    max_batch_tokens: int = 8192
    chunk_size: int = 512
    kv_block_tokens: int = 16
    kv_memory_fraction: float = 0.85
    quantization: str = "fp16"
    seed: int = 7
    slo_ttft_ms: float = 500.0
    slo_e2e_ms: float = 8000.0
    slo_attainment_target: float = 0.99
    burst_multiplier: float = 3.0
    burst_period_s: float = 10.0
    timeline_points: int = 300

    # Prefix-cache scenario. v0.3 intentionally models one reusable shared
    # prefix rather than a full radix tree. Hits share one persistent KV entry.
    prefix_cache_enabled: bool = False
    shared_prefix_tokens: int = 256
    prefix_reuse_fraction: float = 0.0

    # Prefill/decode disaggregation scenario. These fields are ignored for the
    # colocated topology. Interconnect bandwidth is expressed in GB/s.
    prefill_accelerator: str = "L4"
    decode_accelerator: str = "L4"
    prefill_workers: int = 1
    decode_workers: int = 1
    interconnect_gbps: float = 50.0
    transfer_base_ms: float = 0.20

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "SimulationConfig":
        allowed = cls.__dataclass_fields__.keys()
        return cls(**{k: data[k] for k in allowed if k in data})

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class Request:
    request_id: int
    arrival_time: float
    prompt_tokens: int
    output_tokens: int
    deadline_time: float
    remaining_prefill: int
    cached_prefix_tokens: int = 0
    generated_tokens: int = 0
    first_prefill_time: float | None = None
    prefill_complete_time: float | None = None
    transfer_start_time: float | None = None
    transfer_end_time: float | None = None
    first_token_time: float | None = None
    completion_time: float | None = None
    decode_worker_id: int | None = None
    priority: int = 0

    @property
    def context_tokens(self) -> int:
        return self.prompt_tokens + self.generated_tokens

    @property
    def uncached_prompt_tokens(self) -> int:
        return max(0, self.prompt_tokens - self.cached_prefix_tokens)

    @property
    def prefix_cache_hit(self) -> bool:
        return self.cached_prefix_tokens > 0

    @property
    def complete(self) -> bool:
        return self.generated_tokens >= self.output_tokens


@dataclass
class RequestMetrics:
    request_id: int
    arrival_time: float
    prompt_tokens: int
    output_tokens: int
    ttft_ms: float
    e2e_ms: float
    tpot_ms: float
    queue_ms: float
    met_ttft_slo: bool
    met_e2e_slo: bool
    met_all_slos: bool


@dataclass
class TimelinePoint:
    time_s: float
    waiting: int
    prefill_pending: int
    decoding: int
    completed: int
    kv_used_gb: float
    kv_capacity_gb: float
    transfer_pending: int = 0
    decode_ready: int = 0
    prefill_active: int = 0


@dataclass
class SimulationResult:
    config: dict[str, Any]
    provenance: dict[str, Any]
    summary: dict[str, Any]
    latency: dict[str, Any]
    resource: dict[str, Any]
    diagnostics: dict[str, Any] = field(default_factory=dict)
    requests: list[dict[str, Any]] = field(default_factory=list)
    timeline: list[dict[str, Any]] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)