server.py

"""
Continuous Batching Profiler
Real PyTorch inference server with continuous batching.
"""

import csv
import gc
import math
import random
import statistics
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32
RESULTS = Path("results")
RESULTS.mkdir(exist_ok=True)


@dataclass
class Request:
    id: int
    prompt_ids: torch.Tensor
    max_output_tokens: int
    arrival_time: float

    generated_ids: List[int] = field(default_factory=list)
    prefill_done: bool = False
    prefill_pos: int = 0
    first_token_time: float = -1.0
    complete_time: float = -1.0
    past_key_values: Optional[object] = None
    last_tok: Optional[torch.Tensor] = None

    @property
    def prompt_len(self):
        return self.prompt_ids.shape[0]

    def done(self, max_out):
        return len(self.generated_ids) >= max_out

    def ttft_ms(self):
        if self.first_token_time < 0: return -1.0
        return (self.first_token_time - self.arrival_time) * 1000.0

    def latency_ms(self):
        if self.complete_time < 0: return -1.0
        return (self.complete_time - self.arrival_time) * 1000.0


def generate_workload(n_requests, arrival_rate, prompt_len_mean,
                      prompt_len_std, output_len_mean, output_len_std,
                      seed=42):
    rng = random.Random(seed)
    requests = []
    now = 0.0
    for i in range(n_requests):
        inter = -math.log(max(1e-9, 1 - rng.random())) / arrival_rate
        now += inter
        prompt_len = max(8, int(rng.gauss(prompt_len_mean, prompt_len_std)))
        output_len = max(1, int(rng.gauss(output_len_mean, output_len_std)))
        prompt_ids = torch.randint(0, 50000, (prompt_len,))
        requests.append(Request(
            id=i, prompt_ids=prompt_ids,
            max_output_tokens=output_len, arrival_time=now))
    return requests


def cleanup():
    gc.collect()
    if DEVICE == "cuda":
        try: torch.cuda.empty_cache()
        except Exception: pass


@torch.no_grad()
def do_prefill(model, req, start, end):
    """Prefill tokens [start:end]. Returns elapsed seconds."""
    assert end > start, f"Empty chunk: start={start} end={end}"
    chunk = req.prompt_ids[start:end].unsqueeze(0).to(DEVICE)
    t0 = time.perf_counter()
    if req.past_key_values is None:
        out = model(chunk, use_cache=True)
    else:
        out = model(chunk, use_cache=True, past_key_values=req.past_key_values)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0
    req.past_key_values = out.past_key_values
    req.prefill_pos = end
    req.last_tok = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
    del chunk, out
    return elapsed


@torch.no_grad()
def do_decode_step(model, req):
    """One decode step. Returns elapsed seconds."""
    t0 = time.perf_counter()
    out = model(req.last_tok, use_cache=True, past_key_values=req.past_key_values)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0
    req.past_key_values = out.past_key_values
    req.last_tok = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
    req.generated_ids.append(req.last_tok.item())
    del out
    return elapsed


def free_req(req):
    del req.past_key_values, req.last_tok
    req.past_key_values = None
    req.last_tok = None


# ── FCFS ─────────────────────────────────────────────────────────────────────

class FCFSServer:
    def __init__(self, model, max_output_tokens=32):
        self.model = model
        self.max_out = max_output_tokens

    def process(self, requests):
        queue = sorted(requests, key=lambda r: r.arrival_time)
        now = 0.0
        for req in queue:
            now = max(now, req.arrival_time)
            elapsed = do_prefill(self.model, req, 0, req.prompt_len)
            now += elapsed
            req.generated_ids.append(req.last_tok.item())
            req.first_token_time = now
            req.prefill_done = True
            max_out = min(self.max_out, req.max_output_tokens)
            for _ in range(max_out - 1):
                elapsed = do_decode_step(self.model, req)
                now += elapsed
            req.complete_time = now
            free_req(req)
            cleanup()
        return requests


# ── Continuous Batching ───────────────────────────────────────────────────────

class ContinuousBatchingServer:
    def __init__(self, model, max_batch_size=4, max_output_tokens=32):
        self.model   = model
        self.max_bs  = max_batch_size
        self.max_out = max_output_tokens

    def process(self, requests):
        pending   = sorted(requests, key=lambda r: r.arrival_time)
        active    = []
        completed = []
        now = 0.0
        ptr = [0]  # use list to allow mutation in nested scope

        def admit_pending(current_now):
            """Admit requests that have arrived and fit in batch."""
            elapsed_total = 0.0
            while (ptr[0] < len(pending) and
                   pending[ptr[0]].arrival_time <= current_now and
                   len(active) < self.max_bs):
                req = pending[ptr[0]]
                ptr[0] += 1
                e = do_prefill(self.model, req, 0, req.prompt_len)
                elapsed_total += e
                req.generated_ids.append(req.last_tok.item())
                req.first_token_time = current_now + elapsed_total
                req.prefill_done = True
                active.append(req)
            return elapsed_total

        now += admit_pending(now)

        while ptr[0] < len(pending) or active:
            if not active:
                if ptr[0] < len(pending):
                    now = pending[ptr[0]].arrival_time
                    now += admit_pending(now)
                continue

            # Decode one step for all active requests
            max_elapsed = 0.0
            for req in list(active):
                max_out = min(self.max_out, req.max_output_tokens)
                if req.done(max_out):
                    req.complete_time = now
                    active.remove(req)
                    completed.append(req)
                    free_req(req)
                    continue
                e = do_decode_step(self.model, req)
                max_elapsed = max(max_elapsed, e)

            now += max_elapsed if max_elapsed > 0 else 1e-6

            # Check completions
            for req in list(active):
                max_out = min(self.max_out, req.max_output_tokens)
                if req.done(max_out):
                    req.complete_time = now
                    active.remove(req)
                    completed.append(req)
                    free_req(req)

            # Admit new arrivals
            now += admit_pending(now)

        for req in active:
            req.complete_time = now
            completed.append(req)
            free_req(req)

        cleanup()
        return requests


