kostyl-toolkit 0.1.25__py3-none-any.whl → 0.1.27__py3-none-any.whl

kostyl/ml/configs/hyperparams.py

@@ -29,15 +29,11 @@ class Lr(BaseModel):
     @model_validator(mode="after")
     def validate_warmup(self) -> "Lr":
         """Validates the warmup parameters based on use_scheduler."""
-        if (self.warmup_value is None) != (
-            self.warmup_iters_ratio is None
-        ) and self.use_scheduler:
+        if (self.warmup_value is None) != (self.warmup_iters_ratio is None):  # fmt: skip
             raise ValueError(
                 "Both warmup_value and warmup_iters_ratio must be provided or neither"
             )
-        elif (
-            (self.warmup_value is not None) or (self.warmup_iters_ratio is not None)
-        ) and (not self.use_scheduler):
+        if ((self.warmup_value is not None) or (self.warmup_iters_ratio is not None)) and not self.use_scheduler:  # fmt: skip
             logger.warning(
                 "use_scheduler is False, warmup_value and warmup_iters_ratio will be ignored."
             )

kostyl/ml/lightning/callbacks/checkpoint.py

@@ -2,12 +2,16 @@ from datetime import timedelta
 from pathlib import Path
 from shutil import rmtree
 from typing import Literal
+from typing import override
 
+import lightning.pytorch as pl
+import torch.distributed as dist
 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.ml.lightning import KostylLightningModule
 from kostyl.utils import setup_logger
 
 from .registry_uploader import RegistryUploaderCallback
@@ -16,7 +20,7 @@ from .registry_uploader import RegistryUploaderCallback
 logger = setup_logger("callbacks/checkpoint.py")
 
 
