Inherit incumbent topology for runtime validation

src/aituner/store.py

@@ -5,7 +5,15 @@ from dataclasses import replace
 from pathlib import Path
 from typing import Any
 
-from .spec import Proposal, StudySpec, StudyState, TrialSpec, TrialSummary, to_jsonable
+from .spec import ConfigPatch, Proposal, StudySpec, StudyState, TrialSpec, TrialSummary, to_jsonable
+
+
+_TOPOLOGY_FLAG_KEYS = {
+    "tensor-parallel-size",
+    "data-parallel-size",
+    "expert-parallel-size",
+    "enable-expert-parallel",
+}
 
 
 class StudyStore:
@@ -65,6 +73,11 @@ class StudyStore:
         state: StudyState,
         proposal: Proposal,
     ) -> tuple[TrialSpec, StudyState]:
+        proposal = _inherit_incumbent_topology_for_runtime_patch(
+            study=study,
+            state=state,
+            proposal=proposal,
+        )
         trial_id = f"trial-{state.next_trial_index:04d}"
         trial_root = self.study_root(study.study_id) / "trials" / trial_id
         trial_root.mkdir(parents=True, exist_ok=True)
@@ -225,6 +238,47 @@ def _parallel_size_for_proposal(*, study: StudySpec, proposal: Proposal) -> int:
     return _parallel_size_for_config(study=study, flag_patch=proposal.config_patch.flag_patch)
 
 
+def _inherit_incumbent_topology_for_runtime_patch(
+    *,
+    study: StudySpec,
+    state: StudyState,
+    proposal: Proposal,
+) -> Proposal:
+    flag_patch = dict(proposal.config_patch.flag_patch)
+    env_patch = dict(proposal.config_patch.env_patch)
+    if not flag_patch and not env_patch:
+        return proposal
+    if _TOPOLOGY_FLAG_KEYS.intersection(flag_patch):
+        return proposal
+    if not state.best_trial_id:
+        return proposal
+    incumbent = next(
+        (trial for trial in state.trials if trial.trial_id == state.best_trial_id),
+        None,
+    )
+    if incumbent is None or not isinstance(incumbent.config_patch, dict):
+        return proposal
+    incumbent_patch = incumbent.config_patch.get("flag_patch")
+    if not isinstance(incumbent_patch, dict):
+        return proposal
+    inherited_topology = {
+        key: value
+        for key, value in incumbent_patch.items()
+        if key in _TOPOLOGY_FLAG_KEYS and study.engine.base_flags.get(key) != value
+    }
+    if not inherited_topology:
+        return proposal
+    merged_flag_patch = dict(inherited_topology)
+    merged_flag_patch.update(flag_patch)
+    return replace(
+        proposal,
+        config_patch=ConfigPatch(
+            env_patch=env_patch,
+            flag_patch=merged_flag_patch,
+        ),
+    )
+
+
 def _parallel_size_for_trial_id(*, study: StudySpec, study_root: Path, trial_id: str) -> int | None:
     trial_spec_path = study_root / "trials" / trial_id / "trial_spec.json"
     if not trial_spec_path.exists():