[ENH] Fix model caching at refit and checkpointing in v2

.gitignore

@@ -144,6 +144,7 @@ lightning_logs
 *.ckpt
 *.pkl
 .DS_Store
+checkpoints
 
 # data
 pytorch_forecasting/data/*.parquet

pytorch_forecasting/callbacks/artifact_registry.py

@@ -0,0 +1,160 @@
+from pathlib import Path
+from typing import Any
+import warnings
+
+from lightning.pytorch.callbacks import Callback, ModelCheckpoint
+import yaml
+
+
+class _ArtifactRegistry:
+    """Reads/writes artifacts.yaml."""
+
+    @staticmethod
+    def _read(registry_path: Path) -> dict[str, Any]:
+        """Read the yaml file, returns ``{"artifacts": {}}`` if it doesn't exist yet."""
+        if not registry_path.exists():
+            return {"artifacts": {}}
+        with open(registry_path) as f:
+            data = yaml.safe_load(f)
+        return data if data is not None else {"artifacts": {}}
+
+    @staticmethod
+    def write(
+        registry_path: Path, artifacts: dict[str, Any], overwrite: bool = False
+    ) -> None:
+        """Initial write.
+
+        It is used by _save() to save scalers and other "static" artifacts
+        before training starts. The static artifacts are those artifacts which never
+        change during the training process - like the configs, datamodule.metadata,
+        scalers etc. The non-static artifacts changes as the training/prediction
+        process moves forward. Eg of this would be the model checkpoints - if we plan
+        to save "best" model (any model that has better performance based on some
+        metric)
+
+        Parameters
+        ----------
+        registry_path : Path
+           the path of the yaml file where we have to write the artifacts.
+        artifacts : dict[str, Any]
+           A dictionary of the artifacts we want to write.
+           The Keys of the dictionary would be the "type" of artifact - like scaler,
+           configs etc. And the value would be the actual object path to be written.
+        overwrite: bool, default=False
+            Whether to overwrite the artifact.yaml (if present) or not.
+        """
+        registry_path = Path(registry_path)
+        registry_path.parent.mkdir(parents=True, exist_ok=True)
+
+        if registry_path.exists():
+            if not overwrite:
+                raise FileExistsError(
+                    f"{registry_path} already exists. Pass `overwrite=True` to "
+                    "_ArtifactRegistry.write() to replace it, or use .update() "
+                    "if you only want to add/replace specific keys. You can also delete"
+                    "the file if that is not needed."
+                )
+            warnings.warn(
+                f"Overwriting existing {registry_path}. Any keys not present in "
+                "the new `artifacts` dict will be lost. Use .update() instead if "
+                "you want to preserve existing keys."
+            )
+
+        payload = {"artifacts": {k: str(v) for k, v in artifacts.items()}}
+        with open(registry_path, "w") as f:
+            yaml.safe_dump(payload, f)
+
+    @staticmethod
+    def update(registry_path: Path, artifacts: dict[str, Any]) -> None:
+        """Merge-update specific keys.
+        It is used by ArtifactRegistryCallback during
+        training, and by _save() for the manual_checkpoint key. It will never overwrite
+        the artifacts.yaml file, but just update an existing artifact entry.
+
+        Parameters
+        ----------
+        registry_path : Path
+             the path of the yaml file where we have to write the artifacts.
+        artifacts : dict[str, Any]
+            A dictionary where each key is the artifact type (eg.
+            `best_model_checkpoint`) and the value is the path to set/replace
+            for that key. Can contain one or more entries. Existing keys not
+            present in this dict are left untouched.
+        """
+        registry_path = Path(registry_path)
+        if not artifacts:
+            return
+
+        existing = _ArtifactRegistry._read(registry_path)
+        existing.setdefault("artifacts", {})
+        existing["artifacts"].update({k: str(v) for k, v in artifacts.items()})
+
+        registry_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(registry_path, "w") as f:
+            yaml.safe_dump(existing, f)
+
+    @staticmethod
+    def get(registry_path: Path, key: str | None = None) -> dict[str, Any] | None:
+        """Read one key, or the whole registry.
+
+        Parameters
+        -----------
+        registry_path : Path
+            the path of the yaml file from where we have to read the artifacts.
+        key : str
+            the key we want to know.
+
+        Returns
+        -------
+        dict
+            The dictionary if with the key as the key passed and value as the value
+            read, or None with a warning if missing/file doesn't exist yet.
+        key : str, default=None
+            the key we want to know. If None, the entire artifacts
+            dict is returned.
+        """
+        registry_path = Path(registry_path)
+        if not registry_path.exists():
+            warnings.warn(f"{registry_path} does not exist.")
+            return None
+
+        artifacts = _ArtifactRegistry._read(registry_path).get("artifacts", {})
+        if key is None:
+            return artifacts
+
+        value = artifacts.get(key)
+        if value is None:
+            raise KeyError(f"Key '{key}' not found in {registry_path}.")
+
+        return {key: value}
+
+
+class ArtifactRegistryCallback(Callback):
+    """Called every time Lightning's ModelCheckpoint actually persists a file.
+    Re-reads ckpt_cb.best_model_path / .last_model_path (always current,
+    never stale -- Lightning deletes superseded files itself) and writes
+    that into artifacts.yaml as the live truth, not a historical log.
+    """
+
+    def __init__(self, registry_path: Path):
+        self.registry_path = registry_path
+
+    def on_save_checkpoint(self, trainer, pl_module, checkpoint) -> None:
+        ckpt_cb = next(
+            (cb for cb in trainer.callbacks if isinstance(cb, ModelCheckpoint)), None
+        )
+        if ckpt_cb is None:
+            warnings.warn(
+                "ArtifactRegistryCallback found no ModelCheckpoint among "
+                "trainer.callbacks; skipping artifacts.yaml update."
+            )
+            return
+
+        updates: dict[str, Any] = {}
+        if ckpt_cb.best_model_path:
+            updates["best_model_checkpoint"] = ckpt_cb.best_model_path
+        if ckpt_cb.save_last and ckpt_cb.last_model_path:
+            updates["last_model_checkpoint"] = ckpt_cb.last_model_path
+
+        if updates:
+            _ArtifactRegistry.update(self.registry_path, updates)

pytorch_forecasting/data/data_module/_encoder_decoder_data_module.py

@@ -10,6 +10,7 @@
 from pathlib import Path
 import pickle
 from typing import Any, Optional, Union
+import warnings
 from warnings import warn
 
 from lightning.pytorch import LightningDataModule
@@ -803,6 +804,117 @@ def __getitem__(self, idx):
                 y = y.squeeze(-1)
             return x, y
 
+    _ARTIFACT_SPECS = [
+        ("scaler", "_scalers", "scalers"),
+        ("target_normalizer", "_target_normalizer", "scalers"),
+        ("datamodule_metadata", "_metadata", "metadata"),
+    ]
+
+    def save_artifacts(
+        self, artifact_dir: Path, include: list[str] = [], exclude: list[str] = []
+    ):
+        """Save data module artifacts.
+
+        Parameters
+        ----------
+        artifact_dir : Path
+            Path to save artifacts.
+        exclude : list[str], default=None
+            The list of artifacts that need to be excluded from saving.
+
+        Saves
+        -----
+        scalers, target_normalizers, and datamodule's metadata.
+
+        Returns
+        -------
+        dict
+            A dictionary containing artifacts with keys as the "type" of artifact
+            while the values are the paths where they are saved.
+
+        Raises
+        ------
+        UserWarning
+            If some artifact that is to be stored but the datamodule doesnt have them.
+            Eg if the scalers were to be stored but they were not initialized so they
+            are not present in data module's memory.
+        """
+        artifact_dir = Path(artifact_dir)
+        saved_artifacts: dict[str, Path] = {}
+
+        for key, attr, subdir in self._ARTIFACT_SPECS:
+            if key in exclude:
+                continue
+
+            value = getattr(self, attr)
+            if not value:
+                warnings.warn(
+                    f"No {key} found in the datamodule to save. "
+                    f"If you expected {key} to be saved, ensure it is "
+                    "passed to the datamodule constructor."
+                )
+                continue
+
+            target_dir = artifact_dir / subdir
+            target_dir.mkdir(parents=True, exist_ok=True)
+            path = target_dir / f"{key}.pkl"
+            with open(path, "wb") as f:
+                pickle.dump(value, f)
+            saved_artifacts[key] = path
+
+        return saved_artifacts
+
+    def load_artifacts(self, artifacts: dict[str, Any]):
+        """Save data module artifacts.
+
+        Parameters
+        ----------
+        artifacts : dict
+            Dictionary mapping artifact names to their file paths.
+            Expected keys (all optional):
+
+            - ``"scalers"``: Path to pickled scalers dict.
+            - ``"target_normalizer"``: Path to pickled target normalizer.
+            - ``"datamodule_metadata"``: Path to pickled metadata dict.
+
+        exclude : list[str], default=None
+            The list of artifacts that need to be excluded from saving.
+
+        Loads
+        -----
+        scalers : dict
+            Feature scalers loaded from ``artifacts["scalers"]``.
+        target_normalizer : object
+            Target normalizer loaded from ``artifacts["target_normalizer"]``.
+        metadata : dict
+            Metadata loaded from ``artifacts["datamodule_metadata"]``.
+
+        Returns
+        -------
+        dict
+            A dictionary containing artifacts with keys as the "type" of artifact
+            while the values are the paths where they are saved.
+
+        Raises
+        ------
+        FileNotFoundError
+            If a path specified in ``artifacts`` does not exist.
+        UserWarning
+            If a key is present in ``artifacts`` but the file doesn't exist.
+        """
+        for key, attr, _ in self._ARTIFACT_SPECS:
+            path = artifacts.get(key)
+            if path is None:
+                continue
+
+            path = Path(path)
+            if not path.exists():
+                warnings.warn(f"{key} file not found at {path}. Skipping {key} load.")
+                continue
+
+            with open(path, "rb") as f:
+                setattr(self, attr, pickle.load(f))  # noqa: S301
+
     def _create_windows(self, indices: torch.Tensor) -> list[tuple[int, int, int, int]]:
         """Generate sliding windows for training, validation, and testing.
 

pytorch_forecasting/data/data_module/_tslib_data_module.py

@@ -3,6 +3,8 @@
 """
 
 from collections.abc import Callable
+from pathlib import Path
+import pickle
 from typing import Any, Optional
 import warnings
 
@@ -671,6 +673,92 @@ def _create_windows(self, indices: torch.Tensor) -> list[tuple[int, int, int, in
 
         return windows
 
+    def save_artifacts(
+        self, artifact_dir: Path, include: list[str] = [], exclude: list[str] = []
+    ):
+        """Save data module artifacts.
+
+        Parameters
+        ----------
+        artifact_dir : Path
+            Path to save artifacts.
+        exclude : list[str], default=None
+            The list of artifacts that need to be excluded from saving.
+
+        Saves
+        -----
+        datamodule's metadata.
+
+        Returns
+        -------
+        dict
+            A dictionary containing artifacts with keys as the "type" of artifact
+            while the values are the paths where they are saved.
+
+        Raises
+        ------
+        UserWarning
+            If some artifact that is to be stored but the datamodule doesnt have them.
+        """
+        artifact_dir = Path(artifact_dir)
+        saved_artifacts: dict[str, Path] = {}
+
+        if "datamodule_metadata" not in exclude:
+            metadata = self.metadata
+            if metadata:
+                metadata_dir = artifact_dir / "metadata"
+                metadata_dir.mkdir(parents=True, exist_ok=True)
+                metadata_path = metadata_dir / "datamodule_metadata.pkl"
+                with open(metadata_path, "wb") as f:
+                    pickle.dump(metadata, f)
+                saved_artifacts["datamodule_metadata"] = metadata_path
+
+        return saved_artifacts
+
+    def load_artifacts(self, artifacts: dict[str, Any]):
+        """Save data module artifacts.
+
+        Parameters
+        ----------
+        artifacts : dict
+            Dictionary mapping artifact names to their file paths.
+            Expected keys (all optional):
+
+            - ``"datamodule_metadata"``: Path to pickled metadata dict.
+
+        exclude : list[str], default=None
+            The list of artifacts that need to be excluded from saving.
+
+        Loads
+        -----
+        metadata : dict
+            Metadata loaded from ``artifacts["datamodule_metadata"]``.
+
+        Returns
+        -------
+        dict
+            A dictionary containing artifacts with keys as the "type" of artifact
+            while the values are the paths where they are saved.
+
+        Raises
+        ------
+        FileNotFoundError
+            If a path specified in ``artifacts`` does not exist.
+        UserWarning
+            If a key is present in ``artifacts`` but the file doesn't exist.
+        """
+        metadata_path = artifacts.get("datamodule_metadata")
+        if metadata_path is not None:
+            metadata_path = Path(metadata_path)
+            if not metadata_path.exists():
+                warnings.warn(
+                    f"Metadata file not found at {metadata_path}. "
+                    "Skipping metadata load."
+                )
+            else:
+                with open(metadata_path, "rb") as f:
+                    self._metadata = pickle.load(f)  # noqa: S301
+
     def setup(self, stage: str | None = None) -> None:
         """
         Setup the data module by preparing the datasets for training,