-class CustomModelCheckpoint(ModelCheckpoint):
+class ModelCheckpointWithRegistryUploader(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
@@ -226,6 +230,8 @@ class CustomModelCheckpoint(ModelCheckpoint):
 
     def __init__(  # noqa: D107
         self,
+        registry_uploader_callback: RegistryUploaderCallback,
+        uploading_mode: Literal["only-best", "every-checkpoint"] = "only-best",
         dirpath: _PATH | None = None,
         filename: str | None = None,
         monitor: str | None = None,
@@ -241,10 +247,10 @@ class CustomModelCheckpoint(ModelCheckpoint):
         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:
         self.registry_uploader_callback = registry_uploader_callback
-        self._custom_best_model_path = ""
+        self.process_group: dist.ProcessGroup | None = None
+        self.uploading_mode = uploading_mode
         super().__init__(
             dirpath=dirpath,
             filename=filename,
@@ -264,16 +270,30 @@ class CustomModelCheckpoint(ModelCheckpoint):
         )
         return
 
-    @property
-    def best_model_path(self) -> str:
-        """Best model path."""
-        return self._custom_best_model_path
+    @override
+    def setup(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule | KostylLightningModule,
+        stage: str,
+    ) -> None:
+        super().setup(trainer, pl_module, stage)
+        if isinstance(pl_module, KostylLightningModule):
+            self.process_group = pl_module.get_process_group()
+        return
 
-    @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
+    @override
+    def _save_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
+        super()._save_checkpoint(trainer, filepath)
+        if dist.is_initialized():
+            dist.barrier(group=self.process_group)
+        if trainer.is_global_zero and self.registry_uploader_callback is not None:
+            match self.uploading_mode:
+                case "every-checkpoint":
+                    self.registry_uploader_callback.upload_checkpoint(filepath)
+                case "only-best":
+                    if filepath == self.best_model_path:
+                        self.registry_uploader_callback.upload_checkpoint(filepath)
         return
 
 
@@ -282,28 +302,44 @@ def setup_checkpoint_callback(
     ckpt_cfg: CheckpointConfig,
     save_weights_only: bool = True,
     registry_uploader_callback: RegistryUploaderCallback | None = None,
-) -> CustomModelCheckpoint:
+    uploading_mode: Literal["only-best", "every-checkpoint"] | None = None,
+) -> ModelCheckpointWithRegistryUploader | ModelCheckpoint:
     """
-    Sets up a ModelCheckpoint callback for PyTorch Lightning.
+    Create and configure a checkpoint callback for model saving.
 
-    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.
+    Creates the checkpoint directory (removing existing one if present) and returns
+    a configured callback for saving models during training. When registry_uploader_callback
+    is provided, returns an extended version with support for uploading checkpoints to a remote registry.
 
     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.
+        dirpath: Path to the directory for saving checkpoints.
+        ckpt_cfg: Checkpoint configuration (filename, monitor, mode, save_top_k).
+        save_weights_only: If True, only model weights are saved without optimizer and lr-scheduler state.
+            Defaults to True.
+        registry_uploader_callback: Optional callback for uploading checkpoints to a remote registry.
+            Must be specified together with uploading_mode.
+        uploading_mode: Checkpoint upload mode:
+            - "only-best": only the best checkpoint is uploaded
+            - "every-checkpoint": every saved checkpoint is uploaded
+            Must be specified together with registry_uploader_callback.
 
     Returns:
-        ModelCheckpoint: The configured ModelCheckpoint callback instance.
+        ModelCheckpointWithRegistryUploader if registry_uploader_callback is provided,
+        otherwise standard ModelCheckpoint.
+
+    Raises:
+        ValueError: If only one of registry_uploader_callback or uploading_mode is None.
+
+    Note:
+        If the dirpath directory already exists, it will be removed and recreated
+        (only on the main process in distributed training).
 
     """
+    if (registry_uploader_callback is None) != (uploading_mode is None):
+        raise ValueError(
+            "Both registry_uploader_callback and uploading_mode must be provided or neither."
+        )
+
     if dirpath.exists():
         if is_main_process():
             logger.warning(f"Checkpoint directory {dirpath} already exists.")
@@ -313,14 +349,26 @@ def setup_checkpoint_callback(
         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,
-    )
+    if (registry_uploader_callback is not None) and (uploading_mode is not None):
+        checkpoint_callback = ModelCheckpointWithRegistryUploader(
+            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,
+            uploading_mode=uploading_mode,
+        )
+    else:
+        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

kostyl/ml/lightning/callbacks/registry_uploader.py

@@ -1,40 +1,26 @@
 from abc import ABC
 from abc import abstractmethod
 from collections.abc import Callable
-from functools import partial
-from typing import Literal
+from pathlib import Path
 from typing import override
 
 from clearml import OutputModel
 from clearml import Task
-from lightning import Trainer
-from lightning.pytorch.callbacks import Callback
 
 from kostyl.ml.clearml.logging_utils import find_version_in_tags
 from kostyl.ml.clearml.logging_utils import increment_version
-from kostyl.ml.lightning import KostylLightningModule
 from kostyl.utils.logging import setup_logger
 
 
 logger = setup_logger()
 
 
-class RegistryUploaderCallback(Callback, ABC):
+class RegistryUploaderCallback(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:
+    def upload_checkpoint(self, path: str | Path) -> None:
+        """Upload the checkpoint located at the given path to the configured registry backend."""
         raise NotImplementedError
 
 
@@ -50,9 +36,6 @@ class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
         enable_tag_versioning: bool = True,
         label_enumeration: dict[str, int] | None = None,
         config_dict: dict[str, str] | None = None,
-        uploading_frequency: Literal[
-            "after-every-eval", "on-train-end"
-        ] = "on-train-end",
     ) -> None:
         """
         Initializes the ClearMLRegistryUploaderCallback.
@@ -67,9 +50,6 @@ class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
             config_dict: Optional configuration dictionary to associate with the model.
             enable_tag_versioning: Whether to enable versioning in tags. If True,
                 the version tag (e.g., "v1.0") will be automatically incremented or if not present, added as "v1.0".
-            uploading_frequency: When to upload:
-                - "after-every-eval": after each validation phase.
-                - "on-train-end": once at the end of training.
 
         """
         super().__init__()
@@ -82,29 +62,16 @@ class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
         self.config_dict = config_dict
         self.label_enumeration = label_enumeration
         self.verbose = verbose
-        self.uploading_frequency = uploading_frequency
         self.enable_tag_versioning = enable_tag_versioning
 
+        self.best_model_path: str = ""
+
         self._output_model: OutputModel | None = None
         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:
+    def _create_output_model(self) -> OutputModel:
         if self.enable_tag_versioning:
             version = find_version_in_tags(self.output_model_tags)
             if version is None:
@@ -117,13 +84,6 @@ class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
         if "LightningCheckpoint" not in self.output_model_tags:
             self.output_model_tags.append("LightningCheckpoint")
 
-        if self.config_dict is None:
-            config = pl_module.model_config
-            if config is not None:
-                config = config.to_dict()
-        else:
-            config = self.config_dict
-
         return OutputModel(
             task=self.task,
             name=self.output_model_name,
@@ -134,60 +94,29 @@ class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
         )
 
     @override
-    def _upload_best_checkpoint(self, pl_module: "KostylLightningModule") -> None:
-        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)
+    def upload_checkpoint(
+        self,
+        path: str | Path,
+    ) -> None:
+        if isinstance(path, Path):
+            path = str(path)
+        if path == self._last_uploaded_model_path:
+            if self.verbose:
+                logger.info("Model unchanged since last upload")
             return
