wehub-resource-sync
332 lines
12 KiB
Python
332 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Intermediate state tracker for interactive configuration.
|
|
|
|
``InteractiveState`` holds the partially-configured benchmark parameters,
|
|
can be initialized from CLI args or a saved JSON file, and can be
|
|
converted to the ``argparse.Namespace`` that ``run_engine_bench()`` expects.
|
|
"""
|
|
|
|
# Standard
|
|
from typing import Any
|
|
import argparse
|
|
import json
|
|
|
|
# First Party
|
|
from lmcache.cli.commands.bench.engine_bench.interactive.schema import (
|
|
ALL_ITEMS,
|
|
PHASE_GENERAL,
|
|
PHASE_REQUIRED,
|
|
PHASE_WORKLOAD,
|
|
ConfigItem,
|
|
)
|
|
|
|
# Keys that exist on argparse.Namespace but are NOT part of the interactive
|
|
# config item registry (operational flags, handled separately).
|
|
_OUTPUT_KEYS = ("output_dir", "seed", "no_csv", "export_csv", "json", "quiet")
|
|
|
|
# Keys used only during the interactive flow, never serialized or
|
|
# passed to the orchestrator.
|
|
_INTERACTIVE_ONLY_KEYS = {"has_lmcache"}
|
|
|
|
# Keys excluded from exported JSON configs. These are either
|
|
# environment-specific (engine_url, lmcache_url) or interactive-only.
|
|
_EXPORT_EXCLUDED_KEYS = _INTERACTIVE_ONLY_KEYS | {"engine_url", "lmcache_url"}
|
|
|
|
# Mapping from ConfigItem.key to the argparse attribute name when they differ.
|
|
# Most keys match directly; these are the exceptions.
|
|
_KEY_TO_ATTR: dict[str, str] = {
|
|
"kv_cache_volume": "kv_cache_volume",
|
|
"tokens_per_gb_kvcache": "tokens_per_gb_kvcache",
|
|
}
|
|
|
|
# argparse attribute names where the CLI default is None (meaning "not set"),
|
|
# versus attributes where a non-None argparse default is the real default
|
|
# (e.g., kv_cache_volume defaults to 100.0).
|
|
_ARGPARSE_NONE_MEANS_UNSET = {
|
|
"engine_url",
|
|
"workload",
|
|
"model",
|
|
"lmcache_url",
|
|
"tokens_per_gb_kvcache",
|
|
}
|
|
|
|
|
|
class InteractiveState:
|
|
"""Tracks which config items have been set and their values.
|
|
|
|
Keys present in ``_values`` are considered "set". Missing keys are
|
|
"unset" and will either be prompted for or filled with defaults.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._values: dict[str, Any] = {}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Basic accessors
|
|
# ------------------------------------------------------------------
|
|
|
|
def is_set(self, key: str) -> bool:
|
|
return key in self._values
|
|
|
|
def get(self, key: str, default: Any = None) -> Any:
|
|
return self._values.get(key, default)
|
|
|
|
def set(self, key: str, value: Any) -> None:
|
|
self._values[key] = value
|
|
|
|
def unset(self, key: str) -> None:
|
|
"""Remove a key so it counts as unset again (used when stepping back)."""
|
|
self._values.pop(key, None)
|
|
|
|
@property
|
|
def values(self) -> dict[str, Any]:
|
|
return dict(self._values)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Readiness checks
|
|
# ------------------------------------------------------------------
|
|
|
|
def is_ready(self) -> bool:
|
|
"""True when all required items (whose conditions are met) have values."""
|
|
for item in ALL_ITEMS:
|
|
if not item.required:
|
|
continue
|
|
if not self._condition_met(item):
|
|
continue
|
|
if not self.is_set(item.key):
|
|
return False
|
|
return True
|
|
|
|
def get_missing_required(self) -> list[ConfigItem]:
|
|
"""Return phase-1 items that still need user input.
|
|
|
|
Includes both required items that are unset, and non-required
|
|
phase-1 items (like ``lmcache_url``) that are relevant because
|
|
a downstream required item (``tokens_per_gb_kvcache``) is unset.
|
|
"""
|
|
missing: list[ConfigItem] = []
|
|
for item in ALL_ITEMS:
|
|
if item.phase != PHASE_REQUIRED:
|
|
continue
|
|
if not self._condition_met(item):
|
|
continue
|
|
if self.is_set(item.key):
|
|
continue
|
|
if item.required:
|
|
missing.append(item)
|
|
elif item.key in ("has_lmcache", "lmcache_url") and not self.is_set(
|
|
"tokens_per_gb_kvcache"
|
|
):
|
|
# Only ask about LMCache when tokens_per_gb is needed
|
|
missing.append(item)
|
|
return missing
|
|
|
|
def get_general_items(self) -> list[ConfigItem]:
|
|
"""Return phase-2 (general) items that are not yet set."""
|
|
return [
|
|
item
|
|
for item in ALL_ITEMS
|
|
if item.phase == PHASE_GENERAL
|
|
and not self.is_set(item.key)
|
|
and self._condition_met(item)
|
|
]
|
|
|
|
def get_workload_items(self) -> list[ConfigItem]:
|
|
"""Return phase-3 (workload-specific) items whose conditions are met."""
|
|
return [
|
|
item
|
|
for item in ALL_ITEMS
|
|
if item.phase == PHASE_WORKLOAD and self._condition_met(item)
|
|
]
|
|
|
|
def has_unconfigured_general(self) -> bool:
|
|
"""True if there are general items the user hasn't explicitly set."""
|
|
return len(self.get_general_items()) > 0
|
|
|
|
def has_workload_items(self) -> bool:
|
|
"""True if there are workload-specific items to configure."""
|
|
return len(self.get_workload_items()) > 0
|
|
|
|
def workload_items_all_default(self) -> bool:
|
|
"""True if no workload-specific items have been explicitly set."""
|
|
for item in ALL_ITEMS:
|
|
if item.phase != PHASE_WORKLOAD:
|
|
continue
|
|
if not self._condition_met(item):
|
|
continue
|
|
if self.is_set(item.key):
|
|
return False
|
|
return True
|
|
|
|
# ------------------------------------------------------------------
|
|
# Defaults
|
|
# ------------------------------------------------------------------
|
|
|
|
def fill_defaults(self) -> None:
|
|
"""Set all unset items (whose conditions are met) to their defaults."""
|
|
for item in ALL_ITEMS:
|
|
if self.is_set(item.key):
|
|
continue
|
|
if not self._condition_met(item):
|
|
continue
|
|
if item.default is not None:
|
|
self._values[item.key] = item.default
|
|
|
|
# ------------------------------------------------------------------
|
|
# Conversion: CLI args ↔ InteractiveState
|
|
# ------------------------------------------------------------------
|
|
|
|
@classmethod
|
|
def from_cli_args(cls, args: argparse.Namespace) -> "InteractiveState":
|
|
"""Build state from parsed CLI arguments.
|
|
|
|
Only sets values that the user explicitly provided (i.e., not the
|
|
argparse default). This lets us distinguish "user set
|
|
``--kv-cache-volume 100``" from "user didn't touch it".
|
|
"""
|
|
state = cls()
|
|
for item in ALL_ITEMS:
|
|
if item.key in _INTERACTIVE_ONLY_KEYS:
|
|
continue
|
|
attr = _KEY_TO_ATTR.get(item.key, item.key)
|
|
value = getattr(args, attr, None)
|
|
if value is None:
|
|
continue
|
|
# For keys where argparse default is None, any non-None value
|
|
# means the user set it.
|
|
if attr in _ARGPARSE_NONE_MEANS_UNSET:
|
|
state._values[item.key] = value
|
|
continue
|
|
# For keys with real argparse defaults, we can't easily tell
|
|
# if the user typed --kv-cache-volume 100 vs it being the
|
|
# default. We mark it as "set" only if it differs from the
|
|
# schema default. This is imperfect but good enough — the
|
|
# worst case is we re-prompt for a value the user explicitly
|
|
# set to the default.
|
|
if item.default is not None and value == item.default:
|
|
continue
|
|
state._values[item.key] = value
|
|
|
|
# Derive has_lmcache from lmcache_url if provided via CLI
|
|
if state.is_set("lmcache_url"):
|
|
state._values["has_lmcache"] = True
|
|
|
|
return state
|
|
|
|
def to_namespace(self) -> argparse.Namespace:
|
|
"""Convert to ``argparse.Namespace`` compatible with ``run_engine_bench``.
|
|
|
|
Fills defaults for any unset items, then builds the namespace
|
|
with the attribute names that ``parse_args_to_config()`` and
|
|
``create_workload()`` expect.
|
|
"""
|
|
self.fill_defaults()
|
|
ns = argparse.Namespace()
|
|
|
|
# Map state keys to namespace attributes
|
|
for item in ALL_ITEMS:
|
|
if item.key in _INTERACTIVE_ONLY_KEYS:
|
|
continue
|
|
attr = _KEY_TO_ATTR.get(item.key, item.key)
|
|
# Only fall back to schema default if the item's condition is met.
|
|
# This prevents lmcache_url's default from leaking when
|
|
# has_lmcache is not set.
|
|
if item.key in self._values:
|
|
value = self._values[item.key]
|
|
elif self._condition_met(item):
|
|
value = item.default
|
|
else:
|
|
value = None
|
|
setattr(ns, attr, value)
|
|
|
|
# Output settings (not in the interactive registry)
|
|
ns.output_dir = self._values.get("output_dir", ".")
|
|
ns.seed = self._values.get("seed", 42)
|
|
ns.no_csv = self._values.get("no_csv", False)
|
|
ns.json = self._values.get("export_json", False)
|
|
ns.quiet = self._values.get("quiet", False)
|
|
ns.bench_target = "engine"
|
|
|
|
# Ensure format/output attrs exist for create_metrics
|
|
if not hasattr(ns, "format"):
|
|
ns.format = None
|
|
if not hasattr(ns, "output"):
|
|
ns.output = None
|
|
|
|
return ns
|
|
|
|
# ------------------------------------------------------------------
|
|
# Conversion: JSON ↔ InteractiveState
|
|
# ------------------------------------------------------------------
|
|
|
|
def to_json(self) -> dict[str, Any]:
|
|
"""Serialize to a JSON-compatible dict for config export.
|
|
|
|
Excludes environment-specific keys (``engine_url``,
|
|
``lmcache_url``) and interactive-only keys so the exported
|
|
config is portable and works without an LMCache server.
|
|
"""
|
|
self.fill_defaults()
|
|
return {k: v for k, v in self._values.items() if k not in _EXPORT_EXCLUDED_KEYS}
|
|
|
|
@classmethod
|
|
def from_json(cls, data: dict[str, Any]) -> "InteractiveState":
|
|
"""Load from a saved config JSON dict."""
|
|
state = cls()
|
|
for key, value in data.items():
|
|
state._values[key] = value
|
|
return state
|
|
|
|
def save_json(self, path: str) -> None:
|
|
"""Export the current state to a JSON file."""
|
|
with open(path, "w") as f:
|
|
json.dump(self.to_json(), f, indent=2)
|
|
f.write("\n")
|
|
|
|
@classmethod
|
|
def load_json(cls, path: str) -> "InteractiveState":
|
|
"""Load state from a JSON config file."""
|
|
with open(path) as f:
|
|
data = json.load(f)
|
|
return cls.from_json(data)
|
|
|
|
def merge_cli_args(self, args: argparse.Namespace) -> None:
|
|
"""Merge CLI args on top of existing state (CLI args win)."""
|
|
cli_state = InteractiveState.from_cli_args(args)
|
|
self._values.update(cli_state._values)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Summary
|
|
# ------------------------------------------------------------------
|
|
|
|
def summary_lines(self) -> list[tuple[str, str]]:
|
|
"""Return ``(label, value_str)`` pairs for the config summary."""
|
|
lines: list[tuple[str, str]] = []
|
|
for item in ALL_ITEMS:
|
|
if item.key in _INTERACTIVE_ONLY_KEYS:
|
|
continue
|
|
if not self._condition_met(item):
|
|
continue
|
|
value = self._values.get(item.key, item.default)
|
|
if value is None or value == "":
|
|
if item.key == "model":
|
|
display = "(auto-detect)"
|
|
elif item.key == "lmcache_url":
|
|
continue # skip empty lmcache_url
|
|
else:
|
|
display = "(not set)"
|
|
else:
|
|
display = str(value)
|
|
lines.append((item.display_name, display))
|
|
return lines
|
|
|
|
# ------------------------------------------------------------------
|
|
# Internal
|
|
# ------------------------------------------------------------------
|
|
|
|
def _condition_met(self, item: ConfigItem) -> bool:
|
|
"""Check whether an item's condition is satisfied."""
|
|
if item.condition is None:
|
|
return True
|
|
return item.condition(self._values)
|