nshtrainer 1.0.0b13__tar.gz → 1.0.0b14__tar.gz

{nshtrainer-1.0.0b13 → nshtrainer-1.0.0b14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 1.0.0b13
+Version: 1.0.0b14
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b13 → nshtrainer-1.0.0b14}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "1.0.0-beta13"
+version = "1.0.0-beta14"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-1.0.0b13 → nshtrainer-1.0.0b14}/src/nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py

@@ -96,8 +96,8 @@ class OnExceptionCheckpointCallback(_OnExceptionCheckpoint):
     def on_exception(self, trainer: LightningTrainer, *args: Any, **kwargs: Any):
         # Monkey-patch the strategy instance to make the barrier operation a no-op.
         # We do this because `save_checkpoint` calls `barrier`. This is okay in most
-        # cases, but when we want to save a checkpoint in the case of an exception,
-        # `barrier` causes a deadlock. So we monkey-patch the strategy instance to
-        # make the barrier operation a no-op.
+        #   cases, but when we want to save a checkpoint in the case of an exception,
+        #   `barrier` causes a deadlock. So we monkey-patch the strategy instance to
+        #   make the barrier operation a no-op.
         with _monkey_patch_disable_barrier(trainer):
             return super().on_exception(trainer, *args, **kwargs)

{nshtrainer-1.0.0b13 → nshtrainer-1.0.0b14}/src/nshtrainer/configs/__init__.py

@@ -9,19 +9,7 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from nshtrainer import MetricConfig as MetricConfig
     from nshtrainer import TrainerConfig as TrainerConfig
-    from nshtrainer._checkpoint.loader import (
-        BestCheckpointStrategyConfig as BestCheckpointStrategyConfig,
-    )
-    from nshtrainer._checkpoint.loader import (
-        CheckpointLoadingStrategyConfig as CheckpointLoadingStrategyConfig,
-    )
-    from nshtrainer._checkpoint.loader import CheckpointMetadata as CheckpointMetadata
-    from nshtrainer._checkpoint.loader import (
-        LastCheckpointStrategyConfig as LastCheckpointStrategyConfig,
-    )
-    from nshtrainer._checkpoint.loader import (
-        UserProvidedPathCheckpointStrategyConfig as UserProvidedPathCheckpointStrategyConfig,
-    )
+    from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
     from nshtrainer._directory import DirectoryConfig as DirectoryConfig
     from nshtrainer._hf_hub import CallbackConfigBase as CallbackConfigBase
     from nshtrainer._hf_hub import (
@@ -122,9 +110,6 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         CheckpointCallbackConfig as CheckpointCallbackConfig,
     )
-    from nshtrainer.trainer._config import (
-        CheckpointLoadingConfig as CheckpointLoadingConfig,
-    )
     from nshtrainer.trainer._config import (
         CheckpointSavingConfig as CheckpointSavingConfig,
     )
@@ -199,21 +184,13 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).BestCheckpointCallbackConfig
-        if name == "BestCheckpointStrategyConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).BestCheckpointStrategyConfig
         if name == "CSVLoggerConfig":
             return importlib.import_module("nshtrainer.loggers").CSVLoggerConfig
         if name == "CallbackConfigBase":
             return importlib.import_module("nshtrainer._hf_hub").CallbackConfigBase
-        if name == "CheckpointLoadingConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).CheckpointLoadingConfig
         if name == "CheckpointMetadata":
             return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
+                "nshtrainer._checkpoint.metadata"
             ).CheckpointMetadata
         if name == "CheckpointSavingConfig":
             return importlib.import_module(
@@ -317,10 +294,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).LastCheckpointCallbackConfig
-        if name == "LastCheckpointStrategyConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).LastCheckpointStrategyConfig
         if name == "LeakyReLUNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").LeakyReLUNonlinearityConfig
         if name == "LearningRateMonitorConfig":
@@ -403,10 +376,6 @@ else:
             return importlib.import_module("nshtrainer.loggers").TensorboardLoggerConfig
         if name == "TrainerConfig":
             return importlib.import_module("nshtrainer").TrainerConfig
-        if name == "UserProvidedPathCheckpointStrategyConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).UserProvidedPathCheckpointStrategyConfig
         if name == "WandbLoggerConfig":
             return importlib.import_module("nshtrainer.loggers").WandbLoggerConfig
         if name == "WandbUploadCodeCallbackConfig":
@@ -423,10 +392,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).CheckpointCallbackConfig
-        if name == "CheckpointLoadingStrategyConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).CheckpointLoadingStrategyConfig
         if name == "DurationConfig":
             return importlib.import_module("nshtrainer.util.config").DurationConfig
         if name == "LRSchedulerConfig":

