diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index d5149969d13..a5f439f8227 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -57,14 +57,14 @@ overrides: python examples/puzzletron/puzzletron_setup.py --detailed ``` -Every answer is saved atomically. A new invocation starts fresh; resume is always -explicit: +Every answer is saved atomically. Invocations without `--resume` start a campaign; +resuming an existing campaign requires its path: ```bash python examples/puzzletron/puzzletron_setup.py --resume /path/to/campaign ``` -The initial profiles support Nemotron 3 and Qwen 3.5/3.6 dense, MoE, text, and +The supported profiles cover Nemotron 3 and Qwen 3.5/3.6 dense, MoE, text, and multimodal configurations. Unsupported configs exit with detected metadata and point to `.agents/skills/running-puzzletron/SKILL.md` for descriptor onboarding. @@ -75,8 +75,20 @@ with `localhost` for a single local host. ### Setup wizard v2 -The new schema-driven wizard keeps the existing entry point unchanged and adds -local defaults-versus-customize decisions at every section: +The setup wizard v2 offers three guided profiles: + +- **Quick smoke** is the fastest way to verify that the campaign shape is valid. +- **Balanced pruning** is recommended for a first real campaign. +- **High-confidence search** spends more runtime on scoring and sanity checks. + +The selected profile supplies nested pruning and MIP defaults from the detected +model family's `setup_v2_defaults.yaml`. A family file can refine those values +for an exact model geometry, so a small and large model in the same family do +not need to share scoring budgets. Unspecified model values inherit the family +profile, while an explicitly selected defaults file has the highest +default precedence. Setup then asks for the model and dataset, and requires +explicit acceptance or customization of infrastructure-specific worker and +cluster defaults: ```bash python examples/puzzletron/puzzletron_setup_v2.py \ @@ -85,9 +97,17 @@ python examples/puzzletron/puzzletron_setup_v2.py \ The example defaults use only repository-relative values. Copy the file and add site-specific data, scheduler, and container settings before selecting it. -Defaults are loaded only when passed explicitly. Selection prompts have -a visible **← Back** action; text and numeric prompts accept `:back`. Every -accepted answer and the exact navigation frame are saved in +The defaults file is loaded only when passed explicitly and takes precedence +over the selected profile. To expose every per-section and nested setting, use +the advanced flow explicitly: + +```bash +python examples/puzzletron/puzzletron_setup_v2.py --full +``` + +Press **Esc** to go back from any prompt. Selection prompts show a visible +**← Back** action, and text or numeric prompts accept `:back`. +Every accepted answer and the exact navigation frame are saved in `answers_v2.yaml`, so an interrupted session can resume with: ```bash diff --git a/examples/puzzletron/configs/families/nemotron3/setup_v2_defaults.yaml b/examples/puzzletron/configs/families/nemotron3/setup_v2_defaults.yaml new file mode 100644 index 00000000000..8c143a3d286 --- /dev/null +++ b/examples/puzzletron/configs/families/nemotron3/setup_v2_defaults.yaml @@ -0,0 +1,84 @@ +schema_version: 2 + +profiles: + smoke: + pruning: + depth_remove: 1 + depth_importance_samples: 32 + width_importance_samples: 512 + sort_sanity: false + width_sanity: false + slicing_sanity: false + replacement_samples: 32 + bypass: + enabled: false + samples: 64 + mip: + goal_value: 90% + num_solutions: 2 + + balanced: + pruning: + depth_remove: 4 + depth_importance_samples: 128 + width_importance_samples: 32768 + sort_sanity: false + width_sanity: false + slicing_sanity: false + replacement_samples: 128 + bypass: + enabled: true + samples: 4096 + mip: + goal_value: 75% + num_solutions: 8 + + high-confidence: + pruning: + depth_remove: 6 + depth_importance_samples: 512 + width_importance_samples: 65536 + sort_sanity: true + sort_sanity_samples: 512 + width_sanity: true + width_sanity_samples: 512 + slicing_sanity: true + replacement_samples: 512 + bypass: + enabled: true + samples: 8192 + mip: + goal_value: 70% + num_solutions: 16 + +# Model-specific values are seeded from historical campaign configurations. +# Unspecified fields continue to inherit the family profile above. +model_overrides: + nano_30b_a3b_bf16: + match: + num_layers: 52 + facts: + hidden_size: 2688 + intermediate_size: 1856 + num_attention_heads: 32 + num_key_value_heads: 2 + num_experts: 128 + profiles: + smoke: + pruning: + depth_importance_samples: 2 + width_importance_samples: 2 + replacement_samples: 2 + bypass: + enabled: true + mip: + num_solutions: 1 + high-confidence: + pruning: + depth_remove: 5 + depth_importance_samples: 128 + width_importance_samples: 8192 + replacement_samples: 128 + mip: + goal_value: 75% + num_solutions: 5 diff --git a/examples/puzzletron/configs/families/qwen3_5/setup_v2_defaults.yaml b/examples/puzzletron/configs/families/qwen3_5/setup_v2_defaults.yaml new file mode 100644 index 00000000000..b255658245c --- /dev/null +++ b/examples/puzzletron/configs/families/qwen3_5/setup_v2_defaults.yaml @@ -0,0 +1,105 @@ +schema_version: 2 + +profiles: + smoke: + pruning: + depth_remove: 1 + depth_importance_samples: 32 + width_importance_samples: 512 + sort_sanity: false + width_sanity: false + slicing_sanity: false + replacement_samples: 32 + bypass: + enabled: false + samples: 64 + mip: + goal_value: 90% + num_solutions: 2 + + balanced: + pruning: + depth_remove: 4 + depth_importance_samples: 128 + width_importance_samples: 32768 + sort_sanity: false + width_sanity: false + slicing_sanity: false + replacement_samples: 128 + bypass: + enabled: true + samples: 4096 + mip: + goal_value: 75% + num_solutions: 8 + + high-confidence: + pruning: + depth_remove: 6 + depth_importance_samples: 512 + width_importance_samples: 65536 + sort_sanity: true + sort_sanity_samples: 512 + width_sanity: true + width_sanity_samples: 512 + slicing_sanity: true + replacement_samples: 512 + bypass: + enabled: true + samples: 8192 + mip: + goal_value: 70% + num_solutions: 16 + +# Model-specific values are seeded from historical campaign configurations. +# Unspecified fields continue to inherit the family profile above. +model_overrides: + qwen3p5_0p8b: + match: + num_layers: 24 + facts: + hidden_size: 1024 + intermediate_size: 3584 + num_attention_heads: 8 + num_key_value_heads: 2 + profiles: + smoke: + pruning: + width_importance_samples: 8 + replacement_samples: 4 + mip: + goal_value: 85% + + qwen3p5_9b: + match: + num_layers: 32 + facts: + hidden_size: 4096 + intermediate_size: 12288 + num_attention_heads: 16 + num_key_value_heads: 4 + profiles: + high-confidence: + pruning: + depth_importance_samples: 128 + replacement_samples: 128 + bypass: + enabled: false + mip: + goal_value: 75% + + qwen3p6_27b: + match: + num_layers: 64 + facts: + hidden_size: 5120 + intermediate_size: 17408 + num_attention_heads: 24 + num_key_value_heads: 4 + profiles: + balanced: + pruning: + width_importance_samples: 16384 + replacement_samples: 16 + mip: + goal_value: 85% diff --git a/examples/puzzletron/puzzletron_setup_v2.py b/examples/puzzletron/puzzletron_setup_v2.py index 5d96addb185..10bf75cc5d4 100644 --- a/examples/puzzletron/puzzletron_setup_v2.py +++ b/examples/puzzletron/puzzletron_setup_v2.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Create a fully configurable Puzzletron campaign without launching it.""" +"""Create a guided or fully configurable Puzzletron campaign without launching it.""" from __future__ import annotations diff --git a/puzzletron_setup/v2/__init__.py b/puzzletron_setup/v2/__init__.py index 041c4948195..6d0e478f279 100644 --- a/puzzletron_setup/v2/__init__.py +++ b/puzzletron_setup/v2/__init__.py @@ -3,7 +3,7 @@ """Schema-driven Puzzletron campaign setup.""" -from .defaults import DefaultsResolver, ResolvedDefault, load_defaults +from .defaults import DefaultsResolver, ResolvedDefault, load_defaults, validate_defaults from .state import FieldRecord, PromptFrame, WizardState __all__ = [ @@ -13,4 +13,5 @@ "ResolvedDefault", "WizardState", "load_defaults", + "validate_defaults", ] diff --git a/puzzletron_setup/v2/bundle.py b/puzzletron_setup/v2/bundle.py index 532a07720b6..c0eb7cd8720 100644 --- a/puzzletron_setup/v2/bundle.py +++ b/puzzletron_setup/v2/bundle.py @@ -257,10 +257,7 @@ def _bundle_readme( acquisition_command.extend( [ "--subset-rows", - *[ - f"{name}={rows}" - for name, rows in subset_rows.items() - ], + *[f"{name}={rows}" for name, rows in subset_rows.items()], ] ) acquisition_command.extend( @@ -342,14 +339,20 @@ def build_bundles_v2(campaign_dir: Path, state: WizardState) -> BundleResult: (bundle / "dry-run-plan.txt").write_text(dry_run_bundle(bundle)) resolved = { - path: { - "value": record.value, - "requested": record.requested, - "effective": record.effective, - "source": record.source, - } - for path, record in state.records().items() + str(path): dict(record) + for path, record in dict(state.collection("default_resolutions") or {}).items() } + resolved.update( + { + path: { + "value": record.value, + "requested": record.requested, + "effective": record.effective, + "source": record.source, + } + for path, record in state.records().items() + } + ) _write_yaml(temp_root / "resolved_defaults.yaml", resolved) repository = str( state.get_field( diff --git a/puzzletron_setup/v2/cli.py b/puzzletron_setup/v2/cli.py index 42cfe3ce8c9..a012eeae3cc 100644 --- a/puzzletron_setup/v2/cli.py +++ b/puzzletron_setup/v2/cli.py @@ -7,10 +7,13 @@ import argparse from pathlib import Path -from typing import Optional, Sequence +from typing import TYPE_CHECKING from puzzletron_setup import SetupError +if TYPE_CHECKING: + from collections.abc import Sequence + __all__ = ["main"] @@ -28,10 +31,19 @@ def _parser() -> argparse.ArgumentParser: type=Path, help="Explicit versioned defaults YAML; never discovered automatically.", ) + parser.add_argument( + "--full", + action="store_true", + help=( + "Expose every advanced section and nested setting. " + "Without this flag, setup uses a guided profile." + ), + ) return parser -def main(argv: Optional[Sequence[str]] = None) -> int: +def main(argv: Sequence[str] | None = None) -> int: + """Run the setup-v2 command-line interface.""" args = _parser().parse_args(argv) from .wizard import run_wizard_v2 @@ -39,6 +51,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int: campaign = run_wizard_v2( resume=args.resume, defaults_path=args.defaults, + full=args.full, ) except KeyboardInterrupt: target = args.resume or "" diff --git a/puzzletron_setup/v2/defaults.py b/puzzletron_setup/v2/defaults.py index dedd2849183..dddf21211dd 100644 --- a/puzzletron_setup/v2/defaults.py +++ b/puzzletron_setup/v2/defaults.py @@ -30,7 +30,7 @@ from puzzletron_setup import SetupError -__all__ = ["DefaultsResolver", "ResolvedDefault", "load_defaults"] +__all__ = ["DefaultsResolver", "ResolvedDefault", "load_defaults", "validate_defaults"] class _AnyMapping: @@ -125,6 +125,16 @@ def _validate_mapping(value: Any, schema: Any, path: str) -> None: _validate_mapping(item, schema[key], child_path) +def validate_defaults(payload: Mapping[str, Any]) -> dict[str, Any]: + """Validate and isolate one versioned setup-defaults mapping.""" + if payload.get("schema_version") != 1: + raise SetupError( + f"Unsupported defaults schema {payload.get('schema_version')!r}; expected 1." + ) + _validate_mapping(payload, _SCHEMA, "") + return deepcopy(dict(payload)) + + def load_defaults(path: Path | None) -> dict[str, Any]: """Load an explicitly selected versioned defaults file.""" if path is None: @@ -138,12 +148,7 @@ def load_defaults(path: Path | None) -> dict[str, Any]: raise SetupError(f"Cannot read defaults file {resolved}: {error}") from error if not isinstance(payload, Mapping): raise SetupError(f"Defaults file must contain a YAML mapping: {resolved}") - if payload.get("schema_version") != 1: - raise SetupError( - f"Unsupported defaults schema {payload.get('schema_version')!r}; expected 1." - ) - _validate_mapping(payload, _SCHEMA, "") - return deepcopy(dict(payload)) + return validate_defaults(payload) def _lookup(mapping: Mapping[str, Any], dotted_path: str) -> tuple[bool, Any]: @@ -163,13 +168,18 @@ def __init__( *, builtins: Mapping[str, Any] | None = None, model_derived: Mapping[str, Any] | None = None, + preset_defaults: Mapping[str, Any] | None = None, + model_profile_defaults: Mapping[str, Any] | None = None, file_defaults: Mapping[str, Any] | None = None, preserved: Mapping[str, Any] | None = None, ) -> None: - """Build the ordered builtin, model, file, and preserved default layers.""" + """Build the ordered builtin, model, profile, file, and preserved layers.""" + self._resolutions: dict[str, ResolvedDefault] = {} self._default_layers = ( ("builtin", dict(builtins or {})), ("model", dict(model_derived or {})), + ("preset", dict(preset_defaults or {})), + ("model_profile", dict(model_profile_defaults or {})), ("defaults_file", dict(file_defaults or {})), ) self._file_defaults = dict(file_defaults or {}) @@ -193,15 +203,25 @@ def _resolve_layers( def resolve(self, path: str, fallback: Any = None) -> ResolvedDefault: """Return the suggested value, including preserved wizard answers.""" - return self._resolve_layers(self._layers, path, fallback) + resolved = self._resolve_layers(self._layers, path, fallback) + self._resolutions[path] = deepcopy(resolved) + return resolved def resolve_default(self, path: str, fallback: Any = None) -> ResolvedDefault: """Return built-in, model-derived, or explicit-file defaults.""" - return self._resolve_layers(self._default_layers, path, fallback) + resolved = self._resolve_layers(self._default_layers, path, fallback) + self._resolutions[path] = deepcopy(resolved) + return resolved def file_default(self, path: str) -> ResolvedDefault | None: """Return an explicitly supplied file default, if present.""" found, value = _lookup(self._file_defaults, path) if not found: return None - return ResolvedDefault(deepcopy(value), "defaults_file") + resolved = ResolvedDefault(deepcopy(value), "defaults_file") + self._resolutions[path] = deepcopy(resolved) + return resolved + + def resolutions(self) -> Mapping[str, ResolvedDefault]: + """Return every default decision resolved during this wizard run.""" + return deepcopy(self._resolutions) diff --git a/puzzletron_setup/v2/presets.py b/puzzletron_setup/v2/presets.py new file mode 100644 index 00000000000..83eb8362122 --- /dev/null +++ b/puzzletron_setup/v2/presets.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Intent metadata and model-family defaults for guided setup v2.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml + +from puzzletron_setup import SetupError + +from .defaults import validate_defaults + +__all__ = ["QUICK_SETUP_PRESETS", "SetupPreset", "get_setup_preset"] + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +_FAMILY_DEFAULTS_FILENAME = "setup_v2_defaults.yaml" +_SUPPORTED_SCHEMA_VERSIONS = {1, 2} +_MODEL_OVERRIDE_FIELDS = {"match", "profiles"} +_MODEL_MATCH_FIELDS = {"facts", "moe", "num_layers", "num_sublayers"} + + +def _deep_merge(base: Mapping[str, Any], overlay: Mapping[str, Any]) -> dict[str, Any]: + """Return a recursive copy of ``base`` with ``overlay`` applied.""" + merged = deepcopy(dict(base)) + for key, value in overlay.items(): + if isinstance(value, Mapping) and isinstance(merged.get(key), Mapping): + merged[key] = _deep_merge(merged[key], value) + else: + merged[key] = deepcopy(value) + return merged + + +def _matches_inventory(expected: Mapping[str, Any], inventory: Any) -> bool: + """Return whether an inventory contains every value in a declarative selector.""" + for key, expected_value in expected.items(): + if isinstance(inventory, Mapping): + if key not in inventory: + return False + actual_value = inventory[key] + else: + if not hasattr(inventory, key): + return False + actual_value = getattr(inventory, key) + if isinstance(expected_value, Mapping): + if not isinstance(actual_value, Mapping) or not _matches_inventory( + expected_value, actual_value + ): + return False + elif actual_value != expected_value: + return False + return True + + +def _validated_profile(defaults: Any, *, profile: str, path: Path) -> dict[str, Any]: + if not isinstance(defaults, Mapping): + raise SetupError(f"Guided setup profile {profile!r} must be a mapping in {path}.") + try: + validated = validate_defaults({"schema_version": 1, **dict(defaults)}) + except SetupError as error: + raise SetupError(f"Invalid guided setup profile {profile!r} in {path}: {error}") from error + validated.pop("schema_version") + return validated + + +@dataclass(frozen=True) +class SetupPreset: + """One guided setup profile whose tuning is owned by each model family.""" + + name: str + title: str + guidance: str + + @property + def choice_title(self) -> str: + """Render the preset name and its selection guidance.""" + return f"{self.title}: {self.guidance}" + + def resolved_defaults( + self, + family_config: str | Path, + model_inventory: Any | None = None, + ) -> dict[str, Any]: + """Load family defaults and apply a matching model-specific overlay.""" + family_defaults, model_defaults = self.resolved_default_layers( + family_config, + model_inventory, + ) + return _deep_merge(family_defaults, model_defaults) + + def resolved_default_layers( + self, + family_config: str | Path, + model_inventory: Any | None = None, + ) -> tuple[dict[str, Any], dict[str, Any]]: + """Load separate family and matching model layers for provenance.""" + family_path = Path(family_config) + if not family_path.is_absolute(): + family_path = _REPOSITORY_ROOT / family_path + defaults_path = family_path.with_name(_FAMILY_DEFAULTS_FILENAME) + try: + payload = yaml.safe_load(defaults_path.read_text()) or {} + except (OSError, yaml.YAMLError) as error: + raise SetupError( + f"Cannot read guided setup defaults for model family at {defaults_path}: {error}" + ) from error + if not isinstance(payload, Mapping): + raise SetupError(f"Guided setup defaults must contain a YAML mapping: {defaults_path}") + schema_version = payload.get("schema_version") + unknown_fields = set(payload) - {"schema_version", "profiles", "model_overrides"} + if unknown_fields: + fields = ", ".join(sorted(str(field) for field in unknown_fields)) + raise SetupError(f"Unknown guided setup defaults fields in {defaults_path}: {fields}") + if schema_version not in _SUPPORTED_SCHEMA_VERSIONS: + raise SetupError( + f"Unsupported guided setup defaults schema " + f"{schema_version!r} in {defaults_path}; expected 1 or 2." + ) + if schema_version == 1 and "model_overrides" in payload: + raise SetupError(f"Guided setup model overrides require schema 2 in {defaults_path}.") + profiles = payload.get("profiles") + if not isinstance(profiles, Mapping): + raise SetupError(f"Guided setup defaults profiles must be a mapping: {defaults_path}") + defaults = profiles.get(self.name) + if not isinstance(defaults, Mapping): + raise SetupError( + f"Guided setup profile {self.name!r} is not configured in {defaults_path}." + ) + resolved = _validated_profile(defaults, profile=self.name, path=defaults_path) + + model_overrides = payload.get("model_overrides", {}) + if not isinstance(model_overrides, Mapping): + raise SetupError(f"Guided setup model overrides must be a mapping: {defaults_path}") + matches = [] + for model_name, model_override in model_overrides.items(): + if not isinstance(model_override, Mapping): + raise SetupError( + f"Guided setup model override {model_name!r} must be a mapping " + f"in {defaults_path}." + ) + unknown_model_fields = set(model_override) - _MODEL_OVERRIDE_FIELDS + if unknown_model_fields: + fields = ", ".join(sorted(str(field) for field in unknown_model_fields)) + raise SetupError( + f"Unknown fields for guided setup model override {model_name!r} " + f"in {defaults_path}: {fields}" + ) + selector = model_override.get("match") + if not isinstance(selector, Mapping) or not selector: + raise SetupError( + f"Guided setup model override {model_name!r} requires a non-empty " + f"match mapping in {defaults_path}." + ) + unknown_match_fields = set(selector) - _MODEL_MATCH_FIELDS + if unknown_match_fields: + fields = ", ".join(sorted(str(field) for field in unknown_match_fields)) + raise SetupError( + f"Unknown match fields for guided setup model override " + f"{model_name!r} in {defaults_path}: {fields}" + ) + override_profiles = model_override.get("profiles") + if not isinstance(override_profiles, Mapping): + raise SetupError( + f"Guided setup model override {model_name!r} profiles must be a " + f"mapping in {defaults_path}." + ) + known_profiles = {preset.name for preset in QUICK_SETUP_PRESETS} + unknown_profiles = set(override_profiles) - known_profiles + if unknown_profiles: + profiles_list = ", ".join(sorted(str(profile) for profile in unknown_profiles)) + raise SetupError( + f"Unknown profiles for guided setup model override {model_name!r} " + f"in {defaults_path}: {profiles_list}" + ) + validated_overrides = { + profile: _validated_profile(profile_defaults, profile=profile, path=defaults_path) + for profile, profile_defaults in override_profiles.items() + } + if model_inventory is not None and _matches_inventory(selector, model_inventory): + matches.append((str(model_name), validated_overrides)) + + if len(matches) > 1: + names = ", ".join(name for name, _ in matches) + raise SetupError( + f"Model inventory matches multiple guided setup overrides in " + f"{defaults_path}: {names}." + ) + model_defaults = matches[0][1].get(self.name, {}) if matches else {} + return resolved, model_defaults + + +QUICK_SETUP_PRESETS = ( + SetupPreset( + name="smoke", + title="Quick smoke", + guidance="fastest; verifies the campaign shape with minimal scoring", + ), + SetupPreset( + name="balanced", + title="Balanced pruning (recommended)", + guidance="best first real campaign; useful coverage at moderate cost", + ), + SetupPreset( + name="high-confidence", + title="High-confidence search", + guidance="more checks and scoring; choose when extra runtime is acceptable", + ), +) + + +def get_setup_preset(name: str) -> SetupPreset: + """Return a known setup preset or fail with an actionable error.""" + for preset in QUICK_SETUP_PRESETS: + if preset.name == name: + return preset + choices = ", ".join(preset.name for preset in QUICK_SETUP_PRESETS) + raise SetupError(f"Unknown setup preset {name!r}; choose one of: {choices}.") diff --git a/puzzletron_setup/v2/prompts.py b/puzzletron_setup/v2/prompts.py index c5ef38102ff..6abde2c325b 100644 --- a/puzzletron_setup/v2/prompts.py +++ b/puzzletron_setup/v2/prompts.py @@ -6,12 +6,14 @@ from __future__ import annotations from collections import deque -from collections.abc import Sequence from dataclasses import dataclass -from typing import Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol from puzzletron_setup import SetupError +if TYPE_CHECKING: + from collections.abc import Sequence + __all__ = [ "BACK", "InteractiveBackend", @@ -42,6 +44,7 @@ class PromptBackend(Protocol): """Minimal backend used by the navigable wizard session.""" def text(self, message: str, default: str) -> Any: + """Request a text value.""" raise NotImplementedError def select( @@ -50,6 +53,7 @@ def select( choices: Sequence[PromptChoice], default: Any, ) -> Any: + """Request one value from a list of choices.""" raise NotImplementedError def checkbox( @@ -58,6 +62,7 @@ def checkbox( choices: Sequence[PromptChoice], defaults: Sequence[Any], ) -> Any: + """Request multiple values from a list of choices.""" raise NotImplementedError @@ -87,14 +92,36 @@ def _choice_style(questionary: Any) -> Any: ) +def _bind_escape_back(question: Any) -> Any: + """Make Escape return the same sentinel for every interactive widget.""" + key_bindings = question.application.key_bindings + if not hasattr(key_bindings, "add"): + from prompt_toolkit.key_binding import KeyBindings, merge_key_bindings + + escape_bindings = KeyBindings() + question.application.key_bindings = merge_key_bindings([key_bindings, escape_bindings]) + key_bindings = escape_bindings + + @key_bindings.add("escape", eager=True) + def go_back(event): + event.app.exit(result=BACK) + + return question + + class InteractiveBackend: """Questionary-backed prompts with visible Back controls.""" _BACK_TITLE = "← Back" def text(self, message: str, default: str) -> Any: - print(" Type :back to return to the previous question.") - value = str(_answer(_questionary().text(message, default=default))) + """Request text while supporting semantic Back navigation.""" + print(" Press Esc to go back (or type :back).") + question = _bind_escape_back(_questionary().text(message, default=default)) + value = _answer(question) + if value is BACK: + return BACK + value = str(value) return BACK if value.strip().lower() == ":back" else value def select( @@ -103,6 +130,7 @@ def select( choices: Sequence[PromptChoice], default: Any, ) -> Any: + """Request one choice while exposing semantic Back navigation.""" questionary = _questionary() rendered = [ questionary.Choice( @@ -114,11 +142,13 @@ def select( ] rendered.append(questionary.Choice(title=self._BACK_TITLE, value=BACK)) return _answer( - questionary.select( - message, - choices=rendered, - default=default, - style=_choice_style(questionary), + _bind_escape_back( + questionary.select( + message, + choices=rendered, + default=default, + style=_choice_style(questionary), + ) ) ) @@ -128,6 +158,7 @@ def checkbox( choices: Sequence[PromptChoice], defaults: Sequence[Any], ) -> Any: + """Request multiple choices while supporting semantic Back navigation.""" questionary = _questionary() selected = set(defaults) rendered = [ @@ -140,20 +171,18 @@ def checkbox( for choice in choices ] rendered.append(questionary.Separator(f" {self._BACK_TITLE} (press Esc)")) - question = questionary.checkbox( - message, - choices=rendered, - instruction=( - "(Use arrow keys to move, to select, to toggle, " - " to invert, to go back)" - ), - style=_choice_style(questionary), + question = _bind_escape_back( + questionary.checkbox( + message, + choices=rendered, + instruction=( + "(Use arrow keys to move, to select, to toggle, " + " to invert, to go back)" + ), + style=_choice_style(questionary), + ) ) - @question.application.key_bindings.add("escape", eager=True) - def go_back(event): - event.app.exit(result=BACK) - values = _answer(question) if values is BACK: return BACK @@ -164,10 +193,12 @@ class ScriptedBackend: """Deterministic non-interactive backend for embedding and automation.""" def __init__(self, answers: Sequence[Any]) -> None: + """Initialize the backend with deterministic answers.""" self._answers = deque(answers) @property def remaining(self) -> int: + """Return the number of unconsumed answers.""" return len(self._answers) def _next(self) -> Any: @@ -177,6 +208,7 @@ def _next(self) -> Any: return BACK if value == ":back" else value def text(self, message: str, default: str) -> Any: + """Return the next scripted text answer.""" del message, default return self._next() @@ -186,6 +218,7 @@ def select( choices: Sequence[PromptChoice], default: Any, ) -> Any: + """Return the next scripted single-choice answer.""" del message, choices, default return self._next() @@ -195,5 +228,6 @@ def checkbox( choices: Sequence[PromptChoice], defaults: Sequence[Any], ) -> Any: + """Return the next scripted multiple-choice answer.""" del message, choices, defaults return self._next() diff --git a/puzzletron_setup/v2/session.py b/puzzletron_setup/v2/session.py index 1ec540c7640..5802df08a56 100644 --- a/puzzletron_setup/v2/session.py +++ b/puzzletron_setup/v2/session.py @@ -5,14 +5,17 @@ from __future__ import annotations -from collections.abc import Callable, Sequence -from typing import Any +from typing import TYPE_CHECKING, Any from .prompts import BACK, InteractiveBackend, PromptBackend, PromptChoice from .state import PromptFrame, WizardState +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + __all__ = ["BACK", "WizardSession"] + class WizardSession: """Bind prompt interactions to atomic answer and navigation state.""" @@ -20,9 +23,13 @@ def __init__( self, state: WizardState, backend: PromptBackend | None = None, + *, + guided: bool = False, ) -> None: + """Bind wizard state to an interactive or scripted prompt backend.""" self.state = state self.backend = backend or InteractiveBackend() + self.guided = bool(guided) self._section = "campaign" self._collection: str | None = None self._item_id: str | None = None @@ -33,10 +40,12 @@ def __init__( @property def current_frame(self) -> PromptFrame | None: + """Return the most recent prompt frame, if one exists.""" frames = self.state.frames return frames[-1] if frames else None def begin(self, section: str) -> None: + """Begin or resume prompt replay for a wizard section.""" self._section = section self._collection = None self._item_id = None @@ -53,22 +62,26 @@ def enter_collection( item_id: str | None, cursor: int | None, ) -> None: + """Enter one item in a repeatable prompt collection.""" self._collection = collection self._item_id = item_id self._cursor = cursor def leave_collection(self) -> None: + """Leave the active repeatable prompt collection.""" self._collection = None self._item_id = None self._cursor = None def collection_cursor(self, collection: str) -> int | None: + """Return the saved cursor for an active collection.""" frame = self.current_frame if frame is not None and frame.collection == collection: return frame.cursor return None def back(self) -> PromptFrame | None: + """Remove and return the current prompt frame.""" return self.state.pop_frame() def consume_back_target(self) -> PromptFrame | None: @@ -109,6 +122,7 @@ def _ask(self, prompt_id: str, invoke: Callable[[], Any]) -> Any: @staticmethod def describe_default(value: Any, source: str) -> None: + """Print a resolved default and its provenance.""" print(f" Default: {value!r} ({source})") @staticmethod @@ -131,6 +145,7 @@ def text( default: str = "", validate: Callable[[Any], bool | str] | None = None, ) -> Any: + """Request, validate, and persist a text answer.""" while True: value = self._ask(prompt_id, lambda: self.backend.text(message, default)) if value is BACK: @@ -151,6 +166,8 @@ def integer( minimum: int = 0, maximum: int | None = None, ) -> Any: + """Request a bounded integer answer.""" + def validate(value: str) -> bool | str: try: parsed = int(value) @@ -178,6 +195,7 @@ def select( *, default: Any = None, ) -> Any: + """Request one answer from a set of choices.""" rendered = self._choices(choices) if len(rendered) == 1: print(f" {message} {rendered[0].title} (only option)") @@ -194,6 +212,7 @@ def confirm( *, default: bool, ) -> Any: + """Request a yes-or-no answer.""" return self.select( prompt_id, message, @@ -210,6 +229,7 @@ def checkbox( defaults: Sequence[Any] = (), validate: Callable[[Any], bool | str] | None = None, ) -> Any: + """Request and validate multiple selected choices.""" rendered_choices = self._choices(choices) def disabled_verdict(values: Sequence[Any]) -> bool | str: @@ -229,8 +249,7 @@ def disabled_verdict(values: Sequence[Any]) -> bool | str: verdict = validate(selected) if verdict is True: print( - f" {message} {rendered_choices[0].title} " - "(only option, selected automatically)" + f" {message} {rendered_choices[0].title} (only option, selected automatically)" ) return selected while True: diff --git a/puzzletron_setup/v2/state.py b/puzzletron_setup/v2/state.py index 5076775f53f..b8417992480 100644 --- a/puzzletron_setup/v2/state.py +++ b/puzzletron_setup/v2/state.py @@ -54,6 +54,7 @@ def __post_init__(self) -> None: @classmethod def from_dict(cls, payload: Mapping[str, Any]) -> FieldRecord: + """Restore a field record from serialized state.""" return cls( value=payload.get("value"), source=str(payload.get("source", "user")), @@ -65,6 +66,7 @@ def from_dict(cls, payload: Mapping[str, Any]) -> FieldRecord: ) def to_dict(self) -> dict[str, Any]: + """Serialize the field record to plain Python values.""" return _plain(asdict(self)) @@ -80,6 +82,7 @@ class PromptFrame: @classmethod def from_dict(cls, payload: Mapping[str, Any]) -> PromptFrame: + """Restore a prompt frame from serialized state.""" return cls( section=str(payload["section"]), prompt_id=str(payload["prompt_id"]), @@ -103,7 +106,10 @@ def start( campaign_dir: Path, *, defaults_path: Path | None, + setup_mode: str = "full", + preset: str | None = None, ) -> WizardState: + """Create and persist a new setup-v2 campaign state.""" campaign_dir = Path(campaign_dir).expanduser().resolve() if campaign_dir.exists() and any(campaign_dir.iterdir()): raise SetupError( @@ -111,7 +117,7 @@ def start( "Choose a new directory or use --resume." ) campaign_dir.mkdir(parents=True, exist_ok=True) - payload = { + payload: dict[str, Any] = { "schema_version": SCHEMA_VERSION, "wizard_version": WIZARD_VERSION, "defaults_path": ( @@ -119,6 +125,10 @@ def start( if defaults_path is not None else None ), + "setup": { + "mode": str(setup_mode), + "preset": str(preset) if preset is not None else None, + }, "fields": {}, "navigation": {"frames": [], "cursor": None}, "collections": {}, @@ -132,6 +142,7 @@ def start( @classmethod def resume(cls, path: Path) -> WizardState: + """Load a compatible setup-v2 campaign state.""" candidate = Path(path).expanduser().resolve() state_path = candidate / "answers_v2.yaml" if candidate.is_dir() else candidate if not state_path.is_file(): @@ -158,24 +169,61 @@ def resume(cls, path: Path) -> WizardState: @property def campaign_dir(self) -> Path: + """Return the campaign directory containing this state.""" return self.path.parent @property def defaults_path(self) -> Path | None: + """Return the persisted defaults-file path, if configured.""" value = self.payload.get("defaults_path") return Path(str(value)) if value else None + @property + def setup_mode(self) -> str: + """Return the persisted guided or full interaction mode.""" + setup = self.payload.get("setup") + if not isinstance(setup, Mapping): + return "full" + return str(setup.get("mode", "full")) + + @property + def preset(self) -> str | None: + """Return the persisted guided preset name, if any.""" + setup = self.payload.get("setup") + if not isinstance(setup, Mapping): + return None + value = setup.get("preset") + return str(value) if value else None + + def set_setup_mode(self, mode: str) -> None: + """Persist an explicit guided or full mode transition.""" + self.payload.setdefault("setup", {})["mode"] = str(mode) + self.save() + + def set_preset(self, preset: str) -> None: + """Persist a replacement guided profile selection.""" + self.payload.setdefault("setup", {})["preset"] = str(preset) + self.save() + + def set_defaults_path(self, path: Path) -> None: + """Persist an explicitly accepted replacement defaults file.""" + self.payload["defaults_path"] = str(Path(path).expanduser().resolve()) + self.save() + def field(self, path: str) -> FieldRecord: + """Return one required authored field record.""" try: return self._fields[path] except KeyError as error: raise KeyError(f"Unknown setup field: {path}") from error def get_field(self, path: str, default: Any = None) -> Any: + """Return one effective field value or a fallback.""" record = self._fields.get(path) return default if record is None else record.effective def records(self) -> Mapping[str, FieldRecord]: + """Return a copy of all authored field records.""" return dict(self._fields) def set_field( @@ -188,6 +236,7 @@ def set_field( requested: Any = None, effective: Any = None, ) -> FieldRecord: + """Persist a field value and invalidate downstream dependents.""" previous = self._fields.get(path) resolved_effective = value if effective is None else effective changed = previous is None or previous.effective != resolved_effective @@ -207,6 +256,7 @@ def set_field( return record def mark_dependents_stale(self, changed_path: str, *, save: bool = True) -> tuple[str, ...]: + """Mark transitive dependents of a changed field as stale.""" reverse: dict[str, set[str]] = {} for field_path, record in self._fields.items(): for dependency in record.dependencies: @@ -233,6 +283,7 @@ def revalidate( self, validators: Mapping[str, Callable[[Any, WizardState], str | None]], ) -> Mapping[str, str]: + """Revalidate stale fields and return unresolved issues.""" issues: dict[str, str] = {} for path, record in self._fields.items(): if not record.stale: @@ -251,14 +302,17 @@ def revalidate( return issues def collection(self, path: str) -> Any: + """Return a named collection from persisted state.""" return self.payload.setdefault("collections", {}).get(path) def set_collection(self, path: str, value: Any) -> None: + """Persist a named collection.""" self.payload.setdefault("collections", {})[path] = _plain(value) self.save() @property def frames(self) -> tuple[PromptFrame, ...]: + """Return the persisted prompt-navigation stack.""" return tuple( PromptFrame.from_dict(item) for item in self.payload.setdefault("navigation", {}).get("frames", ()) @@ -276,6 +330,7 @@ def answered_frames(self) -> tuple[tuple[PromptFrame, Any, int], ...]: ) def push_frame(self, frame: PromptFrame) -> None: + """Push a prompt frame onto the navigation stack.""" frames = self.payload.setdefault("navigation", {}).setdefault("frames", []) if not frames or frames[-1] != asdict(frame): frames.append(_plain(asdict(frame))) @@ -291,6 +346,7 @@ def answer_frame(self, frame: PromptFrame, value: Any) -> None: self.save() def pop_frame(self) -> PromptFrame | None: + """Pop the active frame and return the new active frame.""" navigation = self.payload.setdefault("navigation", {}) frames = navigation.setdefault("frames", []) if frames: @@ -308,6 +364,7 @@ def truncate_frames(self, count: int) -> None: self.save() def replace_frames(self, frames: Sequence[PromptFrame]) -> None: + """Replace the complete prompt-navigation stack.""" rendered = [_plain(asdict(frame)) for frame in frames] self.payload["navigation"] = { "frames": rendered, @@ -316,11 +373,13 @@ def replace_frames(self, frames: Sequence[PromptFrame]) -> None: self.save() def set_model(self, model: Mapping[str, Any], inventory: Mapping[str, Any]) -> None: + """Persist inspected model metadata and inventory.""" self.payload["model"] = _plain(model) self.payload["inventory"] = _plain(inventory) self.save() def save(self) -> None: + """Atomically write the current state to disk.""" self.path.parent.mkdir(parents=True, exist_ok=True) self.payload["fields"] = { path: record.to_dict() for path, record in sorted(self._fields.items()) diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index c26372ffa3b..f292342051f 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -53,6 +53,7 @@ ) from .parallel_validation import validate_automodel_parallelism, validate_vllm_parallelism from .post_mip import FlowDraft, NodeDraft, PostMIPFlowEditor, recommended_flow +from .presets import QUICK_SETUP_PRESETS, SetupPreset, get_setup_preset from .prompts import BACK, InteractiveBackend, PromptBackend, PromptChoice from .resources import ( ParallelProfile, @@ -67,7 +68,7 @@ __all__ = ["SECTION_BUILDERS", "run_wizard_v2"] BUILTINS = { - "data": {"modality": "text", "layout": "fixed", "sequence_length": 4096}, + "data": {"layout": "fixed", "sequence_length": 4096}, "infrastructure": { "gpus_per_node": 8, "execution_contract": { @@ -244,15 +245,28 @@ def _data_source_choices( explicit_source = normalize_dataset_source(str(explicit.value)) if multimodal_model or explicit_source != _NEMOTRON_VLM_DATA_SOURCE: choices.append(PromptChoice(f"Default — {explicit.value}", _DEFAULT_DATA_SOURCE)) - first_class = [("NVIDIA Puzzle-KD v2 (text)", _PUZZLE_KD_DATA_SOURCE)] + first_class = [ + ( + "NVIDIA Puzzle-KD v2: recommended text pruning dataset", + _PUZZLE_KD_DATA_SOURCE, + ) + ] if multimodal_model: first_class.append( - ("NVIDIA Nemotron-VLM v2 (image-text)", _NEMOTRON_VLM_DATA_SOURCE) + ( + "NVIDIA Nemotron-VLM v2: recommended image-text dataset", + _NEMOTRON_VLM_DATA_SOURCE, + ) ) choices.extend( PromptChoice(title, source) for title, source in first_class if source != explicit_source ) - choices.append(PromptChoice("Custom local path or Hugging Face dataset", _CUSTOM_DATA_SOURCE)) + choices.append( + PromptChoice( + "Custom dataset: choose a local path or Hugging Face dataset", + _CUSTOM_DATA_SOURCE, + ) + ) return choices @@ -314,10 +328,10 @@ def _select_hf_subsets( f"Hugging Face dataset {source} has no selectable subsets with " "known positive row counts and sizes." ) - configured_defaults = resolver.resolve( - "data.subsets", - resolver.resolve("data.acquisition.subsets", None).value, - ).value + configured = resolver.resolve("data.subsets", None) + if configured.value is None: + configured = resolver.resolve("data.acquisition.subsets", None) + configured_defaults = configured.value if isinstance(configured_defaults, str): configured_defaults = [ item.strip() for item in configured_defaults.split(",") if item.strip() @@ -331,6 +345,13 @@ def _select_hf_subsets( else: preferred = [selectable[0].name] selectable_names = {item.name for item in selectable} + unavailable = [name for name in preferred if name not in selectable_names] + if configured_defaults and unavailable: + raise SetupError( + "Configured dataset subsets are missing or unavailable for " + f"{source}: {', '.join(unavailable)}. Choose from: " + f"{', '.join(sorted(selectable_names))}." + ) defaults = [name for name in preferred if name in selectable_names] if not defaults: defaults = [selectable[0].name] @@ -342,23 +363,33 @@ def validate(selected_names: list[str]) -> bool | str: return str(error) return True - selected = session.checkbox( - "data.subsets", - "Dataset subsets:", - [ - PromptChoice( - format_subset_choice(item), - item.name, - disabled=item.disabled_reason, - ) - for item in catalog.subsets - ], - defaults=defaults, - validate=validate, - ) - if selected is BACK: - return BACK - selected_names = [str(name) for name in selected] + if session.guided: + selected_names = list(defaults) + verdict = validate(selected_names) + if verdict is not True: + raise SetupError(str(verdict)) + print( + " Dataset subsets: " + f"{', '.join(selected_names)} (recommended defaults; use --full to customize)" + ) + else: + selected = session.checkbox( + "data.subsets", + "Dataset subsets:", + [ + PromptChoice( + format_subset_choice(item), + item.name, + disabled=item.disabled_reason, + ) + for item in catalog.subsets + ], + defaults=defaults, + validate=validate, + ) + if selected is BACK: + return BACK + selected_names = [str(name) for name in selected] weights = proportional_subset_weights(catalog, selected_names) return catalog, selected_names, weights @@ -374,10 +405,25 @@ def _nested_records(state: WizardState) -> dict[str, Any]: return nested -def _resolver(state: WizardState, defaults_path: Path | None) -> DefaultsResolver: +def _resolver( + state: WizardState, + defaults_path: Path | None, + preset: SetupPreset | None = None, + family_config: str | Path | None = None, + model_inventory: Any | None = None, +) -> DefaultsResolver: + preset_defaults = {} + model_profile_defaults = {} + if preset is not None and family_config is not None: + preset_defaults, model_profile_defaults = preset.resolved_default_layers( + family_config, + model_inventory, + ) return DefaultsResolver( builtins=BUILTINS, model_derived={}, + preset_defaults=preset_defaults, + model_profile_defaults=model_profile_defaults, file_defaults=load_defaults(defaults_path), preserved=_nested_records(state), ) @@ -417,8 +463,12 @@ def _section_action( section: str, summary: str, defaults: Mapping[str, Any], + *, + prompt_in_guided: bool = False, ) -> Any: session.begin(section) + if session.guided and not prompt_in_guided: + return "defaults" print(f"\n[{section}] {summary}") _print_default_decisions(defaults) return session.select( @@ -445,9 +495,7 @@ def _print_default_decisions(defaults: Mapping[str, Any]) -> None: def _plain_review_value(value: Any) -> Any: if isinstance(value, Mapping): - return { - str(key): _plain_review_value(item) for key, item in value.items() - } + return {str(key): _plain_review_value(item) for key, item in value.items()} if isinstance(value, (list, tuple)): return [_plain_review_value(item) for item in value] return value @@ -488,10 +536,7 @@ def _replacement_granularity_choices( def label(name: str, per_width: int, total: int) -> str: if width_count == 1: return f"{name} — {total} options" - return ( - f"{name} — {per_width} options/width, " - f"{total} total across {width_count} widths" - ) + return f"{name} — {per_width} options/width, {total} total across {width_count} widths" return [ ( @@ -643,9 +688,7 @@ def data_section( catalog_loader: Callable[..., HfSubsetCatalog] = discover_hf_subset_catalog, ) -> bool: session.begin("data") - previous_acquisition = _mapping_copy( - session.state.collection("data_acquisition") - ) + previous_acquisition = _mapping_copy(session.state.collection("data_acquisition")) explicit = resolver.file_default("data.source") if explicit is not None and explicit.value: _print_default_decisions({"source": explicit.value}) @@ -653,16 +696,12 @@ def data_section( resolver, multimodal_model=bool(context["model"].inventory.multimodal), ) - default_choice_available = any( - choice.value == _DEFAULT_DATA_SOURCE for choice in choices - ) + default_choice_available = any(choice.value == _DEFAULT_DATA_SOURCE for choice in choices) mode = session.select( "data.source_mode", "Dataset:", choices, - default=( - _DEFAULT_DATA_SOURCE if default_choice_available else _CUSTOM_DATA_SOURCE - ), + default=(_DEFAULT_DATA_SOURCE if default_choice_available else _CUSTOM_DATA_SOURCE), ) if mode is BACK: return False @@ -700,21 +739,44 @@ def data_section( modality_choices = [("Text", "text")] if context["model"].inventory.multimodal: modality_choices.append(("Multimodal", "multimodal")) - suggested_modality = str( - resolver.resolve( - "data.modality", - finding_modality if finding_modality != "unknown" else "text", - ).value + modality_default = resolver.resolve( + "data.modality", + finding_modality if finding_modality != "unknown" else "text", + ) + suggested_modality = str(modality_default.value) + modality_source = ( + "inferred" if modality_default.source == "fallback" else modality_default.source ) - valid_modalities = {value for _, value in modality_choices} - if suggested_modality not in valid_modalities: - suggested_modality = "text" if fixed_modality == "multimodal" and not context["model"].inventory.multimodal: raise SetupError( "NVIDIA Nemotron-VLM v2 requires a multimodal model. " "Choose a multimodal checkpoint or a text dataset." ) - if fixed_modality is None: + valid_modalities = {value for _, value in modality_choices} + if suggested_modality not in valid_modalities: + if modality_source != "inferred" or session.guided: + raise SetupError( + f"Resolved data modality {suggested_modality!r} ({modality_source}) " + "is incompatible with the selected model." + ) + suggested_modality = "text" + modality_source = "inferred" + if ( + fixed_modality is not None + and modality_source != "inferred" + and suggested_modality != fixed_modality + ): + raise SetupError( + f"Configured data modality {suggested_modality!r} ({modality_source}) " + f"conflicts with {source}, which requires {fixed_modality!r}." + ) + if fixed_modality is None and session.guided: + modality = suggested_modality + print( + f" Data modality: {modality} ({modality_source}; {finding_evidence}; " + "use --full to override)" + ) + elif fixed_modality is None: modality = session.select( "data.modality", f"Data modality ({finding_evidence}):", @@ -723,6 +785,7 @@ def data_section( ) if modality is BACK: return False + modality_source = "user" else: modality = fixed_modality print(f" Data modality: {modality} ({finding_evidence})") @@ -737,27 +800,45 @@ def data_section( / f"{session.state.campaign_dir.name}_datasets" / adapter ) - output = _text_field( - session, - resolver, - "data.acquisition.output", - "Local materialization directory:", - str(default_root), - ) - if output is BACK: - return False + if session.guided: + output = _record_default( + session.state, + resolver, + "data.acquisition.output", + str(default_root), + ) + print(f" Local materialization directory: {output}") + else: + output = _text_field( + session, + resolver, + "data.acquisition.output", + "Local materialization directory:", + str(default_root), + ) + if output is BACK: + return False runtime_source = str(Path(str(output)).expanduser().resolve()) seed_default = 408 if adapter == _PUZZLE_KD_ADAPTER else 42 - seed = _integer_field( - session, - resolver, - "data.acquisition.seed", - "Deterministic dataset selection seed:", - seed_default, - minimum=0, - ) - if seed is BACK: - return False + if session.guided: + seed = _record_default( + session.state, + resolver, + "data.acquisition.seed", + seed_default, + ) + print(f" Deterministic dataset selection seed: {seed}") + else: + seed = _integer_field( + session, + resolver, + "data.acquisition.seed", + "Deterministic dataset selection seed:", + seed_default, + minimum=0, + ) + if seed is BACK: + return False acquisition = { "adapter": adapter, "source": source, @@ -784,9 +865,7 @@ def data_section( { "name": name, "num_rows": by_name[name].num_rows, - "num_bytes_original_files": by_name[ - name - ].num_bytes_original_files, + "num_bytes_original_files": by_name[name].num_bytes_original_files, "num_media_shards": by_name[name].num_media_shards, "weight": weights[name], } @@ -813,20 +892,16 @@ def data_section( acquisition.update( subsets=selected_subsets, subset_rows={ - record["name"]: record["num_rows"] - for record in subset_selection["subsets"] + record["name"]: record["num_rows"] for record in subset_selection["subsets"] }, subset_weights={ - record["name"]: record["weight"] - for record in subset_selection["subsets"] + record["name"]: record["weight"] for record in subset_selection["subsets"] }, subset_media_shards=subset_media_shards, revision=subset_selection["revision"], ) if len(subset_media_shards) != len(selected_subsets): - previous_shard_cap = int( - previous_acquisition.get("max_shards_per_subset", 0) - ) + previous_shard_cap = int(previous_acquisition.get("max_shards_per_subset", 0)) if previous_shard_cap > 0: acquisition["max_shards_per_subset"] = previous_shard_cap print( @@ -836,38 +911,54 @@ def data_section( default_layout = str(resolver.resolve("data.layout", "fixed").value) if default_layout == "padded": default_layout = "padded_varlen" - layout = session.select( - "data.layout", - "Dataset layout:", - [ - ("Fixed-length", "fixed"), - ("Packed variable-length", "packed_varlen"), - ("Padded variable-length", "padded_varlen"), - ], - default=default_layout, - ) - if layout is BACK: - return False - sequence = _integer_field( - session, - resolver, - "data.sequence_length", - "Sequence length used by width, depth, bypass, evaluation, and global KD:", - 4096, - ) - if sequence is BACK: - return False + if session.guided: + layout_default = resolver.resolve_default("data.layout", default_layout) + layout = str(layout_default.value) + if layout == "padded": + layout = "padded_varlen" + sequence_default = resolver.resolve_default("data.sequence_length", 4096) + sequence = int(sequence_default.value) + print( + f" Data shape: {layout}, sequence length {sequence} " + "(resolved defaults; use --full to customize)" + ) + else: + layout = session.select( + "data.layout", + "Dataset layout:", + [ + ("Fixed-length", "fixed"), + ("Packed variable-length", "packed_varlen"), + ("Padded variable-length", "padded_varlen"), + ], + default=default_layout, + ) + if layout is BACK: + return False + sequence = _integer_field( + session, + resolver, + "data.sequence_length", + "Sequence length used by width, depth, bypass, evaluation, and global KD:", + 4096, + ) + if sequence is BACK: + return False session.state.set_field("data.source", runtime_source, source=source_kind) session.state.set_field("data.selected_source", source, source=source_kind) session.state.set_field("data.adapter", adapter or "custom", source=source_kind) session.state.set_collection("data_acquisition", acquisition or {}) session.state.set_collection("data_subset_selection", subset_selection or {}) - session.state.set_field("data.modality", modality, source="user") - session.state.set_field("data.layout", layout, source="user") + session.state.set_field("data.modality", modality, source=modality_source) + session.state.set_field( + "data.layout", + layout, + source=layout_default.source if session.guided else "user", + ) session.state.set_field( "data.sequence_length", int(sequence), - source="user", + source=sequence_default.source if session.guided else "user", ) return True @@ -906,12 +997,41 @@ def infrastructure_section( "infrastructure", "Configure the worker contract and cluster facts before stage allocations.", preview, + prompt_in_guided=True, ) if action is BACK: return False if action == "defaults": for path, fallback in paths: _record_default(session.state, resolver, path, fallback) + worker_paths = ( + ( + "infrastructure.execution_contract.repository", + "Repository path on workers:", + WORKER_REPOSITORY_PLACEHOLDER, + ), + ( + "infrastructure.execution_contract.venv", + "Python environment:", + WORKER_VENV_PLACEHOLDER, + ), + ) + if any( + validate_worker_path(str(session.state.get_field(path))) is not True + for path, _, _ in worker_paths + ): + print(" Enter the worker-visible repository and Python environment paths.") + for path, label, fallback in worker_paths: + value = _text_field( + session, + resolver, + path, + label, + fallback, + validate=validate_worker_path, + ) + if value is BACK: + return False commands = resolver.resolve_default("infrastructure.execution_contract.prerun_commands", []) session.state.set_field( "infrastructure.execution_contract.prerun_commands", @@ -1220,6 +1340,47 @@ def _print_parallel_issues(issues) -> None: print(" Choose a different parallel setting.") +def _compatible_default_profile( + session: WizardSession, + registry: ResourceProfileRegistry, + stage_id: str, + model: Any, + *, + node_type: str | None = None, +) -> tuple[ParallelProfile | None, tuple[Any, ...]]: + """Reuse the first compatible profile without opening advanced prompts.""" + last_issues: tuple[Any, ...] = () + for name in registry.names(): + profile = registry.get(name) + issues = tuple( + _profile_compatibility_issues( + session, + profile, + stage_id, + model, + node_type=node_type, + ) + ) + if not issues: + return registry.reuse(name, consumer=stage_id), () + last_issues = issues + if not registry.names(): + profile = ParallelProfile(stage_id) + issues = tuple( + _profile_compatibility_issues( + session, + profile, + stage_id, + model, + node_type=node_type, + ) + ) + if not issues: + return registry.create(profile, consumer=stage_id), () + last_issues = issues + return None, last_issues + + def _profile_prompt( session: WizardSession, registry: ResourceProfileRegistry, @@ -1369,20 +1530,17 @@ def _stage_resource_defaults( profile = registry.get(registry.names()[0]) if registry.names() else ParallelProfile(stage_id) strategy = strategy or CANONICAL_STAGE_STRATEGIES[stage_id] gpus_per_node = int(session.state.get_field("infrastructure.gpus_per_node", 8)) - instances = ( - 1 - if strategy == "single" - else int( - resolver.resolve_default( - f"stages.{stage_id}.instances", - gpus_per_node, - ).value - ) + resolved_instances = resolver.resolve_default( + f"stages.{stage_id}.instances", + gpus_per_node, ) - requested_batch = int(resolver.resolve_default(f"stages.{stage_id}.batch", batch).value) + instances = 1 if strategy == "single" else int(resolved_instances.value) + resolved_batch = resolver.resolve_default(f"stages.{stage_id}.batch", batch) + requested_batch = int(resolved_batch.value) resolution = resolve_batch(requested_batch, profile) return { "instances": instances, + "instances_source": resolved_instances.source, "parallel_profile": profile.name, "parallel": { "tp": profile.tp, @@ -1395,6 +1553,7 @@ def _stage_resource_defaults( }, "requested_batch": resolution.requested, "effective_batch": resolution.effective, + "batch_source": resolved_batch.source, } @@ -1451,14 +1610,19 @@ def _configure_stage_resource( if requested_batch is BACK: return BACK else: - profile = ( - registry.reuse(registry.names()[0], consumer=stage_id) - if registry.names() - else registry.create(ParallelProfile(stage_id), consumer=stage_id) + profile, issues = _compatible_default_profile( + session, + registry, + stage_id, + model, ) - issues = _profile_compatibility_issues(session, profile, stage_id, model) - if issues: + if profile is None: _print_parallel_issues(issues) + if session.guided: + raise SetupError( + f"No configured parallel profile is compatible with {stage_id}. " + "Supply a compatible profile in --defaults or resume with --full." + ) profile = _profile_prompt(session, registry, stage_id, model) if profile is BACK: return BACK @@ -1502,7 +1666,7 @@ def _configure_stage_resource( session.state.set_field( f"stages.{stage_id}.batch", resolution.effective, - source="user" if action == "customize" else "builtin", + source="user" if action == "customize" else str(defaults["batch_source"]), requested=resolution.requested, effective=resolution.effective, dependencies=(f"profiles.{profile.name}",), @@ -2490,9 +2654,7 @@ def serving_workloads_section( result = _serving_workload_prompt( session, "serving_workloads", - default_name=( - "serving-default" if not workloads else f"serving-{len(workloads) + 1}" - ), + default_name=("serving-default" if not workloads else f"serving-{len(workloads) + 1}"), default_workload=default_workload, existing_names=set(workloads), ) @@ -2521,9 +2683,7 @@ def _set_vllm_stage_resource( source: str, ) -> None: allocation_mesh = vllm_topology_to_mesh(topology) - gpus_per_node = int( - session.state.get_field("infrastructure.gpus_per_node", 8) - ) + gpus_per_node = int(session.state.get_field("infrastructure.gpus_per_node", 8)) resources = _mapping_copy(session.state.collection("stage_resources")) resources["vllm_stats"] = { "strategy": "sharded", @@ -2545,9 +2705,7 @@ def _set_vllm_stage_resource( def vllm_section(session: WizardSession, resolver: DefaultsResolver, context: dict) -> bool: enabled_default = bool(resolver.resolve_default("vllm.enabled", False).value) - gpus_per_node = int( - session.state.get_field("infrastructure.gpus_per_node", 8) - ) + gpus_per_node = int(session.state.get_field("infrastructure.gpus_per_node", 8)) resolved_instances = resolver.resolve_default( "stages.vllm_stats.instances", gpus_per_node, @@ -2558,9 +2716,7 @@ def vllm_section(session: WizardSession, resolver: DefaultsResolver, context: di sequence_length=int(session.state.get_field("data.sequence_length", 4096)), ) default_topology = default_measurement_settings["runtime_stats"]["topology"] - workloads = OrderedDict( - _mapping_copy(session.state.collection("serving_workloads")).items() - ) + workloads = OrderedDict(_mapping_copy(session.state.collection("serving_workloads")).items()) if not workloads: workloads["serving-default"] = _default_serving_workload( resolver, @@ -2577,8 +2733,7 @@ def vllm_section(session: WizardSession, resolver: DefaultsResolver, context: di preview = { "enabled": enabled_default, "available_serving_workloads": [ - _serving_workload_label(name, setting) - for name, setting in workloads.items() + _serving_workload_label(name, setting) for name, setting in workloads.items() ], } if enabled_default: @@ -2666,9 +2821,7 @@ def vllm_section(session: WizardSession, resolver: DefaultsResolver, context: di instances: int | None = None while True: unused_workloads = OrderedDict( - (name, setting) - for name, setting in workloads.items() - if name not in measurements + (name, setting) for name, setting in workloads.items() if name not in measurements ) if unused_workloads: workload_choice = session.select( @@ -2794,8 +2947,7 @@ def vllm_section(session: WizardSession, resolver: DefaultsResolver, context: di def _mip_default_search_id(metric: str, value: Any) -> str: suffix = str(value).replace("%", "") slug = "".join( - character if character.isalnum() or character in "_-" else "-" - for character in suffix + character if character.isalnum() or character in "_-" else "-" for character in suffix ).strip("-") return f"{metric}-{slug or 'target'}" @@ -2968,9 +3120,7 @@ def _mip_constraints_prompt( default = _mip_constraint_default(metric, configured_goal) if metric not in {"memory", "runtime"}: prompt = ( - _mip_maximum_prompt - if metric in {"params", "active_params"} - else _mip_bound_prompt + _mip_maximum_prompt if metric in {"params", "active_params"} else _mip_bound_prompt ) bound = prompt( session, @@ -3014,9 +3164,7 @@ def _mip_axis_specs(inventory: Any, pruning: Mapping[str, Any]) -> list[dict[str if axis.axis_id == "hidden_width" or axis.axis_id not in _MIP_AXIS_ALIASES: continue setting = _mapping_copy(configured.get(axis.axis_id)) - values = list( - dict.fromkeys(int(value) for value in setting.get("values") or axis.values) - ) + values = list(dict.fromkeys(int(value) for value in setting.get("values") or axis.values)) if not values or not bool(setting.get("enabled", False)): continue specs.append( @@ -3041,11 +3189,7 @@ def _mip_scenario_domains( embeddings = list(dict.fromkeys(int(value) for value in hidden.get("values") or ())) if not embeddings: teacher_width = next( - ( - int(axis.teacher_value) - for axis in inventory.axes - if axis.axis_id == "hidden_width" - ), + (int(axis.teacher_value) for axis in inventory.axes if axis.axis_id == "hidden_width"), None, ) if teacher_width is None: @@ -3133,21 +3277,17 @@ def _mip_variant_prompt( ) if axes is BACK: return BACK - selected_axes = set(str(axis_id) for axis_id in axes) + selected_axes = {str(axis_id) for axis_id in axes} axis_options: OrderedDict[str, Any] = OrderedDict() for spec in axis_specs: if spec["axis_id"] not in selected_axes: continue values = session.checkbox( - ( - f"mip.search.{search_id}.variant.{variant_id}." - f"axes.{spec['axis_id']}.values" - ), + (f"mip.search.{search_id}.variant.{variant_id}.axes.{spec['axis_id']}.values"), f"Allowed values for {spec['label']}:", [ ( - f"{value}" - + (" (teacher)" if value == spec["teacher_value"] else ""), + f"{value}" + (" (teacher)" if value == spec["teacher_value"] else ""), value, ) for value in spec["values"] @@ -3184,9 +3324,7 @@ def _mip_solution_estimate( else: homogeneous_per_solve = int(homogeneous_keep) homogeneous_label = str(homogeneous_per_solve) - candidate_upper_bound = solve_count * ( - heterogeneous_per_solve + homogeneous_per_solve - ) + candidate_upper_bound = solve_count * (heterogeneous_per_solve + homogeneous_per_solve) return { "concrete_solves": solve_count, "heterogeneous_per_solve": heterogeneous_per_solve, @@ -3248,31 +3386,19 @@ def mip_section(session: WizardSession, resolver: DefaultsResolver, context: dic for name, raw in serving_workloads.items() ) runtime_workloads = OrderedDict( - (name, workloads[name]) - for name in measurements - if name in workloads - ) - default_goal_metric = str( - resolver.resolve_default("mip.goal_metric", "params").value + (name, workloads[name]) for name in measurements if name in workloads ) + default_goal_metric = str(resolver.resolve_default("mip.goal_metric", "params").value) default_goal_value = resolver.resolve_default("mip.goal_value", "75%").value default_objective = str( resolver.resolve_default( "mip.objective", "metrics.cosine_embedding_loss_hidden_states" ).value ) - default_num_solutions = int( - resolver.resolve_default("mip.num_solutions", 8).value - ) - available_depths, available_embeddings = _mip_scenario_domains( - inventory, pruning - ) + default_num_solutions = int(resolver.resolve_default("mip.num_solutions", 8).value) + available_depths, available_embeddings = _mip_scenario_domains(inventory, pruning) teacher_embedding = next( - ( - int(axis.teacher_value) - for axis in inventory.axes - if axis.axis_id == "hidden_width" - ), + (int(axis.teacher_value) for axis in inventory.axes if axis.axis_id == "hidden_width"), max(available_embeddings), ) axis_specs = _mip_axis_specs(inventory, pruning) @@ -3296,9 +3422,7 @@ def mip_section(session: WizardSession, resolver: DefaultsResolver, context: dic default_goal_metric = "params" if default_objective not in {value for _, value in _MIP_RANKING_CHOICES}: default_objective = _MIP_RANKING_CHOICES[0][1] - default_search_id = _mip_default_search_id( - default_goal_metric, default_goal_value - ) + default_search_id = _mip_default_search_id(default_goal_metric, default_goal_value) default_estimate = _mip_solution_estimate( variant_count=1, metric_count=1, @@ -3324,9 +3448,7 @@ def mip_section(session: WizardSession, resolver: DefaultsResolver, context: dic "variants": ["baseline"], "heterogeneous_per_solve": default_num_solutions, "homogeneous_per_solve": 5, - "candidate_origin_upper_bound": default_estimate[ - "candidate_origin_upper_bound" - ], + "candidate_origin_upper_bound": default_estimate["candidate_origin_upper_bound"], }, ) if action is BACK: @@ -3335,9 +3457,7 @@ def mip_section(session: WizardSession, resolver: DefaultsResolver, context: dic if action == "defaults": constraint: Any = {"max": _mip_parse_scalar(str(default_goal_value))} if default_goal_metric in {"memory", "runtime"}: - target_workloads = ( - runtime_workloads if default_goal_metric == "runtime" else workloads - ) + target_workloads = runtime_workloads if default_goal_metric == "runtime" else workloads if not target_workloads: raise SetupError( f"Default MIP constraint {default_goal_metric!r} requires " @@ -3347,7 +3467,7 @@ def mip_section(session: WizardSession, resolver: DefaultsResolver, context: dic else "a serving workload." ) ) - constraint = {"at": {name: constraint for name in target_workloads}} + constraint = {"at": dict.fromkeys(target_workloads, constraint)} session.state.set_collection( "mip_config", { @@ -3356,9 +3476,7 @@ def mip_section(session: WizardSession, resolver: DefaultsResolver, context: dic "runs": { default_search_id: { "constraints": {default_goal_metric: constraint}, - "objectives": [ - {"metric": default_objective, "direction": "minimize"} - ], + "objectives": [{"metric": default_objective, "direction": "minimize"}], "search_space": { "depth": available_depths, "embedding": available_embeddings, @@ -3378,9 +3496,7 @@ def mip_section(session: WizardSession, resolver: DefaultsResolver, context: dic }, }, ) - session.state.set_collection( - "mip_search_estimates", {default_search_id: default_estimate} - ) + session.state.set_collection("mip_search_estimates", {default_search_id: default_estimate}) return True runs: OrderedDict[str, Any] = OrderedDict() @@ -3570,9 +3686,7 @@ def mip_section(session: WizardSession, resolver: DefaultsResolver, context: dic run = { "constraints": constraints, - "objectives": [ - {"metric": metric, "direction": "minimize"} for metric in objectives - ], + "objectives": [{"metric": metric, "direction": "minimize"} for metric in objectives], "search_space": { "depth": depths, "embedding": embeddings, @@ -4166,17 +4280,37 @@ def _configure_dynamic_resources( "sequence_parallel": False, } else: - profile = ( - _profile_prompt( + if customize: + profile = _profile_prompt( session, registry, stage_id, model, node_type=node.node_type, ) - if customize - else registry.reuse(registry.names()[0], consumer=stage_id) - ) + else: + profile, issues = _compatible_default_profile( + session, + registry, + stage_id, + model, + node_type=node.node_type, + ) + if profile is None: + _print_parallel_issues(issues) + if session.guided: + raise SetupError( + "No configured parallel profile is compatible with " + f"{stage_id}. Supply a compatible profile in --defaults " + "or resume with --full." + ) + profile = _profile_prompt( + session, + registry, + stage_id, + model, + node_type=node.node_type, + ) if profile is BACK: return BACK issues = _profile_compatibility_issues( @@ -4253,9 +4387,7 @@ def _acquisition_sample_requirements(state: WizardState) -> tuple[int, int]: validation_requirements.append(int(pruning.get("depth_importance_samples", 1))) if bool(pruning.get("sort_sanity", False)): validation_requirements.append(int(pruning.get("sort_sanity_samples", 1))) - if bool(pruning.get("sort_sanity", False)) and bool( - pruning.get("width_sanity", False) - ): + if bool(pruning.get("sort_sanity", False)) and bool(pruning.get("width_sanity", False)): validation_requirements.append(int(pruning.get("width_sanity_samples", 1))) for raw_flow in _mapping_copy(state.collection("post_mip_flows")).values(): @@ -4282,20 +4414,13 @@ def _apportion_vlm_samples( subset_rows: Mapping[str, Any], total: int, ) -> dict[str, int]: - rows = { - str(name): int(value) - for name, value in subset_rows.items() - } + rows = {str(name): int(value) for name, value in subset_rows.items()} if not rows or any(value <= 0 for value in rows.values()): raise SetupError( - "Cannot infer Nemotron-VLM acquisition without positive selected-subset " - "row counts." + "Cannot infer Nemotron-VLM acquisition without positive selected-subset row counts." ) source_total = sum(rows.values()) - quotas = { - name: total * value // source_total - for name, value in rows.items() - } + quotas = {name: total * value // source_total for name, value in rows.items()} remaining = total - sum(quotas.values()) ranked = sorted( enumerate(rows.items()), @@ -4314,14 +4439,8 @@ def _infer_vlm_shard_cap( num_samples: int, ) -> int: subset_rows = _mapping_copy(acquisition.get("subset_rows")) - subset_media_shards = _mapping_copy( - acquisition.get("subset_media_shards") - ) - missing = [ - name - for name in subset_rows - if int(subset_media_shards.get(name, 0)) <= 0 - ] + subset_media_shards = _mapping_copy(acquisition.get("subset_media_shards")) + missing = [name for name in subset_rows if int(subset_media_shards.get(name, 0)) <= 0] if missing: previous_shard_cap = int(acquisition.get("max_shards_per_subset", 0)) if previous_shard_cap > 0: @@ -4337,9 +4456,7 @@ def _infer_vlm_shard_cap( for name, source_rows in subset_rows.items(): available_shards = int(subset_media_shards[name]) quota = quotas[name] - estimated = ( - quota * available_shards + int(source_rows) - 1 - ) // int(source_rows) + estimated = (quota * available_shards + int(source_rows) - 1) // int(source_rows) estimates.append(min(available_shards, max(1, estimated))) return max(estimates) @@ -4411,28 +4528,106 @@ def output_review_section( else: _record_default(session.state, resolver, "output.result_root", default_root) print("\nEffective setup:") - print( - yaml.safe_dump( - _plain_review_value( - { - "fields": { - path: { - "effective": record.effective, - "requested": record.requested, - "source": record.source, - } - for path, record in session.state.records().items() - }, - "profiles": session.state.collection("parallel_profiles"), - "serving_workloads": session.state.collection("serving_workloads"), - "vllm_measurements": session.state.collection("vllm_measurements"), - "mip": session.state.collection("mip_config"), - "post_mip": session.state.collection("post_mip_flows"), - } - ), - sort_keys=False, + if session.guided: + pruning = _mapping_copy(session.state.collection("pruning")) + mip = _mapping_copy(session.state.collection("mip_config")) + runs = _mapping_copy(mip.get("runs")) + profiles = _mapping_copy(session.state.collection("parallel_profiles")) + axes = { + axis_id: list(_mapping_copy(axis).get("values") or ()) + for axis_id, axis in _mapping_copy(pruning.get("axes")).items() + } + mip_review = { + run_id: { + "constraints": _mapping_copy(_mapping_copy(run).get("constraints")), + "num_solutions": _mapping_copy(_mapping_copy(run).get("solver")).get( + "num_solutions" + ), + } + for run_id, run in runs.items() + } + profile_review = { + name: { + key: _mapping_copy(profile).get(key) + for key in ( + "tp", + "cp", + "pp", + "dp_shard", + "dp_replicate", + "ep", + "sequence_parallel", + ) + } + for name, profile in profiles.items() + } + summary = { + "preset": session.state.preset, + "model": session.state.get_field("model.source"), + "dataset": session.state.get_field("data.selected_source"), + "sequence_length": session.state.get_field("data.sequence_length"), + "pruning": { + "maximum_depth_removed": pruning.get("depth_remove"), + "depth_importance_samples": pruning.get("depth_importance_samples"), + "width_importance_samples": pruning.get("width_importance_samples"), + "width_axes": axes, + "sort_sanity": pruning.get("sort_sanity"), + "width_sanity": pruning.get("width_sanity"), + "slicing_sanity": pruning.get("slicing_sanity"), + "bypass": _mapping_copy(pruning.get("bypass")), + "replacement_samples": pruning.get("replacement_samples"), + }, + "mip_searches": mip_review, + "parallel_profiles": profile_review, + "execution": { + "repository": session.state.get_field( + "infrastructure.execution_contract.repository" + ), + "venv": session.state.get_field("infrastructure.execution_contract.venv"), + "container": session.state.get_field("infrastructure.execution_contract.container"), + "slurm_account": session.state.get_field("infrastructure.runner.slurm.account"), + "interactive_partition": session.state.get_field( + "infrastructure.runner.slurm.partition_interactive" + ), + "batch_partition": session.state.get_field( + "infrastructure.runner.slurm.partition_batch" + ), + "gpus_per_node": session.state.get_field("infrastructure.gpus_per_node"), + }, + "results": session.state.get_field("output.result_root"), + } + print(yaml.safe_dump(summary, sort_keys=False)) + print( + " Nested values and provenance are saved in answers_v2.yaml. " + "Use --full for per-section customization." + ) + print( + " Execution values target the machine or cluster running setup. " + "Use --defaults or --full when that environment needs different values." + ) + else: + print( + yaml.safe_dump( + _plain_review_value( + { + "fields": { + path: { + "effective": record.effective, + "requested": record.requested, + "source": record.source, + } + for path, record in session.state.records().items() + }, + "profiles": session.state.collection("parallel_profiles"), + "serving_workloads": session.state.collection("serving_workloads"), + "vllm_measurements": session.state.collection("vllm_measurements"), + "mip": session.state.collection("mip_config"), + "post_mip": session.state.collection("post_mip_flows"), + } + ), + sort_keys=False, + ) ) - ) generate = session.confirm( "output.generate", "Validate and generate smoke and production bundles?", @@ -4486,23 +4681,55 @@ def output_review_section( def _fresh_state( backend: PromptBackend, defaults_path: Path | None, + *, + full: bool, ) -> WizardState: + if full: + while True: + value = backend.text("Campaign directory:", "") + if value is BACK: + continue + path = Path(str(value)).expanduser() + if str(path): + return WizardState.start( + path, + defaults_path=defaults_path, + setup_mode="full", + ) + while True: + preset = _select_setup_preset(backend) + if preset is BACK: + continue value = backend.text("Campaign directory:", "") if value is BACK: continue path = Path(str(value)).expanduser() if str(path): - return WizardState.start(path, defaults_path=defaults_path) + return WizardState.start( + path, + defaults_path=defaults_path, + setup_mode="quick", + preset=str(preset), + ) + + +def _select_setup_preset( + backend: PromptBackend, + *, + default: str = "balanced", +) -> Any: + return backend.select( + "Setup profile:", + [PromptChoice(item.choice_title, item.name) for item in QUICK_SETUP_PRESETS], + default, + ) def _refresh_legacy_state(state: WizardState) -> None: pruning = deepcopy(state.collection("pruning") or {}) subset_selection = _mapping_copy(state.collection("data_subset_selection")) - subset_records = [ - _mapping_copy(item) - for item in subset_selection.get("subsets") or () - ] + subset_records = [_mapping_copy(item) for item in subset_selection.get("subsets") or ()] serving_workloads = _mapping_copy(state.collection("serving_workloads")) measurements = _mapping_copy(state.collection("vllm_measurements")) first_workload = next( @@ -4517,9 +4744,7 @@ def _refresh_legacy_state(state: WizardState) -> None: "vllm_enabled": bool(measurements), "granularity": measurement.get("granularity", "subblock"), "workload_id": workload_id, - "isl": int( - workload.get("prefill_seq_len", state.get_field("data.sequence_length", 4096)) - ), + "isl": int(workload.get("prefill_seq_len", state.get_field("data.sequence_length", 4096))), "osl": int(workload.get("generation_seq_len", 1024)), "concurrency": int(workload.get("max_num_seqs", 1)), } @@ -4619,9 +4844,7 @@ def _refresh_legacy_state(state: WizardState) -> None: "sequence_length": state.get_field("data.sequence_length", 4096), "subsets": [record["name"] for record in subset_records], "subset_revision": subset_selection.get("revision"), - "subset_weights": { - record["name"]: record["weight"] for record in subset_records - }, + "subset_weights": {record["name"]: record["weight"] for record in subset_records}, "acquisition": deepcopy(state.collection("data_acquisition") or {}), }, "pruning": pruning, @@ -4640,31 +4863,112 @@ def run_wizard_v2( resume: Path | None, defaults_path: Path | None, backend: PromptBackend | None = None, + full: bool = False, ) -> Path: """Run setup v2, save every answer, validate bundles, and never launch jobs.""" backend = backend or InteractiveBackend() - print("Welcome to Puzzletron setup v2 — defaults are local, control is per stage.") + print("Welcome to Puzzletron setup v2.") if resume is None: - state = _fresh_state(backend, defaults_path) + if full: + print(" Full setup enabled: every advanced section is customizable.") + else: + print( + " Guided setup asks only for essential choices and applies nested " + "defaults from a profile." + ) + print(" Use --full only when you need every advanced control.") + state = _fresh_state(backend, defaults_path, full=full) else: state = WizardState.resume(resume) + if full and state.setup_mode != "full": + print( + " Promoting this guided campaign to full setup. Existing model " + "and dataset answers are preserved." + ) + state.set_setup_mode("full") + if state.setup_mode not in {"quick", "full"}: + raise SetupError(f"Unsupported setup mode in {state.path}: {state.setup_mode!r}.") + preset = None + if state.preset: + preset = get_setup_preset(state.preset) + if state.setup_mode == "quick": + if preset is None: + raise SetupError(f"Guided setup state {state.path} does not record a setup preset.") + print(f" Profile: {preset.choice_title}") + elif preset is not None: + print(f" Full setup baseline: {preset.choice_title}") + else: + print(" Full setup enabled: every advanced section is customizable.") selected_defaults = defaults_path or state.defaults_path - resolver = _resolver(state, selected_defaults) - session = WizardSession(state, backend) + if resume is not None and defaults_path is not None: + state.set_defaults_path(defaults_path) + print(f" Persisted replacement defaults file: {state.defaults_path}") + session = WizardSession( + state, + backend, + guided=state.setup_mode == "quick", + ) context: dict[str, Any] = {} if state.payload.get("model", {}).get("source"): saved = state.payload["model"] context["model"] = inspect_model(str(saved["source"])) + family_config = context["model"].inventory.family_config if "model" in context else None + model_inventory = context["model"].inventory if "model" in context else None + resolver = _resolver(state, selected_defaults, preset, family_config, model_inventory) index = 0 while index < len(SECTION_BUILDERS): + if session.guided and index == 2: + print( + "\nApplying the selected profile to advanced pruning, runtime, " + "MIP, and post-MIP settings..." + ) builder = SECTION_BUILDERS[index] completed = builder(session, resolver, context) if completed: + if SECTION_NAMES[index] == "model" and "model" in context: + resolver = _resolver( + state, + selected_defaults, + preset, + context["model"].inventory.family_config, + context["model"].inventory, + ) index += 1 else: target = session.consume_back_target() + if target is None and index == 0 and session.guided: + replacement = _select_setup_preset( + backend, + default=state.preset or "balanced", + ) + if replacement is not BACK: + state.set_preset(str(replacement)) + preset = get_setup_preset(str(replacement)) + family_config = ( + context["model"].inventory.family_config if "model" in context else None + ) + model_inventory = context["model"].inventory if "model" in context else None + resolver = _resolver( + state, + selected_defaults, + preset, + family_config, + model_inventory, + ) + print(f" Profile changed to: {preset.choice_title}") + continue index = SECTION_NAMES.index(target.section) if target is not None else index + state.set_collection( + "default_resolutions", + { + path: { + "value": resolved.value, + "source": resolved.source, + } + for path, resolved in resolver.resolutions().items() + }, + ) _refresh_legacy_state(state) build_bundles_v2(state.campaign_dir, state) return state.campaign_dir diff --git a/tests/unit/torch/puzzletron/test_setup_v2_data.py b/tests/unit/torch/puzzletron/test_setup_v2_data.py index c64269782e0..8340d29fb5f 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_data.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_data.py @@ -17,6 +17,9 @@ from types import SimpleNamespace +import pytest + +from puzzletron_setup import SetupError from puzzletron_setup.v2.bundle import _bundle_readme from puzzletron_setup.v2.defaults import DefaultsResolver from puzzletron_setup.v2.hf_datasets import HfSubsetCatalog, HfSubsetInfo @@ -84,8 +87,7 @@ def _nemotron_catalog(): ("external", 400, 8192, "external media required"), ] entries.extend( - (f"subset_{index:02d}", index + 1, (index + 1) * 1000, None) - for index in range(42) + (f"subset_{index:02d}", index + 1, (index + 1) * 1000, None) for index in range(42) ) return _catalog(_NEMOTRON_VLM_DATA_SOURCE, entries) @@ -97,8 +99,8 @@ def test_data_choices_include_first_class_sources_and_deduplicate_default(): assert [choice.title for choice in choices] == [ f"Default — {_PUZZLE_KD_DATA_SOURCE}", - "NVIDIA Nemotron-VLM v2 (image-text)", - "Custom local path or Hugging Face dataset", + "NVIDIA Nemotron-VLM v2: recommended image-text dataset", + "Custom dataset: choose a local path or Hugging Face dataset", ] @@ -292,6 +294,36 @@ def test_generic_hugging_face_dataset_uses_dynamic_subset_checkbox( assert backend.checkbox_calls[0][2] == ("small",) +def test_guided_explicit_invalid_subset_fails_instead_of_falling_back( + tmp_path, + monkeypatch, +): + state = WizardState.start( + tmp_path / "campaign", + defaults_path=None, + setup_mode="quick", + preset="balanced", + ) + backend = ScriptedBackend([_CUSTOM_DATA_SOURCE, "owner/generic"]) + catalog = _catalog( + "owner/generic", + [("small", 10, 100, None), ("disabled", 20, 200, "media unavailable")], + default="small", + ) + monkeypatch.setattr( + "puzzletron_setup.v2.wizard.infer_dataset_modality", + lambda source: SimpleNamespace(modality="text", evidence="test catalog"), + ) + + with pytest.raises(SetupError, match=r"typo.*Choose from: small"): + data_section( + WizardSession(state, backend, guided=True), + DefaultsResolver(file_defaults={"data": {"subsets": ["typo"]}}), + _context(multimodal=False), + catalog_loader=lambda source, **kwargs: catalog, + ) + + def test_resume_reuses_the_revision_locked_subset_catalog(tmp_path): state = WizardState.start(tmp_path / "campaign", defaults_path=None) answers = [ @@ -408,7 +440,7 @@ def Style(rules): # noqa: N802 - mirrors questionary's public constructor return rules @staticmethod - def checkbox(message, choices, instruction, style): + def checkbox(message, choices, *, instruction, style): assert message == "Subsets:" assert choices[:2] == rendered assert choices[2] == {"separator": " ← Back (press Esc)"} @@ -420,6 +452,10 @@ def checkbox(message, choices, instruction, style): "puzzletron_setup.v2.prompts._questionary", lambda: _Questionary(), ) + monkeypatch.setattr( + "puzzletron_setup.v2.prompts._bind_escape_back", + lambda question: question, + ) selected = InteractiveBackend().checkbox( "Subsets:", diff --git a/tests/unit/torch/puzzletron/test_setup_v2_parallel_prompts.py b/tests/unit/torch/puzzletron/test_setup_v2_parallel_prompts.py index 170ec1bd937..f93ee766c38 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_parallel_prompts.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_parallel_prompts.py @@ -113,8 +113,10 @@ def test_profile_prompt_rejects_incompatible_reuse_and_asks_again( assert "Choose a different parallel setting." in output -def test_default_stage_profile_is_rejected_and_replaced(tmp_path, capsys): - session, backend = _session(tmp_path, ["reuse:good"]) +def test_default_stages_reuse_first_compatible_profile_without_reprompting( + tmp_path, +): + session, backend = _session(tmp_path, []) registry = ResourceProfileRegistry( { "bad": ParallelProfile(name="bad", tp=8), @@ -125,7 +127,7 @@ def test_default_stage_profile_is_rejected_and_replaced(tmp_path, capsys): resolver = DefaultsResolver() model = SimpleNamespace(inventory=MOE_INVENTORY) - result = _configure_stage_resource( + first = _configure_stage_resource( session, resolver, model, @@ -133,11 +135,22 @@ def test_default_stage_profile_is_rejected_and_replaced(tmp_path, capsys): action="defaults", batch_default=8, ) + second = _configure_stage_resource( + session, + resolver, + model, + "replacement_scoring", + action="defaults", + batch_default=8, + ) - assert result.profile.name == "good" + assert first.profile.name == "good" + assert second.profile.name == "good" assert session.state.collection("stage_resources")["width_sanity"]["profile_name"] == "good" + assert ( + session.state.collection("stage_resources")["replacement_scoring"]["profile_name"] == "good" + ) assert backend.remaining == 0 - assert "Choose a different parallel setting." in capsys.readouterr().out def test_serving_prompt_asks_aiperf_inputs_and_boolean_expert_parallel(tmp_path): diff --git a/tests/unit/torch/puzzletron/test_setup_v2_quick.py b/tests/unit/torch/puzzletron/test_setup_v2_quick.py new file mode 100644 index 00000000000..fb9c52adc0f --- /dev/null +++ b/tests/unit/torch/puzzletron/test_setup_v2_quick.py @@ -0,0 +1,837 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import yaml + +import puzzletron_setup.v2.wizard as wizard_module +from puzzletron_setup import SetupError +from puzzletron_setup.inspection import InspectedModel +from puzzletron_setup.profiles import AxisInventory, ModelInventory +from puzzletron_setup.v2.cli import _parser +from puzzletron_setup.v2.defaults import DefaultsResolver +from puzzletron_setup.v2.presets import QUICK_SETUP_PRESETS, get_setup_preset +from puzzletron_setup.v2.prompts import ( + BACK, + InteractiveBackend, + PromptChoice, + ScriptedBackend, + _bind_escape_back, +) +from puzzletron_setup.v2.session import WizardSession +from puzzletron_setup.v2.state import WizardState +from puzzletron_setup.v2.wizard import ( + _CUSTOM_DATA_SOURCE, + _CUSTOM_MODEL_SOURCE, + _fresh_state, + _section_action, + data_section, + depth_section, + infrastructure_section, + output_review_section, +) + +_QWEN_FAMILY_CONFIG = "examples/puzzletron/configs/families/qwen3_5/family.yaml" +_NEMOTRON_FAMILY_CONFIG = "examples/puzzletron/configs/families/nemotron3/family.yaml" + + +def _qwen_inventory( + *, + num_layers, + hidden_size, + intermediate_size, + num_attention_heads, + num_key_value_heads, +): + return SimpleNamespace( + num_layers=num_layers, + facts={ + "hidden_size": hidden_size, + "intermediate_size": intermediate_size, + "num_attention_heads": num_attention_heads, + "num_key_value_heads": num_key_value_heads, + }, + ) + + +def _context(): + return { + "model": SimpleNamespace( + inventory=SimpleNamespace(multimodal=False), + ) + } + + +def test_guided_profiles_explain_cost_and_load_family_defaults(): + assert [preset.name for preset in QUICK_SETUP_PRESETS] == [ + "smoke", + "balanced", + "high-confidence", + ] + assert "recommended" in get_setup_preset("balanced").choice_title.lower() + + resolver = DefaultsResolver( + builtins={"pruning": {"bypass": {"enabled": True}}}, + preset_defaults=get_setup_preset("smoke").resolved_defaults(_QWEN_FAMILY_CONFIG), + ) + + resolved = resolver.resolve_default("pruning.bypass.enabled") + assert resolved.value is False + assert resolved.source == "preset" + + +def test_guided_profile_defaults_are_selected_by_model_family(tmp_path): + families = {} + for family, num_solutions in (("first", 2), ("second", 7)): + family_dir = tmp_path / family + family_dir.mkdir() + family_config = family_dir / "family.yaml" + family_config.write_text("family: {}\n") + (family_dir / "setup_v2_defaults.yaml").write_text( + yaml.safe_dump( + { + "schema_version": 1, + "profiles": {"smoke": {"mip": {"num_solutions": num_solutions}}}, + } + ) + ) + families[family] = family_config + + preset = get_setup_preset("smoke") + + assert preset.resolved_defaults(families["first"])["mip"]["num_solutions"] == 2 + assert preset.resolved_defaults(families["second"])["mip"]["num_solutions"] == 7 + + +def test_same_family_profile_uses_smaller_sample_counts_for_qwen_0p8b(): + qwen_0p8b = _qwen_inventory( + num_layers=24, + hidden_size=1024, + intermediate_size=3584, + num_attention_heads=8, + num_key_value_heads=2, + ) + qwen_9b = _qwen_inventory( + num_layers=32, + hidden_size=4096, + intermediate_size=12288, + num_attention_heads=16, + num_key_value_heads=4, + ) + + smoke = get_setup_preset("smoke") + small_defaults = smoke.resolved_defaults(_QWEN_FAMILY_CONFIG, qwen_0p8b) + large_defaults = smoke.resolved_defaults(_QWEN_FAMILY_CONFIG, qwen_9b) + + assert small_defaults["pruning"]["width_importance_samples"] == 8 + assert small_defaults["pruning"]["replacement_samples"] == 4 + assert large_defaults["pruning"]["width_importance_samples"] == 512 + assert large_defaults["pruning"]["replacement_samples"] == 32 + assert small_defaults["pruning"]["depth_remove"] == 1 + assert large_defaults["pruning"]["depth_remove"] == 1 + + +def test_model_override_changes_only_its_historical_qwen_9b_values(): + qwen_9b = _qwen_inventory( + num_layers=32, + hidden_size=4096, + intermediate_size=12288, + num_attention_heads=16, + num_key_value_heads=4, + ) + + defaults = get_setup_preset("high-confidence").resolved_defaults( + _QWEN_FAMILY_CONFIG, + qwen_9b, + ) + + assert defaults["pruning"]["depth_importance_samples"] == 128 + assert defaults["pruning"]["replacement_samples"] == 128 + assert defaults["pruning"]["bypass"]["enabled"] is False + assert defaults["mip"]["goal_value"] == "75%" + assert defaults["pruning"]["width_importance_samples"] == 65536 + assert defaults["mip"]["num_solutions"] == 16 + + +def test_qwen_27b_balanced_profile_uses_its_historical_campaign_budgets(): + qwen_27b = _qwen_inventory( + num_layers=64, + hidden_size=5120, + intermediate_size=17408, + num_attention_heads=24, + num_key_value_heads=4, + ) + + defaults = get_setup_preset("balanced").resolved_defaults( + _QWEN_FAMILY_CONFIG, + qwen_27b, + ) + + assert defaults["pruning"]["width_importance_samples"] == 16384 + assert defaults["pruning"]["replacement_samples"] == 16 + assert defaults["mip"]["goal_value"] == "85%" + + +def test_nemotron_nano_profiles_use_model_specific_smoke_and_search_budgets(): + nano = SimpleNamespace( + num_layers=52, + facts={ + "hidden_size": 2688, + "intermediate_size": 1856, + "num_attention_heads": 32, + "num_key_value_heads": 2, + "num_experts": 128, + }, + ) + + smoke = get_setup_preset("smoke").resolved_defaults(_NEMOTRON_FAMILY_CONFIG, nano) + high_confidence = get_setup_preset("high-confidence").resolved_defaults( + _NEMOTRON_FAMILY_CONFIG, + nano, + ) + + assert smoke["pruning"]["width_importance_samples"] == 2 + assert smoke["pruning"]["bypass"]["enabled"] is True + assert smoke["mip"]["num_solutions"] == 1 + assert high_confidence["pruning"]["depth_remove"] == 5 + assert high_confidence["pruning"]["width_importance_samples"] == 8192 + assert high_confidence["mip"]["num_solutions"] == 5 + + +def test_explicit_defaults_still_override_model_specific_profile(): + qwen_0p8b = _qwen_inventory( + num_layers=24, + hidden_size=1024, + intermediate_size=3584, + num_attention_heads=8, + num_key_value_heads=2, + ) + resolver = DefaultsResolver( + preset_defaults=get_setup_preset("smoke").resolved_defaults( + _QWEN_FAMILY_CONFIG, + qwen_0p8b, + ), + file_defaults={"pruning": {"width_importance_samples": 64}}, + ) + + resolved = resolver.resolve_default("pruning.width_importance_samples") + + assert resolved.value == 64 + assert resolved.source == "defaults_file" + + +def test_ambiguous_model_specific_defaults_fail_closed(tmp_path): + family_config = tmp_path / "family.yaml" + family_config.write_text("family: {}\n") + (tmp_path / "setup_v2_defaults.yaml").write_text( + yaml.safe_dump( + { + "schema_version": 2, + "profiles": {"smoke": {"mip": {"num_solutions": 2}}}, + "model_overrides": { + "first": { + "match": {"num_layers": 24}, + "profiles": {"smoke": {"mip": {"num_solutions": 3}}}, + }, + "second": { + "match": {"facts": {"hidden_size": 1024}}, + "profiles": {"smoke": {"mip": {"num_solutions": 4}}}, + }, + }, + } + ) + ) + inventory = SimpleNamespace(num_layers=24, facts={"hidden_size": 1024}) + + with pytest.raises(SetupError, match="matches multiple guided setup overrides"): + get_setup_preset("smoke").resolved_defaults(family_config, inventory) + + +def test_guided_profile_defaults_fail_closed_when_family_profile_is_missing(tmp_path): + family_config = tmp_path / "family.yaml" + family_config.write_text("family: {}\n") + (tmp_path / "setup_v2_defaults.yaml").write_text( + yaml.safe_dump({"schema_version": 1, "profiles": {}}) + ) + + with pytest.raises(SetupError, match="profile 'balanced' is not configured"): + get_setup_preset("balanced").resolved_defaults(family_config) + + +@pytest.mark.parametrize("family_config", [_QWEN_FAMILY_CONFIG, _NEMOTRON_FAMILY_CONFIG]) +@pytest.mark.parametrize("preset_name", ["smoke", "balanced", "high-confidence"]) +def test_each_model_family_defines_every_guided_profile(family_config, preset_name): + defaults = get_setup_preset(preset_name).resolved_defaults(family_config) + + assert defaults["pruning"] + assert defaults["mip"] + + +def test_explicit_defaults_override_guided_profile(): + resolver = DefaultsResolver( + preset_defaults={"mip": {"num_solutions": 2}}, + file_defaults={"mip": {"num_solutions": 5}}, + ) + + resolved = resolver.resolve_default("mip.num_solutions") + assert resolved.value == 5 + assert resolved.source == "defaults_file" + assert resolver.resolutions()["mip.num_solutions"] == resolved + + +def test_fresh_guided_state_records_profile_and_cli_full_is_explicit(tmp_path): + campaign = tmp_path / "campaign" + state = _fresh_state( + ScriptedBackend(["balanced", str(campaign)]), + None, + full=False, + ) + + assert state.setup_mode == "quick" + assert state.preset == "balanced" + assert _parser().parse_args([]).full is False + assert _parser().parse_args(["--full"]).full is True + + +def test_back_from_first_guided_section_can_change_profile( + tmp_path, + monkeypatch, +): + state = WizardState.start( + tmp_path / "campaign", + defaults_path=None, + setup_mode="quick", + preset="balanced", + ) + attempts = 0 + + def model_builder(session, resolver, context): + nonlocal attempts + del resolver, context + attempts += 1 + session.begin("model") + if attempts == 1: + return ( + session.select( + "model.source", + "Model:", + [PromptChoice("Custom", "custom"), PromptChoice("Known", "known")], + ) + is not BACK + ) + return True + + monkeypatch.setattr(wizard_module, "SECTION_BUILDERS", (model_builder,)) + monkeypatch.setattr(wizard_module, "SECTION_NAMES", ("model",)) + monkeypatch.setattr(wizard_module, "_refresh_legacy_state", lambda state: None) + monkeypatch.setattr(wizard_module, "build_bundles_v2", lambda campaign, state: None) + + wizard_module.run_wizard_v2( + resume=state.campaign_dir, + defaults_path=None, + backend=ScriptedBackend([BACK, "smoke"]), + ) + + assert WizardState.resume(state.path).preset == "smoke" + assert attempts == 2 + + +def test_resume_full_promotes_guided_state_and_preserves_profile_baseline( + tmp_path, + monkeypatch, +): + state = WizardState.start( + tmp_path / "campaign", + defaults_path=None, + setup_mode="quick", + preset="balanced", + ) + state.payload["model"] = {"source": "saved-model"} + state.save() + resolved_baseline = {} + inspected = SimpleNamespace( + inventory=SimpleNamespace(family_config=_QWEN_FAMILY_CONFIG), + ) + + def capture_baseline(session, resolver, context): + del session, context + resolved = resolver.resolve_default("pruning.depth_remove") + resolved_baseline.update(value=resolved.value, source=resolved.source) + return True + + monkeypatch.setattr(wizard_module, "SECTION_BUILDERS", (capture_baseline,)) + monkeypatch.setattr(wizard_module, "SECTION_NAMES", ("model",)) + monkeypatch.setattr(wizard_module, "inspect_model", lambda source: inspected) + monkeypatch.setattr(wizard_module, "_refresh_legacy_state", lambda state: None) + monkeypatch.setattr(wizard_module, "build_bundles_v2", lambda campaign, state: None) + + wizard_module.run_wizard_v2( + resume=state.campaign_dir, + defaults_path=None, + backend=ScriptedBackend([]), + full=True, + ) + + resumed = WizardState.resume(state.path) + assert resumed.setup_mode == "full" + assert resumed.preset == "balanced" + assert resolved_baseline == {"value": 4, "source": "preset"} + + +def test_resume_replacement_defaults_file_is_persisted( + tmp_path, + monkeypatch, +): + state = WizardState.start( + tmp_path / "campaign", + defaults_path=None, + setup_mode="quick", + preset="balanced", + ) + replacement = tmp_path / "replacement.yaml" + replacement.write_text(yaml.safe_dump({"schema_version": 1})) + monkeypatch.setattr(wizard_module, "SECTION_BUILDERS", ()) + monkeypatch.setattr(wizard_module, "SECTION_NAMES", ()) + monkeypatch.setattr(wizard_module, "_refresh_legacy_state", lambda state: None) + monkeypatch.setattr(wizard_module, "build_bundles_v2", lambda campaign, state: None) + + wizard_module.run_wizard_v2( + resume=state.campaign_dir, + defaults_path=replacement, + backend=ScriptedBackend([]), + ) + + assert WizardState.resume(state.path).defaults_path == replacement.resolve() + + +def test_legacy_state_without_setup_metadata_resumes_in_full_mode(tmp_path): + state = WizardState.start(tmp_path / "campaign", defaults_path=None) + state.payload.pop("setup") + state.save() + + resumed = WizardState.resume(state.path) + + assert resumed.setup_mode == "full" + assert resumed.preset is None + + +def test_guided_data_asks_only_for_source_and_uses_nested_defaults( + tmp_path, + monkeypatch, +): + dataset = tmp_path / "dataset" + dataset.mkdir() + state = WizardState.start( + tmp_path / "campaign", + defaults_path=None, + setup_mode="quick", + preset="balanced", + ) + backend = ScriptedBackend([_CUSTOM_DATA_SOURCE, str(dataset)]) + monkeypatch.setattr( + "puzzletron_setup.v2.wizard.infer_dataset_modality", + lambda source: SimpleNamespace(modality="text", evidence="local fixture"), + ) + + assert data_section( + WizardSession(state, backend, guided=True), + DefaultsResolver( + preset_defaults=get_setup_preset("balanced").resolved_defaults(_QWEN_FAMILY_CONFIG), + file_defaults={ + "data": { + "modality": "text", + "layout": "padded", + "sequence_length": 2048, + } + }, + ), + _context(), + ) + + assert backend.remaining == 0 + assert state.get_field("data.source") == str(dataset.resolve()) + assert state.get_field("data.modality") == "text" + assert state.get_field("data.layout") == "padded_varlen" + assert state.get_field("data.sequence_length") == 2048 + assert state.field("data.modality").source == "defaults_file" + assert state.field("data.layout").source == "defaults_file" + + +def test_guided_data_rejects_an_explicit_modality_incompatible_with_the_model( + tmp_path, + monkeypatch, +): + dataset = tmp_path / "dataset" + dataset.mkdir() + state = WizardState.start( + tmp_path / "campaign", + defaults_path=None, + setup_mode="quick", + preset="balanced", + ) + monkeypatch.setattr( + "puzzletron_setup.v2.wizard.infer_dataset_modality", + lambda source: SimpleNamespace(modality="text", evidence="local fixture"), + ) + + with pytest.raises(SetupError, match="multimodal.*incompatible"): + data_section( + WizardSession( + state, + ScriptedBackend([_CUSTOM_DATA_SOURCE, str(dataset)]), + guided=True, + ), + DefaultsResolver(file_defaults={"data": {"modality": "multimodal"}}), + _context(), + ) + + +def test_guided_review_renders_the_actual_non_parameter_mip_constraint( + tmp_path, + capsys, +): + state = WizardState.start( + tmp_path / "campaign", + defaults_path=None, + setup_mode="quick", + preset="balanced", + ) + state.set_collection( + "mip_config", + { + "runs": { + "memory-search": { + "constraints": {"memory": {"at": {"serving-default": {"max": "24GiB"}}}}, + "solver": {"num_solutions": 4}, + } + } + }, + ) + + assert output_review_section( + WizardSession(state, ScriptedBackend([True]), guided=True), + DefaultsResolver(), + {}, + ) + + output = capsys.readouterr().out + assert "constraints:" in output + assert "memory:" in output + assert "24GiB" in output + assert "parameter_target" not in output + + +def test_guided_wizard_runs_real_sections_and_generates_valid_bundles( + tmp_path, + monkeypatch, +): + campaign = tmp_path / "campaign" + model_path = tmp_path / "model" + dataset = tmp_path / "dataset" + model_path.mkdir() + dataset.mkdir() + inventory = ModelInventory( + family="qwen3_5", + descriptor="qwen3_5_text", + family_config="examples/puzzletron/configs/families/qwen3_5/family.yaml", + model_type="qwen3_5_text", + architectures=("Qwen3_5ForCausalLM",), + multimodal=False, + moe=False, + num_layers=24, + num_sublayers=48, + layer_counts={"full_attention": 6, "linear_attention": 18}, + facts={ + "hidden_size": 1024, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "intermediate_size": 3584, + }, + axes=( + AxisInventory( + axis_id="hidden_width", + label="Hidden width", + teacher_value=1024, + values=(1024, 768), + alignment=256, + ), + ), + ) + inspected = InspectedModel( + source=str(model_path), + requested_revision=None, + resolved_revision=None, + is_local=True, + config={ + "model_type": "qwen3_5_text", + "text_config": { + "num_hidden_layers": 24, + "layer_types": ["linear_attention"] * 18 + ["full_attention"] * 6, + }, + }, + inventory=inventory, + ) + monkeypatch.setattr(wizard_module, "inspect_model", lambda source: inspected) + monkeypatch.setattr( + wizard_module, + "infer_dataset_modality", + lambda source: SimpleNamespace(modality="text", evidence="local fixture"), + ) + backend = ScriptedBackend( + [ + "smoke", + str(campaign), + _CUSTOM_MODEL_SOURCE, + str(model_path), + _CUSTOM_DATA_SOURCE, + str(dataset), + "defaults", + "/worker/modelopt", + "/worker/venv", + True, + ] + ) + + result = wizard_module.run_wizard_v2( + resume=None, + defaults_path=None, + backend=backend, + ) + + assert result == campaign.resolve() + assert backend.remaining == 0 + assert (campaign / "smoke" / "experiment.yaml").is_file() + assert (campaign / "production" / "experiment.yaml").is_file() + assert (campaign / "resolved_defaults.yaml").is_file() + generated = WizardState.resume(campaign) + assert generated.collection("pruning")["depth_remove"] == 1 + assert generated.collection("pruning")["width_importance_samples"] == 8 + assert generated.collection("pruning")["replacement_samples"] == 4 + assert generated.collection("default_resolutions")["pruning.depth_remove"] == { + "value": 1, + "source": "preset", + } + assert generated.collection("default_resolutions")["pruning.width_importance_samples"] == { + "value": 8, + "source": "model_profile", + } + + +def test_guided_section_uses_profile_without_action_prompt(tmp_path): + state = WizardState.start( + tmp_path / "campaign", + defaults_path=None, + setup_mode="quick", + preset="balanced", + ) + backend = ScriptedBackend([]) + session = WizardSession(state, backend, guided=True) + + action = _section_action( + session, + "mip", + "Configure the search.", + {"num_solutions": 8}, + ) + + assert action == "defaults" + assert backend.remaining == 0 + + +def test_guided_infrastructure_prompts_for_unresolved_worker_paths(tmp_path): + state = WizardState.start( + tmp_path / "campaign", + defaults_path=None, + setup_mode="quick", + preset="balanced", + ) + backend = ScriptedBackend(["defaults", "/worker/modelopt", "/worker/venv"]) + + assert infrastructure_section( + WizardSession(state, backend, guided=True), + DefaultsResolver(), + {}, + ) + + assert state.get_field("infrastructure.execution_contract.repository") == "/worker/modelopt" + assert state.get_field("infrastructure.execution_contract.venv") == "/worker/venv" + assert state.get_field("infrastructure.runner.kind") == "slurm" + assert backend.remaining == 0 + + +def test_full_section_keeps_the_existing_customize_prompt(tmp_path): + state = WizardState.start(tmp_path / "campaign", defaults_path=None) + backend = ScriptedBackend(["customize"]) + + action = _section_action( + WizardSession(state, backend), + "mip", + "Configure the search.", + {"num_solutions": 8}, + ) + + assert action == "customize" + assert backend.remaining == 0 + + +def test_back_reasks_the_previous_prompt_with_replay_intact(tmp_path): + state = WizardState.start(tmp_path / "campaign", defaults_path=None) + backend = ScriptedBackend(["one", "two", BACK, "revised"]) + session = WizardSession(state, backend) + session.begin("data") + + assert session.text("data.one", "One:") == "one" + assert session.text("data.two", "Two:") == "two" + assert session.text("data.three", "Three:") is BACK + target = session.consume_back_target() + assert target is not None + assert target.prompt_id == "data.two" + + session.begin("data") + assert session.text("data.one", "One:") == "one" + assert session.text("data.two", "Two:") == "revised" + assert backend.remaining == 0 + + +def test_depth_back_replays_conditional_path_and_zero_removes_resources(tmp_path): + state = WizardState.start(tmp_path / "campaign", defaults_path=None) + state.set_collection("stage_resources", {"depth_importance": {"instances": 8}}) + state.set_collection( + "stage_batches", + {"depth_importance.micro_batch_size": 8}, + ) + backend = ScriptedBackend(["customize", "subblock", 2, BACK, 0]) + session = WizardSession(state, backend) + context = { + "model": SimpleNamespace( + inventory=SimpleNamespace(num_sublayers=8, num_layers=4), + ) + } + + assert not depth_section(session, DefaultsResolver(), context) + target = session.consume_back_target() + assert target is not None + assert target.prompt_id == "pruning.depth_remove" + + assert depth_section(session, DefaultsResolver(), context) + assert state.collection("pruning")["depth_remove"] == 0 + assert "depth_importance" not in state.collection("stage_resources") + assert "depth_importance.micro_batch_size" not in state.collection("stage_batches") + assert backend.remaining == 0 + + +def test_escape_returns_back_for_text_select_and_checkbox(monkeypatch): + questions = [] + + class _Bindings: + def __init__(self): + self.handlers = {} + + def add(self, key, eager=False): + assert eager + + def register(handler): + self.handlers[key] = handler + return handler + + return register + + class _Application: + def __init__(self): + self.key_bindings = _Bindings() + self.result = None + + def exit(self, *, result): + self.result = result + + class _Question: + def __init__(self): + self.application = _Application() + questions.append(self) + + def ask(self): + event = SimpleNamespace(app=self.application) + self.application.key_bindings.handlers["escape"](event) + return self.application.result + + class _Questionary: + @staticmethod + def Style(value): # noqa: N802 - mirrors questionary's public constructor + return value + + @staticmethod + def Choice(**kwargs): # noqa: N802 - mirrors questionary's public constructor + return kwargs + + @staticmethod + def Separator(value): # noqa: N802 - mirrors questionary's public constructor + return value + + @staticmethod + def text(*args, **kwargs): + return _Question() + + @staticmethod + def select(*args, **kwargs): + return _Question() + + @staticmethod + def checkbox(*args, **kwargs): + return _Question() + + monkeypatch.setattr( + "puzzletron_setup.v2.prompts._questionary", + lambda: _Questionary(), + ) + backend = InteractiveBackend() + + assert backend.text("Text:", "") is BACK + assert backend.select("Select:", [PromptChoice("One", 1)], 1) is BACK + assert backend.checkbox("Checkbox:", [PromptChoice("One", 1)], [1]) is BACK + assert len(questions) == 3 + + +def test_escape_binding_supports_a_merged_binding_adapter(monkeypatch): + registered = {} + existing_bindings = object() + + class _Bindings: + def add(self, key, eager=False): + assert eager + + def register(handler): + registered[key] = handler + return handler + + return register + + escape_bindings = _Bindings() + merged_bindings = object() + + key_binding_module = ModuleType("prompt_toolkit.key_binding") + key_binding_module.KeyBindings = lambda: escape_bindings + + def merge_key_bindings(bindings): + assert bindings == [existing_bindings, escape_bindings] + return merged_bindings + + key_binding_module.merge_key_bindings = merge_key_bindings + prompt_toolkit_module = ModuleType("prompt_toolkit") + prompt_toolkit_module.key_binding = key_binding_module + monkeypatch.setitem(sys.modules, "prompt_toolkit", prompt_toolkit_module) + monkeypatch.setitem(sys.modules, "prompt_toolkit.key_binding", key_binding_module) + + application = SimpleNamespace(key_bindings=existing_bindings, result=None) + application.exit = lambda *, result: setattr(application, "result", result) + question = SimpleNamespace(application=application) + + _bind_escape_back(question) + + assert question.application.key_bindings is merged_bindings + registered["escape"](SimpleNamespace(app=application)) + assert application.result is BACK