kostyl-toolkit 0.1.15__tar.gz

kostyl_toolkit-0.1.15/PKG-INFO

@@ -0,0 +1,105 @@
+Metadata-Version: 2.3
+Name: kostyl-toolkit
+Version: 0.1.15
+Summary: Kickass Orchestration System for Training, Yielding & Logging 
+Requires-Dist: case-converter>=1.2.0
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: case-converter>=1.2.0 ; extra == 'ml-core'
+Requires-Dist: clearml[s3]>=2.0.2 ; extra == 'ml-core'
+Requires-Dist: lightning>=2.5.6 ; extra == 'ml-core'
+Requires-Dist: pydantic>=2.12.4 ; extra == 'ml-core'
+Requires-Dist: torch>=2.9.1 ; extra == 'ml-core'
+Requires-Dist: transformers>=4.57.1 ; extra == 'ml-core'
+Requires-Python: >=3.12
+Provides-Extra: ml-core
+Description-Content-Type: text/markdown
+
+# Kostyl Toolkit
+
+Kickass Orchestration System for Training, Yielding & Logging — a batteries-included toolbox that glues PyTorch Lightning, Hugging Face Transformers, and ClearML into a single workflow.
+
+## Overview
+- Rapidly bootstrap Lightning experiments with opinionated defaults (`KostylLightningModule`, custom schedulers, grad clipping and metric formatting).
+- Keep model configs source-controlled via Pydantic mixins, with ClearML syncing out of the box (`ConfigLoadingMixin`, `ClearMLConfigMixin`).
+- Reuse Lightning checkpoints directly inside Transformers models through `LightningCheckpointLoaderMixin`.
+- Ship distributed-friendly utilities (deterministic logging, FSDP helpers, LR scaling, ClearML tag management).
+
+## Installation
+```bash
+# Latest release from PyPI
+pip install kostyl-toolkit
+
+# or with uv
+uv pip install kostyl-toolkit
+```
+
+Development setup:
+```bash
+uv sync                # creates the virtualenv declared in pyproject.toml
+source .venv/bin/activate.fish
+pre-commit install     # optional but recommended
+```
+
+## Quick Start
+```python
+from lightning import Trainer
+from transformers import AutoModelForSequenceClassification
+
+from kostyl.ml_core.configs.hyperparams import HyperparamsConfig
+from kostyl.ml_core.configs.training_params import TrainingParams
+from kostyl.ml_core.lightning.extenstions.custom_module import KostylLightningModule
+
+
+class TextClassifier(KostylLightningModule):
+	def __init__(self, hyperparams: HyperparamsConfig):
+				super().__init__()
+		self.hyperparams = hyperparams  # grad clipping + scheduler knobs
+				self.model = AutoModelForSequenceClassification.from_pretrained(
+						"distilbert-base-uncased",
+						num_labels=2,
+				)
+
+		def training_step(self, batch, batch_idx):
+				outputs = self.model(**batch)
+				self.log("train/loss", outputs.loss)
+				return outputs.loss
+
+train_cfg = TrainingParams.from_file("configs/training.yaml")
+hyperparams = HyperparamsConfig.from_file("configs/hyperparams.yaml")
+
+module = TextClassifier(hyperparams)
+
+trainer = Trainer(**train_cfg.trainer.model_dump())
+trainer.fit(module)
+```
+
+Restoring a plain Transformers model from a Lightning checkpoint:
+```python
+from kostyl.ml_core.lightning.extenstions.pretrained_model import LightningCheckpointLoaderMixin
+
+
+model = LightningCheckpointLoaderMixin.from_lighting_checkpoint(
+		"checkpoints/epoch=03-step=500.ckpt",
+		config_key="config",
+		weights_prefix="model.",
+)
+```
+
+## Components
+- **Configurations** (`kostyl/ml_core/configs`): strongly-typed training, optimizer, and scheduler configs with ClearML syncing helpers.
+- **Lightning Extensions** (`kostyl/ml_core/lightning`): custom LightningModule base class, callbacks, logging bridges, and the checkpoint loader mixin.
+- **Schedulers** (`kostyl/ml_core/schedulers`): extensible LR schedulers (base/composite/cosine) with serialization helpers and on-step logging.
+- **ClearML Utilities** (`kostyl/ml_core/clearml`): tag/version helpers and logging bridges for ClearML Tasks.
+- **Distributed + Metrics Utils** (`kostyl/ml_core/dist_utils.py`, `metrics_formatting.py`): world-size-aware LR scaling, rank-aware metric naming, and per-class formatting.
+- **Logging Helpers** (`kostyl/utils/logging.py`): rank-aware Loguru setup and uniform handling of incompatible checkpoint keys.
+
+## Project Layout
+```
+kostyl/
+	ml_core/
+		configs/                # Pydantic configs + ClearML mixins
+		lightning/              # Lightning module, callbacks, loggers, extensions
+		schedulers/             # Base + composite/cosine schedulers
+		clearml/                # Logging + pulling utilities
+	utils/                    # Dict helpers, logging utilities
+```

