Spaces:
Running
Running
| from __future__ import annotations | |
| import math | |
| import random | |
| from .models import Request, SimulationConfig | |
| def _sample_lognormal(mean: float, cv: float, rng: random.Random, minimum: int = 1) -> int: | |
| if cv <= 1e-9: | |
| return max(minimum, int(round(mean))) | |
| variance_ratio = cv * cv | |
| sigma2 = math.log(1.0 + variance_ratio) | |
| sigma = math.sqrt(sigma2) | |
| mu = math.log(max(mean, 1e-6)) - sigma2 / 2.0 | |
| return max(minimum, int(round(rng.lognormvariate(mu, sigma)))) | |
| def _arrival_times(cfg: SimulationConfig, rng: random.Random) -> list[float]: | |
| rate = max(cfg.request_rate_rps, 1e-9) | |
| arrivals: list[float] = [] | |
| t = 0.0 | |
| if cfg.arrival_process == "constant": | |
| step = 1.0 / rate | |
| while t < cfg.duration_s: | |
| arrivals.append(t) | |
| t += step | |
| return arrivals | |
| if cfg.arrival_process == "bursty": | |
| while t < cfg.duration_s: | |
| phase = int(t // max(cfg.burst_period_s, 0.1)) % 2 | |
| local_rate = rate * (cfg.burst_multiplier if phase else 0.55) | |
| t += rng.expovariate(max(local_rate, 1e-9)) | |
| if t < cfg.duration_s: | |
| arrivals.append(t) | |
| return arrivals | |
| if cfg.arrival_process == "trace": | |
| return [] | |
| if cfg.arrival_process != "poisson": | |
| raise ValueError(f"Unknown arrival process: {cfg.arrival_process}") | |
| while t < cfg.duration_s: | |
| t += rng.expovariate(rate) | |
| if t < cfg.duration_s: | |
| arrivals.append(t) | |
| return arrivals | |
| def _cache_tokens(cfg: SimulationConfig, prompt: int, cache_rng: random.Random) -> int: | |
| if not cfg.prefix_cache_enabled or cfg.shared_prefix_tokens <= 0 or cfg.prefix_reuse_fraction <= 0: | |
| return 0 | |
| if cache_rng.random() >= min(max(cfg.prefix_reuse_fraction, 0.0), 1.0): | |
| return 0 | |
| return min(cfg.shared_prefix_tokens, max(prompt - 1, 0)) | |
| def _from_trace(cfg: SimulationConfig, cache_rng: random.Random) -> list[Request]: | |
| rows: list[tuple[float, int, int]] = [] | |
| for idx, raw in enumerate(cfg.trace_requests): | |
| try: | |
| arrival = float(raw["arrival_time"]) | |
| prompt = int(raw["prompt_tokens"]) | |
| output = int(raw["output_tokens"]) | |
| except (KeyError, TypeError, ValueError) as exc: | |
| raise ValueError( | |
| f"Trace row {idx} must contain numeric arrival_time, prompt_tokens, and output_tokens" | |
| ) from exc | |
| if not math.isfinite(arrival) or arrival < 0: | |
| raise ValueError(f"Trace row {idx} has invalid arrival_time") | |
| if prompt < 1 or output < 1: | |
| raise ValueError(f"Trace row {idx} token counts must be positive") | |
| rows.append((arrival, prompt, output)) | |
| rows.sort(key=lambda row: row[0]) | |
| requests: list[Request] = [] | |
| for request_id, (arrival, prompt, output) in enumerate(rows): | |
| cached = _cache_tokens(cfg, prompt, cache_rng) | |
| requests.append( | |
| Request( | |
| request_id=request_id, | |
| arrival_time=arrival, | |
| prompt_tokens=prompt, | |
| output_tokens=output, | |
| deadline_time=arrival + cfg.slo_e2e_ms / 1000.0, | |
| remaining_prefill=max(0, prompt - cached), | |
| cached_prefix_tokens=cached, | |
| ) | |
| ) | |
| return requests | |
| def generate_workload(cfg: SimulationConfig) -> list[Request]: | |
| rng = random.Random(cfg.seed) | |
| cache_rng = random.Random(cfg.seed ^ 0x5A17CACE) | |
| if cfg.arrival_process == "trace": | |
| return _from_trace(cfg, cache_rng) | |
| requests: list[Request] = [] | |
| for idx, arrival in enumerate(_arrival_times(cfg, rng)): | |
| prompt = _sample_lognormal(cfg.prompt_tokens_mean, cfg.prompt_tokens_cv, rng) | |
| output = _sample_lognormal(cfg.output_tokens_mean, cfg.output_tokens_cv, rng) | |
| cached = _cache_tokens(cfg, prompt, cache_rng) | |
| requests.append( | |
| Request( | |
| request_id=idx, | |
| arrival_time=arrival, | |
| prompt_tokens=prompt, | |
| output_tokens=output, | |
| deadline_time=arrival + cfg.slo_e2e_ms / 1000.0, | |
| remaining_prefill=max(0, prompt - cached), | |
| cached_prefix_tokens=cached, | |
| ) | |
| ) | |
| return requests | |