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

Fix decode harness partial probe handling

Browse Source
This commit is contained in:
Gahow Wang
2026-05-16 14:18:07 +08:00
parent f18765b235
commit 5a879a8592
3 changed files with 180 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
tests/test_core_flow.py
View File

@@ -39,6 +39,7 @@ from aituner.spec import (
from aituner.store import StudyStore
from aituner.trace import load_trace_requests, summarize_window
from aituner.worker import (
_best_feasible_probe_record,
_latency_summary,
_run_one_request,
_replay_requests,
@@ -1518,6 +1519,103 @@ class CoreFlowTests(unittest.TestCase):
"\n".join(context["proposal_rules"]),
)
def test_decode_topology_planner_prefers_dp_redistribution_and_preserves_ep(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)
study_path = _write_study_assets(
tmp_path,
trace_overrides={"request_mode": "decode_only"},
slo_overrides={
"ttft_rule": None,
"tpot_rule": {"kind": "fixed_ms", "threshold_ms": 40},
},
engine_overrides={
"base_flags": {
"host": "127.0.0.1",
"port": 8000,
"enable-expert-parallel": True,
"tensor-parallel-size": 4,
"data-parallel-size": 2,
"expert-parallel-size": 8,
"max-num-seqs": 192,
},
"tunable_flags": [
"tensor-parallel-size",
"data-parallel-size",
"expert-parallel-size",
"max-num-seqs",
],
"topology_constraints": {
"allowed_tensor_parallel_sizes": [1, 2, 4, 8],
"allowed_data_parallel_sizes": [1, 2, 4, 8],
"allowed_expert_parallel_sizes": [1, 2, 4, 8],
"require_tp_dp_product_equals_gpu_count": True,
"require_ep_size_leq_tp_dp_product": True,
"require_ep_size_divides_tp_dp_product": True,
"require_enable_expert_parallel_when_ep_gt_one": True,
},
},
)
result_path = tmp_path / "trial-0001-result.json"
result_path.write_text(
json.dumps(
{
"status": "completed",
"best_request_rate": 0.47,
"best_pass_rate": 0.98,
"probes": [
{
"threshold": 0.04,
"feasible": False,
"payload": {
"request_rate": 0.72,
"pass_rate": 0.3,
"early_stop_reason": "slo_pass_rate_unrecoverable",
"latency_summary": {
"failed_reason_counts": {"tpot_ms>40.0": 80}
},
},
}
],
}
),
encoding="utf-8",
)
study = load_study_spec(study_path)
context = build_harness_context(
study=study,
window_summary={},
state=StudyState(
study_id=study.study_id,
best_trial_id="trial-0001",
best_request_rate=0.47,
best_request_rate_per_gpu=0.05875,
trials=[
TrialSummary(
trial_id="trial-0001",
status="completed",
best_request_rate=0.47,
best_request_rate_per_gpu=0.05875,
best_pass_rate=0.98,
result_path=str(result_path),
config_patch={"env_patch": {}, "flag_patch": {}},
)
],
),
)
action = context["experiment_plan"]["next_action"]
self.assertEqual(action["knob_family"], "topology")
self.assertEqual(
action["config_patch"]["flag_patch"],
{"tensor-parallel-size": 2, "data-parallel-size": 4},
)
proposal = build_harness_guided_proposal(context)
self.assertIsNotNone(proposal)
self.assertEqual(
proposal.config_patch.flag_patch,
{"tensor-parallel-size": 2, "data-parallel-size": 4},
)
def test_prompt_can_disable_harness_for_ablation(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)
@@ -1625,6 +1723,33 @@ class CoreFlowTests(unittest.TestCase):
self.assertIn("data-parallel-size", active)
self.assertIn("max-num-seqs", active)
def test_best_feasible_probe_record_keeps_partial_probe_evidence(self) -> None:
best = _best_feasible_probe_record(
[
{
"threshold": 0.03125,
"request_rate": 0.72,
"pass_rate": 0.3,
"feasible": False,
},
{
"threshold": 0.015625,
"request_rate": 0.3533,
"pass_rate": 0.99,
"feasible": True,
},
{
"threshold": 0.017578125,
"request_rate": 0.3833,
"pass_rate": 0.995,
"feasible": True,
},
]
)
self.assertIsNotNone(best)
self.assertEqual(best["threshold"], 0.017578125)
self.assertEqual(best["request_rate"], 0.3833)
def test_load_study_spec_rejects_mismatched_served_model_name(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)