kostyl_toolkit-0.1.15/README.md

@@ -0,0 +1,89 @@
+# Kostyl Toolkit
+
+Kickass Orchestration System for Training, Yielding & Logging — a batteries-included toolbox that glues PyTorch Lightning, Hugging Face Transformers, and ClearML into a single workflow.
+
+## Overview
+- Rapidly bootstrap Lightning experiments with opinionated defaults (`KostylLightningModule`, custom schedulers, grad clipping and metric formatting).
+- Keep model configs source-controlled via Pydantic mixins, with ClearML syncing out of the box (`ConfigLoadingMixin`, `ClearMLConfigMixin`).
+- Reuse Lightning checkpoints directly inside Transformers models through `LightningCheckpointLoaderMixin`.
+- Ship distributed-friendly utilities (deterministic logging, FSDP helpers, LR scaling, ClearML tag management).
+
+## Installation
+```bash
+# Latest release from PyPI
+pip install kostyl-toolkit
+
+# or with uv
+uv pip install kostyl-toolkit
+```
+
+Development setup:
+```bash
+uv sync                # creates the virtualenv declared in pyproject.toml
+source .venv/bin/activate.fish
+pre-commit install     # optional but recommended
+```
+
+## Quick Start
+```python
+from lightning import Trainer
+from transformers import AutoModelForSequenceClassification
+
+from kostyl.ml_core.configs.hyperparams import HyperparamsConfig
+from kostyl.ml_core.configs.training_params import TrainingParams
+from kostyl.ml_core.lightning.extenstions.custom_module import KostylLightningModule
+
+
+class TextClassifier(KostylLightningModule):
+	def __init__(self, hyperparams: HyperparamsConfig):
+				super().__init__()
+		self.hyperparams = hyperparams  # grad clipping + scheduler knobs
+				self.model = AutoModelForSequenceClassification.from_pretrained(
+						"distilbert-base-uncased",
+						num_labels=2,
+				)
+
+		def training_step(self, batch, batch_idx):
+				outputs = self.model(**batch)
+				self.log("train/loss", outputs.loss)
+				return outputs.loss
+
+train_cfg = TrainingParams.from_file("configs/training.yaml")
+hyperparams = HyperparamsConfig.from_file("configs/hyperparams.yaml")
+
+module = TextClassifier(hyperparams)
+
+trainer = Trainer(**train_cfg.trainer.model_dump())
+trainer.fit(module)
+```
+
+Restoring a plain Transformers model from a Lightning checkpoint:
+```python
+from kostyl.ml_core.lightning.extenstions.pretrained_model import LightningCheckpointLoaderMixin
+
+
+model = LightningCheckpointLoaderMixin.from_lighting_checkpoint(
+		"checkpoints/epoch=03-step=500.ckpt",
+		config_key="config",
+		weights_prefix="model.",
+)
+```
+
+## Components
+- **Configurations** (`kostyl/ml_core/configs`): strongly-typed training, optimizer, and scheduler configs with ClearML syncing helpers.
+- **Lightning Extensions** (`kostyl/ml_core/lightning`): custom LightningModule base class, callbacks, logging bridges, and the checkpoint loader mixin.
+- **Schedulers** (`kostyl/ml_core/schedulers`): extensible LR schedulers (base/composite/cosine) with serialization helpers and on-step logging.
+- **ClearML Utilities** (`kostyl/ml_core/clearml`): tag/version helpers and logging bridges for ClearML Tasks.
+- **Distributed + Metrics Utils** (`kostyl/ml_core/dist_utils.py`, `metrics_formatting.py`): world-size-aware LR scaling, rank-aware metric naming, and per-class formatting.
+- **Logging Helpers** (`kostyl/utils/logging.py`): rank-aware Loguru setup and uniform handling of incompatible checkpoint keys.
+
+## Project Layout
+```
+kostyl/
+	ml_core/
+		configs/                # Pydantic configs + ClearML mixins
+		lightning/              # Lightning module, callbacks, loggers, extensions
+		schedulers/             # Base + composite/cosine schedulers
+		clearml/                # Logging + pulling utilities
+	utils/                    # Dict helpers, logging utilities
+```