-        self._upload_callback = None
 
         if self._output_model is None:
-            self._output_model = self._create_output_model(pl_module)
+            self._output_model = self._create_output_model()
 
         if self.verbose:
-            logger.info(f"Uploading best model from {self._best_model_path}")
+            logger.info(f"Uploading model from {path}")
 
         self._output_model.update_weights(
-            self._best_model_path,
+            path,
             auto_delete_file=False,
             async_enable=False,
         )
-        if self.config_dict is None:
-            config = pl_module.model_config
-            if config is not None:
-                config = config.to_dict()
-        else:
-            config = self.config_dict
-        self._output_model.update_design(config_dict=config)
-
-        self._last_uploaded_model_path = self._best_model_path
-        return
-
-    @override
-    def on_validation_end(
-        self, trainer: Trainer, pl_module: "KostylLightningModule"
-    ) -> None:
-        if self.uploading_frequency != "after-every-eval":
-            return
-        if not trainer.is_global_zero:
-            return
-
-        self._upload_best_checkpoint(pl_module)
-        return
-
-    @override
-    def on_train_end(
-        self, trainer: Trainer, pl_module: "KostylLightningModule"
-    ) -> None:
-        if not trainer.is_global_zero:
-            return
+        self._output_model.update_design(config_dict=self.config_dict)
 
-        self._upload_best_checkpoint(pl_module)
+        self._last_uploaded_model_path = path
         return

kostyl/ml/schedulers/cosine.py

@@ -11,20 +11,18 @@ from .base import BaseScheduler
 class _CosineSchedulerCore(BaseScheduler):
     def __init__(
         self,
-        param_group_field: str,
-        total_iters: int,
+        param_name: str,
+        num_iters: int,
         base_value: float,
         final_value: float,
-        warmup_iters_ratio: float | None = None,
+        warmup_ratio: float | None = None,
         warmup_value: float | None = None,
         freeze_ratio: float | None = None,
     ) -> None:
-        if warmup_iters_ratio is not None:
-            if not (0 < warmup_iters_ratio < 1):
-                raise ValueError(
-                    f"Warmup ratio must be in (0, 1), got {warmup_iters_ratio}."
-                )
-        if (warmup_value is None) != (warmup_iters_ratio is None):
+        if warmup_ratio is not None:
+            if not (0 < warmup_ratio < 1):
+                raise ValueError(f"Warmup ratio must be in (0, 1), got {warmup_ratio}.")
+        if (warmup_value is None) != (warmup_ratio is None):
             raise ValueError(
                 "Both warmup_ratio and warmup_value must be provided or neither."
             )
@@ -32,12 +30,12 @@ class _CosineSchedulerCore(BaseScheduler):
             if not (0 < freeze_ratio < 1):
                 raise ValueError(f"Freeze ratio must be in (0, 1), got {freeze_ratio}.")
 
-        self.param_group_field = param_group_field
-        self.total_iters = total_iters
+        self.param_name = param_name
+        self.num_iters = num_iters
         self.base_value = base_value
         self.final_value = final_value
 
-        self.warmup_iters_ratio = warmup_iters_ratio
+        self.warmup_ratio = warmup_ratio
         self.warmup_value = warmup_value
 
         self.freeze_ratio = freeze_ratio
@@ -49,15 +47,15 @@ class _CosineSchedulerCore(BaseScheduler):
     def _create_scheduler(self) -> None:
         # Create freeze schedule
         if self.freeze_ratio is not None:
