kostyl-toolkit 0.1.22__tar.gz → 0.1.23__tar.gz

{kostyl_toolkit-0.1.22 → kostyl_toolkit-0.1.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.22
+Version: 0.1.23
 Summary: Kickass Orchestration System for Training, Yielding & Logging 
 Requires-Dist: case-converter>=1.2.0
 Requires-Dist: loguru>=0.7.3

{kostyl_toolkit-0.1.22 → kostyl_toolkit-0.1.23}/kostyl/ml/lightning/callbacks/__init__.py

@@ -1,6 +1,6 @@
 from .checkpoint import setup_checkpoint_callback
 from .early_stopping import setup_early_stopping_callback
-from .registry_uploading import ClearMLRegistryUploaderCallback
+from .registry_uploader import ClearMLRegistryUploaderCallback
 
 
 __all__ = [

kostyl_toolkit-0.1.23/kostyl/ml/lightning/callbacks/checkpoint.py

@@ -0,0 +1,327 @@
+from datetime import timedelta
+from pathlib import Path
+from shutil import rmtree
+from typing import Literal
+from typing import cast
+
+from lightning.fabric.utilities.types import _PATH
+from lightning.pytorch.callbacks import ModelCheckpoint
+
+from kostyl.ml.configs import CheckpointConfig
+from kostyl.ml.dist_utils import is_main_process
+from kostyl.utils import setup_logger
+
+from .registry_uploader import RegistryUploaderCallback
+
+
+logger = setup_logger("callbacks/checkpoint.py")
+
+
+class CustomModelCheckpoint(ModelCheckpoint):
+    r"""
+    Save the model after every epoch by monitoring a quantity. Every logged metrics are passed to the
+    :class:`~lightning.pytorch.loggers.logger.Logger` for the version it gets saved in the same directory as the
+    checkpoint.
+
+    After training finishes, use :attr:`best_model_path` to retrieve the path to the
+    best checkpoint file and :attr:`best_model_score` to get its score.
+
+    .. note::
+        When using manual optimization with ``every_n_train_steps``, you should save the model state
+        in your ``training_step`` before the optimizer step if you want the checkpoint to reflect
+        the pre-optimization state. Example:
+
+        .. code-block:: python
+
+            def training_step(self, batch, batch_idx):
+                # ... forward pass, loss calculation, backward pass ...
+
+                # Save model state before optimization
+                if not hasattr(self, 'saved_models'):
+                    self.saved_models = {}
+                self.saved_models[batch_idx] = {
+                    k: v.detach().clone()
+                    for k, v in self.layer.state_dict().items()
+                }
+
+                # Then perform optimization
+                optimizer.zero_grad()
+                self.manual_backward(loss)
+                optimizer.step()
+
+                # Optional: Clean up old states to save memory
+                if batch_idx > 10:  # Keep last 10 states
+                    del self.saved_models[batch_idx - 10]
+
+    Args:
+        dirpath: Directory to save the model file.
+            Example: ``dirpath='my/path/'``.
+
+            .. warning::
+                In a distributed environment like DDP, it's recommended to provide a `dirpath` to avoid race conditions.
+                When using manual optimization with ``every_n_train_steps``, make sure to save the model state
+                in your training loop as shown in the example above.
+
+            Can be remote file paths such as `s3://mybucket/path/` or 'hdfs://path/'
+            (default: ``None``). If dirpath is ``None``, we only keep the ``k`` best checkpoints
+            in memory, and do not save anything to disk.
+
+        filename: Checkpoint filename. Can contain named formatting options to be auto-filled.
+            If no name is provided, it will be ``None`` and the checkpoint will be saved to
+            ``{epoch}``.and if the Trainer uses a logger, the path will also contain logger name and version.
+
+        filename: checkpoint filename. Can contain named formatting options to be auto-filled.
+
+            Example::
+
+                # save any arbitrary metrics like `val_loss`, etc. in name
+                # saves a file like: my/path/epoch=2-val_loss=0.02-other_metric=0.03.ckpt
+                >>> checkpoint_callback = ModelCheckpoint(
+                ...     dirpath='my/path',
+                ...     filename='{epoch}-{val_loss:.2f}-{other_metric:.2f}'
+                ... )
+
+            By default, filename is ``None`` and will be set to ``'{epoch}-{step}'``, where "epoch" and "step" match
+            the number of finished epoch and optimizer steps respectively.
+        monitor: quantity to monitor. By default it is ``None`` which saves a checkpoint only for the last epoch.
+        verbose: verbosity mode. Default: ``False``.
+        save_last: When ``True``, saves a `last.ckpt` copy whenever a checkpoint file gets saved. Can be set to
+            ``'link'`` on a local filesystem to create a symbolic link. This allows accessing the latest checkpoint
+            in a deterministic manner. Default: ``None``.
+        save_top_k: if ``save_top_k == k``,
+            the best k models according to the quantity monitored will be saved.
+            If ``save_top_k == 0``, no models are saved.
+            If ``save_top_k == -1``, all models are saved.
+            Please note that the monitors are checked every ``every_n_epochs`` epochs.
+            If ``save_top_k >= 2`` and the callback is called multiple times inside an epoch, and the filename remains
+            unchanged, the name of the saved file will be appended with a version count starting with ``v1`` to avoid
+            collisions unless ``enable_version_counter`` is set to False. The version counter is unrelated to the top-k
+            ranking of the checkpoint, and we recommend formatting the filename to include the monitored metric to avoid
+            collisions.
+        save_on_exception: Whether to save a checkpoint when an exception is raised. Default: ``False``.
+        mode: one of {min, max}.
+            If ``save_top_k != 0``, the decision to overwrite the current save file is made
+            based on either the maximization or the minimization of the monitored quantity.
+            For ``'val_acc'``, this should be ``'max'``, for ``'val_loss'`` this should be ``'min'``, etc.
+        auto_insert_metric_name: When ``True``, the checkpoints filenames will contain the metric name.
+            For example, ``filename='checkpoint_{epoch:02d}-{acc:02.0f}`` with epoch ``1`` and acc ``1.12`` will resolve
+            to ``checkpoint_epoch=01-acc=01.ckpt``. Is useful to set it to ``False`` when metric names contain ``/``
+            as this will result in extra folders.
+            For example, ``filename='epoch={epoch}-step={step}-val_acc={val/acc:.2f}', auto_insert_metric_name=False``
+        save_weights_only: if ``True``, then only the model's weights will be
+            saved. Otherwise, the optimizer states, lr-scheduler states, etc are added in the checkpoint too.
+        every_n_train_steps: How many training steps to wait before saving a checkpoint. This does not take into account
+            the steps of the current epoch. If ``every_n_train_steps == None or every_n_train_steps == 0``,
+            no checkpoints
+            will be saved during training. Mutually exclusive with ``train_time_interval`` and ``every_n_epochs``.
+
+            .. note::
+                When using with manual optimization, the checkpoint will be saved after the optimizer step by default.
+                To save the model state before the optimizer step, you need to save the model state in your
+                ``training_step`` before calling ``optimizer.step()``. See the class docstring for an example.
+        train_time_interval: Checkpoints are monitored at the specified time interval.
+            For all practical purposes, this cannot be smaller than the amount
+            of time it takes to process a single training batch. This is not
+            guaranteed to execute at the exact time specified, but should be close.
+            This must be mutually exclusive with ``every_n_train_steps`` and ``every_n_epochs``.
+        every_n_epochs: Number of epochs between checkpoints.
+            This value must be ``None`` or non-negative.
+            To disable saving top-k checkpoints, set ``every_n_epochs = 0``.
+            This argument does not impact the saving of ``save_last=True`` checkpoints.
+            If all of ``every_n_epochs``, ``every_n_train_steps`` and
+            ``train_time_interval`` are ``None``, we save a checkpoint at the end of every epoch
+            (equivalent to ``every_n_epochs = 1``).
+            If ``every_n_epochs == None`` and either ``every_n_train_steps != None`` or ``train_time_interval != None``,
+            saving at the end of each epoch is disabled
+            (equivalent to ``every_n_epochs = 0``).
+            This must be mutually exclusive with ``every_n_train_steps`` and ``train_time_interval``.
+            Setting both ``ModelCheckpoint(..., every_n_epochs=V, save_on_train_epoch_end=False)`` and
+            ``Trainer(max_epochs=N, check_val_every_n_epoch=M)``
+            will only save checkpoints at epochs 0 < E <= N
+            where both values for ``every_n_epochs`` and ``check_val_every_n_epoch`` evenly divide E.
+        save_on_train_epoch_end: Whether to run checkpointing at the end of the training epoch.
+            If ``True``, checkpoints are saved at the end of every training epoch.
+            If ``False``, checkpoints are saved at the end of validation.
+            If ``None`` (default), checkpointing behavior is determined based on training configuration.
+            If ``val_check_interval`` is a str, dict, or `timedelta` (time-based), checkpointing is performed after
+            validation.
+            If ``check_val_every_n_epoch != 1``, checkpointing will not be performed at the end of
+            every training epoch. If there are no validation batches of data, checkpointing will occur at the
+            end of the training epoch. If there is a non-default number of validation runs per training epoch
+            (``val_check_interval != 1``), checkpointing is performed after validation.
+        enable_version_counter: Whether to append a version to the existing file name.
+            If ``False``, then the checkpoint files will be overwritten.
+
+    Note:
+        For extra customization, ModelCheckpoint includes the following attributes:
+
+        - ``CHECKPOINT_JOIN_CHAR = "-"``
+        - ``CHECKPOINT_EQUALS_CHAR = "="``
+        - ``CHECKPOINT_NAME_LAST = "last"``
+        - ``FILE_EXTENSION = ".ckpt"``
+        - ``STARTING_VERSION = 1``
+
+        For example, you can change the default last checkpoint name by doing
+        ``checkpoint_callback.CHECKPOINT_NAME_LAST = "{epoch}-last"``
+
+        If you want to checkpoint every N hours, every M train batches, and/or every K val epochs,
+        then you should create multiple ``ModelCheckpoint`` callbacks.
+
+        If the checkpoint's ``dirpath`` changed from what it was before while resuming the training,
+        only ``best_model_path`` will be reloaded and a warning will be issued.
+
+        If you provide a ``filename`` on a mounted device where changing permissions is not allowed (causing ``chmod``
+        to raise a ``PermissionError``), install `fsspec>=2025.5.0`. Then the error is caught, the file's permissions
+        remain unchanged, and the checkpoint is still saved. Otherwise, no checkpoint will be saved and training stops.
+
+    Raises:
+        MisconfigurationException:
+            If ``save_top_k`` is smaller than ``-1``,
+            if ``monitor`` is ``None`` and ``save_top_k`` is none of ``None``, ``-1``, and ``0``, or
+            if ``mode`` is none of ``"min"`` or ``"max"``.
+        ValueError:
+            If ``trainer.save_checkpoint`` is ``None``.
+
+    Example::
+
+        >>> from lightning.pytorch import Trainer
+        >>> from lightning.pytorch.callbacks import ModelCheckpoint
+
+        # saves checkpoints to 'my/path/' at every epoch
+        >>> checkpoint_callback = ModelCheckpoint(dirpath='my/path/')
+        >>> trainer = Trainer(callbacks=[checkpoint_callback])
+
+        # save epoch and val_loss in name
+        # saves a file like: my/path/sample-mnist-epoch=02-val_loss=0.32.ckpt
+        >>> checkpoint_callback = ModelCheckpoint(
+        ...     monitor='val_loss',
+        ...     dirpath='my/path/',
+        ...     filename='sample-mnist-{epoch:02d}-{val_loss:.2f}'
+        ... )
+
+        # save epoch and val_loss in name, but specify the formatting yourself (e.g. to avoid problems with Tensorboard
+        # or Neptune, due to the presence of characters like '=' or '/')
+        # saves a file like: my/path/sample-mnist-epoch02-val_loss0.32.ckpt
+        >>> checkpoint_callback = ModelCheckpoint(
+        ...     monitor='val/loss',
+        ...     dirpath='my/path/',
+        ...     filename='sample-mnist-epoch{epoch:02d}-val_loss{val/loss:.2f}',
+        ...     auto_insert_metric_name=False
+        ... )
+
+        # retrieve the best checkpoint after training
+        >>> checkpoint_callback = ModelCheckpoint(dirpath='my/path/')
+        >>> trainer = Trainer(callbacks=[checkpoint_callback])
+        >>> model = ...  # doctest: +SKIP
+        >>> trainer.fit(model)  # doctest: +SKIP
+        >>> print(checkpoint_callback.best_model_path)  # doctest: +SKIP
+
+    .. tip:: Saving and restoring multiple checkpoint callbacks at the same time is supported under variation in the
+        following arguments:
+
+        *monitor, mode, every_n_train_steps, every_n_epochs, train_time_interval*
+
+        Read more: :ref:`Persisting Callback State <extensions/callbacks_state:save callback state>`
+
+    """  # noqa: D205
+
+    def __init__(  # noqa: D107
+        self,
+        dirpath: _PATH | None = None,
+        filename: str | None = None,
+        monitor: str | None = None,
+        verbose: bool = False,
+        save_last: bool | Literal["link"] | None = None,
+        save_top_k: int = 1,
+        save_on_exception: bool = False,
+        save_weights_only: bool = False,
+        mode: str = "min",
+        auto_insert_metric_name: bool = True,
+        every_n_train_steps: int | None = None,
+        train_time_interval: timedelta | None = None,
+        every_n_epochs: int | None = None,
+        save_on_train_epoch_end: bool | None = None,
+        enable_version_counter: bool = True,
+        registry_uploader_callback: RegistryUploaderCallback | None = None,
+    ) -> None:
+        super().__init__(
+            dirpath=dirpath,
+            filename=filename,
+            monitor=monitor,
+            verbose=verbose,
+            save_last=save_last,
+            save_top_k=save_top_k,
+            save_on_exception=save_on_exception,
+            save_weights_only=save_weights_only,
+            mode=mode,
+            auto_insert_metric_name=auto_insert_metric_name,
+            every_n_train_steps=every_n_train_steps,
+            train_time_interval=train_time_interval,
+            every_n_epochs=every_n_epochs,
+            save_on_train_epoch_end=save_on_train_epoch_end,
+            enable_version_counter=enable_version_counter,
+        )
+        self.registry_uploader_callback = registry_uploader_callback
+        self._custom_best_model_path = cast(str, self.best_model_path)
+        return
+
+    @property
+    def best_model_path(self) -> str:
+        """Best model path."""
+        return self._custom_best_model_path
+
+    @best_model_path.setter
+    def best_model_path(self, value: str) -> None:
+        self._custom_best_model_path = value
+        if self.registry_uploader_callback is not None:
+            self.registry_uploader_callback.best_model_path = value
+        return
+
+
+def setup_checkpoint_callback(
+    dirpath: Path,
+    ckpt_cfg: CheckpointConfig,
+    save_weights_only: bool = True,
+    registry_uploader_callback: RegistryUploaderCallback | None = None,
+) -> CustomModelCheckpoint:
+    """
+    Sets up a ModelCheckpoint callback for PyTorch Lightning.
+
+    This function prepares a checkpoint directory and configures a ModelCheckpoint
+    callback based on the provided configuration. If the directory already exists,
+    it is removed (only by the main process) to ensure a clean start. Otherwise,
+    the directory is created.
+
+    Args:
+        dirpath (Path): The path to the directory where checkpoints will be saved.
+        ckpt_cfg (CheckpointConfig): Configuration object containing checkpoint
+            settings such as filename, save_top_k, monitor, and mode.
+        save_weights_only (bool, optional): Whether to save only the model weights
+            or the full model. Defaults to True.
+        registry_uploader_callback (RegistryUploaderCallback | None, optional):
+            An optional callback for uploading checkpoints to a registry. Defaults to None.
+
+    Returns:
+        ModelCheckpoint: The configured ModelCheckpoint callback instance.
+
+    """
+    if dirpath.exists():
+        if is_main_process():
+            logger.warning(f"Checkpoint directory {dirpath} already exists.")
+            rmtree(dirpath)
+            logger.warning(f"Removed existing checkpoint directory {dirpath}.")
+    else:
+        logger.info(f"Creating checkpoint directory {dirpath}.")
+        dirpath.mkdir(parents=True, exist_ok=True)
+
+    checkpoint_callback = CustomModelCheckpoint(
+        dirpath=dirpath,
+        filename=ckpt_cfg.filename,
+        save_top_k=ckpt_cfg.save_top_k,
+        monitor=ckpt_cfg.monitor,
+        mode=ckpt_cfg.mode,
+        verbose=True,
+        save_weights_only=save_weights_only,
+        registry_uploader_callback=registry_uploader_callback,
+    )
+    return checkpoint_callback

kostyl_toolkit-0.1.22/kostyl/ml/lightning/callbacks/registry_uploading.py → kostyl_toolkit-0.1.23/kostyl/ml/lightning/callbacks/registry_uploader.py

@@ -1,3 +1,7 @@
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import Callable
+from functools import partial
 from typing import Literal
 from typing import override
 
@@ -5,7 +9,6 @@ from clearml import OutputModel
 from clearml import Task
 from lightning import Trainer
 from lightning.pytorch.callbacks import Callback
-from lightning.pytorch.callbacks import ModelCheckpoint
 
 from kostyl.ml.clearml.logging_utils import find_version_in_tags
 from kostyl.ml.clearml.logging_utils import increment_version
@@ -16,13 +19,31 @@ from kostyl.utils.logging import setup_logger
 logger = setup_logger()
 
 
-class ClearMLRegistryUploaderCallback(Callback):
+class RegistryUploaderCallback(Callback, ABC):
+    """Abstract Lightning callback responsible for tracking and uploading the best-performing model checkpoint."""
+
+    @property
+    @abstractmethod
+    def best_model_path(self) -> str:
+        """Return the file system path pointing to the best model artifact produced during training."""
+        raise NotImplementedError
+
+    @best_model_path.setter
+    @abstractmethod
+    def best_model_path(self, value: str) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _upload_best_checkpoint(self, pl_module: "KostylLightningModule") -> None:
+        raise NotImplementedError
+
+
+class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
     """PyTorch Lightning callback to upload the best model checkpoint to ClearML."""
 
     def __init__(
         self,
         task: Task,
-        ckpt_callback: ModelCheckpoint,
         output_model_name: str,
         output_model_tags: list[str] | None = None,
         verbose: bool = True,
@@ -56,7 +77,6 @@ class ClearMLRegistryUploaderCallback(Callback):
             output_model_tags = []
 
         self.task = task
-        self.ckpt_callback = ckpt_callback
         self.output_model_name = output_model_name
         self.output_model_tags = output_model_tags
         self.config_dict = config_dict
@@ -66,7 +86,22 @@ class ClearMLRegistryUploaderCallback(Callback):
         self.enable_tag_versioning = enable_tag_versioning
 
         self._output_model: OutputModel | None = None
-        self._last_best_model_path: str = ""
+        self._last_uploaded_model_path: str = ""
+        self._best_model_path: str = ""
+        self._upload_callback: Callable | None = None
+        return
+
+    @property
+    @override
+    def best_model_path(self) -> str:
+        return self._best_model_path
+
+    @best_model_path.setter
+    @override
+    def best_model_path(self, value: str) -> None:
+        self._best_model_path = value
+        if self._upload_callback is not None:
+            self._upload_callback()
         return
 
     def _create_output_model(self, pl_module: "KostylLightningModule") -> OutputModel:
@@ -98,27 +133,29 @@ class ClearMLRegistryUploaderCallback(Callback):
             label_enumeration=self.label_enumeration,
         )
 
+    @override
     def _upload_best_checkpoint(self, pl_module: "KostylLightningModule") -> None:
-        current_best = self.ckpt_callback.best_model_path
-
-        if not current_best:
-            if self.verbose:
-                logger.info("No best model found yet to upload")
-            return
-
-        if current_best == self._last_best_model_path:
-            if self.verbose:
-                logger.info("Best model unchanged since last upload")
+        if not self._best_model_path or (
+            self._best_model_path == self._last_uploaded_model_path
+        ):
+            if not self._best_model_path:
+                if self.verbose:
+                    logger.info("No best model found yet to upload")
+            elif self._best_model_path == self._last_uploaded_model_path:
+                if self.verbose:
+                    logger.info("Best model unchanged since last upload")
+            self._upload_callback = partial(self._upload_best_checkpoint, pl_module)
             return
+        self._upload_callback = None
 
         if self._output_model is None:
             self._output_model = self._create_output_model(pl_module)
 
         if self.verbose:
-            logger.info(f"Uploading best model from {current_best}")
+            logger.info(f"Uploading best model from {self._best_model_path}")
 
         self._output_model.update_weights(
-            current_best,
+            self._best_model_path,
             auto_delete_file=False,
             async_enable=False,
         )
@@ -130,7 +167,7 @@ class ClearMLRegistryUploaderCallback(Callback):
             config = self.config_dict
         self._output_model.update_design(config_dict=config)
 
-        self._last_best_model_path = current_best
+        self._last_uploaded_model_path = self._best_model_path
         return
 
     @override

{kostyl_toolkit-0.1.22 → kostyl_toolkit-0.1.23}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kostyl-toolkit"
-version = "0.1.22"
+version = "0.1.23"
 description = "Kickass Orchestration System for Training, Yielding & Logging "
 readme = "README.md"
 requires-python = ">=3.12"

kostyl_toolkit-0.1.22/kostyl/ml/lightning/callbacks/checkpoint.py

@@ -1,56 +0,0 @@
-from pathlib import Path
-from shutil import rmtree
-
-from lightning.pytorch.callbacks import ModelCheckpoint
-
-from kostyl.ml.configs import CheckpointConfig
-from kostyl.ml.dist_utils import is_main_process
-from kostyl.utils import setup_logger
-
-
-logger = setup_logger("callbacks/checkpoint.py")
-
-
-def setup_checkpoint_callback(
-    dirpath: Path,
-    ckpt_cfg: CheckpointConfig,
-    save_weights_only: bool = True,
-) -> ModelCheckpoint:
-    """
-    Sets up a ModelCheckpoint callback for PyTorch Lightning.
-
-    This function prepares a checkpoint directory and configures a ModelCheckpoint
-    callback based on the provided configuration. If the directory already exists,
-    it is removed (only by the main process) to ensure a clean start. Otherwise,
-    the directory is created.
-
-    Args:
-        dirpath (Path): The path to the directory where checkpoints will be saved.
-        ckpt_cfg (CheckpointConfig): Configuration object containing checkpoint
-            settings such as filename, save_top_k, monitor, and mode.
-        save_weights_only (bool, optional): Whether to save only the model weights
-            or the full model. Defaults to True.
-
-    Returns:
-        ModelCheckpoint: The configured ModelCheckpoint callback instance.
-
-    """
-    if dirpath.exists():
-        if is_main_process():
-            logger.warning(f"Checkpoint directory {dirpath} already exists.")
-            rmtree(dirpath)
-            logger.warning(f"Removed existing checkpoint directory {dirpath}.")
-    else:
-        logger.info(f"Creating checkpoint directory {dirpath}.")
-        dirpath.mkdir(parents=True, exist_ok=True)
-
-    checkpoint_callback = ModelCheckpoint(
-        dirpath=dirpath,
-        filename=ckpt_cfg.filename,
-        save_top_k=ckpt_cfg.save_top_k,
-        monitor=ckpt_cfg.monitor,
-        mode=ckpt_cfg.mode,
-        verbose=True,
-        save_weights_only=save_weights_only,
-    )
-    return checkpoint_callback