Fix decode harness partial probe handling

src/aituner/harness.py

@@ -1038,6 +1038,7 @@ def _topology_candidate_actions(
         score, factors = _score_topology_candidate(
             top_bottleneck,
             bottleneck_hypotheses,
+            request_mode=study.trace.request_mode,
             current_tp=current_tp,
             current_dp=current_dp,
             candidate_tp=point["tensor-parallel-size"],
@@ -1225,7 +1226,13 @@ def _legal_topology_points(
     else:
         dp_values = [current_dp]
 
-    if constraints is not None and constraints.allowed_expert_parallel_sizes:
+    if (
+        study.trace.request_mode == "decode_only"
+        and current_enable_ep
+        and current_ep > 1
+    ):
+        ep_values = [current_ep]
+    elif constraints is not None and constraints.allowed_expert_parallel_sizes:
         ep_values = sorted(set(constraints.allowed_expert_parallel_sizes))
     elif "expert-parallel-size" in tunable:
         ep_values = sorted({1, current_ep})
@@ -1349,6 +1356,7 @@ def _score_topology_candidate(
     top_bottleneck: str,
     bottleneck_hypotheses: list[dict[str, Any]],
     *,
+    request_mode: str,
     current_tp: int,
     current_dp: int,
     candidate_tp: int,
@@ -1360,6 +1368,15 @@ def _score_topology_candidate(
     relief = 0.0
     if top_bottleneck == "ttft_prefill":
         relief = 0.42 if tp_delta > 0 else 0.05
+    elif top_bottleneck == "decode_tpot" and request_mode == "decode_only":
+        if dp_delta > 0 and candidate_tp <= current_tp:
+            relief = 0.44
+        elif dp_delta > 0:
+            relief = 0.24
+        elif tp_delta > 0 and candidate_dp < current_dp:
+            relief = 0.03
+        else:
+            relief = 0.08
     elif top_bottleneck == "decode_tpot":
         relief = 0.34 if tp_delta > 0 else 0.02
     elif top_bottleneck == "admission_or_queueing":
@@ -1485,6 +1502,12 @@ def _topology_frontier_status(
             "reason": "active_bottleneck_does_not_require_tp_frontier",
             "active_bottleneck": active_bottleneck,
         }
+    if active_bottleneck == "decode_tpot" and study.trace.request_mode == "decode_only":
+        return {
+            **default,
+            "reason": "decode_tpot_uses_topology_redistribution_not_higher_tp_frontier",
+            "active_bottleneck": active_bottleneck,
+        }
 
     flags = _effective_flags_for_item(study, best)
     current_tp = _parse_int_like(flags.get("tensor-parallel-size"), default=1)

src/aituner/worker.py

@@ -209,6 +209,17 @@ def _probe_outcome_details(
     }
 
 
+def _best_feasible_probe_record(probe_history: list[dict[str, Any]]) -> dict[str, Any] | None:
+    feasible = [
+        item
+        for item in probe_history
+        if item.get("feasible") and isinstance(item.get("request_rate"), (int, float))
+    ]
+    if not feasible:
+        return None
+    return max(feasible, key=lambda item: float(item["request_rate"]))
+
+
 def _replay_requests(
     requests: list[TraceRequest],
     *,
@@ -633,6 +644,26 @@ def run_trial(trial_spec_path: Path) -> dict[str, Any]:
             StudyStore.write_json(Path(trial.result_path), result)
             return result
         except Exception as exc:  # noqa: BLE001
+            partial_best = _best_feasible_probe_record(probe_history)
+            if partial_best is not None:
+                result = {
+                    "study_id": trial.study_id,
+                    "trial_id": trial.trial_id,
+                    "status": "completed",
+                    "config_patch": to_jsonable(trial.config_patch),
+                    "best_source": "partial_probe_before_failure",
+                    "best_sampling_u": partial_best.get("threshold"),
+                    "best_request_rate": partial_best.get("request_rate"),
+                    "best_pass_rate": partial_best.get("pass_rate"),
+                    "best_request_count": partial_best.get("request_count"),
+                    "completed_with_probe_failure": True,
+                    "failure_stage": failure_stage,
+                    "failure_reason": str(exc),
+                    "failure_traceback": traceback.format_exc(),
+                    "probes": probe_history,
+                }
+                StudyStore.write_json(Path(trial.result_path), result)
+                return result
             result = {
                 "study_id": trial.study_id,
                 "trial_id": trial.trial_id,