jundot--omlx/omlx/admin/benchmark.py

# SPDX-License-Identifier: Apache-2.0
"""Benchmark execution logic for oMLX admin panel.

Provides single-request and continuous-batching benchmarks with
real-time progress reporting via SSE events.
"""

import asyncio
import json
import logging
import re
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

from pydantic import BaseModel, field_validator

from .external_api import ExternalAPIClient, ExternalEndpointConfig

try:
    import mlx.core as mx

    HAS_MLX = True
except ImportError:
    HAS_MLX = False

logger = logging.getLogger(__name__)

# Module-level storage for active benchmark runs
_benchmark_runs: dict[str, "BenchmarkRun"] = {}

# Valid prompt lengths for single request tests
VALID_PROMPT_LENGTHS = [1024, 4096, 8192, 16384, 32768, 65536, 131072, 200000]

# Valid batch sizes for continuous batching tests
VALID_BATCH_SIZES = [2, 4, 8]


class BenchmarkRequest(BaseModel):
    """Request model for starting a benchmark."""

    model_id: str
    prompt_lengths: list[int]
    generation_length: int = 128
    batch_sizes: list[int] = []
    force_lm_engine: bool = False
    # When set, the benchmark runs against a remote OpenAI-compatible
    # endpoint instead of a local engine and model_id is the remote
    # model name (not validated against the local catalog).
    external: Optional[ExternalEndpointConfig] = None

    @field_validator("prompt_lengths")
    @classmethod
    def validate_prompt_lengths(cls, v: list[int]) -> list[int]:
        if not v:
            raise ValueError("At least one prompt length is required")
        for pl in v:
            if pl not in VALID_PROMPT_LENGTHS:
                raise ValueError(
                    f"Invalid prompt length {pl}. Must be one of {VALID_PROMPT_LENGTHS}"
                )
        return sorted(v)

    @field_validator("batch_sizes")
    @classmethod
    def validate_batch_sizes(cls, v: list[int]) -> list[int]:
        for bs in v:
            if bs not in VALID_BATCH_SIZES:
                raise ValueError(
                    f"Invalid batch size {bs}. Must be one of {VALID_BATCH_SIZES}"
                )
        return sorted(v)


@dataclass
class BenchmarkRun:
    """Tracks the state of a running benchmark.

    SSE delivery model: events are appended to `events` (append-only
    log) under `cond`. Subscribers replay `events` from offset 0 then
    wait on `cond` for new entries. `terminal` is set once the final
    event (`upload_done` / `error`) has been published so subscribers
    know to close their stream rather than wait for a follow-up.
    """

    bench_id: str
    request: BenchmarkRequest
    status: str = "running"  # running, completed, cancelled, error
    events: list[dict] = field(default_factory=list)
    cond: asyncio.Condition = field(default_factory=asyncio.Condition)
    terminal: bool = False
    task: Optional[asyncio.Task] = None
    results: list[dict] = field(default_factory=list)
    error_message: str = ""
    # Experimental flags active when the benchmark started. When non-empty
    # the run's results are not uploaded to omlx.ai community benchmarks
    # because experimental features skew the numbers.
    experimental_features: list[str] = field(default_factory=list)
    # Mirror of the upload SSE events so REST consumers (e.g. native Swift
    # app polling /results) can render leaderboard status without opening
    # the stream. Phases: "idle" → "uploading" → "done" | "skipped". The
    # browser HTML still consumes the SSE stream directly; this is purely
    # additive state that lives alongside it.
    upload_state: dict = field(default_factory=lambda: {
        "phase": "idle",
        "results": [],          # per-context-length: {context_length, id?, url?, duplicate?, error?}
        "total": 0,
        "success_count": 0,
        "failed_count": 0,
        "owner_hash": None,     # display hash, populated on upload_done
        "skipped_reason": None, # e.g. "experimental_features"
        "skipped_features": [],
    })


# Event types that close the SSE stream for a bench run. `done` is NOT
# terminal — it marks "tests finished, upload starting"; the real end of
# stream is `upload_done` (or `error`).
_BENCH_TERMINAL_TYPES = frozenset({"upload_done", "error"})


_EXPERIMENTAL_FEATURE_FLAGS = (
    ("dflash_enabled", "dflash"),
    ("specprefill_enabled", "specprefill"),
    ("turboquant_kv_enabled", "turboquant"),
    ("mtp_enabled", "mtp"),
    ("vlm_mtp_enabled", "vlm_mtp"),
)


def _detect_experimental_features(model_settings: Any) -> list[str]:
    """Return benchmark-skewing model features enabled in settings."""
    return [
        feature
        for attr, feature in _EXPERIMENTAL_FEATURE_FLAGS
        if getattr(model_settings, attr, False)
    ]


def get_run(bench_id: str) -> Optional[BenchmarkRun]:
    """Get a benchmark run by ID."""
    return _benchmark_runs.get(bench_id)


def get_active_run() -> Optional[BenchmarkRun]:
    """Return the currently-running throughput benchmark, if any.

    Discovery surface for clients that need to attach to an in-progress
    run without knowing the bench_id upfront (page refresh, second tab).
    Returns the first run with status == "running"; throughput benches
    are 1-at-a-time so there's never more than one.
    """
    for run in _benchmark_runs.values():
        if run.status == "running":
            return run
    return None