# ── ChunkedPrefill ────────────────────────────────────────────────────────────

class ChunkedPrefillServer:
    def __init__(self, model, chunk_size=64, max_batch_size=4,
                 max_output_tokens=32):
        self.model   = model
        self.chunk   = chunk_size
        self.max_bs  = max_batch_size
        self.max_out = max_output_tokens

    def process(self, requests):
        pending    = sorted(requests, key=lambda r: r.arrival_time)
        prefilling = []
        decoding   = []
        completed  = []
        now = 0.0
        ptr = 0

        def admit_pending():
            nonlocal ptr
            while (ptr < len(pending) and
                   pending[ptr].arrival_time <= now and
                   len(prefilling) + len(decoding) < self.max_bs):
                req = pending[ptr]
                ptr += 1
                prefilling.append(req)

        admit_pending()

        while ptr < len(pending) or prefilling or decoding:
            if not prefilling and not decoding:
                if ptr < len(pending):
                    now = pending[ptr].arrival_time
                    admit_pending()
                continue

            # Prefill one chunk per prefilling request
            for req in list(prefilling):
                start = req.prefill_pos
                end   = min(start + self.chunk, req.prompt_len)
                if start >= req.prompt_len:
                    prefilling.remove(req)
                    continue
                elapsed = do_prefill(self.model, req, start, end)
                now += elapsed
                if end >= req.prompt_len:
                    req.generated_ids.append(req.last_tok.item())
                    req.first_token_time = now
                    req.prefill_done = True
                    prefilling.remove(req)
                    decoding.append(req)

            # Decode one step per decoding request
            max_elapsed = 0.0
            for req in list(decoding):
                max_out = min(self.max_out, req.max_output_tokens)
                if req.done(max_out):
                    req.complete_time = now
                    decoding.remove(req)
                    completed.append(req)
                    free_req(req)
                    continue
                elapsed = do_decode_step(self.model, req)
                max_elapsed = max(max_elapsed, elapsed)

            if max_elapsed > 0:
                now += max_elapsed

            # Check completions
            for req in list(decoding):
                max_out = min(self.max_out, req.max_output_tokens)
                if req.done(max_out):
                    req.complete_time = now
                    decoding.remove(req)
                    completed.append(req)
                    free_req(req)

            admit_pending()

        for req in decoding:
            req.complete_time = now
            completed.append(req)
            free_req(req)

        cleanup()
        return requests


# ── Metrics + Run ─────────────────────────────────────────────────────────────

def compute_metrics(requests, server_name, config):
    completed = [r for r in requests if r.complete_time > 0]
    if not completed: return None

    ttfts     = [r.ttft_ms() for r in completed if r.ttft_ms() >= 0]
    latencies = [r.latency_ms() for r in completed if r.latency_ms() >= 0]
    sim_dur   = max(r.complete_time for r in completed) - \
                min(r.arrival_time for r in requests)
    total_tok = sum(len(r.generated_ids) for r in completed)

    def pct(v, p):
        if not v: return 0.0
        sv = sorted(v)
        return sv[int(p * (len(sv) - 1))]

    return {
        "server":          server_name,
        **config,
        "n_requests":      len(requests),
        "n_completed":     len(completed),
        "ttft_mean_ms":    round(statistics.mean(ttfts), 1) if ttfts else 0,
        "ttft_p50_ms":     round(pct(ttfts, 0.50), 1),
        "ttft_p95_ms":     round(pct(ttfts, 0.95), 1),
        "ttft_p99_ms":     round(pct(ttfts, 0.99), 1),
        "latency_mean_ms": round(statistics.mean(latencies), 1) if latencies else 0,
        "latency_p95_ms":  round(pct(latencies, 0.95), 1),
        "throughput_rps":  round(len(completed) / sim_dur, 4) if sim_dur > 0 else 0,
        "throughput_tps":  round(total_tok / sim_dur, 4) if sim_dur > 0 else 0,
        "sim_duration_s":  round(sim_dur, 3),
    }


