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

Use normalized full config signatures

Browse Source
This commit is contained in:
Gahow Wang
2026-06-26 21:28:10 +08:00
parent 7f50b8b8ea
commit 48911b658b
3 changed files with 233 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
src/aituner/harness.py
View File

@@ -902,16 +902,19 @@ def _harness_proposal_decision(
"expected_effects": [],
}
tested_signatures = {
_config_signature(item.get("config_patch") if isinstance(item, dict) else None)
_effective_config_signature(
study,
item.get("config_patch") if isinstance(item, dict) else None,
)
for item in recent_diagnostics
}
tested_signatures.update(_state_tested_signatures(state))
tested_signatures.update(_state_tested_signatures(study, state))
if experiment_plan is not None:
next_action = experiment_plan.get("next_action")
if isinstance(next_action, dict) and _as_float(next_action.get("score")) >= 0.35:
patch = next_action.get("config_patch")
if isinstance(patch, dict):
signature = _config_signature(patch)
signature = _effective_config_signature(study, patch)
if signature not in tested_signatures:
return {
"should_propose": True,
@@ -976,7 +979,7 @@ def _harness_proposal_decision(
"reason": "no_legal_adjacent_tensor_parallel_probe",
}
flag_patch: dict[str, Any] = {"tensor-parallel-size": next_tp}
signature = _config_signature({"env_patch": {}, "flag_patch": flag_patch})
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": flag_patch})
if signature in tested_signatures:
return {
**default,
@@ -1021,7 +1024,7 @@ def _topology_frontier_proposal(
flag_patch = frontier.get("flag_patch")
if not isinstance(flag_patch, dict):
return {**default, "reason": "topology_frontier_patch_missing"}
signature = _config_signature({"env_patch": {}, "flag_patch": flag_patch})
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": flag_patch})
if signature in tested_signatures:
return {**default, "reason": "topology_frontier_already_tested"}
return {
@@ -1051,10 +1054,13 @@ def _experiment_plan(
bottleneck_hypotheses: list[dict[str, Any]],
) -> dict[str, Any]:
tested_signatures = {
_config_signature(item.get("config_patch") if isinstance(item, dict) else None)
_effective_config_signature(
study,
item.get("config_patch") if isinstance(item, dict) else None,
)
for item in recent_diagnostics
}
tested_signatures.update(_state_tested_signatures(state))
tested_signatures.update(_state_tested_signatures(study, state))
candidates = _candidate_actions(
study,
window_summary,
@@ -1183,7 +1189,7 @@ def _topology_candidate_actions(
if point["tensor-parallel-size"] == current_tp and point["data-parallel-size"] == current_dp:
continue
patch = _topology_patch(study, point)
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": patch})
if signature in tested_signatures:
continue
score, factors = _score_topology_candidate(
@@ -1254,7 +1260,8 @@ def _runtime_candidate_actions(
_next_tp = _next_allowed_tp(study, current_tp=cur_tp, current_dp=cur_dp)
tp_frontier_open = (
_next_tp is not None
and _config_signature(
and _effective_config_signature(
study,
{"env_patch": {}, "flag_patch": {"tensor-parallel-size": _next_tp}}
)
not in tested_signatures
@@ -1276,7 +1283,7 @@ def _runtime_candidate_actions(
mbt_targets.append(("lower_mbt", max(8192, current_mbt // 2)))
for action_id, target in mbt_targets:
patch = {**runtime_base_patch, "max-num-batched-tokens": target}
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": patch})
if signature in tested_signatures:
continue
relief = 0.24 if top_bottleneck == "ttft_prefill" else 0.14
@@ -1332,7 +1339,7 @@ def _runtime_candidate_actions(
mns_targets.append(("raise_max_num_seqs", target))
for action_id, target in mns_targets:
patch = {**runtime_base_patch, "max-num-seqs": target}
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": patch})
if signature in tested_signatures:
continue
if top_bottleneck in {"decode_tpot", "admission_or_queueing"}:
@@ -1388,7 +1395,7 @@ def _runtime_candidate_actions(
"max-num-batched-tokens": mbt_target,
"max-num-seqs": mns_target,
}
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": patch})
if signature not in tested_signatures:
actions.append(
_runtime_action(
@@ -1413,7 +1420,7 @@ def _runtime_candidate_actions(
current = bool(anchor_flags.get("enable-chunked-prefill", False))
if not current:
patch = {**runtime_base_patch, "enable-chunked-prefill": True}
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": patch})
if signature not in tested_signatures:
actions.append(
_runtime_action(
@@ -1444,7 +1451,7 @@ def _runtime_candidate_actions(
)
if target is not None:
patch = {**runtime_base_patch, "gpu-memory-utilization": target}
signature = _config_signature({"env_patch": {}, "flag_patch": patch})
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": patch})
if signature not in tested_signatures:
actions.append(
_runtime_action(
@@ -1557,8 +1564,9 @@ def _has_unmeasured_higher_tp_candidate(
or point["tensor-parallel-size"] <= current_tp
):
continue
signature = _config_signature(
{"env_patch": {}, "flag_patch": _topology_patch(study, point)}
signature = _effective_config_signature(
study,
{"env_patch": {}, "flag_patch": _topology_patch(study, point)},
)
if signature not in tested_signatures:
return True
@@ -1770,7 +1778,7 @@ def _parallel_size_can_vary(study: StudySpec) -> bool:
normalized = _normalized_topology_flags(flags)
if any(normalized.get(key) != point.get(key) for key in point):
continue
signatures.add(_config_signature({"env_patch": {}, "flag_patch": patch}))
signatures.add(_effective_config_signature(study, {"env_patch": {}, "flag_patch": patch}))
return len(signatures) > 1
@@ -1948,8 +1956,8 @@ def _topology_frontier_status(
base_dp = _parse_int_like(study.engine.base_flags.get("data-parallel-size"), default=1)
if current_dp != base_dp:
flag_patch["data-parallel-size"] = current_dp
signature = _config_signature({"env_patch": {}, "flag_patch": flag_patch})
if signature in _state_tested_signatures(state):
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": flag_patch})
if signature in _state_tested_signatures(study, state):
return {
**default,
"reason": "higher_tp_frontier_already_tested",
@@ -1979,9 +1987,9 @@ def _effective_flags_for_item(study: StudySpec, item: dict[str, Any]) -> dict[st
return flags
def _state_tested_signatures(state: StudyState) -> set[str]:
def _state_tested_signatures(study: StudySpec, state: StudyState) -> set[str]:
return {
_config_signature(trial.config_patch)
_effective_config_signature(study, trial.config_patch)
for trial in state.trials
if isinstance(trial.config_patch, dict)
}
@@ -2045,7 +2053,7 @@ def _runtime_refinement_proposal(
"reason": "no_larger_mbt_step_available",
}
flag_patch["max-num-batched-tokens"] = target_mbt
signature = _config_signature({"env_patch": {}, "flag_patch": flag_patch})
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": flag_patch})
if signature in tested_signatures:
return {
**default,
@@ -2643,12 +2651,29 @@ def _parse_float_like(value: Any, *, default: float) -> float:
def _config_signature(config_patch: Any) -> str:
return json.dumps(
_normalized_config_patch(config_patch),
ensure_ascii=False,
sort_keys=True,
separators=(",", ":"),
)
def _effective_config_signature(study: StudySpec, config_patch: Any) -> str:
patch = _normalized_config_patch(config_patch)
payload = {
"env": {**study.engine.base_envs, **patch["env_patch"]},
"flags": {**study.engine.base_flags, **patch["flag_patch"]},
}
return json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
def _normalized_config_patch(config_patch: Any) -> dict[str, dict[str, Any]]:
if not isinstance(config_patch, dict):
config_patch = {}
env_patch = config_patch.get("env_patch")
flag_patch = config_patch.get("flag_patch")
payload = {
return {
"env_patch": env_patch if isinstance(env_patch, dict) else {},
"flag_patch": flag_patch if isinstance(flag_patch, dict) else {},
}
return json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))

11
src/aituner/llm.py
View File

@@ -5,7 +5,7 @@ import time
from pathlib import Path
from typing import TYPE_CHECKING, Any
from .harness import build_harness_context, render_harness_context
from .harness import _effective_config_signature, build_harness_context, render_harness_context
from .http_client import chat_completion, stream_text_completion
from .spec import LLMPolicySpec, Proposal, SpecError, StudySpec, StudyState
@@ -306,7 +306,7 @@ def build_prompt(
json.dumps(launch_failures, ensure_ascii=False, indent=2),
"",
"Tested config signatures:",
json.dumps(_tested_config_signatures(state), ensure_ascii=False, indent=2),
json.dumps(_tested_config_signatures(study, state), ensure_ascii=False, indent=2),
]
return "\n".join(sections)
@@ -402,7 +402,7 @@ def build_prompt(
json.dumps(parallel_candidates, ensure_ascii=False, indent=2),
"",
"Tested config signatures:",
json.dumps(_tested_config_signatures(state), ensure_ascii=False, indent=2),
json.dumps(_tested_config_signatures(study, state), ensure_ascii=False, indent=2),
]
sections.extend(
[
@@ -435,12 +435,12 @@ def build_prompt(
return "\n".join(sections)
def _tested_config_signatures(state: StudyState) -> list[dict[str, Any]]:
def _tested_config_signatures(study: StudySpec, state: StudyState) -> list[dict[str, Any]]:
signatures: list[dict[str, Any]] = []
seen: set[str] = set()
for trial in state.trials:
config_patch = trial.config_patch or {}
signature = json.dumps(config_patch, sort_keys=True, ensure_ascii=False)
signature = _effective_config_signature(study, config_patch)
if signature in seen:
continue
seen.add(signature)
@@ -449,6 +449,7 @@ def _tested_config_signatures(state: StudyState) -> list[dict[str, Any]]:
"trial_id": trial.trial_id,
"status": trial.status,
"best_request_rate_per_gpu": trial.best_request_rate_per_gpu,
"effective_config_signature": signature,
"config_patch": config_patch,
}
)