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

Add harness early stop ablation

Browse Source
This commit is contained in:
Gahow Wang
2026-05-02 08:08:14 +08:00
parent 6d3459c82d
commit 1a3d628268
9 changed files with 837 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
src/aituner/cli.py
View File

@@ -6,9 +6,10 @@ import sys
from pathlib import Path
from .compare import run_compare
from .harness import build_harness_context, build_harness_stop_proposal
from .job import append_job, build_trial_job
from .llm import build_prompt, call_llm_for_proposal, load_capability_profile, parse_proposal_text
from .spec import Proposal, SpecError, load_study_spec
from .spec import Proposal, SpecError, load_study_spec, to_jsonable
from .store import StudyStore
from .trace import load_trace_requests, summarize_window
from .worker import run_trial
@@ -118,16 +119,23 @@ def cmd_study_tune(args: argparse.Namespace) -> int:
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)
window_summary = summarize_window(requests, window)
harness_context = (
build_harness_context(
study=study,
window_summary=window_summary,
state=state,
)
if study.llm.use_harness
else None
)
prompt = build_prompt(
study=study,
window_summary=summarize_window(requests, window),
window_summary=window_summary,
state=state,
capability_profile=capability_profile,
)
@@ -162,18 +170,36 @@ def cmd_study_tune(args: argparse.Namespace) -> int:
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}"
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:
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)
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:
if proposal_name.startswith("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": str(proposal_source) if proposal_source else "llm",
"proposal_source": proposal_source_label,
"stopped": True,
"diagnosis": proposal.diagnosis,
"state_best_trial_id": state.best_trial_id,

200
src/aituner/harness.py
View File

@@ -4,7 +4,25 @@ import json
from pathlib import Path
from typing import Any
from .spec import StudySpec, StudyState, TrialSummary
from .spec import ConfigPatch, Proposal, StudySpec, StudyState, TrialSummary
_TOPOLOGY_KEYS = {
"tensor-parallel-size",
"data-parallel-size",
"expert-parallel-size",
"enable-expert-parallel",
}
_RUNTIME_KEYS = {
"max-num-seqs",
"max-num-batched-tokens",
"block-size",
"gpu-memory-utilization",
"enable-chunked-prefill",
}
_STRONG_INCUMBENT_MIN_GAIN = 1.8
_MIN_POST_INCUMBENT_VALIDATION_TRIALS = 2
_VALIDATION_TRIALS_WITHOUT_FAMILY_COVERAGE = 3
def build_harness_context(
@@ -23,11 +41,39 @@ def build_harness_context(
"workload_lca_profile": _workload_lca_profile(window_summary),
"recent_trial_diagnostics": recent_diagnostics,
"convergence_guard": _convergence_guard(state, recent_diagnostics),
"harness_stop": _harness_stop_decision(state, recent_diagnostics),
"knob_harnesses": _knob_harnesses(study, window_summary, recent_diagnostics),
"proposal_rules": _proposal_rules(),
}
def build_harness_stop_proposal(context: dict[str, Any]) -> Proposal | None:
stop = context.get("harness_stop")
if not isinstance(stop, dict) or not stop.get("should_stop"):
return None
reason = str(stop.get("reason") or "harness_converged")
evidence = stop.get("evidence") if isinstance(stop.get("evidence"), dict) else {}
observation = (
"Harness convergence guard triggered before requesting another proposal: "
f"{reason}."
)
diagnosis = str(evidence.get("summary") or reason)
return Proposal(
observation=observation,
diagnosis=diagnosis,
config_patch=ConfigPatch(env_patch={}, flag_patch={}),
expected_effects=[
"stop without spending another GPU trial",
"preserve the current best observed configuration",
],
why_not_previous_failures=(
"The stop decision is based on completed validation evidence and does not "
"repeat any failed configuration."
),
should_stop=True,
)
def render_harness_context(context: dict[str, Any]) -> str:
return json.dumps(context, ensure_ascii=False, indent=2)
@@ -423,6 +469,158 @@ def _convergence_guard(
}
def _harness_stop_decision(
state: StudyState,
recent_diagnostics: list[dict[str, Any]],
) -> dict[str, Any]:
guard = _convergence_guard(state, recent_diagnostics)
if guard["should_stop_if_no_harness_can_justify_a_new_adjacent_probe"]:
return {
"should_stop": True,
"reason": guard["reason"],
"evidence": {
"summary": "The convergence guard fired and no further adjacent probe is required.",
"convergence_guard": guard,
},
}
validation = _validation_exhausted_guard(state, recent_diagnostics)
if validation["exhausted"]:
return {
"should_stop": True,
"reason": validation["reason"],
"evidence": validation,
}
return {
"should_stop": False,
"reason": "continue_harness_guided_search",
"evidence": {
"summary": "No deterministic harness stop condition is satisfied.",
"convergence_guard": guard,
"validation_exhausted": validation,
},
}
def _validation_exhausted_guard(
state: StudyState,
recent_diagnostics: list[dict[str, Any]],
) -> dict[str, Any]:
default = {
"exhausted": False,
"reason": "validation_not_exhausted",
"summary": "Validation probes are not sufficient to stop yet.",
"incumbent_trial_id": state.best_trial_id,
"incumbent_gain_vs_baseline": None,
"validation_trial_ids": [],
"validation_families": [],
}
if not state.best_trial_id or not isinstance(state.best_request_rate_per_gpu, (int, float)):
return default
completed = [
item
for item in recent_diagnostics
if item.get("status") == "completed"
and isinstance(item.get("best_request_rate_per_gpu"), (int, float))
]
if not completed:
return default
baseline_rate = _as_float(completed[0].get("best_request_rate_per_gpu"))
incumbent_rate = _as_float(state.best_request_rate_per_gpu)
if baseline_rate <= 0 or incumbent_rate <= 0:
return default
gain = incumbent_rate / baseline_rate
if gain < _STRONG_INCUMBENT_MIN_GAIN:
return {
**default,
"reason": "incumbent_gain_not_large_enough_for_validation_stop",
"incumbent_gain_vs_baseline": gain,
}
best_index = next(
(
index
for index, item in enumerate(recent_diagnostics)
if item.get("trial_id") == state.best_trial_id
),
None,
)
if best_index is None:
return {
**default,
"reason": "incumbent_not_in_recent_harness_history",
"incumbent_gain_vs_baseline": gain,
}
after_best = [
item
for item in recent_diagnostics[best_index + 1 :]
if item.get("status") in {"completed", "failed"}
]
if len(after_best) < _MIN_POST_INCUMBENT_VALIDATION_TRIALS:
return {
**default,
"reason": "need_at_least_two_post_incumbent_validation_trials",
"incumbent_gain_vs_baseline": gain,
"validation_trial_ids": [str(item.get("trial_id")) for item in after_best],
}
if any(isinstance(item.get("best_request_rate_per_gpu"), (int, float)) for item in after_best):
return {
**default,
"reason": "post_incumbent_validation_found_feasible_candidate",
"incumbent_gain_vs_baseline": gain,
"validation_trial_ids": [str(item.get("trial_id")) for item in after_best],
}
families: set[str] = set()
for item in after_best:
families.update(_validation_families(item))
has_topology = "topology" in families
has_runtime = bool(families & {"runtime", "max-num-seqs", "max-num-batched-tokens"})
enough_evidence = (
len(after_best) >= _VALIDATION_TRIALS_WITHOUT_FAMILY_COVERAGE
or (has_topology and has_runtime)
)
if not enough_evidence:
return {
**default,
"reason": "post_incumbent_validation_has_not_covered_enough_families",
"incumbent_gain_vs_baseline": gain,
"validation_trial_ids": [str(item.get("trial_id")) for item in after_best],
"validation_families": sorted(families),
}
return {
"exhausted": True,
"reason": "post_incumbent_validation_exhausted",
"summary": (
"A strong incumbent was followed by validation probes across nearby "
"topology/runtime families, and none produced a feasible candidate."
),
"incumbent_trial_id": state.best_trial_id,
"incumbent_gain_vs_baseline": gain,
"validation_trial_ids": [str(item.get("trial_id")) for item in after_best],
"validation_families": sorted(families),
}
def _validation_families(item: dict[str, Any]) -> set[str]:
config_patch = item.get("config_patch")
if not isinstance(config_patch, dict):
return set()
flag_patch = config_patch.get("flag_patch")
if not isinstance(flag_patch, dict):
return set()
families: set[str] = set()
if any(key in flag_patch for key in _TOPOLOGY_KEYS):
families.add("topology")
for key in _RUNTIME_KEYS:
if key in flag_patch:
families.add("runtime")
families.add(key)
if not families and flag_patch:
families.add("other")
return families
def _strong_incumbent_guard(
state: StudyState,
recent_diagnostics: list[dict[str, Any]],

46
src/aituner/llm.py
View File

@@ -312,23 +312,45 @@ def build_prompt(
"",
"Tested config signatures:",
json.dumps(_tested_config_signatures(state), ensure_ascii=False, indent=2),
"",
"Harnesses:",
render_harness_context(
build_harness_context(
study=study,
window_summary=window_summary,
state=state,
)
),
"",
]
if study.llm.use_harness:
sections.extend(
[
"",
"Harnesses:",
render_harness_context(
build_harness_context(
study=study,
window_summary=window_summary,
state=state,
)
),
"",
]
)
else:
sections.extend(
[
"",
"Harnesses:",
"Disabled by llm.use_harness=false for ablation.",
"",
]
)
sections.extend(
[
"The primary cross-topology comparison metric is request_rate_per_gpu, not raw request_rate.",
"The proposal should beat the incumbent on request_rate_per_gpu under the 95%+ SLO target.",
"The evaluator uses the best feasible sampling_u from the same tp_dp_product group when it exists.",
"If a tp_dp_product group has no history yet, the evaluator starts from the study's original search.low and runs a full binary search for that group.",
"Do not assume a configuration with fewer GPUs should inherit the global incumbent sampling_u.",
"Follow the active harness. Prefer stop over a weak exploratory proposal once a good incumbent has converged.",
]
(
"Follow the active harness. Prefer stop over a weak exploratory proposal once a good incumbent has converged."
if study.llm.use_harness
else "For this ablation, reason from the raw study stack, trial history, launch failures, and tested config signatures without harness hints."
),
]
)
return "\n".join(sections)

6
src/aituner/spec.py
View File

@@ -576,6 +576,7 @@ class LLMPolicySpec:
endpoint: LLMEndpointSpec | None
system_prompt: str
max_history_trials: int
use_harness: bool = True
@classmethod
def from_dict(cls, data: Mapping[str, Any] | None) -> "LLMPolicySpec":
@@ -593,6 +594,11 @@ class LLMPolicySpec:
max_history_trials=_require_int(
payload.get("max_history_trials", 8), context="llm.max_history_trials"
),
use_harness=(
_require_bool(payload.get("use_harness"), context="llm.use_harness")
if payload.get("use_harness") is not None
else True
),
)