def run_one(model, cls, kw, workload, server_name):
    requests = generate_workload(
        n_requests     = workload["n_requests"],
        arrival_rate   = workload["arrival_rate"],
        prompt_len_mean= workload["prompt_len_mean"],
        prompt_len_std = workload.get("prompt_len_std", 32),
        output_len_mean= workload["output_len_mean"],
        output_len_std = workload.get("output_len_std", 8),
        seed           = workload.get("seed", 42),
    )
    server = cls(model, **kw)
    wall_t0 = time.perf_counter()
    server.process(requests)
    wall_s = time.perf_counter() - wall_t0

    cfg = {
        "arrival_rate": workload["arrival_rate"],
        "prompt_len":   workload["prompt_len_mean"],
        "output_len":   workload["output_len_mean"],
        **{k: v for k, v in kw.items() if k != "model"},
    }
    row = compute_metrics(requests, server_name, cfg)
    if row:
        print(f"    wall={wall_s:.1f}s  "
              f"TTFT mean={row['ttft_mean_ms']:.0f}ms "
              f"p95={row['ttft_p95_ms']:.0f}ms  "
              f"tput={row['throughput_rps']:.2f}rps")
    return row


def save_csv(path, rows):
    if not rows: return
    all_keys = []
    seen = set()
    for row in rows:
        for k in row.keys():
            if k not in seen:
                all_keys.append(k)
                seen.add(k)
    with open(path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=all_keys, extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            writer.writerow({k: row.get(k, "") for k in all_keys})
    print(f"\nSaved {len(rows)} rows -> {path}")


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    print("=" * 72)
    print("CONTINUOUS BATCHING PROFILER")
    print(f"Device: {DEVICE}  Dtype: {DTYPE}")
    if DEVICE == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")
    print("=" * 72)

    print("\nLoading gpt2...")
    tok = AutoTokenizer.from_pretrained("gpt2")
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        "gpt2", torch_dtype=DTYPE).to(DEVICE).eval()
    n = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"  {n:.0f}M params on {DEVICE}")

    all_rows = []
    base = {
        "n_requests": 10, "arrival_rate": 2.0,
        "prompt_len_mean": 128, "prompt_len_std": 32,
        "output_len_mean": 32,  "output_len_std": 8,
        "seed": 42,
    }

    # Exp 1: Main comparison
    print("\n" + "=" * 72)
    print("EXP 1: FCFS vs ContinuousBatching vs ChunkedPrefill")
    print("=" * 72)

    for name, cls, kw in [
        ("FCFS",
         FCFSServer, {"max_output_tokens": 32}),
        ("ContBatch_bs4",
         ContinuousBatchingServer, {"max_batch_size": 4, "max_output_tokens": 32}),
        ("ChunkedPrefill_c64_bs4",
         ChunkedPrefillServer, {"chunk_size": 64, "max_batch_size": 4, "max_output_tokens": 32}),
        ("ChunkedPrefill_c32_bs4",
         ChunkedPrefillServer, {"chunk_size": 32, "max_batch_size": 4, "max_output_tokens": 32}),
    ]:
        print(f"\n  [{name}]")
        row = run_one(model, cls, kw, base, name)
        if row: all_rows.append(row)

    # Exp 2: Arrival rate sweep
    print("\n" + "=" * 72)
    print("EXP 2: Arrival rate sweep")
    print("=" * 72)

    for rate in [1.0, 2.0, 4.0, 8.0]:
        wl = {**base, "arrival_rate": rate, "n_requests": 12}
        print(f"\n  rate={rate}")
        for name, cls, kw in [
            ("FCFS", FCFSServer, {"max_output_tokens": 32}),
            ("ChunkedPrefill_c64_bs4", ChunkedPrefillServer,
             {"chunk_size": 64, "max_batch_size": 4, "max_output_tokens": 32}),
        ]:
            row = run_one(model, cls, kw, wl, name)
            if row: all_rows.append(row)

    # Exp 3: Chunk size sweep
    print("\n" + "=" * 72)
    print("EXP 3: Chunk size sweep")
    print("=" * 72)

    for chunk in [16, 32, 64, 128]:
        print(f"\n  chunk_size={chunk}")
        row = run_one(model, ChunkedPrefillServer,
                      {"chunk_size": chunk, "max_batch_size": 4, "max_output_tokens": 32},
                      base, f"ChunkedPrefill_c{chunk}")
        if row: all_rows.append(row)

    save_csv(RESULTS / "batching_results.csv", all_rows)

    print("\n" + "=" * 72)
    print("FINAL SUMMARY")
    print("=" * 72)
    print(f"\n{'Server':<28s} {'TTFT mean':>10s} {'TTFT p95':>10s} "
          f"{'RPS':>8s} {'TPS':>8s}")
    print("-" * 66)
    for r in all_rows:
        print(f"{r['server']:<28s} "
              f"{r['ttft_mean_ms']:>10.0f} "
              f"{r['ttft_p95_ms']:>10.0f} "
              f"{r['throughput_rps']:>8.3f} "
              f"{r['throughput_tps']:>8.1f}")


if __name__ == "__main__":
    main()