-            freeze_iters = int(self.total_iters * self.freeze_ratio)
+            freeze_iters = int(self.num_iters * self.freeze_ratio)
             freeze_schedule = np.zeros(freeze_iters, dtype=np.float64)
         else:
             freeze_iters = 0
             freeze_schedule = np.array([], dtype=np.float64)
 
         # Create linear warmup schedule
-        if self.warmup_iters_ratio is not None and self.warmup_value is not None:
-            warmup_iters = int(self.total_iters * self.warmup_iters_ratio)
+        if self.warmup_ratio is not None and self.warmup_value is not None:
+            warmup_iters = int(self.num_iters * self.warmup_ratio)
             warmup_schedule = np.linspace(
                 self.warmup_value, self.base_value, warmup_iters, dtype=np.float64
             )
@@ -65,7 +63,7 @@ class _CosineSchedulerCore(BaseScheduler):
             warmup_iters = 0
             warmup_schedule = np.array([], dtype=np.float64)
 
-        cosine_annealing_iters = self.total_iters - warmup_iters - freeze_iters
+        cosine_annealing_iters = self.num_iters - warmup_iters - freeze_iters
         if cosine_annealing_iters <= 0:
             raise ValueError("Cosine annealing iters must be > 0.")
 
@@ -80,9 +78,9 @@ class _CosineSchedulerCore(BaseScheduler):
             (freeze_schedule, warmup_schedule, schedule)
         )
 
-        if len(self.scheduler_values) != self.total_iters:
+        if len(self.scheduler_values) != self.num_iters:
             raise ValueError(
-                f"Scheduler length ({len(self.scheduler_values)}) does not match total_iters ({self.total_iters})."
+                f"Scheduler length ({len(self.scheduler_values)}) does not match num_iters ({self.num_iters})."
             )
         return
 
@@ -100,7 +98,7 @@ class _CosineSchedulerCore(BaseScheduler):
         if len(self.scheduler_values) == 0:
             self._create_scheduler()
 
-        if it >= self.total_iters:
+        if it >= self.num_iters:
             value: float = self.final_value
         else:
             value: float = self.scheduler_values[it]
@@ -109,20 +107,20 @@ class _CosineSchedulerCore(BaseScheduler):
 
     @override
     def current_value(self) -> dict[str, float]:
-        return {self.param_group_field: self.current_value_}
+        return {self.param_name: self.current_value_}
 
 
 class CosineScheduler(_CosineSchedulerCore):
-    """Implements a cosine scheduler for adjusting parameter values in torch.optim.Optimizer."""
+    """Applies a cosine schedule to an optimizer param-group field."""
 
     def __init__(
         self,
         optimizer: torch.optim.Optimizer,
         param_group_field: str,
-        total_iters: int,
+        num_iters: int,
         base_value: float,
         final_value: float,
-        warmup_iters_ratio: float | None = None,
+        warmup_ratio: float | None = None,
         warmup_value: float | None = None,
         freeze_ratio: float | None = None,
         multiplier_field: str | None = None,
@@ -131,21 +129,21 @@ class CosineScheduler(_CosineSchedulerCore):
         ignore_if_field: str | None = None,
     ) -> None:
         """
-        Initialize the scheduler with optimizer and scheduling parameters.
+        Configure cosine scheduling for matching optimizer groups.
 
         Args:
-            optimizer: PyTorch optimizer to schedule parameters for.
-            param_group_field: Name of the parameter group field to modify (e.g., 'lr', 'weight_decay').
-            total_iters: Total number of iterations for the scheduling.
-            base_value: Initial value for the parameter.
-            final_value: Final value for the parameter at the end of scheduling.
-            warmup_iters_ratio: Ratio of total iterations to use for warmup phase. Defaults to None.
-            warmup_value: Value to use during warmup phase. Defaults to None.
-            freeze_ratio: Ratio of total iterations to freeze parameter updates. Defaults to None.
-            multiplier_field: Field name for multiplier values in parameter groups. Defaults to None.
-            skip_if_zero: Whether to skip scheduling if the parameter value is zero. Defaults to False.
-            apply_if_field: Field name that must be present to apply scheduling. Defaults to None.
-            ignore_if_field: Field name that when present causes scheduling to be ignored. Defaults to None.
+            optimizer: Optimizer whose param groups are updated in-place.
+            param_group_field: Name of the field that receives the scheduled value.
+            num_iters: Number of scheduler iterations before clamping at ``final_value``.
+            base_value: Value used on the first cosine step (after warmup/freeze).
+            final_value: Value approached as iterations progress.
+            warmup_ratio: Optional fraction of iterations to linearly ramp from ``warmup_value`` to ``base_value``.
+            warmup_value: Starting value for the warmup ramp.
+            freeze_ratio: Optional fraction of iterations to keep the value frozen at zero at the beginning.
+            multiplier_field: Optional per-group multiplier applied to the scheduled value.
+            skip_if_zero: Leave groups untouched when their target field equals zero.
+            apply_if_field: Require this flag to be present in a param group before updating.
+            ignore_if_field: Skip groups that declare this flag.
 
         """
         self.apply_if_field = apply_if_field