nshtrainer-1.0.0b14/src/nshtrainer/configs/_checkpoint/__init__.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
+    from nshtrainer._checkpoint.metadata import EnvironmentConfig as EnvironmentConfig
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "CheckpointMetadata":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.metadata"
+            ).CheckpointMetadata
+        if name == "EnvironmentConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.metadata"
+            ).EnvironmentConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+# Submodule exports
+from . import metadata as metadata

{nshtrainer-1.0.0b13 → nshtrainer-1.0.0b14}/src/nshtrainer/configs/trainer/__init__.py

@@ -18,9 +18,6 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         CheckpointCallbackConfig as CheckpointCallbackConfig,
     )
-    from nshtrainer.trainer._config import (
-        CheckpointLoadingConfig as CheckpointLoadingConfig,
-    )
     from nshtrainer.trainer._config import (
         CheckpointSavingConfig as CheckpointSavingConfig,
     )
@@ -91,10 +88,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).CallbackConfigBase
-        if name == "CheckpointLoadingConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).CheckpointLoadingConfig
         if name == "CheckpointSavingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
@@ -180,5 +173,4 @@ else:
 
 # Submodule exports
 from . import _config as _config
-from . import checkpoint_connector as checkpoint_connector
 from . import trainer as trainer

{nshtrainer-1.0.0b13 → nshtrainer-1.0.0b14}/src/nshtrainer/configs/trainer/_config/__init__.py

@@ -17,9 +17,6 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         CheckpointCallbackConfig as CheckpointCallbackConfig,
     )
-    from nshtrainer.trainer._config import (
-        CheckpointLoadingConfig as CheckpointLoadingConfig,
-    )
     from nshtrainer.trainer._config import (
         CheckpointSavingConfig as CheckpointSavingConfig,
     )
@@ -91,10 +88,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).CallbackConfigBase
-        if name == "CheckpointLoadingConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).CheckpointLoadingConfig
         if name == "CheckpointSavingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"

{nshtrainer-1.0.0b13 → nshtrainer-1.0.0b14}/src/nshtrainer/trainer/_config.py

@@ -32,7 +32,6 @@ from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.strategies.strategy import Strategy
 from typing_extensions import TypedDict, TypeVar, override
 
-from .._checkpoint.loader import CheckpointLoadingConfig
 from .._directory import DirectoryConfig
 from .._hf_hub import HuggingFaceHubConfig
 from ..callbacks import (
@@ -493,12 +492,6 @@ class TrainerConfig(C.Config):
     ckpt_path: Literal["none"] | str | Path | None = None
     """Path to a checkpoint to load and resume training from. If ``"none"``, will not load a checkpoint."""
 
-    checkpoint_loading: CheckpointLoadingConfig | Literal["auto", "none"] = "auto"
-    """Checkpoint loading configuration options.
-    `"auto"` will automatically determine the best checkpoint loading strategy based on the provided.
-    `"none"` will disable checkpoint loading.
-    """
-
     checkpoint_saving: CheckpointSavingConfig = CheckpointSavingConfig()
     """Checkpoint saving configuration options."""
 

{nshtrainer-1.0.0b13 → nshtrainer-1.0.0b14}/src/nshtrainer/trainer/trainer.py

@@ -29,7 +29,6 @@ from ._config import (
     TrainerConfig,
 )
 from ._runtime_callback import RuntimeTrackerCallback, Stage
-from .checkpoint_connector import _CheckpointConnector
 from .signal_connector import _SignalConnector
 
 log = logging.getLogger(__name__)
@@ -314,9 +313,6 @@ class Trainer(LightningTrainer):
         # Replace the signal connector with our own.
         self._signal_connector = _SignalConnector(self)
 
-        # Replace the checkpoint connector with our own.
-        self._checkpoint_connector = _CheckpointConnector(self)
-
         # Print out the log dir, so that we can easily find it in the logs.
         if log_dir := self.log_dir:
             log_dir = str(Path(log_dir).resolve())
@@ -441,19 +437,13 @@ class Trainer(LightningTrainer):
     ):
         filepath = Path(filepath)
 
-        # List of files that we should upload to HF
-        written_files: list[Path] = [filepath]
-
         super().save_checkpoint(filepath, weights_only, storage_options)
 
         # Save the checkpoint metadata
         metadata_path = None
         if self.hparams.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
-            if (
-                metadata_path := _write_checkpoint_metadata(self, filepath)
-            ) is not None:
-                written_files.append(metadata_path)
+            metadata_path = _write_checkpoint_metadata(self, filepath)
 
         # Call the `on_checkpoint_saved` method on all callbacks
         from .. import _callback

