Files
2026-07-13 12:24:33 +08:00

419 lines
15 KiB
Python

# SPDX-License-Identifier: Apache-2.0
"""Centralized config item definitions for interactive configuration.
Each ``ConfigItem`` declaratively describes one configurable parameter:
its key, display name, description, input type, default, and when it
should be shown. The ``ALL_ITEMS`` list is the single source of truth
for descriptions, ordering, and defaults.
"""
# Standard
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any
# ---------------------------------------------------------------------------
# Phases
# ---------------------------------------------------------------------------
PHASE_REQUIRED = 1
PHASE_GENERAL = 2
PHASE_WORKLOAD = 3
# ---------------------------------------------------------------------------
# ConfigItem
# ---------------------------------------------------------------------------
@dataclass
class ConfigItem:
"""Declarative description of a single configurable parameter.
Attributes:
key: State dict key (matches argparse attr name, e.g., ``"engine_url"``).
display_name: Heading shown in the prompt.
description: One-sentence explanation shown below the heading.
input_type: One of ``"text"``, ``"int"``, ``"float"``, ``"bool"``,
``"choice"``.
default: Default value. ``None`` means required (no default).
required: If True, this item must have a value before the benchmark
can start.
choices: For ``"choice"`` type — list of ``(value, description)`` tuples.
condition: Callable ``(state_dict) -> bool`` that determines whether
this item should be shown. ``None`` means always shown.
phase: Which interactive phase this item belongs to.
"""
key: str
display_name: str
description: str
input_type: str # "text", "int", "float", "bool", "choice"
default: Any = None
required: bool = False
choices: list[tuple[str, str]] = field(default_factory=list)
condition: Callable[[dict[str, Any]], bool] | None = None
phase: int = PHASE_GENERAL
# ---------------------------------------------------------------------------
# Condition helpers
# ---------------------------------------------------------------------------
def _has_lmcache(state: dict[str, Any]) -> bool:
"""Show this item only when the user said they have LMCache."""
return bool(state.get("has_lmcache"))
def _no_lmcache_url(state: dict[str, Any]) -> bool:
"""Show this item only when lmcache_url is not set."""
return not state.get("lmcache_url")
def _workload_is(name: str) -> Callable[[dict[str, Any]], bool]:
"""Return a condition that checks the workload value."""
def check(state: dict[str, Any]) -> bool:
return state.get("workload") == name
return check
# ---------------------------------------------------------------------------
# ALL_ITEMS — the centralized registry
# ---------------------------------------------------------------------------
ALL_ITEMS: list[ConfigItem] = [
# ── Phase 1: Required ─────────────────────────────────────────────
ConfigItem(
key="engine_url",
display_name="Engine URL",
description=(
"URL of the inference engine. Enter just a port (e.g. 8000) to "
"use http://localhost:8000. "
"Set OPENAI_API_KEY env var if authentication is needed."
),
input_type="text",
default="http://localhost:8000",
required=True,
phase=PHASE_REQUIRED,
),
ConfigItem(
key="workload",
display_name="Workload",
description="The type of benchmark workload to run.",
input_type="choice",
default=None,
required=True,
choices=[
(
"long-doc-permutator",
"Query the same set of long documents with different orders",
),
("long-doc-qa", "Repeated Q&A over long documents (tests KV cache reuse)"),
("multi-round-chat", "Multi-turn chat with stateful sessions"),
(
"prefix-suffix-tuner",
"Two-pass sequential workload demonstrating tiered KV cache reuse",
),
("random-prefill", "Prefill-only requests fired simultaneously"),
],
phase=PHASE_REQUIRED,
),
ConfigItem(
key="has_lmcache",
display_name="LMCache Server",
description=(
"Do you have a running LMCache server? "
"It can auto-detect KV cache size information."
),
input_type="bool",
default=True,
required=False,
phase=PHASE_REQUIRED,
),
ConfigItem(
key="lmcache_url",
display_name="LMCache Server URL",
description=(
"URL of the running LMCache HTTP server. Enter just a port "
"(e.g. 8080) to use http://localhost:8080."
),
input_type="text",
default="http://localhost:8080",
required=False,
condition=_has_lmcache,
phase=PHASE_REQUIRED,
),
ConfigItem(
key="tokens_per_gb_kvcache",
display_name="Tokens per GB KV cache",
description=(
"How many tokens fit in 1 GB of KV cache for your model.\n"
" If using vLLM, look for these lines in the startup log:\n"
' "Available KV cache memory: XX.XX GiB"\n'
' "GPU KV cache size: XXX,XXX tokens"\n'
" Then compute: tokens_per_gb = "
"GPU_KV_cache_tokens / Available_KV_cache_GiB"
),
input_type="int",
default=None,
required=True,
condition=_no_lmcache_url,
phase=PHASE_REQUIRED,
),
# ── Phase 2: General ──────────────────────────────────────────────
ConfigItem(
key="model",
display_name="Model name",
description=(
"The model served by the engine. "
"Leave empty to auto-detect from the engine."
),
input_type="text",
default="",
phase=PHASE_GENERAL,
),
ConfigItem(
key="kv_cache_volume",
display_name="KV cache volume (GB)",
description="Target active KV cache size for the benchmark.",
input_type="float",
default=100.0,
phase=PHASE_GENERAL,
),
ConfigItem(
key="ignore_eos",
display_name="Ignore EOS",
description=(
"Force generation to run for the full output length by ignoring "
"the model's EOS token (vLLM extension). Makes decode throughput "
"reproducible."
),
input_type="bool",
default=False,
phase=PHASE_GENERAL,
),
# ── Phase 3: long-doc-permutator ─────────────────────────────────
ConfigItem(
key="ldp_num_contexts",
display_name="Number of contexts",
description="Number of unique context documents to generate.",
input_type="int",
default=5,
condition=_workload_is("long-doc-permutator"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="ldp_context_length",
display_name="Context length (tokens)",
description="Token length of each context document.",
input_type="int",
default=5000,
condition=_workload_is("long-doc-permutator"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="ldp_system_prompt_length",
display_name="System prompt length (tokens)",
description="Token length of the shared system prompt. Use 0 for none.",
input_type="int",
default=1000,
condition=_workload_is("long-doc-permutator"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="ldp_num_permutations",
display_name="Number of permutations",
description="Distinct permutations to send. Capped at N! (N = num_contexts).",
input_type="int",
default=10,
condition=_workload_is("long-doc-permutator"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="ldp_num_inflight_requests",
display_name="Max inflight requests",
description="Maximum concurrent in-flight requests.",
input_type="int",
default=1,
condition=_workload_is("long-doc-permutator"),
phase=PHASE_WORKLOAD,
),
# ── Phase 3: long-doc-qa ──────────────────────────────────────────
ConfigItem(
key="ldqa_document_length",
display_name="Document length (tokens)",
description="Token length of each synthetic document.",
input_type="int",
default=10000,
condition=_workload_is("long-doc-qa"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="ldqa_query_per_document",
display_name="Queries per document",
description="Number of questions asked per document.",
input_type="int",
default=2,
condition=_workload_is("long-doc-qa"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="ldqa_shuffle_policy",
display_name="Shuffle policy",
description="How benchmark requests are ordered.",
input_type="choice",
default="random",
choices=[
("random", "Shuffle all (doc, query) pairs randomly"),
("tile", "Process queries round by round across all documents"),
],
condition=_workload_is("long-doc-qa"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="ldqa_num_inflight_requests",
display_name="Max inflight requests",
description="Maximum concurrent in-flight requests.",
input_type="int",
default=3,
condition=_workload_is("long-doc-qa"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="ldqa_max_output_length",
display_name="Max output length (tokens)",
description="Max tokens to generate per benchmark query.",
input_type="int",
default=128,
condition=_workload_is("long-doc-qa"),
phase=PHASE_WORKLOAD,
),
# ── Phase 3: multi-round-chat ─────────────────────────────────────
ConfigItem(
key="mrc_shared_prompt_length",
display_name="System prompt length (tokens)",
description="Token length of the system prompt per session.",
input_type="int",
default=2000,
condition=_workload_is("multi-round-chat"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="mrc_chat_history_length",
display_name="Chat history length (tokens)",
description="Token length of pre-filled conversation history.",
input_type="int",
default=10000,
condition=_workload_is("multi-round-chat"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="mrc_user_input_length",
display_name="User input length (tokens)",
description="Tokens per user query in each round.",
input_type="int",
default=50,
condition=_workload_is("multi-round-chat"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="mrc_output_length",
display_name="Output length (tokens)",
description="Max tokens to generate per response.",
input_type="int",
default=200,
condition=_workload_is("multi-round-chat"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="mrc_qps",
display_name="Queries per second",
description="Target request dispatch rate.",
input_type="float",
default=1.0,
condition=_workload_is("multi-round-chat"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="mrc_duration",
display_name="Duration (seconds)",
description="How long the benchmark runs.",
input_type="float",
default=60.0,
condition=_workload_is("multi-round-chat"),
phase=PHASE_WORKLOAD,
),
# ── Phase 3: prefix-suffix-tuner ──────────────────────────────────
ConfigItem(
key="psf_context_length",
display_name="Context length (tokens)",
description="Total tokens per request (prefix + breaker + suffix).",
input_type="int",
default=8000,
condition=_workload_is("prefix-suffix-tuner"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="psf_prefix_ratio",
display_name="Prefix ratio",
description=(
"Fraction of context-length used by the prefix. Must be in "
"(0.0, 1.0). The remainder (minus a 32-token breaker) is the "
"shared suffix."
),
input_type="float",
default=0.8,
condition=_workload_is("prefix-suffix-tuner"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="psf_thrash",
display_name="Target tier size (GB)",
description=(
"Size in GB of the KV-cache tier to overflow. The prefix pool "
"is sized to slightly more than this, so every pass-2 request "
"misses the targeted tier. Use the L0 (HBM) size for vanilla "
"vLLM, or the L1 (LMCache DRAM) size for tiered baselines."
),
input_type="float",
default=20.0,
condition=_workload_is("prefix-suffix-tuner"),
phase=PHASE_WORKLOAD,
),
# ── Phase 3: random-prefill ───────────────────────────────────────
ConfigItem(
key="rp_request_length",
display_name="Request length (tokens)",
description="Token length of each prefill request.",
input_type="int",
default=10000,
condition=_workload_is("random-prefill"),
phase=PHASE_WORKLOAD,
),
ConfigItem(
key="rp_num_requests",
display_name="Number of requests",
description="Total prefill requests to fire simultaneously.",
input_type="int",
default=50,
condition=_workload_is("random-prefill"),
phase=PHASE_WORKLOAD,
),
]
def get_items_by_phase(phase: int) -> list[ConfigItem]:
"""Return all items belonging to a given phase."""
return [item for item in ALL_ITEMS if item.phase == phase]
def get_item(key: str) -> ConfigItem:
"""Look up a ConfigItem by key. Raises KeyError if not found."""
for item in ALL_ITEMS:
if item.key == key:
return item
raise KeyError(f"No ConfigItem with key {key!r}")