@@ -154,14 +152,15 @@ class CosineScheduler(_CosineSchedulerCore):
         self.multiplier_field = multiplier_field
         self.skip_if_zero = skip_if_zero
         super().__init__(
-            param_group_field=param_group_field,
-            total_iters=total_iters,
+            param_name=param_group_field,
+            num_iters=num_iters,
             base_value=base_value,
             final_value=final_value,
-            warmup_iters_ratio=warmup_iters_ratio,
+            warmup_ratio=warmup_ratio,
             warmup_value=warmup_value,
             freeze_ratio=freeze_ratio,
         )
+        self.param_group_field = param_group_field
         return
 
     @override
@@ -194,14 +193,7 @@ class CosineScheduler(_CosineSchedulerCore):
 
 
 class CosineParamScheduler(_CosineSchedulerCore):
-    """
-    CosineParamScheduler adjusts a parameter value using a cosine annealing scheduler.
-
-    This class provides a mechanism to schedule the value of a parameter over a
-    predefined number of iterations. It supports linear warm-up and optional freezing
-    periods before the cosine annealing wave begins. The scheduler can be used to
-    gradually transition a parameter value from a starting value to a final value.
-    """
+    """Standalone cosine scheduler for non-optimizer parameters."""
 
     @override
     def step(self, it: int) -> float:

kostyl/ml/schedulers/linear.py

