kostyl-toolkit 0.1.24__tar.gz → 0.1.26__tar.gz

{kostyl_toolkit-0.1.24 → kostyl_toolkit-0.1.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.24
+Version: 0.1.26
 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.24 → kostyl_toolkit-0.1.26}/kostyl/ml/configs/hyperparams.py

@@ -23,21 +23,17 @@ class Lr(BaseModel):
         default=None, gt=0, lt=1, validate_default=False
     )
     warmup_value: float | None = Field(default=None, gt=0, validate_default=False)
-    base_value: float
+    start_value: float
     final_value: float | None = Field(default=None, gt=0, validate_default=False)
 
     @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."
             )
@@ -60,7 +56,7 @@ class WeightDecay(BaseModel):
     """Weight decay hyperparameters configuration."""
 
     use_scheduler: bool = False
-    base_value: float
+    start_value: float
     final_value: float | None = None
 
     @model_validator(mode="after")

{kostyl_toolkit-0.1.24 → kostyl_toolkit-0.1.26}/kostyl/ml/lightning/callbacks/checkpoint.py

@@ -2,13 +2,16 @@ from datetime import timedelta
 from pathlib import Path
 from shutil import rmtree
 from typing import Literal
-from typing import cast
+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
@@ -17,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
@@ -227,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,
@@ -242,8 +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.process_group: dist.ProcessGroup | None = None
+        self.uploading_mode = uploading_mode
         super().__init__(
             dirpath=dirpath,
             filename=filename,
@@ -261,20 +268,32 @@ class CustomModelCheckpoint(ModelCheckpoint):
             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
+    @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
 
 
@@ -283,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.")
@@ -314,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_toolkit-0.1.26/kostyl/ml/lightning/callbacks/registry_uploader.py

@@ -0,0 +1,122 @@
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import Callable
+from pathlib import Path
+from typing import override
+
+from clearml import OutputModel
+from clearml import Task
+
+from kostyl.ml.clearml.logging_utils import find_version_in_tags
+from kostyl.ml.clearml.logging_utils import increment_version
+from kostyl.utils.logging import setup_logger
+
+
+logger = setup_logger()
+
+
+class RegistryUploaderCallback(ABC):
+    """Abstract Lightning callback responsible for tracking and uploading the best-performing model checkpoint."""
+
+    @abstractmethod
+    def upload_checkpoint(self, path: str | Path) -> None:
+        """Upload the checkpoint located at the given path to the configured registry backend."""
+        raise NotImplementedError
+
+
+class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
+    """PyTorch Lightning callback to upload the best model checkpoint to ClearML."""
+
+    def __init__(
+        self,
+        task: Task,
+        output_model_name: str,
+        output_model_tags: list[str] | None = None,
+        verbose: bool = True,
+        enable_tag_versioning: bool = True,
+        label_enumeration: dict[str, int] | None = None,
+        config_dict: dict[str, str] | None = None,
+    ) -> None:
+        """
+        Initializes the ClearMLRegistryUploaderCallback.
+
+        Args:
+            task: ClearML task.
+            ckpt_callback: ModelCheckpoint instance used by Trainer.
+            output_model_name: Name for the ClearML output model.
+            output_model_tags: Tags for the output model.
+            verbose: Whether to log messages.
+            label_enumeration: Optional mapping of label names to integer IDs.
+            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".
+
+        """
+        super().__init__()
+        if output_model_tags is None:
+            output_model_tags = []
+
+        self.task = task
+        self.output_model_name = output_model_name
+        self.output_model_tags = output_model_tags
+        self.config_dict = config_dict
+        self.label_enumeration = label_enumeration
+        self.verbose = verbose
+        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._upload_callback: Callable | None = None
+        return
+
+    def _create_output_model(self) -> OutputModel:
+        if self.enable_tag_versioning:
+            version = find_version_in_tags(self.output_model_tags)
+            if version is None:
+                self.output_model_tags.append("v1.0")
+            else:
+                new_version = increment_version(version)
+                self.output_model_tags.remove(version)
+                self.output_model_tags.append(new_version)
+
+        if "LightningCheckpoint" not in self.output_model_tags:
+            self.output_model_tags.append("LightningCheckpoint")
+
+        return OutputModel(
+            task=self.task,
+            name=self.output_model_name,
+            framework="PyTorch",
+            tags=self.output_model_tags,
+            config_dict=None,
+            label_enumeration=self.label_enumeration,
+        )
+
+    @override
+    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
+
+        if self._output_model is None:
+            self._output_model = self._create_output_model()
+
+        if self.verbose:
+            logger.info(f"Uploading model from {path}")
+
+        self._output_model.update_weights(
+            path,
+            auto_delete_file=False,
+            async_enable=False,
+        )
+        self._output_model.update_design(config_dict=self.config_dict)
+
+        self._last_uploaded_model_path = path
+        return

{kostyl_toolkit-0.1.24 → kostyl_toolkit-0.1.26}/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,
-        base_value: float,
+        param_name: str,
+        num_iters: int,
+        start_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,46 +30,46 @@ 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.base_value = base_value
+        self.param_name = param_name
+        self.num_iters = num_iters
+        self.start_value = start_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
 
         self.scheduler_values: npt.NDArray[np.float64] = np.array([], dtype=np.float64)
-        self.current_value_ = self.base_value
+        self.current_value_ = self.start_value
         return
 
     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
+                self.warmup_value, self.start_value, warmup_iters, dtype=np.float64
             )
         else:
             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.")
 
         # Create cosine schedule
         iters = np.arange(cosine_annealing_iters)
