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

Add study tune loop and smoke configs

Browse Source
This commit is contained in:
Gahow Wang
2026-04-04 22:29:59 +08:00
parent 7b7eaafd78
commit f192c741ed
8 changed files with 387 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
src/aituner/cli.py
View File

@@ -105,6 +105,76 @@ def cmd_study_ingest(args: argparse.Namespace) -> int:
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 1)
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)
if not proposal_files and study.llm.endpoint is None:
raise SpecError("No proposal files provided and study.llm.endpoint is not configured")
executed: list[dict[str, object]] = []
for idx in range(max_trials):
state = store.load_state(study.study_id)
window, requests = load_trace_requests(study, study_spec_path=spec_path)
prompt = build_prompt(
study=study,
window_summary=summarize_window(requests, window),
state=state,
capability_profile=capability_profile,
)
prompt_name = f"prompt-{state.next_trial_index:04d}"
store.write_prompt(study.study_id, prompt_name, prompt)
if proposal_files:
proposal_source = proposal_files[idx]
proposal_text = proposal_source.read_text(encoding="utf-8")
proposal_name = proposal_source.stem
else:
proposal_source = None
proposal_text = call_llm_for_proposal(policy=study.llm, prompt=prompt)
proposal_name = f"proposal-{state.next_trial_index:04d}"
proposal = parse_proposal_text(proposal_text, study)
store.write_proposal(study.study_id, proposal_name, 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": 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,
}
)
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,
},
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))
@@ -154,6 +224,13 @@ def build_parser() -> argparse.ArgumentParser:
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.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")