diff --git a/puzzletron_setup/v2/bundle.py b/puzzletron_setup/v2/bundle.py index d364f0fabd2..d4c47ee28e5 100644 --- a/puzzletron_setup/v2/bundle.py +++ b/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", diff --git a/puzzletron_setup/v2/defaults.py b/puzzletron_setup/v2/defaults.py index 04d8da55673..7718e2ff2f2 100644 --- a/puzzletron_setup/v2/defaults.py +++ b/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 diff --git a/puzzletron_setup/v2/resolved.py b/puzzletron_setup/v2/resolved.py index bcb7776bdcc..e42fb8d7aca 100644 --- a/puzzletron_setup/v2/resolved.py +++ b/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) diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index 3003fd3fdcf..3c0a1935178 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/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"] diff --git a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py index f15fab8d76d..8d7f95de990 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py +++ b/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})