def create_run(request: BenchmarkRequest) -> BenchmarkRun:
    """Create and register a new benchmark run."""
    bench_id = f"bench-{uuid.uuid4().hex[:12]}"
    run = BenchmarkRun(bench_id=bench_id, request=request)
    _benchmark_runs[bench_id] = run
    return run


def cleanup_old_runs(max_runs: int = 10) -> None:
    """Remove old completed runs to prevent memory leaks."""
    completed = [
        (bid, r)
        for bid, r in _benchmark_runs.items()
        if r.status in ("completed", "cancelled", "error")
    ]
    if len(completed) > max_runs:
        for bid, _ in completed[:-max_runs]:
            del _benchmark_runs[bid]



def _generate_prompt(tokenizer: Any, target_tokens: int) -> str:
    """Generate a prompt with exactly target_tokens tokens.

    Uses a unique UUID prefix to prevent SSD cache hits from previous sessions.
    """
    unique_prefix = f"BENCH-{uuid.uuid4().hex} "
    filler = (
        "The quick brown fox jumps over the lazy dog. "
        "In the realm of artificial intelligence, large language models "
        "have demonstrated remarkable capabilities across diverse tasks. "
    )

    # Build a large enough text
    text = unique_prefix + filler * (target_tokens // 10 + 1)
    tokens = tokenizer.encode(text)

    if len(tokens) < target_tokens:
        # Need more tokens, repeat more
        text = unique_prefix + filler * (target_tokens // 5 + 1)
        tokens = tokenizer.encode(text)

    # Truncate to exact target length
    tokens = tokens[:target_tokens]
    return tokenizer.decode(tokens)


# The external path has no tokenizer, so prompt lengths are approximated by
# repeating this filler (~30 tokens per repetition in common BPE vocabs).
# Results report the endpoint's actual usage.prompt_tokens.
_EXTERNAL_FILLER = (
    "The quick brown fox jumps over the lazy dog. "
    "In the realm of artificial intelligence, large language models "
    "have demonstrated remarkable capabilities across diverse tasks. "
)
_EXTERNAL_FILLER_TOKENS = 30


def _generate_external_prompt(target_tokens: int) -> str:
    """Generate an approximately target_tokens-long prompt without a tokenizer.

    Uses a unique UUID prefix so remote prefix caches cannot skew results.
    """
    unique_prefix = f"BENCH-{uuid.uuid4().hex} "
    repeats = max(1, target_tokens // _EXTERNAL_FILLER_TOKENS)
    return unique_prefix + _EXTERNAL_FILLER * repeats


def _compute_single_metrics(
    prompt_tokens: int,
    completion_tokens: int,
    start_time: float,
    first_token_time: float,
    end_time: float,
    peak_memory: int,
    cached_tokens: int,
    prefill_duration_s: float | None = None,
    generation_duration_s: float | None = None,
    generation_measured: bool = True,
) -> dict:
    """Compute all metrics for a single request benchmark."""
    ttft_s = first_token_time - start_time
    prefill_duration = (
        prefill_duration_s if prefill_duration_s is not None else ttft_s
    )
    gen_duration = (
        generation_duration_s
        if generation_duration_s is not None
        else end_time - first_token_time
    )
    e2e_duration = end_time - start_time

    ttft_ms = ttft_s * 1000
    if generation_measured and completion_tokens > 1 and gen_duration > 0:
        tpot_ms = (gen_duration / (completion_tokens - 1)) * 1000
        gen_tps = completion_tokens / gen_duration
    else:
        tpot_ms = 0.0
        gen_tps = 0.0
    processing_tps = prompt_tokens / max(prefill_duration, 1e-9)
    total_throughput = (prompt_tokens + completion_tokens) / max(e2e_duration, 1e-9)

    return {
        "ttft_ms": round(ttft_ms, 1),
        "tpot_ms": round(tpot_ms, 2),
        "gen_tps": round(gen_tps, 1),
        "processing_tps": round(processing_tps, 1),
        "e2e_latency_s": round(e2e_duration, 3),
        "total_throughput": round(total_throughput, 1),
        "peak_memory_bytes": peak_memory,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "cached_tokens": cached_tokens,
    }


def _get_batch_benchmark_core(engine: Any) -> Any | None:
    """Return the scheduler core when this engine supports batch benchmarks."""
    engine_core = getattr(engine, "_engine", None)
    if engine_core is None:
        return None
    if not callable(getattr(engine_core, "add_request", None)):
        return None
    if not callable(getattr(engine_core, "stream_outputs", None)):
        return None
    return engine_core


async def _send_event(run: BenchmarkRun, event: dict) -> None:
    """Append an event to the run's log and wake any subscribers.

    Sets `run.terminal` when the event ends the stream so subscribers
    can return rather than wait for an event that will never come.
    """
    async with run.cond:
        run.events.append(event)
        if event.get("type") in _BENCH_TERMINAL_TYPES:
            run.terminal = True
        run.cond.notify_all()


async def _run_single_test(
    engine: Any,
    prompt: str,
    max_tokens: int,
    pp_len: int,
) -> dict:
    """Run a single request benchmark test and return metrics."""
    # Reset peak memory tracking
    try:
        mx.reset_peak_memory()
    except Exception:
        pass

    start_time = time.perf_counter()
    first_token_time = None
    last_generated_token_time = None
    last_output = None
    prev_completion_tokens = 0

    async for output in engine.stream_generate(
        prompt=prompt,
        max_tokens=max_tokens,
        temperature=0.0,
        top_p=1.0,
    ):
        # Detect first generated token via completion_tokens count,
        # not new_text. Some models (e.g. Harmony/gpt-oss) produce
        # protocol tokens that don't yield visible new_text.
        completion_delta = output.completion_tokens - prev_completion_tokens
        if completion_delta > 0:
            generated_at = getattr(output, "generated_at", None)
            generated_until = getattr(output, "generated_until", None)
            output_first_token_time = (
                float(generated_at) if generated_at is not None else time.perf_counter()
            )
            if first_token_time is None:
                first_token_time = output_first_token_time
            if generated_until is not None:
                last_generated_token_time = float(generated_until)
            elif completion_delta == 1:
                last_generated_token_time = output_first_token_time
        prev_completion_tokens = output.completion_tokens
        last_output = output

    end_time = time.perf_counter()

    if first_token_time is None:
        first_token_time = end_time

    # Get peak memory
    try:
        peak_memory = mx.get_peak_memory()
    except Exception:
        peak_memory = 0

    prompt_tokens = last_output.prompt_tokens if last_output else 0
    completion_tokens = last_output.completion_tokens if last_output else 0
    cached_tokens = last_output.cached_tokens if last_output else 0

    if cached_tokens > 0:
        logger.warning(
            f"Benchmark test pp{pp_len} had {cached_tokens} cached tokens "
            f"(expected 0). Results may not reflect true prefill performance."
        )

    prefill_duration_s = None
    generation_duration_s = None
    producer_generation_duration_s = None
    metric_completion_tokens = completion_tokens
    if first_token_time is not None and last_generated_token_time is not None:
        measured_duration = last_generated_token_time - first_token_time
        if measured_duration > 0:
            producer_generation_duration_s = measured_duration
    if last_output is not None:
        prompt_tps = float(getattr(last_output, "prompt_tps", 0.0) or 0.0)
        if prompt_tps > 0 and prompt_tokens > 0:
            prefill_duration_s = prompt_tokens / prompt_tps

        canvas_tps = float(getattr(last_output, "diffusion_canvas_tps", 0.0) or 0.0)
        canvas_tokens = int(getattr(last_output, "diffusion_canvas_tokens", 0) or 0)
        if canvas_tps > 0 and canvas_tokens > 0:
            metric_completion_tokens = canvas_tokens
            generation_duration_s = canvas_tokens / canvas_tps
        else:
            generation_tps = float(getattr(last_output, "generation_tps", 0.0) or 0.0)
            if generation_tps > 0 and completion_tokens > 0:
                generation_duration_s = completion_tokens / generation_tps

    if generation_duration_s is None:
        generation_duration_s = producer_generation_duration_s

    generation_measured = generation_duration_s is not None

    return _compute_single_metrics(
        prompt_tokens=prompt_tokens,
        completion_tokens=metric_completion_tokens,
        start_time=start_time,
        first_token_time=first_token_time,
        end_time=end_time,
        peak_memory=peak_memory,
        cached_tokens=cached_tokens,
        prefill_duration_s=prefill_duration_s,
        generation_duration_s=generation_duration_s,
        generation_measured=generation_measured,
    )


async def _run_batch_test(
    engine: Any,
    prompts: list[str],
    prompt_tokens: int,
    max_tokens: int,
    batch_size: int,
) -> dict:
    """Run a continuous batching benchmark test.

    Submits batch_size concurrent requests via the engine core and measures
    aggregate throughput including pp TPS and tg TPS.

    Args:
        prompts: List of prompts (one per request). For same-prompt tests,
                 all entries are identical. For different-prompt tests, each
                 has a unique UUID prefix.
        prompt_tokens: Number of prompt tokens per request (for pp TPS calc).
    """
    from ..request import SamplingParams

    engine_core = _get_batch_benchmark_core(engine)
    if engine_core is None:
        raise ValueError("Engine does not support batch benchmarks")

    sampling_params = SamplingParams(
        max_tokens=max_tokens,
        temperature=0.0,
        top_p=1.0,
    )

    async def _single_request(prompt: str) -> dict:
        """Run a single request within the batch."""
        start = time.perf_counter()
        first_token = None
        tokens = 0
        prev_tokens = 0

        request_id = await engine_core.add_request(
            prompt=prompt,
            sampling_params=sampling_params,
        )

        async for output in engine_core.stream_outputs(request_id):
            if first_token is None and output.completion_tokens > prev_tokens:
                first_token = time.perf_counter()
            prev_tokens = output.completion_tokens
            if output.finished:
                tokens = output.completion_tokens

        end = time.perf_counter()
        if first_token is None:
            first_token = end

        return {
            "ttft_s": first_token - start,
            "first_token_abs": first_token,
            "end_abs": end,
            "completion_tokens": tokens,
        }

    # Submit all requests concurrently
    wall_start = time.perf_counter()
    results = await asyncio.gather(
        *[_single_request(prompts[i]) for i in range(batch_size)]
    )
    wall_end = time.perf_counter()

    # Aggregate metrics
    total_gen_tokens = sum(r["completion_tokens"] for r in results)
    total_prompt_tokens = prompt_tokens * batch_size
    wall_time = wall_end - wall_start
    avg_ttft_ms = (sum(r["ttft_s"] for r in results) / batch_size) * 1000

    # pp TPS: total prompt tokens / time until ALL requests finish prefill
    max_first_token = max(r["first_token_abs"] for r in results)
    prefill_wall_time = max_first_token - wall_start
    pp_tps = total_prompt_tokens / max(prefill_wall_time, 1e-9)

    # tg TPS: total generated tokens / generation wall time
    # Generation starts when the last request finishes prefill
    gen_wall_time = wall_end - max_first_token
    tg_tps = total_gen_tokens / max(gen_wall_time, 1e-9)

    return {
        "pp_tps": round(pp_tps, 1),
        "tg_tps": round(tg_tps, 1),
        "avg_ttft_ms": round(avg_ttft_ms, 1),
        "e2e_latency_s": round(wall_time, 3),
        "total_gen_tokens": total_gen_tokens,
        "batch_size": batch_size,
    }


async def _run_external_single_test(
    client: ExternalAPIClient,
    prompt: str,
    max_tokens: int,
) -> dict:
    """Run a single-request benchmark against an external endpoint.

    Token counts come from the endpoint's streamed usage payload, never
    from counting SSE chunks (providers batch multiple tokens per chunk).
    Prefill duration is not observable remotely, so pp TPS falls back to
    prompt_tokens / TTFT (network latency included). Peak memory is not
    measurable for a remote host.
    """
    stats = await client.stream_chat_completion(
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
        temperature=0.0,
    )
    gen_duration = stats.last_content_time - stats.first_content_time
    metrics = _compute_single_metrics(
        prompt_tokens=stats.prompt_tokens,
        completion_tokens=stats.completion_tokens,
        start_time=stats.start_time,
        first_token_time=stats.first_content_time,
        end_time=stats.end_time,
        peak_memory=0,
        cached_tokens=stats.cached_tokens,
        prefill_duration_s=None,
        generation_duration_s=gen_duration if gen_duration > 0 else None,
        generation_measured=gen_duration > 0,
    )
    metrics["peak_memory_bytes"] = None
    return metrics


async def _run_external_batch_test(
    client: ExternalAPIClient,
    prompts: list[str],
    max_tokens: int,
    batch_size: int,
) -> dict:
    """Run a concurrent-requests benchmark against an external endpoint.

    Mirrors _run_batch_test aggregation, with actual per-request token
    counts taken from each stream's usage payload.
    """
    wall_start = time.perf_counter()
    stats_list = await asyncio.gather(*[
        client.stream_chat_completion(
            messages=[{"role": "user", "content": prompt}],
            max_tokens=max_tokens,
            temperature=0.0,
        )
        for prompt in prompts
    ])
    wall_end = time.perf_counter()

    total_gen_tokens = sum(s.completion_tokens for s in stats_list)
    total_prompt_tokens = sum(s.prompt_tokens for s in stats_list)
    wall_time = wall_end - wall_start
    avg_ttft_ms = (
        sum(s.first_content_time - s.start_time for s in stats_list) / batch_size
    ) * 1000

    # pp TPS: total prompt tokens / time until ALL requests emit content
    max_first_token = max(s.first_content_time for s in stats_list)
    pp_tps = total_prompt_tokens / max(max_first_token - wall_start, 1e-9)

    # tg TPS: total generated tokens / generation wall time
    tg_tps = total_gen_tokens / max(wall_end - max_first_token, 1e-9)

    return {
        "pp_tps": round(pp_tps, 1),
        "tg_tps": round(tg_tps, 1),
        "avg_ttft_ms": round(avg_ttft_ms, 1),
        "e2e_latency_s": round(wall_time, 3),
        "total_gen_tokens": total_gen_tokens,
        "batch_size": batch_size,
    }


OMLX_AI_API_URL = "https://omlx.ai/api/benchmarks"

# Quantization patterns to strip from model directory names
_QUANT_SUFFIXES = re.compile(
    r"[-_](2bit|3bit|4bit|6bit|8bit|fp16|bf16|fp32|MXFP4|NVFP4)$", re.IGNORECASE
)
_MLX_SUFFIXES = re.compile(r"[-_]?MLX[-_]?", re.IGNORECASE)


def _detect_quantization(model_path: str) -> str:
    """Detect model quantization from config.json or directory name.

    Fallback chain: config.json → directory name → "unknown"
    """
    config_path = Path(model_path) / "config.json"
    if config_path.exists():
        try:
            with open(config_path) as f:
                config = json.load(f)
            qconfig = config.get("quantization_config", {})
            bits = qconfig.get("bits")
            if bits is not None:
                return f"{bits}bit"
        except Exception:
            pass

    # Fallback: extract from directory name
    dirname = Path(model_path).name
    match = re.search(
        r"(2bit|3bit|4bit|6bit|8bit|fp16|bf16|MXFP4|NVFP4)", dirname, re.IGNORECASE
    )
    if match:
        return match.group(1).lower()

    return "unknown"


def _clean_model_name(model_id: str, quantization: str) -> str:
    """Clean model directory name for display as model_name.

    Strips quantization suffixes and MLX markers.
    e.g. "Qwen3-30B-A3B-4bit" → "Qwen3-30B-A3B"
    """
    name = model_id
    name = _QUANT_SUFFIXES.sub("", name)
    name = _MLX_SUFFIXES.sub("", name)
    return name.strip("-_ ")


def _sanitize_upload_error(resp: Any) -> str:
    """Extract a user-presentable error string from a failed upload response.

    Avoids dumping raw HTML bodies (e.g. Cloudflare's "Just a moment..."
    challenge interstitial) into the dashboard's red-x error column.
    Detects CF mitigation specifically so users get actionable context
    instead of a 5KB markup blob.

    Resolution order:
    1. Cloudflare challenge — header ``cf-mitigated: challenge`` is
       authoritative; a body sniff for "just a moment" / "cf-chl" covers
       edge transports that strip the header.
    2. JSON envelope — the omlx.ai API's normal error shape; extract
       ``error`` / ``detail`` / ``message`` if present, truncated.
    3. Plain-text body — short responses only; HTML-looking bodies are
       collapsed to a one-line "non-JSON response (N bytes)" hint.
    4. Fallback to the bare HTTP status code.
    """
    headers = getattr(resp, "headers", {}) or {}
    cf_mitigated = str(headers.get("cf-mitigated", "")).lower()
    body = getattr(resp, "text", "") or ""
    status = getattr(resp, "status_code", "?")

    body_head = body[:512].lower()
    if cf_mitigated == "challenge" or "just a moment" in body_head or "cf-chl" in body_head:
        return (
            f"Upload blocked by Cloudflare (HTTP {status}). "
            f"This is a server-side issue with omlx.ai — retry later or "
            f"report it to the maintainer."
        )

    try:
        data = resp.json()
        msg = data.get("error") or data.get("detail") or data.get("message")
        if msg:
            return str(msg)[:300]
    except Exception:
        pass

    text = body.strip()
    if "<" in text and ">" in text:
        return f"HTTP {status} — unexpected non-JSON response ({len(body)} bytes)"
    return text[:300] or f"HTTP {status}"


async def _upload_to_omlx_ai(run: BenchmarkRun, engine_pool: Any) -> None:
    """Upload benchmark results to omlx.ai community benchmarks.

    Sends each single-request result as a separate submission,
    grouped by submission_group. Upload failures don't affect
    the benchmark run status.
    """
    import requests

    from .._version import __version__
    from ..utils.hardware import (
        compute_owner_hash,
        get_chip_name,
        get_gpu_core_count,
        get_io_platform_uuid,
        get_os_version,
        get_total_memory_gb,
        parse_chip_info,
    )

    # Skip upload when experimental features were active during the run.
    # These features skew throughput and would pollute the community
    # leaderboard if mixed in unmarked.
    if run.experimental_features:
        run.upload_state["phase"] = "skipped"
        run.upload_state["skipped_reason"] = "experimental_features"
        run.upload_state["skipped_features"] = list(run.experimental_features)
        await _send_event(run, {
            "type": "upload_skipped",
            "reason": "experimental_features",
            "features": list(run.experimental_features),
        })
        logger.info(
            f"Benchmark upload skipped: experimental features active: "
            f"{run.experimental_features}"
        )
        return

    run.upload_state["phase"] = "uploading"
    await _send_event(run, {
        "type": "progress",
        "phase": "upload",
        "message": "Uploading to community benchmarks...",
        "current": 0,
        "total": 0,
    })

    # Collect hardware info
    chip_string = get_chip_name()
    chip_name, chip_variant = parse_chip_info(chip_string)
    memory_gb = round(get_total_memory_gb())
    gpu_cores = get_gpu_core_count()
    os_version = get_os_version()
    omlx_version = __version__

    # Compute owner_hash
    owner_hash_full = None
    owner_hash_display = None
    io_uuid = get_io_platform_uuid()
    if io_uuid:
        owner_hash_full = compute_owner_hash(io_uuid, chip_name, gpu_cores, memory_gb)
        # Display hash is without the verify character
        owner_hash_display = owner_hash_full[:-1]

    # Get model info
    entry = engine_pool.get_entry(run.request.model_id)
    model_path = entry.model_path if entry else ""
    quantization = _detect_quantization(model_path)
    model_name = _clean_model_name(run.request.model_id, quantization)

    # Generate submission group
    submission_group = str(uuid.uuid4())

    # Collect single results and batch results
    single_results = [r for r in run.results if r.get("test_type") == "single"]
    uploadable_single_results = [
        r for r in single_results if float(r.get("gen_tps", 0.0) or 0.0) > 0.0
    ]
    skipped_count = len(single_results) - len(uploadable_single_results)
    batch_results = [r for r in run.results if r.get("test_type") == "batch"]

    # Build batching_results from batch data
    batching_results = []
    pp1024_single = next(
        (r for r in single_results if r.get("pp") == 1024), None
    )
    if (
        pp1024_single
        and float(pp1024_single.get("gen_tps", 0.0) or 0.0) > 0.0
        and batch_results
    ):
        baseline_tps = pp1024_single["gen_tps"]
        batching_results.append({
            "batch_size": 1,
            "tg_tps": baseline_tps,
            "speedup": 1.0,
        })
        for br in batch_results:
            speedup = round(br["tg_tps"] / baseline_tps, 2) if baseline_tps > 0 else 1.0
            batching_results.append({
                "batch_size": br["batch_size"],
                "tg_tps": br["tg_tps"],
                "speedup": speedup,
            })

    success_count = 0
    failed_count = 0

    if skipped_count:
        logger.info(
            f"Benchmark upload skipped {skipped_count} result(s) without "
            f"measurable generation throughput"
        )

    for result in uploadable_single_results:
        context_length = result["pp"]
        peak_mem_gb = None
        if result.get("peak_memory_bytes") and result["peak_memory_bytes"] > 0:
            peak_mem_gb = round(result["peak_memory_bytes"] / (1024**3), 2)

        payload = {
            "chip_name": chip_name,
            "chip_variant": chip_variant,
            "memory_gb": memory_gb,
            "gpu_cores": gpu_cores,
            "omlx_version": omlx_version,
            "os_version": os_version,
            "model_name": model_name,
            "quantization": quantization,
            "context_length": context_length,
            "pp_tps": result["processing_tps"],
            "tg_tps": result["gen_tps"],
            "ttft_ms": result.get("ttft_ms"),
            "peak_memory_gb": peak_mem_gb,
            "submission_group": submission_group,
        }

        if owner_hash_full:
            payload["owner_hash"] = owner_hash_full

        # Attach batching_results only to the first submission (lowest context_length)
        if (
            context_length == uploadable_single_results[0]["pp"]
            and batching_results
        ):
            payload["batching_results"] = batching_results

        try:
            resp = await asyncio.to_thread(
                requests.post,
                OMLX_AI_API_URL,
                json=payload,
                timeout=15,
            )

            if resp.status_code == 201:
                data = resp.json()
                success_count += 1
                result_dict = {
                    "context_length": context_length,
                    "id": data.get("id"),
                    "url": data.get("url"),
                }
                run.upload_state["results"].append(result_dict)
                await _send_event(run, {
                    "type": "upload",
                    "data": result_dict,
                })
            elif resp.status_code == 409:
                data = resp.json()
                success_count += 1  # Duplicate is still ok
                result_dict = {
                    "context_length": context_length,
                    "id": data.get("existing_id"),
                    "url": data.get("existing_url"),
                    "duplicate": True,
                }
                run.upload_state["results"].append(result_dict)
                await _send_event(run, {
                    "type": "upload",
                    "data": result_dict,
                })
            else:
                failed_count += 1
                error_msg = _sanitize_upload_error(resp)
                result_dict = {
                    "context_length": context_length,
                    "error": error_msg,
                }
                run.upload_state["results"].append(result_dict)
                await _send_event(run, {
                    "type": "upload",
                    "data": result_dict,
                })
                # Surface the sanitized message to ops; the full body
                # (truncated) goes to debug so it can still be retrieved
                # from the log file if needed.
                logger.warning(
                    f"Benchmark upload failed for pp{context_length}: "
                    f"{resp.status_code} {error_msg}"
                )
                if (resp.text or "")[:1] not in ("{", "["):
                    logger.debug(
                        "Benchmark upload non-JSON body (truncated): %r",
                        (resp.text or "")[:500],
                    )

        except Exception as e:
            failed_count += 1
            result_dict = {
                "context_length": context_length,
                "error": str(e),
            }
            run.upload_state["results"].append(result_dict)
            await _send_event(run, {
                "type": "upload",
                "data": result_dict,
            })
            logger.warning(f"Benchmark upload error for pp{context_length}: {e}")

    run.upload_state["phase"] = "done"
    run.upload_state["total"] = len(uploadable_single_results)
    run.upload_state["success_count"] = success_count
    run.upload_state["failed_count"] = failed_count
    run.upload_state["skipped_count"] = skipped_count
    run.upload_state["owner_hash"] = owner_hash_display
    await _send_event(run, {
        "type": "upload_done",
        "data": {
            "owner_hash": owner_hash_display,
            "total": len(uploadable_single_results),
            "success": success_count,
            "failed": failed_count,
            "skipped": skipped_count,
        },
    })

    logger.info(
        f"Benchmark upload complete: {success_count}/"
        f"{len(uploadable_single_results)} succeeded, skipped={skipped_count}"
    )


async def run_benchmark(run: BenchmarkRun, engine_pool: Any) -> None:
    """Execute a complete benchmark run.

    Phases:
    1. Unload all loaded models
    2. Load the target model
    3. Run single request tests
    4. Run batch tests
    5. Unload the benchmark model
    """
    request = run.request
    if request.external is not None:
        await _run_external_benchmark(run)
        return
    total_tests = len(request.prompt_lengths) + len(request.batch_sizes)
    current_test = 0
    overall_start = time.perf_counter()

    try:
        # Snapshot experimental flags at run start. Settings can change mid-run,
        # and the produced numbers are tied to whatever was active when
        # generation actually ran.
        model_settings = None
        sm = getattr(engine_pool, "_settings_manager", None)
        if sm is not None:
            try:
                model_settings = sm.get_settings(request.model_id)
                run.experimental_features.extend(
                    _detect_experimental_features(model_settings)
                )
            except Exception as e:
                logger.warning(
                    f"Benchmark: failed to read experimental flags for "
                    f"{request.model_id}: {e}"
                )

        # Phase 1: Unload all loaded models
        loaded_ids = engine_pool.get_loaded_model_ids()
        if loaded_ids:
            await _send_event(run, {
                "type": "progress",
                "phase": "unload",
                "message": f"Unloading {len(loaded_ids)} model(s)...",
                "current": 0,
                "total": total_tests,
            })
            for model_id in loaded_ids:
                try:
                    await engine_pool._unload_engine(model_id)
                    logger.info(f"Benchmark: unloaded {model_id}")
                except Exception as e:
                    logger.warning(f"Benchmark: failed to unload {model_id}: {e}")

        # Phase 2: Load the target model
        await _send_event(run, {
            "type": "progress",
            "phase": "load",
            "message": f"Loading {request.model_id}...",
            "current": 0,
            "total": total_tests,
        })
        # VLM MTP requires VLMBatchedEngine (which has set_vlm_mtp_drafter),
        # so don't force LM-only loading when VLM MTP is enabled.
        vlm_mtp_active = (
            model_settings is not None
            and getattr(model_settings, "vlm_mtp_enabled", False)
            and getattr(model_settings, "vlm_mtp_draft_model", None)
        )
        force_lm = True if request.force_lm_engine else not vlm_mtp_active
        engine = await engine_pool.get_engine(
            request.model_id,
            force_lm=force_lm,
        )
        logger.info(f"Benchmark: loaded {request.model_id}")

        # Generate prompts for all needed lengths
        tokenizer = engine.tokenizer
        prompts: dict[int, str] = {}
        for pp_len in request.prompt_lengths:
            prompts[pp_len] = _generate_prompt(tokenizer, pp_len)

        # Ensure pp1024 prompt exists for batch tests
        if request.batch_sizes and 1024 not in prompts:
            prompts[1024] = _generate_prompt(tokenizer, 1024)

        # Warmup: run a short request to trigger JIT compilation,
        # Metal shader compilation, and KV cache initialization.
        # Without this, the first real benchmark test absorbs all
        # one-time overhead and shows artificially low pp TPS.
        await _send_event(run, {
            "type": "progress",
            "phase": "warmup",
            "message": "Warming up (JIT compile)...",
            "current": 0,
            "total": total_tests,
        })
        warmup_prompt = _generate_prompt(tokenizer, 32)
        warmup_max_tokens = (
            request.generation_length
            if getattr(engine, "is_diffusion_model", False)
            else 8
        )
        async for _ in engine.stream_generate(
            prompt=warmup_prompt, max_tokens=warmup_max_tokens, temperature=0.0
        ):
            pass
        logger.info("Benchmark: warmup complete")

        # Phase 3: Single request tests
        single_pp1024_gen_tps = None

        for pp_len in request.prompt_lengths:
            current_test += 1
            await _send_event(run, {
                "type": "progress",
                "phase": "single",
                "message": f"Single: pp{pp_len}/tg{request.generation_length}",
                "current": current_test,
                "total": total_tests,
            })

            metrics = await _run_single_test(
                engine=engine,
                prompt=prompts[pp_len],
                max_tokens=request.generation_length,
                pp_len=pp_len,
            )

            result = {
                "test_type": "single",
                "pp": pp_len,
                "tg": request.generation_length,
                **metrics,
            }
            run.results.append(result)

            await _send_event(run, {"type": "result", "data": result})

            # Store pp1024 gen_tps for speedup calculation
            if pp_len == 1024:
                single_pp1024_gen_tps = metrics["gen_tps"]

        # Phase 4: Batch tests
        # Each request has a unique UUID prefix (no cache hits)
        max_batch = max(request.batch_sizes) if request.batch_sizes else 0
        batch_prompts = [_generate_prompt(tokenizer, 1024) for _ in range(max_batch)]

        # Skip batch tests for engines without scheduler core (e.g. VLM/Diffusion)
        batch_core = _get_batch_benchmark_core(engine)
        if request.batch_sizes and batch_core is None:
            logger.info(
                "Batch test skipped: engine does not support concurrent batching"
            )
            current_test += len(request.batch_sizes)

        for batch_size in request.batch_sizes if batch_core is not None else []:
            current_test += 1
            await _send_event(run, {
                "type": "progress",
                "phase": "batch",
                "message": f"Batch {batch_size}x: pp1024/tg{request.generation_length}",
                "current": current_test,
                "total": total_tests,
            })

            batch_metrics = await _run_batch_test(
                engine=engine,
                prompts=batch_prompts[:batch_size],
                prompt_tokens=1024,
                max_tokens=request.generation_length,
                batch_size=batch_size,
            )

            result = {
                "test_type": "batch",
                "pp": 1024,
                "tg": request.generation_length,
                **batch_metrics,
            }
            run.results.append(result)
            await _send_event(run, {"type": "result", "data": result})

        # Phase 5: Unload benchmark model
        await _send_event(run, {
            "type": "progress",
            "phase": "cleanup",
            "message": f"Unloading {request.model_id}...",
            "current": total_tests,
            "total": total_tests,
        })
        try:
            await engine_pool._unload_engine(request.model_id)
            logger.info(f"Benchmark: unloaded {request.model_id} after benchmark")
        except Exception as e:
            logger.warning(f"Benchmark: failed to unload {request.model_id}: {e}")

        # Done
        overall_duration = time.perf_counter() - overall_start
        run.status = "completed"
        await _send_event(run, {
            "type": "done",
            "summary": {
                "model_id": request.model_id,
                "total_time": round(overall_duration, 1),
                "total_tests": total_tests,
            },
        })

        # Upload results to omlx.ai (failures don't affect benchmark status)
        try:
            await _upload_to_omlx_ai(run, engine_pool)
        except Exception as e:
            logger.warning(f"Benchmark upload to omlx.ai failed: {e}")
            await _send_event(run, {
                "type": "upload_done",
                "data": {
                    "owner_hash": None,
                    "total": 0,
                    "success": 0,
                    "failed": 0,
                    "error": str(e),
                },
            })

    except asyncio.CancelledError:
        run.status = "cancelled"
        await _send_event(run, {
            "type": "error",
            "message": "Benchmark cancelled by user",
        })
        # Try to unload the model on cancellation
        try:
            await engine_pool._unload_engine(request.model_id)
        except Exception:
            pass

    except Exception as e:
        logger.error(f"Benchmark error: {e}", exc_info=True)
        run.status = "error"
        run.error_message = str(e)
        await _send_event(run, {
            "type": "error",
            "message": str(e),
        })
        # Try to unload the model on error
        try:
            await engine_pool._unload_engine(request.model_id)
        except Exception:
            pass


async def _run_external_benchmark(run: BenchmarkRun) -> None:
    """Execute a benchmark run against an external OpenAI-compatible endpoint.

    No local model phases (unload/load/JIT warmup) and no community
    upload — external numbers measure someone else's hardware.
    """
    request = run.request
    total_tests = len(request.prompt_lengths) + len(request.batch_sizes)
    current_test = 0
    overall_start = time.perf_counter()
    client = ExternalAPIClient(request.external)

    try:
        # Warmup doubles as preflight: fail fast on bad URL/key and on
        # endpoints that do not return streamed usage (hard requirement
        # for accurate token counts) before any long test runs.
        await _send_event(run, {
            "type": "progress",
            "phase": "warmup",
            "message": "Warming up external endpoint...",
            "current": 0,
            "total": total_tests,
        })
        await client.stream_chat_completion(
            messages=[{"role": "user", "content": _generate_external_prompt(32)}],
            max_tokens=8,
            temperature=0.0,
        )
        logger.info("Benchmark: external endpoint warmup complete")

        # Single request tests
        for pp_len in request.prompt_lengths:
            current_test += 1
            await _send_event(run, {
                "type": "progress",
                "phase": "single",
                "message": f"Single: pp{pp_len}/tg{request.generation_length}",
                "current": current_test,
                "total": total_tests,
            })

            metrics = await _run_external_single_test(
                client=client,
                prompt=_generate_external_prompt(pp_len),
                max_tokens=request.generation_length,
            )

            result = {
                "test_type": "single",
                "pp": pp_len,
                "tg": request.generation_length,
                **metrics,
            }
            run.results.append(result)
            await _send_event(run, {"type": "result", "data": result})

        # Batch tests: concurrent requests with unique pp1024 prompts
        for batch_size in request.batch_sizes:
            current_test += 1
            await _send_event(run, {
                "type": "progress",
                "phase": "batch",
                "message": f"Batch {batch_size}x: pp1024/tg{request.generation_length}",
                "current": current_test,
                "total": total_tests,
            })

            batch_metrics = await _run_external_batch_test(
                client=client,
                prompts=[_generate_external_prompt(1024) for _ in range(batch_size)],
                max_tokens=request.generation_length,
                batch_size=batch_size,
            )

            result = {
                "test_type": "batch",
                "pp": 1024,
                "tg": request.generation_length,
                **batch_metrics,
            }
            run.results.append(result)
            await _send_event(run, {"type": "result", "data": result})

        # Done
        overall_duration = time.perf_counter() - overall_start
        run.status = "completed"
        await _send_event(run, {
            "type": "done",
            "summary": {
                "model_id": request.model_id,
                "total_time": round(overall_duration, 1),
                "total_tests": total_tests,
            },
        })

        # External results measure remote hardware — never upload them to
        # the omlx.ai community leaderboard. Mirrors the experimental-
        # features skip so REST pollers see the same upload_state shape.
        run.upload_state["phase"] = "skipped"
        run.upload_state["skipped_reason"] = "external_endpoint"
        await _send_event(run, {
            "type": "upload_skipped",
            "reason": "external_endpoint",
            "features": [],
        })

    except asyncio.CancelledError:
        run.status = "cancelled"
        await _send_event(run, {
            "type": "error",
            "message": "Benchmark cancelled by user",
        })
    except Exception as e:
        logger.error(f"External benchmark error: {e}", exc_info=True)
        run.status = "error"
        run.error_message = str(e)
        await _send_event(run, {
            "type": "error",
            "message": str(e),
        })
    finally:
        await client.aclose()