kostyl_toolkit-0.1.15/kostyl/ml/clearml/dataset_utils.py

@@ -0,0 +1,58 @@
+from collections.abc import Collection
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import as_completed
+from pathlib import Path
+
+from clearml import Dataset as ClearMLDataset
+
+
+def collect_clearml_datasets(
+    datasets_mapping: dict[str, str],
+) -> dict[str, ClearMLDataset]:
+    """
+    Collect ClearML datasets by dataset ID.
+
+    Args:
+        datasets_mapping: Mapping where keys are human-readable names and values
+            are ClearML dataset IDs.
+
+    Returns:
+        A mapping of dataset names to fetched `ClearMLDataset` instances.
+
+    """
+    datasets_list = {}
+    for name, dataset_id in datasets_mapping.items():
+        clearml_dataset = ClearMLDataset.get(dataset_id, alias=name)
+        datasets_list[name] = clearml_dataset
+    return datasets_list
+
+
+def download_clearml_datasets(datasets: Collection[ClearMLDataset]) -> None:
+    """
+    Download all ClearML datasets in parallel.
+
+    Args:
+        datasets: Collection of initialized `ClearMLDataset` instances to download.
+
+    """
+    with ThreadPoolExecutor() as executor:
+        futures = [executor.submit(ds.get_local_copy) for ds in datasets]
+        for future in as_completed(futures):
+            future.result()
+    return
+
+
+def get_datasets_paths(datasets_mapping: dict[str, ClearMLDataset]) -> dict[str, Path]:
+    """
+    Return local filesystem paths for ClearML datasets.
+
+    Args:
+        datasets_mapping: Mapping of dataset names to initialized
+            `ClearMLDataset` instances.
+
+    Returns:
+        Mapping of dataset names to local `Path` objects pointing to the
+        downloaded dataset copies.
+
+    """
+    return {name: Path(ds.get_local_copy()) for name, ds in datasets_mapping.items()}

kostyl_toolkit-0.1.15/kostyl/ml/clearml/logging_utils.py

@@ -0,0 +1,46 @@
+import re
+
+from kostyl.utils import setup_logger
+
+
+logger = setup_logger(name="clearml_logging_utils.py", fmt="only_message")
+
+
+def increment_version(s: str) -> str:
+    """
+    Increments the minor part of a version string.
+
+    Examples:
+        v1.00 -> v1.01
+        v2.99 -> v2.100
+        v.3.009 -> v.3.010
+
+    """
+    s = s.strip()
+    m = re.fullmatch(r"v(\.?)(\d+)\.(\d+)", s)
+    if not m:
+        raise ValueError(f"Invalid version format: {s!r}. Expected 'vX.Y' or 'v.X.Y'.")
+
+    vdot, major_str, minor_str = m.groups()
+    major = int(major_str)
+    minor = int(minor_str) + 1
+
+    # preserve leading zeros based on original width, length may increase (99 -> 100)
+    minor_out = str(minor).zfill(len(minor_str))
+    prefix = f"v{vdot}"  # 'v' or 'v.'
+    return f"{prefix}{major}.{minor_out}"
+
+
+def find_version_in_tags(tags: list[str]) -> str | None:
+    """
+    Finds the first version tag in the list of tags.
+
+    Note:
+        Version tags must be in the format 'vX.Y' or 'v.X.Y' (an optional dot after 'v' is supported).
+
+    """
+    version_pattern = re.compile(r"^v(\.?)(\d+)\.(\d+)$")
+    for tag in tags:
+        if version_pattern.match(tag):
+            return tag
+    return None

