Merge pull request #25 from ben-nowacki/main

Bug fixes with on-disk result storage

Author: ben-nowacki

modularml/core/experiment/experiment.py

@@ -569,6 +569,7 @@ def _execute_training(
         phase: TrainPhase,
         *,
         _artifact_dir: Path | None = None,
+        _callback_dir: Path | None = None,
         _execution_dir: Path | None = None,
         _metric_dir: Path | None = None,
         show_sampler_progress: bool = True,
@@ -620,6 +621,7 @@ def _execute_training(
         res = TrainResults(
             label=phase.label,
             _artifact_dir=_artifact_dir,
+            _callback_dir=_callback_dir,
             _execution_dir=_execution_dir,
             _metric_dir=_metric_dir,
         )
@@ -684,6 +686,7 @@ def _execute_evaluation(
         phase: EvalPhase,
         *,
         _artifact_dir: Path | None = None,
+        _callback_dir: Path | None = None,
         _execution_dir: Path | None = None,
         _metric_dir: Path | None = None,
         show_eval_progress: bool = False,
@@ -719,6 +722,7 @@ def _execute_evaluation(
         res = EvalResults(
             label=phase.label,
             _artifact_dir=_artifact_dir,
+            _callback_dir=_callback_dir,
             _execution_dir=_execution_dir,
             _metric_dir=_metric_dir,
         )
@@ -742,6 +746,7 @@ def _execute_fit(
         phase: FitPhase,
         *,
         _artifact_dir: Path | None = None,
+        _callback_dir: Path | None = None,
         _execution_dir: Path | None = None,
         _metric_dir: Path | None = None,
     ) -> FitResults:
@@ -765,6 +770,7 @@ def _execute_fit(
         res = FitResults(
             label=phase.label,
             _artifact_dir=_artifact_dir,
+            _callback_dir=_callback_dir,
             _execution_dir=_execution_dir,
             _metric_dir=_metric_dir,
         )
@@ -784,8 +790,7 @@ def _execute_phase_with_meta(
         self,
         phase: TrainPhase | EvalPhase | FitPhase,
         *,
-        _path_suffix: Path | None = None,
-        _run_idx: int | None = None,
+        phase_dir: Path | None = None,
         **kwargs,
     ) -> tuple[PhaseResults, PhaseExecutionMeta]:
         """
@@ -808,9 +813,9 @@ def _execute_phase_with_meta(
             if exp_dir is None and self._exp_checkpointing is not None:
                 exp_dir = self._exp_checkpointing.directory
             if exp_dir is not None:
-                phase_dir = exp_dir / phase.label
-                phase_dir.mkdir(parents=True, exist_ok=True)
-                phase.checkpointing._directory = phase_dir
+                ckpt_dir = exp_dir / phase.label
+                ckpt_dir.mkdir(parents=True, exist_ok=True)
+                phase.checkpointing._directory = ckpt_dir
 
         # Skip callbacks and checkpointing when inside a callback
         run_hooks = not self._in_callback
@@ -837,23 +842,8 @@ def _execute_phase_with_meta(
             self._save_experiment_checkpoint(label=phase.label)
 
         # ------------------------------------------------
-        # Compute phase-specific storage directories
+        # Compute phase-specific storage directories from the caller-supplied phase_dir
         # ------------------------------------------------
-        if _path_suffix is not None:
-            # Called from within a group; suffix already contains the run prefix
-            phase_dir = self._results_config.phase_dir(_path_suffix / phase.label)
-        elif _run_idx is not None:
-            # Top-level call from run_phase; prefix with run index for stable ordering
-            phase_dir = self._results_config.phase_dir(f"{_run_idx}_{phase.label}")
-        elif (
-            self._active_phase_dir is not None
-            and self._results_config.results_dir is not None
-        ):
-            # Called from preview_phase during callback execution → nest under callbacks/
-            phase_dir = self._active_phase_dir / "callbacks" / phase.label
-        else:
-            phase_dir = None  # pure preview with no disk storage
-
         cfg = self._results_config
         phase_execution_dir = (
             phase_dir / "execution_data"
@@ -870,6 +860,11 @@ def _execute_phase_with_meta(
             if phase_dir is not None and cfg.save_artifacts
             else None
         )
+        phase_callback_dir = (
+            phase_dir / "callbacks"
+            if phase_dir is not None and cfg.save_execution
+            else None
+        )
 
         # Track active phase dir so nested callback previews nest under callbacks/
         prev_active_phase_dir = self._active_phase_dir
@@ -892,6 +887,7 @@ def _execute_phase_with_meta(
                 phase_res: TrainResults = self._execute_training(
                     phase,
                     _artifact_dir=phase_artifact_dir,
+                    _callback_dir=phase_callback_dir,
                     _execution_dir=phase_execution_dir,
                     _metric_dir=phase_metric_dir,
                     **{k: v for k, v in kwargs.items() if k in train_keys},
@@ -901,6 +897,7 @@ def _execute_phase_with_meta(
                 phase_res: EvalResults = self._execute_evaluation(
                     phase,
                     _artifact_dir=phase_artifact_dir,
+                    _callback_dir=phase_callback_dir,
                     _execution_dir=phase_execution_dir,
                     _metric_dir=phase_metric_dir,
                     **{k: v for k, v in kwargs.items() if k in eval_keys},
@@ -909,6 +906,7 @@ def _execute_phase_with_meta(
                 phase_res: FitResults = self._execute_fit(
                     phase,
                     _artifact_dir=phase_artifact_dir,
+                    _callback_dir=phase_callback_dir,
                     _execution_dir=phase_execution_dir,
                     _metric_dir=phase_metric_dir,
                 )
@@ -957,8 +955,7 @@ def _execute_group_with_meta(
         self,
         group: PhaseGroup,
         *,
-        _path_suffix: Path | None = None,
-        _run_idx: int | None = None,
+        group_dir: Path | None = None,
         **kwargs,
     ) -> tuple[PhaseGroupResults, PhaseGroupExecutionMeta]:
         """
@@ -1001,15 +998,6 @@ def _execute_group_with_meta(
         # - construct result container
         # - run each phase in order
         # ------------------------------------------------
-        if _path_suffix is not None:
-            # Nested group; suffix already contains the run prefix
-            group_suffix = _path_suffix / group.label
-        elif _run_idx is not None:
-            # Top-level call from run_group; prefix with run index for stable ordering
-            group_suffix = Path(f"{_run_idx}_{group.label}")
-        else:
-            group_suffix = Path(group.label)
-
         group_results = PhaseGroupResults(label=group.label)
         group_meta = PhaseGroupExecutionMeta(
             label=group.label,
@@ -1019,9 +1007,12 @@ def _execute_group_with_meta(
         for element in group.all:
             if isinstance(element, ExperimentPhase):
                 # Run phase with meta tracking
+                element_dir = (
+                    group_dir / element.label if group_dir is not None else None
+                )
                 phase_res, phase_meta = self._execute_phase_with_meta(
                     phase=element,
-                    _path_suffix=group_suffix,
+                    phase_dir=element_dir,
                     **kwargs,
                 )
 
@@ -1032,9 +1023,10 @@ def _execute_group_with_meta(
 
             elif isinstance(element, PhaseGroup):
                 # Run group with meta tracking
+                sub_dir = group_dir / element.label if group_dir is not None else None
                 sub_res, sub_meta = self._execute_group_with_meta(
                     group=element,
-                    _path_suffix=group_suffix,
+                    group_dir=sub_dir,
                     **kwargs,
                 )
 
@@ -1137,9 +1129,12 @@ def run_phase(
 
         # Run phase and record phase-level meta data
         try:
+            run_dir = self._results_config.phase_dir(
+                f"{len(self._history)}_{phase.label}",
+            )
             res, meta = self._execute_phase_with_meta(
                 phase=phase,
-                _run_idx=len(self._history),
+                phase_dir=run_dir,
                 **kwargs,
             )
         except Exception:
@@ -1196,9 +1191,12 @@ def run_group(
 
         # Run group and record phase-level meta data
         try:
+            run_dir = self._results_config.phase_dir(
+                f"{len(self._history)}_{group.label}",
+            )
             res, meta = self._execute_group_with_meta(
                 group=group,
-                _run_idx=len(self._history),
+                group_dir=run_dir,
                 **kwargs,
             )
         except Exception:
@@ -1344,10 +1342,11 @@ def preview_phase(
         needs_restore = _phase_mutates_state(phase)
         state = self.get_state() if needs_restore else None
 
-        # Execute phase with checkpointing disabled
+        # Execute phase with checkpointing disabled (no disk storage for previews)
         with self.disable_checkpointing():
             res, _ = self._execute_phase_with_meta(
                 phase=phase,
+                phase_dir=None,
                 **kwargs,
             )
 
@@ -1385,10 +1384,11 @@ def preview_group(
         needs_restore = _phase_mutates_state(group)
         state = self.get_state() if needs_restore else None
 
-        # Execute group with checkpointing disabled
+        # Execute group with checkpointing disabled (no disk storage for previews)
         with self.disable_checkpointing():
             res, _ = self._execute_group_with_meta(
                 group=group,
+                group_dir=None,
                 **kwargs,
             )
 

modularml/core/experiment/results/artifact_store.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-import re
+import pickle
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, ClassVar
@@ -40,10 +40,15 @@ def __init__(
     def artifact(self) -> Any:
         """The artifact object. Transparently loads from disk if serialized as a Path."""
         if isinstance(self._artifact, Path):
-            import pickle
-
             with self._artifact.open("rb") as f:
-                return pickle.load(f)
+                payload = pickle.load(f)
+
+            # New format: dict payload containing the artifact under "artifact" key
+            if isinstance(payload, dict) and "artifact" in payload:
+                return payload["artifact"]
+
+            # Legacy format: file contained just the artifact object
+            return payload
         return self._artifact
 
     def __repr__(self) -> str:
@@ -85,6 +90,7 @@ class ArtifactStore:
     def __init__(self, location: Path | None = None) -> None:
         self._location = location
         self._entries: dict[str, list[ArtifactEntry]] = {}
+        self._count: int = 0
 
     # ================================================
     # Writing
@@ -109,17 +115,18 @@ def log(
                 The batch index. Defaults to None (epoch-level artifact).
 
         """
-        stored = artifact
         if self._location is not None:
-            stored = self._save_to_disk(
+            artifact_value = self._save_to_disk(
                 name=name,
                 artifact=artifact,
                 epoch_idx=epoch_idx,
                 batch_idx=batch_idx,
             )
+        else:
+            artifact_value = artifact
         entry = ArtifactEntry(
             name=name,
-            artifact=stored,
+            artifact=artifact_value,
             epoch_idx=epoch_idx,
             batch_idx=batch_idx,
         )
@@ -135,15 +142,20 @@ def _save_to_disk(
         epoch_idx: int,
         batch_idx: int | None,
     ) -> Path:
-        """Serialize ``artifact`` to disk and return the file path."""
+        """Serialize artifact and metadata to disk, returning the file path."""
         import pickle
 
-        batch_str = f"_b{batch_idx}" if batch_idx is not None else ""
-        filename = f"{name}_e{epoch_idx}{batch_str}.pkl"
-        filepath = Path(self._location) / filename  # type: ignore[arg-type]
+        filepath = Path(self._location) / f"{self._count}.pkl"  # type: ignore[arg-type]
         filepath.parent.mkdir(parents=True, exist_ok=True)
+        payload = {
+            "name": name,
+            "epoch_idx": epoch_idx,
+            "batch_idx": batch_idx,
+            "artifact": artifact,
+        }
         with filepath.open("wb") as f:
-            pickle.dump(artifact, f)
+            pickle.dump(payload, f)
+        self._count += 1
         return filepath
 
     # ================================================
@@ -202,35 +214,31 @@ def from_directory(cls, location: Path) -> ArtifactStore:
         """
         Reconstruct a store index from existing pickle files in ``location``.
 
-        Globs for ``*.pkl`` files and parses ``{name}_e{epoch}`` or
-        ``{name}_e{epoch}_b{batch}`` from filenames.
-
         Args:
             location (Path): Directory containing serialized artifact files.
 
         Returns:
             ArtifactStore: Store with ``_entries`` populated as ``Path`` references.
 
         """
+        import pickle
+
         store = cls(location=location)
-        # Matches: {name}_e{epoch}.pkl  or  {name}_e{epoch}_b{batch}.pkl
-        pattern = re.compile(r"^(.+)_e(\d+)(?:_b(\d+))?\.pkl$")
-        for filepath in sorted(Path(location).glob("*.pkl")):
-            m = pattern.match(filepath.name)
-            if m is None:
-                continue
-            name = m.group(1)
-            epoch_idx = int(m.group(2))
-            batch_idx = int(m.group(3)) if m.group(3) is not None else None
+        pkl_files = sorted(Path(location).glob("*.pkl"), key=lambda p: int(p.stem))
+        for filepath in pkl_files:
+            with filepath.open("rb") as f:
+                payload = pickle.load(f)
+            name = payload["name"]
             entry = ArtifactEntry(
                 name=name,
-                artifact=filepath,
-                epoch_idx=epoch_idx,
-                batch_idx=batch_idx,
+                artifact=filepath,  # lazy-load
+                epoch_idx=payload["epoch_idx"],
+                batch_idx=payload["batch_idx"],
             )
             if name not in store._entries:
                 store._entries[name] = []
             store._entries[name].append(entry)
+        store._count = len(pkl_files)
         return store
 
     # ================================================