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

Add initial config preflight review

Browse Source
This commit is contained in:
Gahow Wang
2026-07-01 11:12:58 +08:00
parent 1b8f5a3af1
commit 46b477f48e
5 changed files with 488 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

157
tests/test_core_flow.py
View File

@@ -40,7 +40,14 @@ from aituner.lca import (
resolve_length_mode,
similarity_report,
)
from aituner.llm import _extract_response_text, build_prompt, parse_proposal_text, validate_proposal
from aituner.llm import (
_extract_response_text,
build_initial_config_review_prompt,
build_prompt,
parse_initial_config_review_text,
parse_proposal_text,
validate_proposal,
)
from aituner.search import ThresholdProbe, binary_search_max_feasible
from aituner.slo import RequestOutcome, evaluate_request, summarize_evaluations
from aituner.spec import (
@@ -266,6 +273,62 @@ class CoreFlowTests(unittest.TestCase):
self.assertIn("knob_harnesses", prompt)
self.assertTrue(study_root.exists())
def test_initial_config_review_schema_prompt_and_parse(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)
study_path = _write_study_assets(tmp_path)
payload = json.loads(study_path.read_text(encoding="utf-8"))
payload["llm"]["initial_config_review"] = {"mode": "warn"}
study_path.write_text(json.dumps(payload), encoding="utf-8")
study = load_study_spec(study_path)
self.assertEqual(study.llm.initial_config_review.mode, "warn")
window, requests = load_trace_requests(study, study_spec_path=study_path)
prompt = build_initial_config_review_prompt(
study=study,
window_summary=summarize_window(requests, window),
capability_profile={"prefill": "profile"},
workload_profile=build_study_workload_profile(study, requests, window),
)
self.assertIn("pre-flight review", prompt)
self.assertIn("knob_descriptors", prompt)
self.assertIn("minimal_repair_patch", prompt)
review = parse_initial_config_review_text(
json.dumps(
{
"verdict": "risky",
"issues": [
{
"knob": "max-num-seqs",
"mechanism": "admission_capacity",
"reason": "low admission capacity may throttle concurrency",
"severity": "high",
}
],
"minimal_repair_patch": {
"env_patch": {},
"flag_patch": {"max-num-seqs": 64},
},
"do_not_change": ["tensor-parallel-size"],
"confidence": 0.8,
"requires_harness_validation": True,
}
),
study,
)
self.assertEqual(review["verdict"], "risky")
self.assertEqual(review["minimal_repair_patch"]["flag_patch"], {"max-num-seqs": 64})
self.assertEqual(review["do_not_change"], ["tensor-parallel-size"])
bad_payload = dict(payload)
bad_payload["llm"] = dict(payload["llm"])
bad_payload["llm"]["initial_config_review"] = {"mode": "repair"}
bad_path = tmp_path / "bad-study.json"
bad_path.write_text(json.dumps(bad_payload), encoding="utf-8")
with self.assertRaisesRegex(SpecError, "llm.initial_config_review.mode"):
load_study_spec(bad_path)
def test_search_auto_high_schema_is_backward_compatible(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
study_path = _write_study_assets(
@@ -7810,6 +7873,98 @@ class CoreFlowTests(unittest.TestCase):
self.assertEqual(state.trials[0].config_patch, {"env_patch": {}, "flag_patch": {}})
self.assertEqual(state.trials[1].config_patch["flag_patch"], {"max-num-seqs": 64})
def test_cli_tune_records_warn_initial_config_review_without_repairing_baseline(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)
study_path = _write_study_assets(tmp_path)
payload = json.loads(study_path.read_text(encoding="utf-8"))
payload["llm"]["initial_config_review"] = {"mode": "warn"}
payload["llm"]["endpoint"] = {
"provider": "custom",
"base_url": "http://llm.example/v1",
"wire_api": "chat.completions",
"model": "test-model",
"api_key_env": "OPENAI_API_KEY",
}
study_path.write_text(json.dumps(payload), encoding="utf-8")
store_root = tmp_path / "store"
def fake_run_trial(trial_spec_path: Path) -> dict[str, object]:
payload = json.loads(trial_spec_path.read_text(encoding="utf-8"))
trial_root = Path(payload["artifact_dir"])
result = {
"study_id": payload["study_id"],
"trial_id": payload["trial_id"],
"status": "completed",
"best_sampling_u": 0.25,
"best_request_rate": 1.0,
"best_pass_rate": 1.0,
"best_request_count": 2,
"probes": [],
}
(trial_root / "result.json").write_text(json.dumps(result), encoding="utf-8")
return result
audit_payload = json.dumps(
{
"verdict": "risky",
"issues": [
{
"knob": "max-num-seqs",
"mechanism": "admission_capacity",
"reason": "initial admission cap may be too low",
"severity": "medium",
}
],
"minimal_repair_patch": {
"env_patch": {},
"flag_patch": {"max-num-seqs": 64},
},
"do_not_change": ["tensor-parallel-size"],
"confidence": 0.7,
"requires_harness_validation": True,
}
)
buffer = io.StringIO()
with mock.patch("aituner.cli.run_trial", side_effect=fake_run_trial):
with mock.patch(
"aituner.cli.call_llm_for_initial_config_review",
return_value=audit_payload,
) as audit_mock:
with contextlib.redirect_stdout(buffer):
exit_code = cli_main(
[
"study",
"tune",
"--spec",
str(study_path),
"--store-root",
str(store_root),
"--max-trials",
"1",
]
)
self.assertEqual(exit_code, 0)
audit_mock.assert_called_once()
summary = json.loads(buffer.getvalue())
self.assertEqual(summary["preflight_audit"]["status"], "completed")
self.assertFalse(summary["preflight_audit"]["repair_applied"])
store = StudyStore(store_root)
state = store.load_state("study-1")
self.assertEqual(state.next_trial_index, 2)
self.assertEqual(state.trials[0].config_patch, {"env_patch": {}, "flag_patch": {}})
audit_dir = store.study_root("study-1") / "preflight_audits"
audit = json.loads((audit_dir / "initial-config-0001.json").read_text(encoding="utf-8"))
self.assertEqual(audit["status"], "completed")
self.assertEqual(audit["review"]["verdict"], "risky")
self.assertEqual(
audit["review"]["minimal_repair_patch"]["flag_patch"],
{"max-num-seqs": 64},
)
self.assertTrue((audit_dir / "initial-config-0001.prompt.txt").exists())
self.assertTrue((audit_dir / "initial-config-0001.raw.txt").exists())
def test_cli_tune_stops_when_baseline_is_all_infeasible(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)