kostyl_toolkit-0.1.15/kostyl/ml/clearml/pulling_utils.py

@@ -0,0 +1,91 @@
+from pathlib import Path
+from typing import Any
+from typing import cast
+
+from clearml import InputModel
+from clearml import Task
+from transformers import AutoModel
+from transformers import AutoTokenizer
+from transformers import PreTrainedModel
+from transformers import PreTrainedTokenizerBase
+
+from kostyl.ml.lightning.extenstions.pretrained_model import (
+    LightningCheckpointLoaderMixin,
+)
+
+
+def get_tokenizer_from_clearml(
+    model_id: str, task: Task | None = None, ignore_remote_overrides: bool = True
+) -> PreTrainedTokenizerBase:
+    """
+    Retrieve a Hugging Face tokenizer stored in a ClearML.
+
+    Args:
+        model_id (str): The ClearML InputModel identifier that holds the tokenizer artifacts.
+        task (Task | None, optional): An optional ClearML Task used to associate and sync
+            the model. Defaults to None.
+        ignore_remote_overrides (bool, optional): Whether to ignore remote hyperparameter
+            overrides when connecting the ClearML task. Defaults to True.
+
+    Returns:
+        PreTrainedTokenizerBase: The instantiated tokenizer loaded from the local copy
+            of the referenced ClearML InputModel.
+
+    """
+    clearml_tokenizer = InputModel(model_id=model_id)
+    if task is not None:
+        clearml_tokenizer.connect(task, ignore_remote_overrides=ignore_remote_overrides)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        clearml_tokenizer.get_local_copy(raise_on_error=True)
+    )
+    return tokenizer
+
+
+def get_model_from_clearml[
+    TModel: PreTrainedModel | LightningCheckpointLoaderMixin | AutoModel
+](
+    model_id: str,
+    model: type[TModel],
+    task: Task | None = None,
+    ignore_remote_overrides: bool = True,
+    **kwargs: Any,
+) -> TModel:
+    """
+    Retrieve a pretrained model from ClearML and instantiate it using the appropriate loader.
+
+    Args:
+        model_id: Identifier of the ClearML input model to retrieve.
+        model: The model class that implements either PreTrainedModel or LightningCheckpointLoaderMixin.
+        task: Optional ClearML task used to resolve the input model reference. If provided, the input model
+            will be connected to this task, with remote overrides optionally ignored.
+        ignore_remote_overrides: When connecting the input model to the provided task, determines whether
+            remote configuration overrides should be ignored.
+        **kwargs: Additional keyword arguments to pass to the model loading method.
+
+    Returns:
+        An instantiated model loaded either from a ClearML package directory or a Lightning checkpoint.
+
+    """
+    input_model = InputModel(model_id=model_id)
+
+    if task is not None:
+        input_model.connect(task, ignore_remote_overrides=ignore_remote_overrides)
+
+    local_path = Path(input_model.get_local_copy(raise_on_error=True))
+
+    if local_path.is_dir() and input_model._is_package():
+        model_instance = model.from_pretrained(local_path, **kwargs)
+    elif local_path.suffix == ".ckpt":
+        if not issubclass(model, LightningCheckpointLoaderMixin):
+            raise ValueError(
+                f"Model class {model.__name__} is not compatible with Lightning checkpoints."
+            )
+        model_instance = model.from_lighting_checkpoint(local_path, **kwargs)
+    else:
+        raise ValueError(
+            f"Unsupported model format for path: {local_path}. "
+            "Expected a ClearML package directory or a .ckpt file."
+        )
+    model_instance = cast(TModel, model_instance)
+    return model_instance

