gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add L-C-A workload profile metric and CLI profile commands

Browse Source 
Implement the paper's 10-dimensional L-C-A workload feature vector
(RobustScaler-normalized, sim=exp(-||dz||)) in lca.py, and wire it into
`aituner profile window` / `aituner profile similarity`. Covered by tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-15 14:02:24 +08:00
parent 984eb1f325
commit 27d1c8fa92
3 changed files with 770 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

212
src/aituner/cli.py
View File

@@ -3,6 +3,7 @@ from __future__ import annotations
import argparse
import json
import sys
from dataclasses import replace
from pathlib import Path
from .compare import run_compare
@@ -12,8 +13,20 @@ from .harness import (
build_harness_stop_proposal,
)
from .job import append_job, build_trial_job
from .lca import (
build_workload_profile,
resolve_length_mode,
similarity_report,
)
from .llm import build_prompt, call_llm_for_proposal, load_capability_profile, parse_proposal_text
from .spec import Proposal, SpecError, load_study_spec, to_jsonable
from .spec import (
Proposal,
SpecError,
StudySpec,
load_structured_file,
load_study_spec,
to_jsonable,
)
from .store import StudyStore
from .trace import load_trace_requests, summarize_window
from .worker import run_trial
@@ -422,6 +435,159 @@ def cmd_compare_run(args: argparse.Namespace) -> int:
return 0
def _resolve_profile_gpu_count(args: argparse.Namespace, study: StudySpec) -> int:
gpu_count = args.gpu_count
if gpu_count is None:
gpu_count = study.hardware.gpu_count
if gpu_count <= 0:
raise SpecError("--gpu-count must be > 0.")
return int(gpu_count)
def _load_profile_study_spec(spec_path: Path) -> StudySpec:
payload = dict(load_structured_file(spec_path))
llm_payload = dict(payload.get("llm") or {})
llm_payload.pop("endpoint", None)
payload["llm"] = llm_payload
return StudySpec.from_dict(payload)
def _profile_current_study_window(args: argparse.Namespace) -> dict[str, object]:
spec_path = Path(args.spec).resolve()
study = _load_profile_study_spec(spec_path)
mode = resolve_length_mode(
request_mode=study.trace.request_mode,
length_mode=args.length_mode,
)
window, requests = load_trace_requests(study, study_spec_path=spec_path)
profile = build_workload_profile(
requests,
window,
gpu_count=_resolve_profile_gpu_count(args, study),
length_mode=mode,
)
return {
"profile": profile.to_dict(),
"source": {
"study_spec_path": str(spec_path),
"window_id": study.trace.window_id,
},
}
def _resolve_windows_path_for_profile(study: StudySpec, *, study_spec_path: Path) -> Path:
path = Path(study.trace.windows_path)
if not path.is_absolute():
path = (study_spec_path.parent / path).resolve()
return path
def _load_profile_windows(
study: StudySpec,
*,
study_spec_path: Path,
) -> list[dict[str, object]]:
windows_path = _resolve_windows_path_for_profile(study, study_spec_path=study_spec_path)
payload = json.loads(windows_path.read_text(encoding="utf-8"))
raw_windows = payload.get("windows") if isinstance(payload, dict) else payload
if not isinstance(raw_windows, list):
raise SpecError(f"windows payload must contain a list: {windows_path}")
return [
{str(key): value for key, value in item.items()}
for item in raw_windows
if isinstance(item, dict)
]
def _selected_profile_windows(
args: argparse.Namespace,
study: StudySpec,
*,
study_spec_path: Path,
) -> list[dict[str, object]]:
windows = _load_profile_windows(study, study_spec_path=study_spec_path)
window_ids = set(args.window_id or [])
selected: list[dict[str, object]] = []
for item in windows:
window_id = str(item.get("window_id") or "").strip()
if not window_id:
continue
if window_ids and window_id not in window_ids:
continue
if not window_ids and not args.all:
if window_id != study.trace.window_id:
continue
trace_type = str(item.get("trace_type") or "").strip()
if args.trace_type and trace_type != args.trace_type:
continue
date_value = str(item.get("date") or "").strip()
if args.date_from and date_value and date_value < args.date_from:
continue
if args.date_to and date_value and date_value > args.date_to:
continue
if args.slot_token and str(item.get("slot_token") or "").strip() != args.slot_token:
continue
selected.append(item)
selected.sort(
key=lambda item: (
str(item.get("date") or ""),
str(item.get("slot_token") or ""),
str(item.get("window_id") or ""),
)
)
if args.limit is not None:
selected = selected[: args.limit]
if not selected:
raise SpecError("No trace windows selected for profile similarity.")
return selected
def cmd_profile_window(args: argparse.Namespace) -> int:
print(json.dumps(_profile_current_study_window(args), ensure_ascii=False, indent=2))
return 0
def cmd_profile_similarity(args: argparse.Namespace) -> int:
spec_path = Path(args.spec).resolve()
study = _load_profile_study_spec(spec_path)
mode = resolve_length_mode(
request_mode=study.trace.request_mode,
length_mode=args.length_mode,
)
gpu_count = _resolve_profile_gpu_count(args, study)
profiles = []
selected = _selected_profile_windows(args, study, study_spec_path=spec_path)
for item in selected:
window_id = str(item["window_id"])
window_study = replace(study, trace=replace(study.trace, window_id=window_id))
window, requests = load_trace_requests(window_study, study_spec_path=spec_path)
profiles.append(
build_workload_profile(
requests,
window,
gpu_count=gpu_count,
length_mode=mode,
)
)
print(
json.dumps(
{
"source": {
"study_spec_path": str(spec_path),
"selected_window_count": len(profiles),
"length_mode": mode,
"gpu_count": gpu_count,
},
"profiles": [profile.to_dict() for profile in profiles],
"similarity": similarity_report(profiles),
},
ensure_ascii=False,
indent=2,
)
)
return 0
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="AITuner CLI")
subparsers = parser.add_subparsers(dest="command", required=True)
@@ -490,6 +656,50 @@ def build_parser() -> argparse.ArgumentParser:
compare_run.add_argument("--output-root")
compare_run.set_defaults(func=cmd_compare_run)
profile = subparsers.add_parser("profile")
profile_sub = profile.add_subparsers(dest="profile_command", required=True)
profile_window = profile_sub.add_parser("window")
profile_window.add_argument("--spec", required=True)
profile_window.add_argument(
"--length-mode",
default="auto",
choices=["auto", "total", "input", "output"],
help="Token length basis for the L vector. auto uses output for decode_only and total otherwise.",
)
profile_window.add_argument(
"--gpu-count",
type=int,
help="GPU denominator for per-GPU arrival rate. Defaults to hardware.gpu_count.",
)
profile_window.set_defaults(func=cmd_profile_window)
profile_similarity = profile_sub.add_parser("similarity")
profile_similarity.add_argument("--spec", required=True)
profile_similarity.add_argument("--window-id", action="append")
profile_similarity.add_argument("--trace-type")
profile_similarity.add_argument("--date-from")
profile_similarity.add_argument("--date-to")
profile_similarity.add_argument("--slot-token")
profile_similarity.add_argument("--limit", type=int)
profile_similarity.add_argument(
"--all",
action="store_true",
help="Profile all windows selected by filters. Without this or --window-id, only the study window is used.",
)
profile_similarity.add_argument(
"--length-mode",
default="auto",
choices=["auto", "total", "input", "output"],
help="Token length basis for the L vector. auto uses output for decode_only and total otherwise.",
)
profile_similarity.add_argument(
"--gpu-count",
type=int,
help="GPU denominator for per-GPU arrival rate. Defaults to hardware.gpu_count.",
)
profile_similarity.set_defaults(func=cmd_profile_similarity)
return parser