Gahow Wang
a8f903498d
Phase 4 of the two-stop work. The harness already pre-empts the LLM with deterministic stops and guided probes, but an LLM-originated should_stop could still end the loop while the validator saw remaining opportunity. Add harness._stop_authority, exposed as context["stop_authority"], whose `authorized` mirrors the deterministic harness stop decision and whose `opportunity_remains` flags an open topology frontier or a high-value planned candidate. In study tune, an LLM-originated should_stop is now honored only when the validator authorizes it; an unauthorized stop is vetoed (bounded budget) so the loop cannot converge prematurely on the agent's say-so. File- and harness-originated stops are unaffected, and the stop reason chain is recorded. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
758 lines
29 KiB
Python
758 lines
29 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from dataclasses import replace
|
|
from pathlib import Path
|
|
|
|
from .compare import run_compare
|
|
from .harness import (
|
|
build_harness_context,
|
|
build_harness_guided_proposal,
|
|
build_harness_stop_proposal,
|
|
)
|
|
from .job import append_job, build_trial_job
|
|
from .lca import (
|
|
build_study_workload_profile,
|
|
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,
|
|
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
|
|
|
|
|
|
def _is_empty_config_patch(proposal: Proposal) -> bool:
|
|
return not proposal.config_patch.env_patch and not proposal.config_patch.flag_patch
|
|
|
|
|
|
def _latency_percentiles(summary: object, metric: str) -> dict[str, float]:
|
|
if not isinstance(summary, dict):
|
|
return {}
|
|
payload = summary.get(metric)
|
|
if not isinstance(payload, dict):
|
|
return {}
|
|
selected: dict[str, float] = {}
|
|
for key in ("mean", "p50", "p95", "p99"):
|
|
value = payload.get(key)
|
|
if isinstance(value, (int, float)):
|
|
selected[key] = float(value)
|
|
return selected
|
|
|
|
|
|
def _format_latency_percentiles(metric: str, values: dict[str, float]) -> str:
|
|
if not values:
|
|
return ""
|
|
ordered = ", ".join(
|
|
f"{key}={values[key]:.3f}"
|
|
for key in ("mean", "p50", "p95", "p99")
|
|
if key in values
|
|
)
|
|
return f"{metric}({ordered})"
|
|
|
|
|
|
def _baseline_all_infeasible_stop(result: dict[str, object]) -> tuple[str, dict[str, object]] | None:
|
|
if result.get("status") != "completed":
|
|
return None
|
|
if isinstance(result.get("best_request_rate"), (int, float)):
|
|
return None
|
|
probes = result.get("probes")
|
|
if not isinstance(probes, list) or not probes:
|
|
return None
|
|
if any(isinstance(probe, dict) and probe.get("feasible") for probe in probes):
|
|
return None
|
|
|
|
diagnostics = result.get("all_infeasible_diagnostics")
|
|
if not isinstance(diagnostics, dict):
|
|
diagnostics = {}
|
|
lowest_rate = diagnostics.get("request_rate")
|
|
lowest_threshold = diagnostics.get("threshold")
|
|
pass_rate = diagnostics.get("pass_rate")
|
|
early_stop_reason = str(diagnostics.get("early_stop_reason") or "").strip()
|
|
latency_summary = diagnostics.get("latency_summary")
|
|
ttft = _latency_percentiles(latency_summary, "ttft_ms")
|
|
tpot = _latency_percentiles(latency_summary, "tpot_ms")
|
|
details: dict[str, object] = {
|
|
"lowest_sampled_request_rate": lowest_rate,
|
|
"lowest_sampling_u": lowest_threshold,
|
|
"lowest_probe_pass_rate": pass_rate,
|
|
"early_stop_reason": early_stop_reason,
|
|
"lowest_probe_latency_ms": {
|
|
"ttft": ttft,
|
|
"tpot": tpot,
|
|
},
|
|
"lowest_probe_latency_summary": latency_summary if isinstance(latency_summary, dict) else {},
|
|
}
|
|
pieces = [
|
|
"Baseline configuration has no feasible probe under the current SLO.",
|
|
"Stopping tuning because even the lowest sampled request rate did not meet the target pass rate.",
|
|
]
|
|
if isinstance(lowest_rate, (int, float)):
|
|
pieces.append(f"lowest_sampled_request_rate={float(lowest_rate):.6g}")
|
|
if isinstance(lowest_threshold, (int, float)):
|
|
pieces.append(f"lowest_sampling_u={float(lowest_threshold):.6g}")
|
|
if isinstance(pass_rate, (int, float)):
|
|
pieces.append(f"lowest_probe_pass_rate={float(pass_rate):.6g}")
|
|
if early_stop_reason:
|
|
pieces.append(f"early_stop_reason={early_stop_reason}")
|
|
for item in (
|
|
_format_latency_percentiles("lowest_probe_ttft_ms", ttft),
|
|
_format_latency_percentiles("lowest_probe_tpot_ms", tpot),
|
|
):
|
|
if item:
|
|
pieces.append(item)
|
|
return " ".join(pieces), details
|
|
|
|
|
|
def _study_source_path(study_root: Path) -> Path:
|
|
return Path((study_root / "study_spec.source").read_text(encoding="utf-8").strip())
|
|
|
|
|
|
def cmd_study_init(args: argparse.Namespace) -> int:
|
|
spec_path = Path(args.spec).resolve()
|
|
study = load_study_spec(spec_path)
|
|
store = StudyStore(Path(args.store_root) if args.store_root else None)
|
|
root = store.init_study(spec_path=spec_path, study=study)
|
|
print(root)
|
|
return 0
|
|
|
|
|
|
def cmd_study_prompt(args: argparse.Namespace) -> int:
|
|
store = StudyStore(Path(args.store_root) if args.store_root else None)
|
|
study_root = Path(args.study_root).resolve()
|
|
source_path = _study_source_path(study_root)
|
|
study = load_study_spec(source_path)
|
|
state = store.load_state(study.study_id)
|
|
capability_profile = load_capability_profile(study, study_spec_path=source_path)
|
|
window, requests = load_trace_requests(study, study_spec_path=source_path)
|
|
prompt = build_prompt(
|
|
study=study,
|
|
window_summary=summarize_window(requests, window),
|
|
state=state,
|
|
capability_profile=capability_profile,
|
|
workload_profile=build_study_workload_profile(study, requests, window),
|
|
)
|
|
prompt_name = args.prompt_name or f"prompt-{state.next_trial_index:04d}"
|
|
path = store.write_prompt(study.study_id, prompt_name, prompt)
|
|
print(path)
|
|
return 0
|
|
|
|
|
|
def cmd_study_llm_propose(args: argparse.Namespace) -> int:
|
|
store = StudyStore(Path(args.store_root) if args.store_root else None)
|
|
study_root = Path(args.study_root).resolve()
|
|
source_path = _study_source_path(study_root)
|
|
study = load_study_spec(source_path)
|
|
state = store.load_state(study.study_id)
|
|
capability_profile = load_capability_profile(study, study_spec_path=source_path)
|
|
window, requests = load_trace_requests(study, study_spec_path=source_path)
|
|
prompt = build_prompt(
|
|
study=study,
|
|
window_summary=summarize_window(requests, window),
|
|
state=state,
|
|
capability_profile=capability_profile,
|
|
workload_profile=build_study_workload_profile(study, requests, window),
|
|
)
|
|
proposal_text = call_llm_for_proposal(
|
|
policy=study.llm,
|
|
prompt=prompt,
|
|
use_harness=study.llm.use_harness,
|
|
)
|
|
proposal = parse_proposal_text(proposal_text, study)
|
|
name = args.proposal_name or f"proposal-{state.next_trial_index:04d}"
|
|
path = store.write_proposal(study.study_id, name, proposal)
|
|
print(path)
|
|
return 0
|
|
|
|
|
|
def cmd_study_register_proposal(args: argparse.Namespace) -> int:
|
|
store = StudyStore(Path(args.store_root) if args.store_root else None)
|
|
study_root = Path(args.study_root).resolve()
|
|
source_path = _study_source_path(study_root)
|
|
study = load_study_spec(source_path)
|
|
proposal = parse_proposal_text(Path(args.proposal_file).read_text(encoding="utf-8"), study)
|
|
name = args.proposal_name or Path(args.proposal_file).stem
|
|
path = store.write_proposal(study.study_id, name, proposal)
|
|
print(path)
|
|
return 0
|
|
|
|
|
|
def cmd_study_emit_job(args: argparse.Namespace) -> int:
|
|
store = StudyStore(Path(args.store_root) if args.store_root else None)
|
|
study_root = Path(args.study_root).resolve()
|
|
source_path = _study_source_path(study_root)
|
|
study = load_study_spec(source_path)
|
|
state = store.load_state(study.study_id)
|
|
proposal_text = Path(args.proposal_file).read_text(encoding="utf-8")
|
|
proposal = parse_proposal_text(proposal_text, study)
|
|
trial, _ = store.materialize_trial(study=study, state=state, proposal=proposal)
|
|
repo_root = Path(__file__).resolve().parents[2]
|
|
job = build_trial_job(study=study, trial=trial, repo_root=repo_root)
|
|
append_job(Path(args.jobs_file).resolve(), job)
|
|
print(trial.trial_id)
|
|
return 0
|
|
|
|
|
|
def cmd_study_ingest(args: argparse.Namespace) -> int:
|
|
store = StudyStore(Path(args.store_root) if args.store_root else None)
|
|
study_root = Path(args.study_root).resolve()
|
|
study_id = study_root.name
|
|
state = store.ingest_trial_results(study_id)
|
|
print(json.dumps({"best_trial_id": state.best_trial_id, "best_request_rate": state.best_request_rate}))
|
|
return 0
|
|
|
|
|
|
def cmd_study_tune(args: argparse.Namespace) -> int:
|
|
spec_path = Path(args.spec).resolve()
|
|
study = load_study_spec(spec_path)
|
|
store = StudyStore(Path(args.store_root) if args.store_root else None)
|
|
study_root = store.init_study(spec_path=spec_path, study=study)
|
|
capability_profile = load_capability_profile(study, study_spec_path=spec_path)
|
|
proposal_files = [Path(item).resolve() for item in (args.proposal_file or [])]
|
|
max_trials = args.max_trials or (len(proposal_files) if proposal_files else 2)
|
|
if max_trials <= 0:
|
|
raise SpecError("max_trials must be positive")
|
|
if proposal_files and max_trials > len(proposal_files):
|
|
max_trials = len(proposal_files)
|
|
executed: list[dict[str, object]] = []
|
|
stop_vetoes = 0
|
|
max_llm_stop_vetoes = 1
|
|
for idx in range(max_trials):
|
|
state = store.load_state(study.study_id)
|
|
if state.tuning_stop_reason:
|
|
executed.append(
|
|
{
|
|
"trial_id": None,
|
|
"stopped": True,
|
|
"reason": state.tuning_stop_reason,
|
|
"diagnosis": state.tuning_stop_diagnosis,
|
|
"details": state.tuning_stop_details,
|
|
"state_best_trial_id": state.best_trial_id,
|
|
"state_best_request_rate": state.best_request_rate,
|
|
}
|
|
)
|
|
break
|
|
if state.next_trial_index > max_trials:
|
|
break
|
|
window, requests = load_trace_requests(study, study_spec_path=spec_path)
|
|
window_summary = summarize_window(requests, window)
|
|
workload_profile = build_study_workload_profile(study, requests, window)
|
|
harness_context = (
|
|
build_harness_context(
|
|
study=study,
|
|
window_summary=window_summary,
|
|
state=state,
|
|
workload_profile=workload_profile,
|
|
)
|
|
if study.llm.use_harness
|
|
else None
|
|
)
|
|
prompt = build_prompt(
|
|
study=study,
|
|
window_summary=window_summary,
|
|
state=state,
|
|
capability_profile=capability_profile,
|
|
workload_profile=workload_profile,
|
|
)
|
|
prompt_name = f"prompt-{state.next_trial_index:04d}"
|
|
store.write_prompt(study.study_id, prompt_name, prompt)
|
|
|
|
if (
|
|
not proposal_files
|
|
and not args.skip_baseline
|
|
and state.next_trial_index == 1
|
|
and not state.trials
|
|
):
|
|
proposal_source = None
|
|
proposal_name = "baseline-0001"
|
|
proposal_text = json.dumps(
|
|
{
|
|
"observation": "Evaluate the study's initial engine configuration before LLM-guided edits.",
|
|
"diagnosis": "Baseline trial aligned with the AITuner evaluate-then-search loop.",
|
|
"config_patch": {"env_patch": {}, "flag_patch": {}},
|
|
"expected_effects": [
|
|
"establish incumbent performance",
|
|
"provide bottleneck evidence for harness-guided proposals",
|
|
],
|
|
"why_not_previous_failures": "No config changes are applied.",
|
|
"should_stop": False,
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
elif proposal_files:
|
|
proposal_index = state.next_trial_index - 1
|
|
if proposal_index >= len(proposal_files):
|
|
break
|
|
proposal_source = proposal_files[proposal_index]
|
|
proposal_text = proposal_source.read_text(encoding="utf-8")
|
|
proposal_name = proposal_source.stem
|
|
else:
|
|
proposal_source = None
|
|
stop_proposal = (
|
|
build_harness_stop_proposal(harness_context)
|
|
if harness_context is not None
|
|
else None
|
|
)
|
|
if stop_proposal is not None:
|
|
proposal_text = json.dumps(to_jsonable(stop_proposal), ensure_ascii=False)
|
|
proposal_name = f"harness-stop-{state.next_trial_index:04d}"
|
|
else:
|
|
guided_proposal = (
|
|
build_harness_guided_proposal(harness_context)
|
|
if harness_context is not None
|
|
else None
|
|
)
|
|
if guided_proposal is not None:
|
|
proposal_text = json.dumps(
|
|
to_jsonable(guided_proposal),
|
|
ensure_ascii=False,
|
|
)
|
|
proposal_name = f"harness-proposal-{state.next_trial_index:04d}"
|
|
else:
|
|
if study.llm.endpoint is None:
|
|
raise SpecError(
|
|
"No proposal files provided, study.llm.endpoint is not configured, "
|
|
"and the harness stop guard did not fire."
|
|
)
|
|
proposal_text = call_llm_for_proposal(
|
|
policy=study.llm,
|
|
prompt=prompt,
|
|
use_harness=study.llm.use_harness,
|
|
)
|
|
proposal_name = f"proposal-{state.next_trial_index:04d}"
|
|
raw_proposal_path = store.study_root(study.study_id) / "proposals" / f"{proposal_name}.raw.txt"
|
|
raw_proposal_path.write_text(proposal_text, encoding="utf-8")
|
|
proposal = parse_proposal_text(proposal_text, study)
|
|
store.write_proposal(study.study_id, proposal_name, proposal)
|
|
if proposal.should_stop:
|
|
is_harness_stop = proposal_name.startswith("harness-stop-")
|
|
is_llm_stop = not is_harness_stop and proposal_source is None
|
|
stop_authority = (
|
|
harness_context.get("stop_authority")
|
|
if isinstance(harness_context, dict)
|
|
else None
|
|
)
|
|
authorized = stop_authority is None or bool(stop_authority.get("authorized"))
|
|
# Stop-B authority: the deterministic validator overrides an
|
|
# LLM-originated stop. Veto an unauthorized stop (bounded) so the
|
|
# loop does not converge prematurely on the agent's say-so alone.
|
|
if is_llm_stop and not authorized and stop_vetoes < max_llm_stop_vetoes:
|
|
stop_vetoes += 1
|
|
executed.append(
|
|
{
|
|
"trial_id": None,
|
|
"proposal_name": proposal_name,
|
|
"proposal_source": "llm",
|
|
"stop_vetoed": True,
|
|
"reason": "validator_did_not_authorize_stop",
|
|
"validator_reason": (
|
|
stop_authority.get("reason") if stop_authority else None
|
|
),
|
|
"diagnosis": proposal.diagnosis,
|
|
}
|
|
)
|
|
continue
|
|
if is_harness_stop:
|
|
proposal_source_label = "harness"
|
|
else:
|
|
proposal_source_label = str(proposal_source) if proposal_source else "llm"
|
|
executed.append(
|
|
{
|
|
"trial_id": None,
|
|
"proposal_name": proposal_name,
|
|
"proposal_source": proposal_source_label,
|
|
"stopped": True,
|
|
"stop_authorized_by": (
|
|
"validator"
|
|
if (is_harness_stop or authorized)
|
|
else "llm_after_veto_budget"
|
|
),
|
|
"diagnosis": proposal.diagnosis,
|
|
"state_best_trial_id": state.best_trial_id,
|
|
"state_best_request_rate": state.best_request_rate,
|
|
}
|
|
)
|
|
break
|
|
is_auto_baseline = (
|
|
not proposal_files
|
|
and not args.skip_baseline
|
|
and state.next_trial_index == 1
|
|
and not state.trials
|
|
and _is_empty_config_patch(proposal)
|
|
)
|
|
trial, _ = store.materialize_trial(study=study, state=state, proposal=proposal)
|
|
trial_spec_path = Path(trial.artifact_dir) / "trial_spec.json"
|
|
result = run_trial(trial_spec_path)
|
|
state = store.ingest_trial_results(study.study_id)
|
|
executed.append(
|
|
{
|
|
"trial_id": trial.trial_id,
|
|
"proposal_name": proposal_name,
|
|
"proposal_source": (
|
|
"harness"
|
|
if proposal_name.startswith("harness-proposal-")
|
|
else str(proposal_source) if proposal_source else "llm"
|
|
),
|
|
"best_sampling_u": result.get("best_sampling_u"),
|
|
"best_request_rate": result.get("best_request_rate"),
|
|
"best_pass_rate": result.get("best_pass_rate"),
|
|
"state_best_trial_id": state.best_trial_id,
|
|
"state_best_request_rate": state.best_request_rate,
|
|
}
|
|
)
|
|
if is_auto_baseline:
|
|
stop = _baseline_all_infeasible_stop(result)
|
|
if stop is not None:
|
|
diagnosis, details = stop
|
|
state.tuning_stop_reason = "baseline_all_infeasible"
|
|
state.tuning_stop_diagnosis = diagnosis
|
|
state.tuning_stop_details = details
|
|
store.save_state(state)
|
|
executed.append(
|
|
{
|
|
"trial_id": None,
|
|
"stopped": True,
|
|
"reason": state.tuning_stop_reason,
|
|
"diagnosis": diagnosis,
|
|
"details": details,
|
|
"state_best_trial_id": state.best_trial_id,
|
|
"state_best_request_rate": state.best_request_rate,
|
|
}
|
|
)
|
|
break
|
|
|
|
final_state = store.load_state(study.study_id)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"study_root": str(study_root),
|
|
"executed_trials": executed,
|
|
"best_trial_id": final_state.best_trial_id,
|
|
"best_request_rate": final_state.best_request_rate,
|
|
"tuning_stop_reason": final_state.tuning_stop_reason,
|
|
"tuning_stop_diagnosis": final_state.tuning_stop_diagnosis,
|
|
"tuning_stop_details": final_state.tuning_stop_details,
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
)
|
|
return 0
|
|
|
|
|
|
def cmd_worker_run_trial(args: argparse.Namespace) -> int:
|
|
result = run_trial(Path(args.trial_spec).resolve())
|
|
print(json.dumps(result))
|
|
return 0
|
|
|
|
|
|
def cmd_compare_run(args: argparse.Namespace) -> int:
|
|
summary = run_compare(
|
|
Path(args.spec).resolve(),
|
|
output_root=Path(args.output_root).resolve() if args.output_root else None,
|
|
)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"compare_id": summary["compare_id"],
|
|
"compare_root": summary["compare_root"],
|
|
"window_count": summary["aggregate"]["window_count"],
|
|
"wins": summary["aggregate"]["wins"],
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
)
|
|
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)
|
|
|
|
study = subparsers.add_parser("study")
|
|
study_sub = study.add_subparsers(dest="study_command", required=True)
|
|
|
|
init = study_sub.add_parser("init")
|
|
init.add_argument("--spec", required=True)
|
|
init.add_argument("--store-root")
|
|
init.set_defaults(func=cmd_study_init)
|
|
|
|
prompt = study_sub.add_parser("prompt")
|
|
prompt.add_argument("--study-root", required=True)
|
|
prompt.add_argument("--store-root")
|
|
prompt.add_argument("--prompt-name")
|
|
prompt.set_defaults(func=cmd_study_prompt)
|
|
|
|
llm_propose = study_sub.add_parser("llm-propose")
|
|
llm_propose.add_argument("--study-root", required=True)
|
|
llm_propose.add_argument("--store-root")
|
|
llm_propose.add_argument("--proposal-name")
|
|
llm_propose.set_defaults(func=cmd_study_llm_propose)
|
|
|
|
register = study_sub.add_parser("register-proposal")
|
|
register.add_argument("--study-root", required=True)
|
|
register.add_argument("--store-root")
|
|
register.add_argument("--proposal-file", required=True)
|
|
register.add_argument("--proposal-name")
|
|
register.set_defaults(func=cmd_study_register_proposal)
|
|
|
|
emit = study_sub.add_parser("emit-job")
|
|
emit.add_argument("--study-root", required=True)
|
|
emit.add_argument("--store-root")
|
|
emit.add_argument("--proposal-file", required=True)
|
|
emit.add_argument("--jobs-file", required=True)
|
|
emit.set_defaults(func=cmd_study_emit_job)
|
|
|
|
ingest = study_sub.add_parser("ingest")
|
|
ingest.add_argument("--study-root", required=True)
|
|
ingest.add_argument("--store-root")
|
|
ingest.set_defaults(func=cmd_study_ingest)
|
|
|
|
tune = study_sub.add_parser("tune")
|
|
tune.add_argument("--spec", required=True)
|
|
tune.add_argument("--store-root")
|
|
tune.add_argument("--proposal-file", action="append")
|
|
tune.add_argument("--max-trials", type=int)
|
|
tune.add_argument(
|
|
"--skip-baseline",
|
|
action="store_true",
|
|
help="Do not automatically evaluate the initial config before LLM proposals.",
|
|
)
|
|
tune.set_defaults(func=cmd_study_tune)
|
|
|
|
worker = subparsers.add_parser("worker")
|
|
worker_sub = worker.add_subparsers(dest="worker_command", required=True)
|
|
run = worker_sub.add_parser("run-trial")
|
|
run.add_argument("--trial-spec", required=True)
|
|
run.set_defaults(func=cmd_worker_run_trial)
|
|
|
|
compare = subparsers.add_parser("compare")
|
|
compare_sub = compare.add_subparsers(dest="compare_command", required=True)
|
|
compare_run = compare_sub.add_parser("run")
|
|
compare_run.add_argument("--spec", required=True)
|
|
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
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
try:
|
|
return int(args.func(args))
|
|
except SpecError as exc:
|
|
print(str(exc), file=sys.stderr)
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|