kostyl_toolkit-0.1.15/kostyl/ml/configs/__init__.py

@@ -0,0 +1,30 @@
+from .base_model import KostylBaseModel
+from .hyperparams import HyperparamsConfig
+from .hyperparams import Lr
+from .hyperparams import Optimizer
+from .hyperparams import WeightDecay
+from .training_settings import CheckpointConfig
+from .training_settings import DataConfig
+from .training_settings import DDPStrategyConfig
+from .training_settings import EarlyStoppingConfig
+from .training_settings import FSDP1StrategyConfig
+from .training_settings import LightningTrainerParameters
+from .training_settings import SingleDeviceStrategyConfig
+from .training_settings import TrainingSettings
+
+
+__all__ = [
+    "CheckpointConfig",
+    "DDPStrategyConfig",
+    "DataConfig",
+    "EarlyStoppingConfig",
+    "FSDP1StrategyConfig",
+    "HyperparamsConfig",
+    "KostylBaseModel",
+    "LightningTrainerParameters",
+    "Lr",
+    "Optimizer",
+    "SingleDeviceStrategyConfig",
+    "TrainingSettings",
+    "WeightDecay",
+]

kostyl_toolkit-0.1.15/kostyl/ml/configs/base_model.py

@@ -0,0 +1,143 @@
+from pathlib import Path
+from typing import Self
+from typing import TypeVar
+
+from caseconverter import pascalcase
+from caseconverter import snakecase
+from clearml import Task
+from pydantic import BaseModel as PydanticBaseModel
+
+from kostyl.utils.dict_manipulations import convert_to_flat_dict
+from kostyl.utils.dict_manipulations import flattened_dict_to_nested
+from kostyl.utils.fs import load_config
+
+
+TConfig = TypeVar("TConfig", bound=PydanticBaseModel)
+
+
+class BaseModelWithConfigLoading(PydanticBaseModel):
+    """Pydantic class providing basic configuration loading functionality."""
+
+    @classmethod
+    def from_file(
+        cls: type[Self],  # pyright: ignore
+        path: str | Path,
+    ) -> Self:
+        """
+        Create an instance of the class from a configuration file.
+
+        Args:
+            cls_: The class type to instantiate.
+            path (str | Path): Path to the configuration file.
+
+        Returns:
+            An instance of the class created from the configuration file.
+
+        """
+        config = load_config(path)
+        instance = cls.model_validate(config)
+        return instance
+
+    @classmethod
+    def from_dict(
+        cls: type[Self],  # pyright: ignore
+        state_dict: dict,
+    ) -> Self:
+        """
+        Creates an instance from a dictionary.
+
+        Args:
+            cls_: The class type to instantiate.
+            state_dict (dict): A dictionary representing the state of the
+                class that must be validated and used for initialization.
+
+        Returns:
+            An initialized instance of the class based on the
+                provided state dictionary.
+
+        """
+        instance = cls.model_validate(state_dict)
+        return instance
+
+
+class BaseModelWithClearmlSyncing(BaseModelWithConfigLoading):
+    """Pydantic class providing ClearML configuration loading and syncing functionality."""
+
+    @classmethod
+    def connect_as_file(
+        cls: type[Self],  # pyright: ignore
+        task: Task,
+        path: str | Path,
+        alias: str | None = None,
+    ) -> Self:
+        """
+        Connects the configuration file to a ClearML task and creates an instance of the class from it.
+
+        This method connects the specified configuration file to the given ClearML task for version control and monitoring,
+        then loads and validates the configuration to the class.
+
+        Args:
+            cls: The class type to instantiate.
+            task: The ClearML Task object to connect the configuration to.
+            path: Path to the configuration file (supports YAML format).
+            alias: Optional alias for the configuration in ClearML. Defaults to PascalCase of the class name if None.
+
+        Returns:
+            An instance of the class created from the connected configuration file.
+
+        """
+        if isinstance(path, Path):
+            str_path = str(path)
+        else:
+            str_path = path
+
+        name = alias if alias is not None else pascalcase(cls.__name__)
+        connected_path = task.connect_configuration(str_path, name=pascalcase(name))
+
+        if not isinstance(connected_path, str):
+            connected_path_str = str(connected_path)
+        else:
+            connected_path_str = connected_path
+
+        model = cls.from_file(path=connected_path_str)
+        return model
+
+    @classmethod
+    def connect_as_dict(
+        cls: type[Self],  # pyright: ignore
+        task: Task,
+        path: str | Path,
+        alias: str | None = None,
+    ) -> Self:
+        """
+        Connects configuration from a file as a dictionary to a ClearML task and creates an instance of the class.
+
+        This class method loads configuration from a file as a dictionary, flattens and sync them with ClearML
+        task parameters. Then it creates an instance of the class using the synced dictionary.
+
+        Args:
+            cls: The class type of the model to be created (must be a TRetuningModel subclass).
+            task: The ClearML task to connect the configuration to.
+            path: Path to the configuration file to load parameters from.
+            alias: Optional alias name for the configuration. If None, uses snake_case of class name.
+
+        Returns:
+            An instance of the specified class created from the loaded configuration.
+
+        """
+        name = alias if alias is not None else snakecase(cls.__name__)
+
+        config = load_config(path)
+
+        flattened_config = convert_to_flat_dict(config)
+        task.connect(flattened_config, name=pascalcase(name))
+        config = flattened_dict_to_nested(flattened_config)
+
+        model = cls.from_dict(state_dict=config)
+        return model
+
+
+class KostylBaseModel(BaseModelWithClearmlSyncing):
+    """A Pydantic model class with basic configuration loading functionality."""
+
+    pass