@@ -0,0 +1,153 @@
+from typing import Any
+from typing import override
+
+import numpy as np
+import numpy.typing as npt
+import torch
+
+from .base import BaseScheduler
+
+
+class _LinearScheduleBase(BaseScheduler):
+    def __init__(
+        self,
+        param_name: str,
+        num_iters: int,
+        base_value: float,
+        final_value: float,
+    ) -> None:
+        self.param_name = param_name
+        self.num_iters = num_iters
+        self.base_value = base_value
+        self.final_value = final_value
+
+        self.scheduler_values: npt.NDArray[np.float64] = np.array([], dtype=np.float64)
+        self.current_value_ = self.base_value
+        return
+
+    def _create_scheduler(self) -> None:
+        self.scheduler_values = np.linspace(
+            self.base_value, self.final_value, num=self.num_iters, dtype=np.float64
+        )
+        if len(self.scheduler_values) != self.num_iters:
+            raise ValueError(
+                f"Scheduler length ({len(self.scheduler_values)}) does not match total_iters ({self.num_iters})."
+            )
+        return
+
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        super().load_state_dict(state_dict)
+        self.scheduler_values = np.array([], dtype=np.float64)
+        return
+
+    @override
+    def step(self, it: int) -> None | float:
+        raise NotImplementedError
+
+    def _get_value(self, it: int) -> float:
+        if len(self.scheduler_values) == 0:
+            self._create_scheduler()
+
+        if it >= self.num_iters:
+            value: float = self.final_value
+        else:
+            value: float = self.scheduler_values[it]
+        self.current_value_ = value
+        return value
+
+    @override
+    def current_value(self) -> dict[str, float]:
+        return {self.param_name: self.current_value_}
+
+
+class LinearScheduler(_LinearScheduleBase):
+    """Implements a linear scheduler for adjusting parameter values in torch.optim.Optimizer."""
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        param_group_field: str,
+        num_iters: int,
+        base_value: float,
+        final_value: float,
+        multiplier_field: str | None = None,
+        skip_if_zero: bool = False,
+        apply_if_field: str | None = None,
+        ignore_if_field: str | None = None,
+    ) -> None:
+        """
+        Configure which optimizer groups get a linear value schedule.
+
+        Args:
+            optimizer: Optimizer whose param groups are updated in-place.
+            param_group_field: Name of the field that receives the scheduled value.
+            num_iters: Number of scheduler iterations before clamping at ``final_value``.
+            base_value: Value used on the first iteration.
+            final_value: Value used once ``num_iters`` iterations are consumed.
+            multiplier_field: Optional per-group multiplier applied to the scheduled value.
+            skip_if_zero: Leave groups untouched when their target field equals zero.
+            apply_if_field: Require this flag to be present in a param group before updating.
+            ignore_if_field: Skip groups that declare this flag.
+
+        """
+        self.apply_if_field = apply_if_field
+        self.ignore_if_field = ignore_if_field
+        self.optimizer = optimizer
+        self.multiplier_field = multiplier_field
+        self.skip_if_zero = skip_if_zero
+        super().__init__(
+            param_name=param_group_field,
+            num_iters=num_iters,
+            base_value=base_value,
+            final_value=final_value,
+        )
+        self.param_group_field = param_group_field
+        return
+
+    @override
+    def step(self, it: int) -> None:
+        value = self._get_value(it)
+        for pg in self.optimizer.param_groups:
+            if self.param_group_field not in pg:
+                raise ValueError(
+                    f"Parameter group field '{self.param_group_field}' not found in optimizer parameter groups."
+                )
+
+            if (self.apply_if_field is not None) and (self.apply_if_field not in pg):
+                continue
+
+            if (self.ignore_if_field is not None) and (self.ignore_if_field in pg):
+                continue
+
+            if self.skip_if_zero and pg[self.param_group_field] == 0:
+                continue
+
+            if self.multiplier_field is not None:
+                if self.multiplier_field not in pg:
+                    multiplier = 1.0
+                else:
+                    multiplier = pg[self.multiplier_field]
+                pg[self.param_group_field] = value * multiplier
+            else:
+                pg[self.param_group_field] = value
+        return
+
+
+class LinearParamScheduler(_LinearScheduleBase):
+    """LinearParamScheduler adjusts a parameter value using a linear scheduler."""
+
+    @override
+    def step(self, it: int) -> float:
+        """
+        Computes the value corresponding to the given iteration step.
+
+        Args:
+            it: The current iteration index used for value computation.
+
+        Returns:
+            The computed value for the provided iteration step as a float.
+
+        """
+        value = self._get_value(it)
+        return value

{kostyl_toolkit-0.1.25.dist-info → kostyl_toolkit-0.1.27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.25
+Version: 0.1.27
 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.25.dist-info → kostyl_toolkit-0.1.27.dist-info}/RECORD

@@ -6,14 +6,14 @@ kostyl/ml/clearml/logging_utils.py,sha256=GBjIIZbH_itd5sj7XpvxjkyZwxxGOpEcQ3BiWa
 kostyl/ml/clearml/pulling_utils.py,sha256=cNa_-_5LHjNVYi9btXBrfl5sPvI6BAAlIFidtpKu310,4078
 kostyl/ml/configs/__init__.py,sha256=IetcivbqYGutowLqxdKp7QR4tkXKBr4m8t4Zkk9jHZU,911
 kostyl/ml/configs/base_model.py,sha256=Eofn14J9RsjpVx_J4rp6C19pDDCANU4hr3JtX-d0FpQ,4820
-kostyl/ml/configs/hyperparams.py,sha256=OqN7mEj3zc5MTqBPCZL3Lcd2VCTDLo_K0yvhRWGfhCs,2924
+kostyl/ml/configs/hyperparams.py,sha256=2S_VEZ07RWquNFSWjHBb3OUpBlTznbUpFSchzMpSBOc,2879
 kostyl/ml/configs/training_settings.py,sha256=Sq2tiRuwkbmi9zKDG2JghZLXo5DDt_eQqN_KYJSdcTY,2509
 kostyl/ml/dist_utils.py,sha256=Onf0KHVLA8oeUgZTcTdmR9qiM22f2uYLoNwgLbMGJWk,3495
 kostyl/ml/lightning/__init__.py,sha256=-F3JAyq8KU1d-nACWryGu8d1CbvWbQ1rXFdeRwfE2X8,175
 kostyl/ml/lightning/callbacks/__init__.py,sha256=enexQt3octktsTiEYHltSF_24CM-NeFEVFimXiavGiY,296