nshtrainer-1.0.0b13/src/nshtrainer/_checkpoint/loader.py

@@ -1,387 +0,0 @@
-from __future__ import annotations
-
-import logging
-from collections.abc import Iterable, Sequence
-from dataclasses import dataclass
-from pathlib import Path
-from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias, overload
-
-import nshconfig as C
-from lightning.pytorch.trainer.states import TrainerFn
-from typing_extensions import assert_never
-
-from ..metrics._config import MetricConfig
-from .metadata import CheckpointMetadata
-
-if TYPE_CHECKING:
-    from ..trainer import Trainer
-    from ..trainer._config import TrainerConfig
-
-log = logging.getLogger(__name__)
-
-
-class BestCheckpointStrategyConfig(C.Config):
-    name: Literal["best"] = "best"
-
-    metric: MetricConfig | None = None
-    """The metric to use for selecting the best checkpoint. If `None`, the primary metric will be used."""
-
-    additional_candidates: Iterable[Path] = []
-    """Additional checkpoint candidates to consider when selecting the last checkpoint."""
-
-
-class UserProvidedPathCheckpointStrategyConfig(C.Config):
-    name: Literal["user_provided_path"] = "user_provided_path"
-
-    path: Path
-    """The path to the checkpoint to load."""
-
-    on_error: Literal["warn", "raise"] = "warn"
-    """The behavior when the checkpoint does not belong to the current run.
-
-    - `warn`: Log a warning and skip the checkpoint.
-    - `raise`: Raise an error.
-    """
-
-
-class LastCheckpointStrategyConfig(C.Config):
-    name: Literal["last"] = "last"
-
-    criterion: Literal["global_step", "runtime"] = "global_step"
-    """The criterion to use for selecting the last checkpoint.
-
-    - `global_step`: The checkpoint with the highest global step will be selected.
-    - `runtime`: The checkpoint with the highest runtime will be selected.
-    """
-
-    additional_candidates: Iterable[Path] = []
-    """Additional checkpoint candidates to consider when selecting the last checkpoint."""
-
-
-CheckpointLoadingStrategyConfig: TypeAlias = Annotated[
-    BestCheckpointStrategyConfig
-    | LastCheckpointStrategyConfig
-    | UserProvidedPathCheckpointStrategyConfig,
-    C.Field(discriminator="name"),
-]
-
-
-class CheckpointLoadingConfig(C.Config):
-    strategies: Sequence[CheckpointLoadingStrategyConfig]
-    """The strategies to use for loading checkpoints.
-
-    The order of the strategies determines the priority of the strategies.
-    The first strategy that resolves a checkpoint will be used.
-    """
-
-    include_hpc: bool
-    """Whether to include checkpoints from HPC pre-emption."""
-
-    @classmethod
-    def none(cls, include_hpc: bool = False):
-        return cls(strategies=[], include_hpc=include_hpc)
-
-    @classmethod
-    def _auto_train(cls, ckpt: Literal["best", "last", "none"] | str | Path | None):
-        if ckpt is None:
-            ckpt = "last"
-        match ckpt:
-            case "best":
-                return cls(
-                    strategies=[BestCheckpointStrategyConfig()],
-                    include_hpc=True,
-                )
-            case "last":
-                return cls(
-                    strategies=[LastCheckpointStrategyConfig()],
-                    include_hpc=True,
-                )
-            case "none":
-                return cls.none()
-            case Path() | str():
-                ckpt = Path(ckpt)
-                return cls(
-                    strategies=[
-                        LastCheckpointStrategyConfig(additional_candidates=[ckpt]),
-                        UserProvidedPathCheckpointStrategyConfig(path=ckpt),
-                    ],
-                    include_hpc=True,
-                )
-            case _:
-                assert_never(ckpt)
-
-    @classmethod
-    def _auto_eval(cls, ckpt: Literal["best", "last", "none"] | str | Path | None):
-        if ckpt is None:
-            log.warn("No checkpoint specified for evaluation. Defaulting to `last`.")
-            ckpt = "last"
-
-        match ckpt:
-            case "best":
-                return cls(
-                    strategies=[BestCheckpointStrategyConfig()],
-                    include_hpc=False,
-                )
-            case "last":
-                return cls(
-                    strategies=[LastCheckpointStrategyConfig()],
-                    include_hpc=False,
-                )
-            case "none":
-                return cls.none(include_hpc=False)
-            case Path() | str():
-                ckpt = Path(ckpt)
-                return cls(
-                    strategies=[UserProvidedPathCheckpointStrategyConfig(path=ckpt)],
-                    include_hpc=False,
-                )
-            case _:
-                assert_never(ckpt)
-
-    @classmethod
-    def auto(
-        cls,
-        ckpt: Literal["best", "last", "none"] | str | Path | None,
-        trainer_mode: TrainerFn,
-    ):
-        """
-        Automatically create a CheckpointLoadingConfig based on the provided checkpoint option and trainer mode.
-
-        This method provides a convenient way to generate a checkpoint loading configuration
-        tailored to different training and evaluation scenarios.
-
-        Parameters:
-        -----------
-        ckpt : Literal["best", "last", "none"] | str | Path | None
-            Specifies the checkpoint loading preference:
-            - "best": Use the best checkpoint based on the primary metric.
-            - "last": Use the most recent checkpoint.
-            - str or Path: Path to a specific checkpoint file.
-            - None: Defaults to "last" for training, raises an error for evaluation.
-
-        trainer_mode : TrainerFn
-            The mode in which the trainer is operating. This affects how the configuration is created.
-            - TrainerFn.FITTING: Used for training scenarios.
-            - TrainerFn.VALIDATING, TrainerFn.TESTING, TrainerFn.PREDICTING: Used for evaluation scenarios.
-
-        Returns:
-        --------
-        CheckpointLoadingConfig
-            A configuration object for checkpoint loading based on the given parameters.
-
-        Behavior:
-        ---------
-        1. For training (TrainerFn.FITTING):
-        - Includes HPC pre-emption checkpoints.
-        - If ckpt is None, defaults to "last".
-        - For "best" or "last", creates a single-strategy configuration that loads the best or last checkpoint.
-        - For a specific path, creates a two-strategy configuration:
-            a) Tries to load the checkpoint as the last checkpoint.
-            b) Falls back to loading it as a user-provided path.
-
-        2. For evaluation (VALIDATING, TESTING, PREDICTING):
-        - Does not include HPC pre-emption checkpoints.
-        - Requires ckpt to be specified (raises ValueError if None).
-        - Creates a single-strategy configuration based on the ckpt value.
-
-        Raises:
-        -------
-        ValueError
-            If ckpt is None during evaluation modes.
-
-        Examples:
-        ---------
-        # Training mode, use last checkpoint
-        config = CheckpointLoadingConfig.auto("last", TrainerFn.FITTING)
-
-        # Evaluation mode, use best checkpoint
-        config = CheckpointLoadingConfig.auto("best", TrainerFn.TESTING)
-
-        # Training mode, use specific checkpoint
-        config = CheckpointLoadingConfig.auto("/path/to/checkpoint.ckpt", TrainerFn.FITTING)
-
-        Notes:
-        ------
-        - The method internally calls _auto_train or _auto_eval based on the trainer_mode.
-        - The resulting configuration always includes strategies as a sequence, even if there's only one strategy.
-        """
-        # Implementation remains the same...
-        match trainer_mode:
-            case TrainerFn.FITTING:
-                return cls._auto_train(ckpt)
-            case TrainerFn.VALIDATING | TrainerFn.TESTING | TrainerFn.PREDICTING:
-                return cls._auto_eval(ckpt)
-            case _:
-                assert_never(trainer_mode)
-
-
-@dataclass
-class _CkptCandidate:
-    meta: CheckpointMetadata
-    meta_path: Path
-
-    @property
-    def ckpt_path(self):
-        return self.meta_path.with_name(self.meta.checkpoint_filename)
-
-
-@overload
-def _load_ckpt_meta(
-    path: Path,
-    trainer_config: TrainerConfig,
-    on_error: Literal["warn"] = "warn",
-) -> _CkptCandidate | None: ...
-@overload
-def _load_ckpt_meta(
-    path: Path,
-    trainer_config: TrainerConfig,
-    on_error: Literal["raise"],
-) -> _CkptCandidate: ...
-def _load_ckpt_meta(
-    path: Path,
-    trainer_config: TrainerConfig,
-    on_error: Literal["warn", "raise"] = "warn",
-):
-    meta = CheckpointMetadata.from_file(path)
-    if trainer_config.id != meta.run_id:
-        error_msg = f"Skipping checkpoint {path} because it belongs to a different run"
-        match on_error:
-            case "warn":
-                log.warning(error_msg)
-            case "raise":
-                raise ValueError(error_msg)
-            case _:
-                assert_never(on_error)
-        return None
-    return _CkptCandidate(meta, path)
-
-
-def _checkpoint_candidates(trainer: Trainer, *, include_hpc: bool = True):
-    # Load the checkpoint directory, and throw if it doesn't exist.
-    # This indicates a non-standard setup, and we don't want to guess
-    # where the checkpoints are.
-    ckpt_dir = trainer.hparams.directory.resolve_subdirectory(
-        trainer.hparams.id, "checkpoint"
-    )
-    if not ckpt_dir.is_dir():
-        raise FileNotFoundError(
-            f"Checkpoint directory {ckpt_dir} not found. "
-            "Please ensure that the checkpoint directory exists."
-        )
-
-    # Load all checkpoints in the directory.
-    # We can do this by looking for metadata files.
-    for path in ckpt_dir.glob(f"*{CheckpointMetadata.PATH_SUFFIX}"):
-        if (meta := _load_ckpt_meta(path, trainer.hparams)) is not None:
-            yield meta
-
-    # If we have a pre-empted checkpoint, load it
-    if include_hpc and (hpc_path := trainer._checkpoint_connector._hpc_resume_path):
-        hpc_meta_path = Path(hpc_path).with_suffix(CheckpointMetadata.PATH_SUFFIX)
-        if (meta := _load_ckpt_meta(hpc_meta_path, trainer.hparams)) is not None:
-            yield meta
-
-
-def _additional_candidates(
-    additional_candidates: Iterable[Path], trainer_config: TrainerConfig
-):
-    for path in additional_candidates:
-        if (
-            meta := _load_ckpt_meta(
-                path.with_suffix(CheckpointMetadata.PATH_SUFFIX), trainer_config
-            )
-        ) is None:
-            continue
-        yield meta
-
-
-def _resolve_checkpoint(config: CheckpointLoadingConfig, trainer: Trainer):
-    # We lazily load the checkpoint candidates to avoid loading them
-    # if they are not needed.
-    _ckpt_candidates: list[_CkptCandidate] | None = None
-
-    def ckpt_candidates():
-        nonlocal _ckpt_candidates, trainer
-
-        if _ckpt_candidates is None:
-            _ckpt_candidates = list(
-                _checkpoint_candidates(trainer, include_hpc=config.include_hpc)
-            )
-        return _ckpt_candidates
-
-    # Iterate over the strategies and try to resolve the checkpoint.
-    for strategy in config.strategies:
-        match strategy:
-            case UserProvidedPathCheckpointStrategyConfig():
-                meta = _load_ckpt_meta(
-                    strategy.path.with_suffix(CheckpointMetadata.PATH_SUFFIX),
-                    trainer.hparams,
-                    on_error=strategy.on_error,
-                )
-                if meta is None:
-                    continue
-                return meta.ckpt_path
-            case BestCheckpointStrategyConfig():
-                candidates = [
-                    *ckpt_candidates(),
-                    *_additional_candidates(
-                        strategy.additional_candidates, trainer.hparams
-                    ),
-                ]
-                if not candidates:
-                    log.warning(
-                        "No checkpoint candidates found for `best` checkpoint strategy."
-                    )
-                    continue
-
-                if (
-                    metric := strategy.metric or trainer.hparams.primary_metric
-                ) is None:
-                    log.warning(
-                        "No metric specified for `best` checkpoint strategy, "
-                        "and no primary metric is set in the configuration. "
-                        "Skipping strategy."
-                    )
-                    continue
-
-                # Find the best checkpoint based on the metric.
-                def metric_value(ckpt: _CkptCandidate):
-                    assert metric is not None
-                    if (
-                        value := ckpt.meta.metrics.get(metric.validation_monitor)
-                    ) is None:
-                        raise ValueError(
-                            f"Metric {metric.validation_monitor} not found in checkpoint metadata. "
-                            f"Available metrics: {ckpt.meta.metrics.keys()}"
-                        )
-                    return value
-
-                best_candidate = metric.best(candidates, key=metric_value)
-                return best_candidate.ckpt_path
-            case LastCheckpointStrategyConfig():
-                candidates = [
-                    *ckpt_candidates(),
-                    *_additional_candidates(
-                        strategy.additional_candidates, trainer.hparams
-                    ),
-                ]
-                if not candidates:
-                    log.warning(
-                        "No checkpoint candidates found for `last` checkpoint strategy."
-                    )
-                    continue
-
-                # Find the last checkpoint based on the criterion.
-                def criterion_value(ckpt: _CkptCandidate):
-                    match strategy.criterion:
-                        case "global_step":
-                            return ckpt.meta.global_step
-                        case "runtime":
-                            return ckpt.meta.training_time.total_seconds()
-                        case _:
-                            assert_never(strategy.criterion)
-
-                last_candidate = max(candidates, key=criterion_value)
-                return last_candidate.ckpt_path
-            case _:
-                assert_never(strategy)