kostyl_toolkit-0.1.15/kostyl/ml/configs/hyperparams.py

@@ -0,0 +1,84 @@
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import model_validator
+
+from kostyl.utils.logging import setup_logger
+
+from .base_model import KostylBaseModel
+
+
+logger = setup_logger(fmt="only_message")
+
+
+class Optimizer(BaseModel):
+    """Optimizer hyperparameters configuration."""
+
+    adamw_beta1: float = 0.9
+    adamw_beta2: float = 0.999
+
+
+class Lr(BaseModel):
+    """Learning rate hyperparameters configuration."""
+
+    use_scheduler: bool = False
+    warmup_iters_ratio: float | None = Field(
+        default=None, gt=0, lt=1, validate_default=False
+    )
+    warmup_value: float | None = Field(default=None, gt=0, validate_default=False)
+    base_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:
+            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):
+            logger.warning(
+                "use_scheduler is False, warmup_value and warmup_iters_ratio will be ignored."
+            )
+            self.warmup_value = None
+            self.warmup_iters_ratio = None
+        return self
+
+    @model_validator(mode="after")
+    def validate_final_value(self) -> "Lr":
+        """Validates the final_value based on use_scheduler."""
+        if self.use_scheduler and (self.final_value is None):
+            raise ValueError("If use_scheduler is True, final_value must be provided.")
+        if (not self.use_scheduler) and (self.final_value is not None):
+            logger.warning("use_scheduler is False, final_value will be ignored.")
+            self.final_value = None
+        return self
+
+
+class WeightDecay(BaseModel):
+    """Weight decay hyperparameters configuration."""
+
+    use_scheduler: bool = False
+    base_value: float
+    final_value: float | None = None
+
+    @model_validator(mode="after")
+    def validate_final_value(self) -> "WeightDecay":
+        """Validates the final_value based on use_scheduler."""
+        if self.use_scheduler and self.final_value is None:
+            raise ValueError("If use_scheduler is True, final_value must be provided.")
+        if not self.use_scheduler and self.final_value is not None:
+            logger.warning("use_scheduler is False, final_value will be ignored.")
+        return self
+
+
+class HyperparamsConfig(KostylBaseModel):
+    """Model training hyperparameters configuration."""
+
+    grad_clip_val: float | None = Field(default=None, gt=0, validate_default=False)
+    optimizer: Optimizer = Optimizer()
+    lr: Lr
+    weight_decay: WeightDecay