kostyl-toolkit 0.1.18__py3-none-any.whl → 0.1.20__py3-none-any.whl

kostyl/ml/dist_utils.py

@@ -86,14 +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."""
-    return _get_rank() == 0
+    rank = get_rank()
+    if rank != 0:
+        return False
+    return True

kostyl/ml/lightning/callbacks/registry_uploading.py

@@ -26,93 +26,113 @@ class ClearMLRegistryUploaderCallback(Callback):
         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:
         """
-        Initialize the callback.
+        Initializes the ClearMLRegistryUploaderCallback.
 
         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".
+            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:
-        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)
-
+    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()
 
-        output_model = OutputModel(
+        return 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)
+    def _upload_best_checkpoint(self, pl_module: "KostylLightningModule") -> None:
+        current_best = self.ckpt_callback.best_model_path
 
-        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:
+        if not current_best:
+            if self.verbose:
                 logger.info("No best model found yet to upload")
-        else:
+            return
+
+        if current_best == self._last_best_model_path:
             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,
-            )
+                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_epoch_end(
-        self, trainer: Trainer, pl_module: KostylLightningModule
+    def on_validation_end(
+        self, trainer: Trainer, pl_module: "KostylLightningModule"
     ) -> None:
-        if (not trainer.is_global_zero) or (
-            self.uploading_frequency != "after-every-eval"
-        ):
+        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:
+    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/ml/lightning/loggers/tb_logger.py

@@ -7,7 +7,7 @@ from kostyl.ml.dist_utils import is_main_process
 from kostyl.utils.logging import setup_logger
 
 
-logger = setup_logger("callbacks/tb_logger.py", add_rank=True)
+logger = setup_logger()
 
 
 def setup_tb_logger(

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.18.dist-info → kostyl_toolkit-0.1.20.dist-info}/METADATA

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

@@ -8,17 +8,17 @@ kostyl/ml/configs/__init__.py,sha256=IetcivbqYGutowLqxdKp7QR4tkXKBr4m8t4Zkk9jHZU
 kostyl/ml/configs/base_model.py,sha256=Eofn14J9RsjpVx_J4rp6C19pDDCANU4hr3JtX-d0FpQ,4820
 kostyl/ml/configs/hyperparams.py,sha256=CaVNEvpW4LvlHhLsbe2FockIGI1mJufCqjH298nYgKE,2971
 kostyl/ml/configs/training_settings.py,sha256=0cyKF6EuTv6KgXC1g0Oy9zbwnMkDP4uXTJJO1TRQ0aY,2556
-kostyl/ml/dist_utils.py,sha256=G8atjzkRbXZZiZh9rdEYBmeXqX26rJdDDovft2n6xiU,3201
+kostyl/ml/dist_utils.py,sha256=Onf0KHVLA8oeUgZTcTdmR9qiM22f2uYLoNwgLbMGJWk,3495
 kostyl/ml/lightning/__init__.py,sha256=-F3JAyq8KU1d-nACWryGu8d1CbvWbQ1rXFdeRwfE2X8,175
 kostyl/ml/lightning/callbacks/__init__.py,sha256=Vd-rozY4T9Prr3IMqbliXxj6sC6y9XsovHQqRwzc2HI,297
 kostyl/ml/lightning/callbacks/checkpoint.py,sha256=FooGeeUz6TtoXQglpcK16NWAmSX3fbu6wntRtK3a_Io,1936
 kostyl/ml/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
-kostyl/ml/lightning/callbacks/registry_uploading.py,sha256=AAohqTGigWMGLHt-yITAznIVsaDVXR_jisLDzFNzlu8,4275
+kostyl/ml/lightning/callbacks/registry_uploading.py,sha256=jJdSoFIkTcGLnZIKBzbAkt9MAgpZURLCQdd0DwAC5gk,4659
 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=ZOKtrVl095cwvI43wAz-Xdzu4l0v0lHH2mfh4WXwxKQ,5059
 kostyl/ml/lightning/loggers/__init__.py,sha256=e51dszaoJbuzwBkbdugmuDsPldoSO4yaRgmZUg1Bdy0,71
-kostyl/ml/lightning/loggers/tb_logger.py,sha256=I0WQV3SVVIds9kHhAYnN3SXodX64jfTk7_5gH-j3uzA,932
+kostyl/ml/lightning/loggers/tb_logger.py,sha256=j02HK5ue8yzXXV8FWKmmXyHkFlIxgHx-ahHWk_rFCZs,893
 kostyl/ml/lightning/steps_estimation.py,sha256=fTZ0IrUEZV3H6VYlx4GYn56oco56mMiB7FO9F0Z7qc4,1511
 kostyl/ml/metrics_formatting.py,sha256=w0rTz61z0Um_d2pomYLvcQFcZX_C-KolZcIPRsa1efE,1421
 kostyl/ml/params_groups.py,sha256=nUyw5d06Pvy9QPiYtZzLYR87xwXqJLxbHthgQH8oSCM,3583
@@ -29,7 +29,7 @@ kostyl/ml/schedulers/cosine.py,sha256=jufULVHn_L_ZZEc3ZTG3QCY_pc0jlAMH5Aw496T31j
 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=3MvfDPArZhwakHu5nMlp_LpOsWg0E0SP26y41clsBtA,5232
-kostyl_toolkit-0.1.18.dist-info/WHEEL,sha256=3id4o64OvRm9dUknh3mMJNcfoTRK08ua5cU6DFyVy-4,79
-kostyl_toolkit-0.1.18.dist-info/METADATA,sha256=SPdsj3qr-Rdi7Rr0yblo51avTETIxA3irGHIJt2sW6Q,4269
-kostyl_toolkit-0.1.18.dist-info/RECORD,,
+kostyl/utils/logging.py,sha256=CQmPZ1x7yiVz56OkK6IZCtdoHs_Owo7fxJ03oOct-Qc,5782
+kostyl_toolkit-0.1.20.dist-info/WHEEL,sha256=3id4o64OvRm9dUknh3mMJNcfoTRK08ua5cU6DFyVy-4,79
+kostyl_toolkit-0.1.20.dist-info/METADATA,sha256=HaAS7ZSdaBNwoyKc6juh3nTDfEgQBPbWtMhAZg7Xg3s,4269
+kostyl_toolkit-0.1.20.dist-info/RECORD,,