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

Improve harness incumbent follow-up search

Browse Source
This commit is contained in:
Gahow Wang
2026-06-20 05:37:15 +08:00
parent b3156a382a
commit 5257fbc1a2
3 changed files with 316 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
src/aituner/harness.py
View File

@@ -1192,6 +1192,7 @@ def _runtime_candidate_actions(
tunable = set(study.engine.tunable_flags)
anchor_flags = _effective_flags_for_item(study, anchor)
topology_patch = _preserve_topology_patch(study, anchor_flags)
runtime_base_patch = {**topology_patch, **_preserve_runtime_patch(study, anchor_flags)}
actions: list[dict[str, Any]] = []
cur_tp = _parse_int_like(anchor_flags.get("tensor-parallel-size"), default=1)
@@ -1224,7 +1225,7 @@ def _runtime_candidate_actions(
elif top_bottleneck == "decode_tpot" and current_mbt > 8192:
mbt_targets.append(("lower_mbt", max(8192, current_mbt // 2)))
for action_id, target in mbt_targets:
patch = {**topology_patch, "max-num-batched-tokens": target}
patch = {**runtime_base_patch, "max-num-batched-tokens": target}
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
if signature in tested_signatures:
continue
@@ -1264,7 +1265,7 @@ def _runtime_candidate_actions(
)
mns_targets.append(("raise_max_num_seqs", raise_target))
for action_id, target in mns_targets:
patch = {**topology_patch, "max-num-seqs": target}
patch = {**runtime_base_patch, "max-num-seqs": target}
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
if signature in tested_signatures:
continue
@@ -1289,7 +1290,7 @@ def _runtime_candidate_actions(
if "enable-chunked-prefill" in tunable and top_bottleneck == "ttft_prefill":
current = bool(anchor_flags.get("enable-chunked-prefill", False))
if not current:
patch = {**topology_patch, "enable-chunked-prefill": True}
patch = {**runtime_base_patch, "enable-chunked-prefill": True}
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
if signature not in tested_signatures:
actions.append(
@@ -1320,7 +1321,7 @@ def _runtime_candidate_actions(
if 0.0 < current_gmu < _GMU_SAFE_CEILING:
target = round(min(_GMU_SAFE_CEILING, current_gmu + _GMU_STEP), 4)
if target > current_gmu:
patch = {**topology_patch, "gpu-memory-utilization": target}
patch = {**runtime_base_patch, "gpu-memory-utilization": target}
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
if signature not in tested_signatures:
actions.append(
@@ -1488,6 +1489,18 @@ def _preserve_topology_patch(study: StudySpec, flags: dict[str, Any]) -> dict[st
return patch
def _preserve_runtime_patch(study: StudySpec, flags: dict[str, Any]) -> dict[str, Any]:
patch: dict[str, Any] = {}
tunable = set(study.engine.tunable_flags)
base = study.engine.base_flags
for key in _RUNTIME_KEYS:
if key not in tunable or key not in flags:
continue
if flags.get(key) != base.get(key):
patch[key] = flags[key]
return patch
def _normalized_topology_flags(flags: dict[str, Any]) -> dict[str, Any]:
return {
"tensor-parallel-size": _parse_int_like(
@@ -1762,11 +1775,15 @@ def _runtime_refinement_proposal(
best_flags = best_patch.get("flag_patch")
if not isinstance(best_flags, dict):
best_flags = {}
best_tp = _parse_int_like(best_flags.get("tensor-parallel-size"), default=1)
best_effective_flags = _effective_flags_for_item(study, best)
best_tp = _parse_int_like(best_effective_flags.get("tensor-parallel-size"), default=1)
if best_tp <= 1:
return default
tunable = set(study.engine.tunable_flags)
flag_patch: dict[str, Any] = {"tensor-parallel-size": best_tp}
flag_patch = {
**_preserve_topology_patch(study, best_effective_flags),
**_preserve_runtime_patch(study, best_effective_flags),
}
if "enable-chunked-prefill" in tunable:
flag_patch["enable-chunked-prefill"] = True
if "max-num-batched-tokens" not in tunable:
@@ -1801,7 +1818,7 @@ def _runtime_refinement_proposal(
"config_patch": {"env_patch": {}, "flag_patch": flag_patch},
"expected_effects": [
"preserve the incumbent topology",
"increase batching headroom without also raising memory pressure",
"increase batching headroom without dropping measured runtime gains",
],
"incumbent_trial_id": best.get("trial_id"),
}
@@ -1989,12 +2006,18 @@ def _validation_exhausted_guard(
"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):
improving_trials = [
item
for item in after_best
if isinstance(item.get("best_request_rate_per_gpu"), (int, float))
and float(item["best_request_rate_per_gpu"]) > incumbent_rate
]
if improving_trials:
return {
**default,
"reason": "post_incumbent_validation_found_feasible_candidate",
"reason": "post_incumbent_validation_found_improving_candidate",
"incumbent_gain_vs_baseline": gain,
"validation_trial_ids": [str(item.get("trial_id")) for item in after_best],
"validation_trial_ids": [str(item.get("trial_id")) for item in improving_trials],
}
families: set[str] = set()
@@ -2020,7 +2043,7 @@ def _validation_exhausted_guard(
"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."
"topology/runtime families, and none improved request_rate_per_gpu."
),
"incumbent_trial_id": state.best_trial_id,
"incumbent_gain_vs_baseline": gain,

42
src/aituner/worker.py
View File

@@ -781,20 +781,28 @@ def run_trial(trial_spec_path: Path) -> dict[str, Any]:
best = primary_search.best_feasible_payload
best_source = "primary_search"
fallback_search = None
skipped_lower_range_fallback = False
lower_range_fallback_skip_reason = ""
original_search_low = float(study.search.low)
inherited_search_floor = float(trial.search.low)
if best is None and inherited_search_floor > original_search_low:
fallback_search = binary_search_max_feasible(
low=original_search_low,
high=inherited_search_floor,
tolerance=trial.search.tolerance,
max_probes=trial.search.max_probes,
evaluator=evaluator,
)
if fallback_search.best_feasible_payload is not None:
search_for_best = fallback_search
best = fallback_search.best_feasible_payload
best_source = "lower_range_fallback"
if trial.search.inherit_incumbent_floor:
skipped_lower_range_fallback = True
lower_range_fallback_skip_reason = (
"primary_search_above_incumbent_floor_all_infeasible"
)
else:
fallback_search = binary_search_max_feasible(
low=original_search_low,
high=inherited_search_floor,
tolerance=trial.search.tolerance,
max_probes=trial.search.max_probes,
evaluator=evaluator,
)
if fallback_search.best_feasible_payload is not None:
search_for_best = fallback_search
best = fallback_search.best_feasible_payload
best_source = "lower_range_fallback"
def serialize_probe(probe: ThresholdProbe[ProbePayload]) -> dict[str, Any]:
return {
@@ -826,7 +834,7 @@ def run_trial(trial_spec_path: Path) -> dict[str, Any]:
"best_request_count": best.request_count if best is not None else None,
"probes": [serialize_probe(probe) for probe in all_probes],
}
if fallback_search is not None:
if fallback_search is not None or skipped_lower_range_fallback:
result["primary_search"] = {
"low": inherited_search_floor,
"high": trial.search.high,
@@ -838,6 +846,16 @@ def run_trial(trial_spec_path: Path) -> dict[str, Any]:
else None,
"probes": [serialize_probe(probe) for probe in primary_search.probes],
}
if skipped_lower_range_fallback:
result["lower_range_fallback"] = {
"triggered": False,
"skipped": True,
"reason": lower_range_fallback_skip_reason,
"low": original_search_low,
"high": inherited_search_floor,
"probes": [],
}
if fallback_search is not None:
result["lower_range_fallback"] = {
"triggered": True,
"low": original_search_low,