-kostyl/ml/lightning/callbacks/checkpoint.py,sha256=1gk5-NjsMXe5cZP0OgNcoc9KUTzRDTHIokVEDr74sjI,16740
+kostyl/ml/lightning/callbacks/checkpoint.py,sha256=KNwNVB2TFh2dcn133NbeTo5ul0jgiPYCeA-8NQ7U_mw,18951
 kostyl/ml/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
-kostyl/ml/lightning/callbacks/registry_uploader.py,sha256=ksoh02dzIde4E_GaZykfiOgfSjZti-IJt_i61enem3s,6779
+kostyl/ml/lightning/callbacks/registry_uploader.py,sha256=pIZHzHVANO_VsxPIbYhS8SwgZFHL341mP2HJnQ4iMFs,4216
 kostyl/ml/lightning/extenstions/__init__.py,sha256=OY6QGv1agYgqqKf1xJBrxgp_i8FunVfPzYezfaRrGXU,182
 kostyl/ml/lightning/extenstions/custom_module.py,sha256=nB5jW7cqRD1tyh-q5LD2EtiFQwFkLXpnS9Yu6c5xMRg,5987
 kostyl/ml/lightning/extenstions/pretrained_model.py,sha256=QJGr2UvYJcU2Gy2w8z_cEvTodjv7hGdd2PPPfdOI-Mw,4017
@@ -25,11 +25,12 @@ kostyl/ml/params_groups.py,sha256=nUyw5d06Pvy9QPiYtZzLYR87xwXqJLxbHthgQH8oSCM,35
 kostyl/ml/schedulers/__init__.py,sha256=bxXbsU_WYnVbhvNNnuI7cOAh2Axz7D25TaleBTZhYfc,197
 kostyl/ml/schedulers/base.py,sha256=9M2iOoOVSRojR_liPX1qo3Nn4iMXSM5ZJuAFWZTulUk,1327
 kostyl/ml/schedulers/composite.py,sha256=ee4xlMDMMtjKPkbTF2ue9GTr9DuGCGjZWf11mHbi6aE,2387
-kostyl/ml/schedulers/cosine.py,sha256=jufULVHn_L_ZZEc3ZTG3QCY_pc0jlAMH5Aw496T31jo,8203
+kostyl/ml/schedulers/cosine.py,sha256=t74_ByT22L5NQKpnBVU9UGzBVx1ZM2GTylb9ct3_PVg,7627
+kostyl/ml/schedulers/linear.py,sha256=62mYEfd_2cQjOWrd0Vl5_sFeEokBKYmx496szhY04aU,5159
 kostyl/utils/__init__.py,sha256=hkpmB6c5pr4Ti5BshOROebb7cvjDZfNCw83qZ_FFKMM,240
 kostyl/utils/dict_manipulations.py,sha256=e3vBicID74nYP8lHkVTQc4-IQwoJimrbFELy5uSF6Gk,1073
 kostyl/utils/fs.py,sha256=gAQNIU4R_2DhwjgzOS8BOMe0gZymtY1eZwmdgOdDgqo,510
 kostyl/utils/logging.py,sha256=Vye0u4-yeOSUc-f03gpQbxSktTbFiilTWLEVr00ZHvc,5796
-kostyl_toolkit-0.1.25.dist-info/WHEEL,sha256=z-mOpxbJHqy3cq6SvUThBZdaLGFZzdZPtgWLcP2NKjQ,79
-kostyl_toolkit-0.1.25.dist-info/METADATA,sha256=Go9dF8W4vQ4HzQpzRCLa-i9NsYq5o0J34NtvGWZnEvA,4269
-kostyl_toolkit-0.1.25.dist-info/RECORD,,
+kostyl_toolkit-0.1.27.dist-info/WHEEL,sha256=ZyFSCYkV2BrxH6-HRVRg3R9Fo7MALzer9KiPYqNxSbo,79
+kostyl_toolkit-0.1.27.dist-info/METADATA,sha256=kg7Y2CJqhAI-3--rIKsPlarm1Ukk6jQLJpW2ZBvysI8,4269
+kostyl_toolkit-0.1.27.dist-info/RECORD,,

{kostyl_toolkit-0.1.25.dist-info → kostyl_toolkit-0.1.27.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: uv 0.9.15
+Generator: uv 0.9.18
 Root-Is-Purelib: true
 Tag: py3-none-any