Fix resolved configuration edge cases

puzzletron_setup/v2/bundle.py

@@ -27,7 +27,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
-import yaml
+import yaml  # type: ignore[import-untyped, unused-ignore]
 
 from puzzletron_setup import SetupError
 from puzzletron_setup.bundle import (
@@ -235,8 +235,15 @@ def _render_experiment_v2(
     if config.post_mip_flows_configured:
         rendered["post_mip"] = {"flows": deepcopy(flows)}
 
+    batch_mirrors = {
+        "pruning.micro_batch_size": "data.calibration.micro_batch_size",
+        "replacement_scoring.micro_batch_size": "data.replacement_scoring.micro_batch_size",
+    }
     for dotted, value in config.stage_batches.items():
         _set_dotted(rendered, str(dotted), value)
+        mirrored = batch_mirrors.get(str(dotted))
+        if mirrored is not None:
+            _set_dotted(rendered, mirrored, value)
     parallel_paths = {
         "depth_importance": "depth_importance.automodel.parallel",
         "width_importance": "pruning.automodel.parallel",

puzzletron_setup/v2/defaults.py

@@ -24,7 +24,7 @@
 from pathlib import Path
 from typing import Any
 
-import yaml
+import yaml  # type: ignore[import-untyped, unused-ignore]
 
 from puzzletron_setup import WORKER_REPOSITORY_PLACEHOLDER, WORKER_VENV_PLACEHOLDER, SetupError
 

puzzletron_setup/v2/resolved.py

@@ -602,15 +602,7 @@ def resolve_campaign_config(state: WizardState) -> ResolvedCampaignConfig:
     """Resolve one mutable wizard state into an immutable campaign snapshot."""
     payload = deepcopy(state.payload)
     collections = _mapping(payload.get("collections"))
-    field_records = {
-        str(path): {
-            "value": deepcopy(record.value),
-            "source": str(record.source),
-            "requested": deepcopy(record.requested),
-            "effective": deepcopy(record.effective),
-        }
-        for path, record in state.records().items()
-    }
+    field_records = {str(path): record.to_dict() for path, record in state.records().items()}
 
     def effective(path: str, default: Any = _USE_BUILTIN_DEFAULT) -> Any:
         record = field_records.get(path)

puzzletron_setup/v2/wizard.py

@@ -588,11 +588,10 @@ def data_section(
                 "from the completed stage configuration."
             )
         else:
-            if not selected_subsets or subset_selection is None or acquisition is None:
+            if not selected_subsets or subset_selection is None:
                 raise SetupError(
                     "Nemotron-VLM v2 requires at least one selectable hosted-media subset."
                 )
-            assert subset_selection is not None
             subset_media_shards = {
                 record["name"]: record["num_media_shards"]
                 for record in subset_selection["subsets"]

tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py

@@ -21,6 +21,7 @@
 from copy import deepcopy
 from dataclasses import FrozenInstanceError
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pytest
 import yaml
@@ -35,6 +36,9 @@
 from puzzletron_setup.v2.resolved import resolve_campaign_config
 from puzzletron_setup.v2.state import WizardState
 
+if TYPE_CHECKING:
+    from puzzletron_setup.bundle import BundleValidation
+
 REPOSITORY_ROOT = Path(__file__).resolve().parents[4]
 
 
@@ -345,6 +349,22 @@ def test_resolved_sections_take_precedence_over_compatibility_overrides(tmp_path
     assert experiment["pruning"]["automodel"]["parallel"]["tp"] == 2
 
 
+def test_stage_batches_update_shared_data_sections(tmp_path: Path) -> None:
+    state = _campaign_state(tmp_path)
+    state.set_collection(
+        "stage_batches",
+        {
+            "pruning.micro_batch_size": 7,
+            "replacement_scoring.micro_batch_size": 5,
+        },
+    )
+
+    experiment = render_experiment_v2(state, "production")
+
+    assert experiment["data"]["calibration"]["micro_batch_size"] == 7
+    assert experiment["data"]["replacement_scoring"]["micro_batch_size"] == 5
+
+
 def test_runner_compatibility_override_is_applied_to_resolved_runner(tmp_path: Path) -> None:
     state = _campaign_state(tmp_path)
 
@@ -452,10 +472,9 @@ def test_build_freezes_one_snapshot_before_rendering_both_budgets(
     monkeypatch,
 ) -> None:
     state = _campaign_state(tmp_path)
-
     real_validate_bundle = bundle_module.validate_bundle
 
-    def mutate_state_after_smoke(path: Path):
+    def mutate_state_after_smoke(path: Path) -> BundleValidation:
         if path.name == "smoke":
             state.set_field("stages.width_importance.batch", 99, source="user")
             state.set_collection("stage_batches", {"pruning.micro_batch_size": 99})