diff --git a/src/aituner/harness.py b/src/aituner/harness.py index d2fcafe..d774f83 100644 --- a/src/aituner/harness.py +++ b/src/aituner/harness.py @@ -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, diff --git a/src/aituner/worker.py b/src/aituner/worker.py index 23b1948..b7c6cc3 100644 --- a/src/aituner/worker.py +++ b/src/aituner/worker.py @@ -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, diff --git a/tests/test_core_flow.py b/tests/test_core_flow.py index e6a4444..9b3395f 100644 --- a/tests/test_core_flow.py +++ b/tests/test_core_flow.py @@ -998,6 +998,76 @@ class CoreFlowTests(unittest.TestCase): self.assertIsNotNone(proposal) self.assertTrue(proposal.should_stop) + def test_harness_stop_after_non_improving_feasible_validation_is_exhausted(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + study_path = _write_study_assets(tmp_path) + study = load_study_spec(study_path) + state = StudyState( + study_id=study.study_id, + best_trial_id="trial-0002", + best_parallel_size=8, + best_sampling_u=0.02, + best_request_rate=2.4, + best_request_rate_per_gpu=0.3, + trials=[ + TrialSummary( + trial_id="trial-0001", + status="completed", + parallel_size=8, + best_request_rate=0.8, + best_request_rate_per_gpu=0.1, + config_patch={"env_patch": {}, "flag_patch": {}}, + ), + TrialSummary( + trial_id="trial-0002", + status="completed", + parallel_size=8, + best_request_rate=2.4, + best_request_rate_per_gpu=0.3, + config_patch={ + "env_patch": {}, + "flag_patch": { + "tensor-parallel-size": 2, + "data-parallel-size": 4, + }, + }, + ), + TrialSummary( + trial_id="trial-0003", + status="completed", + parallel_size=8, + best_request_rate=2.0, + best_request_rate_per_gpu=0.25, + config_patch={ + "env_patch": {}, + "flag_patch": { + "tensor-parallel-size": 1, + "data-parallel-size": 8, + }, + }, + ), + TrialSummary( + trial_id="trial-0004", + status="completed", + parallel_size=8, + best_request_rate=2.1, + best_request_rate_per_gpu=0.2625, + config_patch={ + "env_patch": {}, + "flag_patch": {"max-num-seqs": 160}, + }, + ), + ], + ) + context = build_harness_context( + study=study, + window_summary={"prompt_tokens_p95": 2048}, + state=state, + ) + self.assertTrue(context["harness_stop"]["should_stop"]) + self.assertEqual(context["harness_stop"]["reason"], "post_incumbent_validation_exhausted") + def test_harness_does_not_stop_immediately_after_strong_incumbent(self) -> None: with tempfile.TemporaryDirectory() as tmp: tmp_path = Path(tmp) @@ -1318,6 +1388,100 @@ class CoreFlowTests(unittest.TestCase): }, ) + def test_harness_runtime_refinement_preserves_incumbent_runtime_knobs(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + study_path = _write_study_assets( + tmp_path, + engine_overrides={ + "tunable_flags": [ + "tensor-parallel-size", + "gpu-memory-utilization", + "max-num-seqs", + "enable-chunked-prefill", + "max-num-batched-tokens", + ], + "topology_constraints": { + "allowed_tensor_parallel_sizes": [1, 2, 4], + "allowed_tp_dp_products": [1, 2, 4], + }, + }, + ) + study = load_study_spec(study_path) + result_path = tmp_path / "trial-0002.json" + result_path.write_text( + json.dumps( + { + "status": "completed", + "best_sampling_u": 0.098, + "best_request_rate": 3.3, + "best_pass_rate": 0.97, + "probes": [ + { + "threshold": 0.098, + "feasible": True, + "payload": { + "request_count": 100, + "pass_rate": 0.97, + "request_rate": 3.3, + "early_stopped": False, + "early_stop_reason": "", + "latency_summary": {"failed_reason_counts": {}}, + }, + } + ], + } + ), + encoding="utf-8", + ) + state = StudyState( + study_id=study.study_id, + best_trial_id="trial-0002", + best_request_rate=3.3, + best_request_rate_per_gpu=0.825, + trials=[ + TrialSummary( + trial_id="trial-0001", + status="completed", + best_request_rate=2.5, + best_request_rate_per_gpu=0.625, + config_patch={"env_patch": {}, "flag_patch": {"tensor-parallel-size": 4}}, + ), + TrialSummary( + trial_id="trial-0002", + status="completed", + best_request_rate=3.3, + best_request_rate_per_gpu=0.825, + result_path=str(result_path), + config_patch={ + "env_patch": {}, + "flag_patch": { + "tensor-parallel-size": 4, + "gpu-memory-utilization": 0.92, + "max-num-seqs": 48, + }, + }, + ), + ], + ) + context = build_harness_context( + study=study, + window_summary={"prompt_tokens_p99": 8100}, + state=state, + ) + proposal = build_harness_guided_proposal(context) + self.assertIsNotNone(proposal) + self.assertEqual( + proposal.config_patch.flag_patch, + { + "tensor-parallel-size": 4, + "gpu-memory-utilization": 0.92, + "max-num-seqs": 48, + "enable-chunked-prefill": True, + "max-num-batched-tokens": 16384, + }, + ) + def test_harness_raises_gpu_mem_util_on_settled_decode_bound_incumbent(self) -> None: """Regression for the coverage gap that let the naive baseline beat the harness: a settled TP incumbent that is decode_tpot-bound must get a gpu-memory-utilization @@ -3511,6 +3675,94 @@ class CoreFlowTests(unittest.TestCase): [0.25, 0.375], ) + def test_run_trial_skips_fallback_below_incumbent_floor(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["search"]["max_probes"] = 2 + payload["search"]["inherit_incumbent_floor"] = True + study_path.write_text(json.dumps(payload), encoding="utf-8") + study = load_study_spec(study_path) + store = StudyStore(tmp_path / ".aituner" / "studies") + store.init_study(spec_path=study_path, study=study) + state = StudyState( + study_id=study.study_id, + best_trial_id="trial-0001", + best_parallel_size=1, + best_sampling_u=0.5, + best_request_rate=2.0, + best_request_rate_per_gpu=2.0, + next_trial_index=2, + best_by_parallel_size={ + "1": { + "trial_id": "trial-0001", + "parallel_size": 1, + "best_sampling_u": 0.5, + "best_request_rate": 2.0, + "best_request_rate_per_gpu": 2.0, + } + }, + trials=[], + ) + proposal = Proposal.from_dict( + { + "observation": "runtime patch", + "diagnosis": "primary range all infeasible", + "config_patch": {"env_patch": {}, "flag_patch": {"max-num-seqs": 2}}, + "expected_effects": ["measure"], + } + ) + trial, _ = store.materialize_trial(study=study, state=state, proposal=proposal) + self.assertEqual(trial.search.low, 0.5) + self.assertTrue(trial.search.inherit_incumbent_floor) + + def fake_replay(requests, **kwargs): + return ( + [ + RequestOutcome( + request_id=request.row_id, + success=True, + ttft_ms=10000.0, + tpot_ms=1000.0, + prompt_tokens=request.prompt_tokens_hint, + completion_tokens=request.completion_tokens_hint, + ) + for request in requests + ], + False, + "", + ) + + process = mock.Mock() + process.poll.return_value = 0 + with mock.patch("aituner.worker.subprocess.Popen", return_value=process): + with mock.patch("aituner.worker._wait_for_server_or_exit", return_value=None): + with mock.patch("aituner.worker._terminate_process_tree", return_value=None): + with mock.patch("aituner.worker._replay_requests", side_effect=fake_replay): + result = run_trial(Path(trial.artifact_dir) / "trial_spec.json") + + self.assertEqual(result["status"], "completed") + self.assertIsNone(result["best_request_rate"]) + self.assertEqual(result["best_source"], "primary_search") + self.assertEqual(result["primary_search"]["low"], 0.5) + self.assertIsNone(result["primary_search"]["best_request_rate"]) + self.assertEqual( + [probe["threshold"] for probe in result["primary_search"]["probes"]], + [0.75, 0.625], + ) + self.assertEqual(result["lower_range_fallback"]["triggered"], False) + self.assertEqual(result["lower_range_fallback"]["skipped"], True) + self.assertEqual(result["lower_range_fallback"]["probes"], []) + self.assertEqual( + result["lower_range_fallback"]["reason"], + "primary_search_above_incumbent_floor_all_infeasible", + ) + self.assertEqual( + result["all_infeasible_diagnostics"]["threshold"], + 0.625, + ) + def test_materialize_trial_does_not_mutate_input_state_trials(self) -> None: with tempfile.TemporaryDirectory() as tmp: tmp_path = Path(tmp)