Implement mode-agnostic refactor

Author: bknutson0

.gitignore

@@ -3,7 +3,7 @@
 uv.lock
 /data/*/**
 /out/*/*/**
-/tmp/
+/old/*
 
 # Ignore Python cache files
 __pycache__

README.md

@@ -2,10 +2,9 @@
 
 ## Configuration & Trials
 
-Configuration lives under `cfg/<mode>/...` where `<mode>` is one of:
-
-- `sweep` – hyperparameter exploration (multiple non-seed dimensions allowed; seed must be scalar)
-- `experiment` – focused evaluation (only the seed may vary; other hyperparameters must be scalar)
+Configuration lives under `cfg/<mode>/...`.
+The `<mode>` segment is only a folder namespace (e.g., `sweep`, `experiment`, `sweep_1`).
+It does not affect expansion semantics.
 
 Directory pattern for a leaf config:
 
@@ -21,33 +20,27 @@ Expansion writes resolved per‑trial configurations to:
 out/log/<mode>/<dataset>/<model>/<trainer>/trial_###/cfg.yaml
 ```
 
-Trials are the Cartesian product of list‑valued hyperparameters excluding those on a structural allowlist (currently only `batch_metrics`). Special handling:
-
-- `seed:` list -> becomes a sweep over seeds (each trial gets one scalar value)
-- `seed:` scalar -> required if no list sweep is desired
-
-Mode rules enforced during expansion:
-
-1. sweep mode: `seed` must be scalar; any other list hyperparameters are allowed.
-2. experiment mode: only `seed` may be list-valued; all other hyperparameters must be scalar (aside from structural allowlist like `batch_metrics`).
+Trials are the Cartesian product of list‑valued hyperparameters, excluding logging lists
+(currently only `batch_metrics`). `seed` may be scalar or list; if a list, it is just another
+dimension.
 
 Re‑running expansion is idempotent: resolved trial `cfg.yaml` files are regenerated (so parent edits propagate). A trial is considered complete when `batch_log.csv` exists in the same `trial_###` directory.
 
 Summary:
 
 1. Hierarchical inheritance (leaf overrides parents).
-2. Lists -> sweep axes (except structural allowlist).
-3. A list-form `seed` defines a seed axis; else scalar `seed` required.
-4. Mode-specific constraints (see above) validated at expansion time.
-5. Resolved configs live only under `out/log/...` and drive all training & analysis.
+2. Lists -> sweep axes (except logging lists like `batch_metrics`).
+3. `seed` can be scalar or a list (a list becomes an axis like any other).
+4. `mode` is only a folder namespace and does not affect expansion behavior.
+5. Resolved configs live under `out/log/...` and drive all training & analysis.
 
 Programmatic expansion + execution:
 
 ```python
 from src.run import run_trials
-# All modes (sweep + experiment) across all datasets
+# All configs across all datasets
 run_trials()
-# Only sweep mode for mnist
+# Only configs under cfg/sweep for mnist
 run_trials(datasets=["mnist"], modes=["sweep"])
 ```
 

run/experiment.py

@@ -1,14 +1,21 @@
 import sys
 from collections.abc import Sequence
 
-from src.run import run_trials
+from src.run import discover_modes, run_trials
 
 
-def main(datasets: Sequence[str] | None = None) -> None:  # pragma: no cover - thin wrapper
-    """Run all pending experiment trials for given datasets (or all)."""
-    run_trials(datasets=datasets, modes=['experiment'])
+def main(datasets: Sequence[str] | None = None, *, dry_run: bool = False) -> None:  # pragma: no cover - thin wrapper
+    """Run all pending trials under cfg modes containing 'experiment'."""
+    exp_like = discover_modes('experiment')
+    run_trials(datasets=datasets, modes=exp_like, dry_run=dry_run)
 
 
 if __name__ == '__main__':  # pragma: no cover
-    ds = sys.argv[1:] if len(sys.argv) > 1 else None
-    main(ds)
+    # usage: uv run -m run.experiment [dataset ...] [--dry-run]
+    args = sys.argv[1:]
+    dry = False
+    if '--dry-run' in args:
+        dry = True
+        args = [a for a in args if a != '--dry-run']
+    ds = args if args else None
+    main(ds, dry_run=dry)

run/sweep.py

@@ -1,16 +1,19 @@
 import sys
 from collections.abc import Sequence
 
-from src.run import run_trials
+from src.run import discover_modes, run_trials
 
 
 def main(datasets: Sequence[str] | None = None, *, dry_run: bool = False) -> None:  # pragma: no cover - thin wrapper
-    """Run all pending sweep trials for given datasets (or all)."""
-    run_trials(datasets=datasets, modes=['sweep'], dry_run=dry_run)
+    """Run all pending trials under cfg modes containing 'sweep'."""
+    # Discover modes with 'sweep' in their name (mode-agnostic semantics)
+    sweep_like = discover_modes('sweep')
+    run_trials(datasets=datasets, modes=sweep_like, dry_run=dry_run)
 
 
 if __name__ == '__main__':  # pragma: no cover
     # usage: uv run -m run.sweep [dataset ...] [--dry-run]
+    # 'sweep' is a folder namespace only; expansion is mode-agnostic.
     args = sys.argv[1:]
     dry = False
     if '--dry-run' in args:

src/run.py

@@ -1,7 +1,11 @@
 """Training orchestration for uniform config expansion.
 
-Each leaf `cfg.yaml` under `cfg/<mode>/` is expanded (respecting mode rules) and
-written to `out/log/<mode>/<dataset>/<model>/<trainer>/trial_###/cfg.yaml`.
+Each leaf `cfg.yaml` under `cfg/<mode>/` is expanded and written to
+`out/log/<mode>/<dataset>/<model>/<trainer>/trial_###/cfg.yaml`.
+
+Notes:
+- The `<mode>` segment is just a folder namespace (e.g., `sweep`, `experiment`, `sweep_1`).
+    It does not affect expansion semantics.
 """
 
 from collections.abc import Sequence
@@ -12,6 +16,19 @@
 from src.utils.cfg import create_trials, has_cfg_been_run
 
 
+def discover_modes(contains: str | Sequence[str] | None = None) -> list[str]:
+    """Return list of cfg/<mode> folder names, optionally filtered by substring(s)."""
+    cfg_root = Path('cfg')
+    modes = [p.name for p in cfg_root.iterdir() if p.is_dir()] if cfg_root.exists() else []
+    if contains is None:
+        return modes
+    if isinstance(contains, str):
+        key = contains.lower()
+        return [m for m in modes if key in m.lower()]
+    keys = {s.lower() for s in contains}
+    return [m for m in modes if any(k in m.lower() for k in keys)]
+
+
 def run_trials(
     datasets: Sequence[str] | None = None,
     modes: Sequence[str] | None = None,
@@ -36,12 +53,12 @@ def run_trials(
     if not cfg_root.exists():
         return {'trials_total': 0, 'trials_skipped': 0}
 
-    # Allowed values and defaults (datasets -> registry keys; modes -> canonical list)
+    # Allowed datasets come from the registry; modes are discovered dynamically
     allowed_datasets = list(DATASETS.keys())
-    allowed_modes = ['sweep', 'experiment']
+    discovered_modes = discover_modes()
 
-    # Resolve selected modes: default to all; otherwise intersect with allowed
-    requested_modes: list[str] = allowed_modes if modes is None else [m for m in allowed_modes if m in set(modes)]
+    # Resolve selected modes: default to all discovered; otherwise intersect
+    requested_modes: list[str] = discovered_modes if modes is None else [m for m in discovered_modes if m in set(modes)]
 
     # Map selected modes to existing cfg/<mode> directories
     mode_dirs = [cfg_root / m for m in requested_modes if (cfg_root / m).exists()]

src/utils/cfg.py

@@ -10,7 +10,6 @@
 from src.models.registry import MODELS
 from src.trainers.registry import TRAINERS
 
-# Configuration constants (simplified)
 DERIVED_CFG_KEYS = {'input_shape', 'num_classes'}
 REQUIRED_KEYS = {
     'dataset',
@@ -20,7 +19,7 @@
     'base_lr',
     'batch_size',
     'epochs',
-    'seed',  # a default seed (can be overridden by sweep dimension `seed` list)
+    'seed',
     'num_workers',
 }
 KEY_ORDER = [
@@ -40,7 +39,6 @@
 # Expansion constants
 LOGGING_LIST_KEYS = {'batch_metrics'}  # list-valued keys not treated as sweep dimensions
 HASH_IGNORE_KEYS = DERIVED_CFG_KEYS | LOGGING_LIST_KEYS
-SEED_KEY = 'seed'  # may be scalar or list (list => sweep over seeds)
 
 
 @dataclass
@@ -110,7 +108,7 @@ def cfg_hash(cfg: dict) -> str:
     configuration, with mapping keys sorted and without ignored keys.
     """
     norm = normalize_resolved_cfg(cfg)
-    data = json.dumps(norm, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False)
+    data = json.dumps(norm, sort_keys=True, separators=(',', ':'), ensure_ascii=False, allow_nan=False)
     return hashlib.sha256(data.encode('utf-8')).hexdigest()
 
 
@@ -120,7 +118,8 @@ def _extract_path_info(path: Path) -> tuple[str, str, str, str]:
     Expected pattern (after first 'cfg' segment):
         cfg/<mode>/<dataset>/<model>/<trainer>/.../cfg.yaml
 
-    <mode> must currently be one of {'sweep', 'experiment'}.
+    The `<mode>` segment is a free-form folder namespace (e.g., 'sweep',
+    'experiment', 'sweep_1'). Expansion semantics do not depend on this value.
     """
     parts = list(path.parts)
     try:
@@ -131,17 +130,9 @@ def _extract_path_info(path: Path) -> tuple[str, str, str, str]:
         raise ValueError(f'Invalid cfg path (no cfg segment): {path}') from e
 
     tail = parts[cfg_idx + 1 :]
-    # Legacy directory prefixes ('tune', 'experiments') are no longer supported and not skipped.
-    # 'sweep' and 'experiment' are treated as real dataset names now, so we do not strip anything.
-    if tail and tail[0] in {'tune', 'experiments'}:
-        raise ValueError(
-            f"Legacy path prefix no longer supported (remove 'tune/' or 'experiments/' from cfg path): {path}"
-        )
     if len(tail) < 4:
         raise ValueError(f'Insufficient path depth for mode/dataset/model/trainer: {path}')
     mode, dataset, model, trainer = tail[:4]
-    if mode not in {'sweep', 'experiment'}:
-        raise ValueError(f"Unsupported mode '{mode}' in cfg path (expected 'sweep' or 'experiment'): {path}")
     return mode, dataset, model, Path(trainer).stem
 
 
@@ -181,27 +172,17 @@ def _dimension_items(cfg: dict) -> list[tuple[str, list]]:
     return dims
 
 
-def expand_cfg_from_dict(cfg: dict, mode: str) -> list[TrialSpec]:
+def expand_cfg_from_dict(cfg: dict) -> list[TrialSpec]:
     """Expand a (possibly merged) cfg dict into trial specs.
 
-    For 'sweep' mode: all list-valued keys (except structural allowlist) become
-    dimensions. 'seed' must be scalar.
-    For 'experiment' mode: only 'seed' may be list-valued (creating per-seed trials).
+    Unified behavior (mode-agnostic):
+    - Every list-valued key at the top level (except those in LOGGING_LIST_KEYS)
+      is treated as a sweep dimension.
+    - 'seed' may be scalar or list; if list, it is simply another dimension.
+    - If no list-valued keys are present, produce a single TrialSpec.
     """
-    dims: list[tuple[str, list]] = []
-    for k, v in cfg.items():
-        if isinstance(v, list) and k not in LOGGING_LIST_KEYS:
-            if mode == 'sweep':
-                if k == SEED_KEY:
-                    raise ValueError("sweep cfg.yaml must not contain list-valued 'seed'; provide a single scalar seed")
-                dims.append((k, v))
-            elif mode == 'experiment':
-                if k == SEED_KEY:
-                    dims.append((k, v))
-                else:  # pragma: no cover - defensive, normally enforced earlier
-                    raise ValueError('experiment cfg.yaml may only vary seed; offending list key: ' + k)
-            else:  # pragma: no cover - future mode types
-                raise ValueError(f'Unsupported mode for expansion: {mode}')
+    # Identify dimensions using the shared helper for a single source of truth
+    dims = _dimension_items(cfg)
 
     if not dims:
         return [TrialSpec(idx=1, assignments={})]
@@ -215,9 +196,9 @@ def expand_cfg_from_dict(cfg: dict, mode: str) -> list[TrialSpec]:
 def create_trials(base_cfg_path: str | Path) -> list[Path]:
     """Write resolved trial cfgs beneath out/log/<mode>/<dataset>/<model>/<trainer>/trial_###/cfg.yaml.
 
-    Mode-specific rules:
-    - sweep: if 'seed' provided it must be scalar (not list); other list hyperparams allowed.
-    - experiment: only 'seed' may be list-valued (aside from structural allowlist).
+    Expansion uses unified rules (see `expand_cfg_from_dict`): all list-valued
+    top-level keys except logging lists become sweep dimensions. Resolved cfgs
+    must not contain any remaining list-valued hyperparameters.
     """
     path = Path(base_cfg_path)
     if path.name != 'cfg.yaml':
@@ -234,7 +215,7 @@ def create_trials(base_cfg_path: str | Path) -> list[Path]:
     # properly swept. Previously only the leaf cfg file was inspected which
     # caused inherited lists (e.g. batch_size) to remain unresolved in the
     # written trial cfgs.
-    trials = expand_cfg_from_dict(merged_base, mode)
+    trials = expand_cfg_from_dict(merged_base)
     # Root output directory for this cfg tree
     out_root = Path('out') / 'log' / mode / dataset / model / trainer
     out_root.mkdir(parents=True, exist_ok=True)
@@ -269,21 +250,19 @@ def _parse_idx(p: Path) -> int:
         resolved = merged_base.copy()
         resolved.update(spec.assignments)  # apply concrete assignments
         # Ensure that any list-valued dimension keys have been resolved to scalars
-        for k, v in list(resolved.items()):
-            if isinstance(v, list) and k not in LOGGING_LIST_KEYS:
-                # If still a list here, it was not selected as a dimension (e.g. leftover due to mode rules)
-                # For experiment mode only 'seed' may remain list; others are invalid.
-                if mode == 'experiment' and k == SEED_KEY:
-                    continue
-                raise ValueError(
-                    f'Unresolved list-valued hyperparameter {k!r} remained in resolved cfg. Check expansion rules.'
-                )
+        unresolved_dims = _dimension_items(resolved)
+        if unresolved_dims:
+            # With unified expansion, no list-valued hyperparameters should remain.
+            names = [name for name, _ in unresolved_dims]
+            raise ValueError(
+                f'Unresolved list-valued hyperparameters remained in resolved cfg: {names}. Check expansion rules.'
+            )
         # Compute identity hash for de-duplication
         new_hash = cfg_hash(resolved)
 
         # If this configuration already exists anywhere under out_root, reuse that path
         if new_hash in hash_to_dir:
-            existing_path = (hash_to_dir[new_hash] / 'cfg.yaml')
+            existing_path = hash_to_dir[new_hash] / 'cfg.yaml'
             written.append(existing_path)
             continue
 
@@ -356,23 +335,14 @@ def get_output_path(cfg_path: str | Path) -> Path:
     raise ValueError('Only resolved trial cfg paths are supported.')
 
 
-def _validate_expansion_rules(_: dict) -> None:  # pragma: no cover - retained for potential future rules
-    return
-
-
 def has_cfg_been_run(cfg_path: str | Path) -> tuple[bool, str]:
     """A trial is considered run if its cfg exists and batch_log.csv is present."""
     try:
         cfg_path = Path(cfg_path)
         if not cfg_path.exists():
             return False, 'Resolved cfg missing'
-        with cfg_path.open() as f:
-            resolved = yaml.safe_load(f) or {}
         if not (cfg_path.parent / 'batch_log.csv').exists():
             return False, 'No batch_log.csv'
-        # Lightweight validation of registries
-        if not cfgs_equal(resolved, resolved):  # always true, placeholder for future diff
-            return False, 'Internal mismatch'
         return True, 'Outputs present'
-    except Exception as e:  # noqa: BLE001
+    except Exception as e:
         return False, f'Error checking cfg: {e}'

tests/integration/test_dynamic_mode_discovery.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+
+from _pytest.monkeypatch import MonkeyPatch
+
+from src.run import run_trials
+
+
+def test_dynamic_mode_discovery_includes_custom_mode(tmp_path: Path, monkeypatch: MonkeyPatch) -> None:
+    """run_trials should discover cfg/<mode> subfolders (e.g., 'sweep_1') and plan runs."""
+    # Minimal config under a custom mode folder
+    base_cfg = tmp_path / 'cfg/sweep_1/mnist/simple-conv/jfb/cfg.yaml'
+    base_cfg.parent.mkdir(parents=True, exist_ok=True)
+    base_cfg.write_text(
+        'model: simple-conv\n'
+        'trainer: jfb\n'
+        'optimizer: sgd\n'
+        'batch_size: 2\n'
+        'epochs: 1\n'
+        'seed: 0\n'
+        'num_workers: 0\n'
+        'base_lr: 0.01\n'
+    )
+
+    # Ensure cwd is temp
+    cwd = Path.cwd()
+    monkeypatch.chdir(tmp_path)
+
+    # Dry-run planning should count this trial when filtering by the discovered mode
+    summary = run_trials(datasets=['mnist'], modes=['sweep_1'], dry_run=True)
+    assert summary['trials_planned'] >= 1
+    assert summary['trials_run'] == 0
+
+    # Without specifying modes, discovered modes should include 'sweep_1'
+    summary2 = run_trials(datasets=['mnist'], dry_run=True)
+    assert summary2['trials_planned'] >= 1
+    assert summary2['trials_run'] == 0
+
+    monkeypatch.chdir(cwd)