-        schedule = self.final_value + 0.5 * (self.base_value - self.final_value) * (
+        schedule = self.final_value + 0.5 * (self.start_value - self.final_value) * (
             1 + np.cos(np.pi * iters / len(iters))
         )
 
@@ -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,
-        base_value: float,
+        num_iters: int,
+        start_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``.
+            start_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 ``start_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,
-            base_value=base_value,
+            param_name=param_group_field,
+            num_iters=num_iters,
+            start_value=start_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_toolkit-0.1.26/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,
+        start_value: float,
+        final_value: float,
+    ) -> None:
+        self.param_name = param_name
+        self.num_iters = num_iters
+        self.start_value = start_value
+        self.final_value = final_value
+
+        self.scheduler_values: npt.NDArray[np.float64] = np.array([], dtype=np.float64)
+        self.current_value_ = self.start_value
+        return
+
+    def _create_scheduler(self) -> None:
+        self.scheduler_values = np.linspace(
+            self.start_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,
+        start_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``.
+            start_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,
+            start_value=start_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.24 → kostyl_toolkit-0.1.26}/pyproject.toml

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

kostyl_toolkit-0.1.24/kostyl/ml/lightning/callbacks/registry_uploader.py

@@ -1,193 +0,0 @@
-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
-
-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):
-    """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,
-        output_model_name: str,
-        output_model_tags: list[str] | None = None,
-        verbose: bool = True,
-        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.
-
-        Args:
-            task: ClearML task.
-            ckpt_callback: ModelCheckpoint instance used by Trainer.
-            output_model_name: Name for the ClearML output model.
-            output_model_tags: Tags for the output model.
-            verbose: Whether to log messages.
-            label_enumeration: Optional mapping of label names to integer IDs.
-            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__()
-        if output_model_tags is None:
-            output_model_tags = []
-
-        self.task = task
-        self.output_model_name = output_model_name
-        self.output_model_tags = output_model_tags
-        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._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:
-        if self.enable_tag_versioning:
-            version = find_version_in_tags(self.output_model_tags)
-            if version is None:
-                self.output_model_tags.append("v1.0")
-            else:
-                new_version = increment_version(version)
-                self.output_model_tags.remove(version)
-                self.output_model_tags.append(new_version)
-
-        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,
-            framework="PyTorch",
-            tags=self.output_model_tags,
-            config_dict=None,
-            label_enumeration=self.label_enumeration,
-        )
-
-    @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)
-            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 {self._best_model_path}")
-
-        self._output_model.update_weights(
-            self._best_model_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._upload_best_checkpoint(pl_module)
-        return