kostyl-toolkit 0.1.19__tar.gz → 0.1.20__tar.gz

{kostyl_toolkit-0.1.19 → kostyl_toolkit-0.1.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.19
+Version: 0.1.20
 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.19 → kostyl_toolkit-0.1.20}/kostyl/ml/dist_utils.py

@@ -86,22 +86,22 @@ def scale_lrs_by_world_size(
     return lrs
 
 
-def _get_rank() -> int:
+def get_rank() -> int:
+    """Gets the rank of the current process in a distributed setting."""
     if dist.is_initialized():
-        rank = dist.get_rank()
-    else:
-        rank = int(os.environ.get("RANK", 0))
-    return rank
+        return dist.get_rank()
+    if "RANK" in os.environ:
+        return int(os.environ["RANK"])
+    if "SLURM_PROCID" in os.environ:
+        return int(os.environ["SLURM_PROCID"])
+    if "LOCAL_RANK" in os.environ:
+        return int(os.environ["LOCAL_RANK"])
+    return 0
 
 
 def is_main_process() -> bool:
     """Checks if the current process is the main process (rank 0) in a distributed setting."""
-    if dist.is_initialized():
-        return dist.get_rank() == 0
-    if "RANK" in os.environ:
-        return int(os.environ["RANK"]) == 0
-    if "SLURM_PROCID" in os.environ:
-        return int(os.environ["SLURM_PROCID"]) == 0
-    if "LOCAL_RANK" in os.environ:
-        return int(os.environ["LOCAL_RANK"]) == 0
+    rank = get_rank()
+    if rank != 0:
+        return False
     return True

kostyl_toolkit-0.1.20/kostyl/ml/lightning/callbacks/registry_uploading.py

@@ -0,0 +1,138 @@
+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 lightning.pytorch.callbacks import ModelCheckpoint
+
+from kostyl.ml.clearml.logging_utils import find_version_in_tags
+from kostyl.ml.clearml.logging_utils import increment_version
+from kostyl.ml.lightning import KostylLightningModule
+from kostyl.utils.logging import setup_logger
+
+
+logger = setup_logger()
+
+
+class ClearMLRegistryUploaderCallback(Callback):
+    """PyTorch Lightning callback to upload the best model checkpoint to ClearML."""
+
+    def __init__(
+        self,
+        task: Task,
+        ckpt_callback: ModelCheckpoint,
+        output_model_name: str,
+        output_model_tags: list[str] | None = None,
+        verbose: bool = True,
+        enable_tag_versioning: bool = True,
+        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.
+            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.ckpt_callback = ckpt_callback
+        self.output_model_name = output_model_name
+        self.output_model_tags = output_model_tags
+        self.verbose = verbose
+        self.uploading_frequency = uploading_frequency
+        self.enable_tag_versioning = enable_tag_versioning
+
+        self._output_model: OutputModel | None = None
+        self._last_best_model_path: str = ""
+        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")
+        config = pl_module.model_config
+        if config is not None:
+            config = config.to_dict()
+
+        return OutputModel(
+            task=self.task,
+            name=self.output_model_name,
+            framework="PyTorch",
+            tags=self.output_model_tags,
+            config_dict=config,
+        )
+
+    def _upload_best_checkpoint(self, pl_module: "KostylLightningModule") -> None:
+        current_best = self.ckpt_callback.best_model_path
+
+        if not current_best:
+            if self.verbose:
+                logger.info("No best model found yet to upload")
+            return
+
+        if current_best == self._last_best_model_path:
+            if self.verbose:
+                logger.info("Best model unchanged since last upload")
+            return
+
+        if self._output_model is None:
+            self._output_model = self._create_output_model(pl_module)
+
+        if self.verbose:
+            logger.info(f"Uploading best model from {current_best}")
+
+        self._output_model.update_weights(
+            current_best,
+            auto_delete_file=False,
+            async_enable=False,
+        )
+
+        self._last_best_model_path = current_best
+        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

{kostyl_toolkit-0.1.19 → kostyl_toolkit-0.1.20}/kostyl/utils/logging.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 import inspect
-import os
 import sys
 import uuid
+from collections import namedtuple
 from copy import deepcopy
 from functools import partialmethod
 from pathlib import Path
@@ -13,7 +13,6 @@ from typing import Literal
 from typing import cast
 
 from loguru import logger as _base_logger
-from torch.nn.modules.module import _IncompatibleKeys
 
 
 if TYPE_CHECKING:
@@ -27,6 +26,9 @@ else:
 
 try:
     import torch.distributed as dist
+    from torch.nn.modules.module import (
+        _IncompatibleKeys,  # pyright: ignore[reportAssignmentType]
+    )
 except Exception:
 
     class _Dummy:
@@ -38,7 +40,24 @@ except Exception:
         def is_initialized() -> bool:
             return False
 
+        @staticmethod
+        def get_rank() -> int:
+            return 0
+
+    class _IncompatibleKeys(
+        namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"]),
+    ):
+        __slots__ = ()
+
+        def __repr__(self) -> str:
+            if not self.missing_keys and not self.unexpected_keys:
+                return "<All keys matched successfully>"
+            return super().__repr__()
+
+        __str__ = __repr__
+
     dist = _Dummy()
+    _IncompatibleKeys = _IncompatibleKeys
 
 _once_lock = Lock()
 _once_keys: set[tuple[str, str]] = set()
@@ -106,7 +125,7 @@ def setup_logger(
             add_rank = False
 
     if add_rank:
-        rank = int(os.environ.get("RANK", "0"))
+        rank = dist.get_rank()
         channel = f"rank:{rank} - {base}"
     else:
         channel = base

{kostyl_toolkit-0.1.19 → kostyl_toolkit-0.1.20}/pyproject.toml

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

kostyl_toolkit-0.1.19/kostyl/ml/lightning/callbacks/registry_uploading.py

@@ -1,118 +0,0 @@
-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 lightning.pytorch.callbacks import ModelCheckpoint
-
-from kostyl.ml.clearml.logging_utils import find_version_in_tags
-from kostyl.ml.clearml.logging_utils import increment_version
-from kostyl.ml.lightning import KostylLightningModule
-from kostyl.utils.logging import setup_logger
-
-
-logger = setup_logger()
-
-
-class ClearMLRegistryUploaderCallback(Callback):
-    """PyTorch Lightning callback to upload the best model checkpoint to ClearML."""
-
-    def __init__(
-        self,
-        task: Task,
-        ckpt_callback: ModelCheckpoint,
-        output_model_name: str,
-        output_model_tags: list[str] | None = None,
-        verbose: bool = True,
-        uploading_frequency: Literal[
-            "after-every-eval", "on-train-end"
-        ] = "on-train-end",
-    ) -> None:
-        """
-        Initialize the callback.
-
-        Args:
-            task (Task): The ClearML task object.
-            ckpt_callback (ModelCheckpoint): The model checkpoint callback.
-            output_model_name (str): The name for the output model.
-            output_model_tags (list[str] | None, optional): Tags for the output model. Defaults to None, which is converted to an empty list.
-            verbose (bool, optional): Whether to log verbose messages. Defaults to True.
-            uploading_frequency (Literal["after-every-eval", "on-train-end"]): Frequency of uploading the model. Defaults to "on-train-end".
-
-        """
-        super().__init__()
-        if output_model_tags is None:
-            output_model_tags = []
-        self.task = task
-        self.ckpt_callback = ckpt_callback
-        self.output_model_name = output_model_name
-        self.output_model_tags = output_model_tags
-        self.verbose = verbose
-        self.uploading_frequency = uploading_frequency
-
-        self._output_model: OutputModel | None = None
-        self._last_best_model_path: str = ""
-        return
-
-    def _create_output_model(self, pl_module: KostylLightningModule) -> OutputModel:
-        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)
-
-        config = pl_module.model_config
-        if config is not None:
-            config = config.to_dict()
-
-        output_model = OutputModel(
-            task=self.task,
-            name=self.output_model_name,
-            framework="PyTorch",
-            tags=self.output_model_tags,
-            config_dict=config,
-        )
-        return output_model
-
-    def _upload_best_checkpoint(self, pl_module: KostylLightningModule) -> None:
-        if self._output_model is None:
-            self._output_model = self._create_output_model(pl_module)
-
-        if self.ckpt_callback.best_model_path == self._last_best_model_path:
-            if self.verbose and (self._last_best_model_path != ""):
-                logger.info("Best model unchanged since last upload")
-            elif self.verbose:
-                logger.info("No best model found yet to upload")
-        else:
-            if self.verbose:
-                logger.info(
-                    f"Uploading best model from {self.ckpt_callback.best_model_path}"
-                )
-            self._output_model.update_weights(
-                self.ckpt_callback.best_model_path,
-                auto_delete_file=False,
-                async_enable=False,
-            )
-        return
-
-    @override
-    def on_validation_epoch_end(
-        self, trainer: Trainer, pl_module: KostylLightningModule
-    ) -> None:
-        if (not trainer.is_global_zero) or (
-            self.uploading_frequency != "after-every-eval"
-        ):
-            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