kostyl-toolkit 0.1.36__py3-none-any.whl → 0.1.37__py3-none-any.whl

kostyl/ml/dist_utils.py

@@ -4,38 +4,61 @@ from typing import Literal
 
 import torch.distributed as dist
 
+from kostyl.utils.logging import KostylLogger
 from kostyl.utils.logging import setup_logger
 
 
-logger = setup_logger(add_rank=True)
+module_logger = setup_logger()
 
 
-def log_dist(msg: str, how: Literal["only-zero-rank", "world"]) -> None:
+def log_dist(
+    msg: str,
+    logger: KostylLogger | None = None,
+    level: Literal["info", "warning", "error", "warning_once", "debug"] = "info",
+    log_scope: Literal["only-zero-rank", "world"] = "world",
+    group: dist.ProcessGroup | None = None,
+) -> None:
     """
     Log a message in a distributed environment based on the specified verbosity level.
 
     Args:
         msg (str): The message to log.
-        how (Literal["only-zero-rank", "world"]): The verbosity level for logging.
+        log_scope (Literal["only-zero-rank", "world"]): The verbosity level for logging.
             - "only-zero-rank": Log only from the main process (rank 0).
             - "world": Log from all processes in the distributed environment.
+        logger (KostylLogger | None): The logger instance to use. If None, the module logger is used.
+        level (Literal["info", "warning", "error", "warning_once", "debug"]): The logging level.
+        group (dist.ProcessGroup | None): Optional process group used to determine ranks. Defaults to the global process group.
 
     """
-    match how:
-        case _ if not dist.is_initialized():
-            logger.warning_once(
-                "Distributed logging requested but torch.distributed is not initialized."
-            )
-            logger.info(msg)
+    if logger is None:
+        logger = module_logger
+
+    log_attr = getattr(logger, level, None)
+    if log_attr is None:
+        raise ValueError(f"Invalid logging level: {level}")
+
+    if not dist.is_initialized():
+        module_logger.warning_once(
+            "Distributed process group is not initialized; logging from all ranks."
+        )
+        log_attr(msg)
+        return
+
+    match log_scope:
         case "only-zero-rank":
-            if is_main_process():
-                logger.info(msg)
+            if group is None:
+                module_logger.debug(
+                    "No process group provided; assuming global group for rank check."
+                )
+                group = dist.group.WORLD
+            group_rank = dist.get_rank(group=group)
+            if dist.get_global_rank(group=group, group_rank=group_rank) == 0:  # pyright: ignore[reportArgumentType]
+                log_attr(msg)
         case "world":
-            logger.info(msg)
+            log_attr(msg)
         case _:
-            logger.warning_once(
-                f"Invalid logging verbosity level requested: {how}. Message not logged."
-            )
+            raise ValueError(f"Invalid logging verbosity level: {log_scope}")
     return
 
 
@@ -44,7 +67,7 @@ def scale_lrs_by_world_size(
     group: dist.ProcessGroup | None = None,
     config_name: str = "",
     inv_scale: bool = False,
-    verbose: Literal["only-zero-rank", "world"] | None = None,
+    verbose_level: Literal["only-zero-rank", "world"] | None = None,
 ) -> dict[str, float]:
     """
     Scale learning-rate configuration values to match the active distributed world size.
@@ -58,7 +81,7 @@ def scale_lrs_by_world_size(
             the target world size. Defaults to the global process group.
         config_name (str): Human-readable identifier included in log messages.
         inv_scale (bool): If True, use the inverse square-root scale factor.
-        verbose (Literal["only-zero-rank", "world"] | None): Verbosity level for logging scaled values.
+        verbose_level (Literal["only-zero-rank", "world"] | None): Verbosity level for logging scaled values.
             - "only-zero-rank": Log only from the main process (rank 0).
             - "world": Log from all processes in the distributed environment.
             -  None: No logging.
@@ -77,31 +100,30 @@ def scale_lrs_by_world_size(
     for name, value in lrs.items():
         old_value = value
         new_value = value * scale
-        if verbose is not None:
+        if verbose_level is not None:
             log_dist(
                 f"New {config_name} lr {name.upper()}: {new_value}; OLD: {old_value}",
-                verbose,
+                log_scope=verbose_level,
+                group=group,
             )
         lrs[name] = new_value
     return lrs
 
 
-def get_rank() -> int:
-    """Gets the rank of the current process in a distributed setting."""
-    if dist.is_initialized():
-        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"])
+def get_local_rank(group: dist.ProcessGroup | None = None) -> int:
+    """Gets the local rank of the current process in a distributed setting."""
+    if dist.is_initialized() and group is not None:
+        return dist.get_rank(group=group)
+    if "SLURM_LOCALID" in os.environ:
+        return int(os.environ["SLURM_LOCALID"])
     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."""
-    rank = get_rank()
+def is_local_zero_rank() -> bool:
+    """Checks if the current process is the main process (rank 0) for the local node in a distributed setting."""
+    rank = get_local_rank()
     if rank != 0:
         return False
     return True

kostyl/ml/lightning/callbacks/checkpoint.py

@@ -10,7 +10,7 @@ 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.dist_utils import is_local_zero_rank
 from kostyl.ml.lightning import KostylLightningModule
 from kostyl.ml.registry_uploader import RegistryUploaderCallback
 from kostyl.utils import setup_logger
@@ -339,7 +339,7 @@ def setup_checkpoint_callback(
         )
 
     if dirpath.exists():
-        if is_main_process():
+        if is_local_zero_rank():
             logger.warning(f"Checkpoint directory {dirpath} already exists.")
             if remove_folder_if_exists:
                 rmtree(dirpath)

kostyl/ml/lightning/extensions/custom_module.py

@@ -26,11 +26,6 @@ module_logger = setup_logger(fmt="only_message")
 class KostylLightningModule(L.LightningModule):
     """Custom PyTorch Lightning Module with logging, checkpointing, and distributed training utilities."""
 
-    @property
-    def process_group(self) -> ProcessGroup | None:
-        """Returns the data parallel process group for distributed training."""
-        return self.get_process_group()
-
     def get_process_group(self) -> ProcessGroup | None:
         """
         Retrieves the data parallel process group for distributed training.

kostyl/ml/lightning/extensions/pretrained_model.py

@@ -12,12 +12,12 @@ from kostyl.utils.logging import setup_logger
 logger = setup_logger("LightningPretrainedModelMixin", fmt="only_message")
 
 
-class LightningCheckpointLoaderMixin(PreTrainedModel):
+class LightningCheckpointLoaderMixin:
     """A mixin class for loading pretrained models from PyTorch Lightning checkpoints."""
 
     @classmethod
-    def from_lightning_checkpoint[TModelInstance: LightningCheckpointLoaderMixin](  # noqa: C901
-        cls: type[TModelInstance],
+    def from_lightning_checkpoint[TModelInstance: PreTrainedModel](  # noqa: C901
+        cls: type[TModelInstance],  # pyright: ignore[reportGeneralTypeIssues]
         checkpoint_path: str | Path,
         config_key: str = "config",
         weights_prefix: str | None = "model.",
@@ -78,7 +78,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             mmap=True,
         )
 
-        # 1. Восстанавливаем конфиг
+        # Load config
         config_cls = cast(type[PretrainedConfig], cls.config_class)
         config_dict = checkpoint_dict[config_key]
         config_dict.update(kwargs)
@@ -91,6 +91,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
 
         raw_state_dict: dict[str, torch.Tensor] = checkpoint_dict["state_dict"]
 
+        # Handle weights prefix
         if weights_prefix:
             if not weights_prefix.endswith("."):
                 weights_prefix = weights_prefix + "."
@@ -117,6 +118,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
         else:
             state_dict = raw_state_dict
 
+        # Instantiate model and load state dict
         model = cls.from_pretrained(
             pretrained_model_name_or_path=None,
             config=config,

kostyl/ml/lightning/loggers/tb_logger.py

@@ -3,7 +3,7 @@ from shutil import rmtree
 
 from lightning.pytorch.loggers import TensorBoardLogger
 
-from kostyl.ml.dist_utils import is_main_process
+from kostyl.ml.dist_utils import is_local_zero_rank
 from kostyl.utils.logging import setup_logger
 
 
@@ -15,7 +15,7 @@ def setup_tb_logger(
 ) -> TensorBoardLogger:
     """Sets up a TensorBoardLogger for PyTorch Lightning."""
     if runs_dir.exists():
-        if is_main_process():
+        if is_local_zero_rank():
             logger.warning(f"TensorBoard log directory {runs_dir} already exists.")
             rmtree(runs_dir)
             logger.warning(f"Removed existing TensorBoard log directory {runs_dir}.")

kostyl/ml/lightning/utils.py

@@ -14,7 +14,7 @@ TRAINING_STRATEGIES = (
     FSDP1StrategyConfig | DDPStrategyConfig | SingleDeviceStrategyConfig
 )
 
-logger = setup_logger(add_rank=True)
+logger = setup_logger()
 
 
 def estimate_total_steps(

kostyl/utils/logging.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import inspect
+import os
 import sys
 import uuid
 from collections import namedtuple
@@ -18,32 +19,18 @@ from loguru import logger as _base_logger
 if TYPE_CHECKING:
     from loguru import Logger
 
-    class CustomLogger(Logger):  # noqa: D101
+    class KostylLogger(Logger):  # noqa: D101
         def log_once(self, level: str, message: str, *args, **kwargs) -> None: ...  # noqa: ANN003, D102
         def warning_once(self, message: str, *args, **kwargs) -> None: ...  # noqa: ANN003, D102
 else:
-    CustomLogger = type(_base_logger)
+    KostylLogger = type(_base_logger)
 
 try:
-    import torch.distributed as dist
     from torch.nn.modules.module import (
         _IncompatibleKeys,  # pyright: ignore[reportAssignmentType]
     )
 except Exception:
 
-    class _Dummy:
-        @staticmethod
-        def is_available() -> bool:
-            return False
-
-        @staticmethod
-        def is_initialized() -> bool:
-            return False
-
-        @staticmethod
-        def get_rank() -> int:
-            return 0
-
     class _IncompatibleKeys(
         namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"]),
     ):
@@ -56,14 +43,13 @@ except Exception:
 
         __str__ = __repr__
 
-    dist = _Dummy()
     _IncompatibleKeys = _IncompatibleKeys
 
 _once_lock = Lock()
 _once_keys: set[tuple[str, str]] = set()
 
 
-def _log_once(self: CustomLogger, level: str, message: str, *args, **kwargs) -> None:  # noqa: ANN003
+def _log_once(self: KostylLogger, level: str, message: str, *args, **kwargs) -> None:  # noqa: ANN003
     key = (message, level)
 
     with _once_lock:
@@ -75,7 +61,7 @@ def _log_once(self: CustomLogger, level: str, message: str, *args, **kwargs) ->
     return
 
 
-_base_logger = cast(CustomLogger, _base_logger)
+_base_logger = cast(KostylLogger, _base_logger)
 _base_logger.log_once = _log_once  # pyright: ignore[reportAttributeAccessIssue]
 _base_logger.warning_once = partialmethod(_log_once, "WARNING")  # pyright: ignore[reportAttributeAccessIssue]
 
@@ -91,44 +77,83 @@ _DEFAULT_FMT = "<level>{level: <8}</level> {time:HH:mm:ss.SSS} [{extra[channel]}
 _ONLY_MESSAGE_FMT = "<level>{message}</level>"
 _PRESETS = {"default": _DEFAULT_FMT, "only_message": _ONLY_MESSAGE_FMT}
 
+KOSTYL_LOG_LEVEL = os.getenv("KOSTYL_LOG_LEVEL", "INFO")
+
 
 def setup_logger(
     name: str | None = None,
     fmt: Literal["default", "only_message"] | str = "only_message",
-    level: str = "INFO",
-    add_rank: bool | None = None,
+    level: str | None = None,
     sink=sys.stdout,
     colorize: bool = True,
     serialize: bool = False,
-) -> CustomLogger:
+) -> KostylLogger:
     """
-    Returns a bound logger with its own sink and formatting.
+    Creates and configures a logger with custom formatting and output.
+
+    The function automatically removes the default sink on first call and creates
+    an isolated logger with a unique identifier for message filtering.
+
+    Args:
+        name (str | None, optional): Logger channel name. If None, automatically
+            uses the calling function's filename. Defaults to None.
+        fmt (Literal["default", "only_message"] | str, optional): Log message format.
+            Available presets:
+            - "default": includes level, time, and channel
+            - "only_message": outputs only the message itself
+            Custom format strings are also supported. Defaults to "only_message".
+        level (str | None, optional): Logging level (TRACE, DEBUG, INFO, SUCCESS,
+            WARNING, ERROR, CRITICAL). If None, uses the KOSTYL_LOG_LEVEL environment
+            variable or "INFO" by default. Defaults to None.
+        sink: Output object for logs (file, sys.stdout, sys.stderr, etc.).
+            Defaults to sys.stdout.
+        colorize (bool, optional): Enable colored output formatting.
+            Defaults to True.
+        serialize (bool, optional): Serialize logs to JSON format.
+            Defaults to False.
+
+    Returns:
+        CustomLogger: Configured logger instance with additional methods
+            log_once() and warning_once().
+
+    Example:
+        >>> # Basic usage with automatic name detection
+        >>> logger = setup_logger()
+        >>> logger.info("Hello World")
 
-    Note: If name=None, the caller's filename (similar to __file__) is used automatically.
+        >>> # With custom name and level
+        >>> logger = setup_logger(name="MyApp", level="DEBUG")
+
+        >>> # With custom format
+        >>> logger = setup_logger(
+        ...     name="API",
+        ...     fmt="{level} | {time:YYYY-MM-DD HH:mm:ss} | {message}"
+        ... )
 
-    Format example: "{level} {time:MM-DD HH:mm:ss} [{extra[channel]}] {message}"
     """
     global _DEFAULT_SINK_REMOVED
     if not _DEFAULT_SINK_REMOVED:
         _base_logger.remove()
         _DEFAULT_SINK_REMOVED = True
 
-    if name is None:
-        base = _caller_filename()
-    else:
-        base = name
+    if level is None:
+        if KOSTYL_LOG_LEVEL not in {
+            "TRACE",
+            "DEBUG",
+            "INFO",
+            "SUCCESS",
+            "WARNING",
+            "ERROR",
+            "CRITICAL",
+        }:
+            level = "INFO"
+        else:
+            level = KOSTYL_LOG_LEVEL
 
-    if (add_rank is None) or add_rank:
-        try:
-            add_rank = dist.is_available() and dist.is_initialized()
-        except Exception:
-            add_rank = False
-
-    if add_rank:
-        rank = dist.get_rank()
-        channel = f"rank:{rank} - {base}"
+    if name is None:
+        channel = _caller_filename()
     else:
-        channel = base
+        channel = name
 
     if fmt in _PRESETS:
         fmt = _PRESETS[fmt]
@@ -146,7 +171,7 @@ def setup_logger(
         filter=lambda r: r["extra"].get("logger_id") == logger_id,
     )
     logger = _base_logger.bind(logger_id=logger_id, channel=channel)
-    return cast(CustomLogger, logger)
+    return cast(KostylLogger, logger)
 
 
 def log_incompatible_keys(
@@ -154,22 +179,12 @@ def log_incompatible_keys(
     incompatible_keys: _IncompatibleKeys
     | tuple[list[str], list[str]]
     | dict[str, list[str]],
-    model_specific_msg: str = "",
+    postfix_msg: str = "",
 ) -> None:
     """
     Logs warnings for incompatible keys encountered during model loading or state dict operations.
 
     Note: If incompatible_keys is of an unsupported type, an error message is logged and the function returns early.
-
-    Args:
-        logger (Logger): The logger instance used to output warning messages.
-        incompatible_keys (_IncompatibleKeys | tuple[list[str], list[str]] | dict[str, list[str]]): An object containing lists of missing and unexpected keys.
-        model_specific_msg (str, optional): A custom message to append to the log output, typically
-            indicating the model or context. Defaults to an empty string.
-
-    Returns:
-        None
-
     """
     incompatible_keys_: dict[str, list[str]] = {}
     match incompatible_keys:
@@ -192,5 +207,5 @@ def log_incompatible_keys(
             return
 
     for name, keys in incompatible_keys_.items():
-        logger.warning(f"{name} {model_specific_msg}: {', '.join(keys)}")
+        logger.warning(f"{name} {postfix_msg}: {', '.join(keys)}")
     return

{kostyl_toolkit-0.1.36.dist-info → kostyl_toolkit-0.1.37.dist-info}/METADATA

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

@@ -9,17 +9,17 @@ kostyl/ml/configs/base_model.py,sha256=Eofn14J9RsjpVx_J4rp6C19pDDCANU4hr3JtX-d0F
 kostyl/ml/configs/hyperparams.py,sha256=lvtbvOFEoTBAJug7FR35xMQdPLgDQjRoP2fyDP-jD7E,3305
 kostyl/ml/configs/training_settings.py,sha256=wT9CHuLaKrLwonsc87Ee421EyFis_c9fqOgn9bSClm8,2747
 kostyl/ml/data_processing_utils.py,sha256=jjEjV0S0wREgZkzg27ip0LpI8cQqkwe2QwATmAqm9-g,3832
-kostyl/ml/dist_utils.py,sha256=Onf0KHVLA8oeUgZTcTdmR9qiM22f2uYLoNwgLbMGJWk,3495
+kostyl/ml/dist_utils.py,sha256=lK9_aAh9L1SvvXWzcWiBoFjczfDiKzEpcno5csImAYQ,4635
 kostyl/ml/lightning/__init__.py,sha256=R36PImjVvzBF9t_z9u6RYVnUFJJ-sNDUOdboWUojHmM,173
 kostyl/ml/lightning/callbacks/__init__.py,sha256=EnKkNwwNDZnEqKRlpY4FVrqP88ECPF6nlT2bSLUIKRk,194
-kostyl/ml/lightning/callbacks/checkpoint.py,sha256=COW7WErj4EMxJNMn97WQO-G2A3LbI6GQOCpIZu3Cblk,19060
+kostyl/ml/lightning/callbacks/checkpoint.py,sha256=HI17gu-GxnfXUchflWBTwly7cCYnlpKcshuR-TgD6s4,19066
 kostyl/ml/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
 kostyl/ml/lightning/extensions/__init__.py,sha256=OY6QGv1agYgqqKf1xJBrxgp_i8FunVfPzYezfaRrGXU,182
-kostyl/ml/lightning/extensions/custom_module.py,sha256=iQrnPz-WTmRfvLo94C5fQc2Qwa1IpHtUh1sCpVwTSFM,6602
-kostyl/ml/lightning/extensions/pretrained_model.py,sha256=eRfQBzAjVernHl9A4PP5uTLvjjmcNKPdTu7ABFLq7HI,5196
+kostyl/ml/lightning/extensions/custom_module.py,sha256=qYffgPwIB_ePwK_MIaRruuDxPKJZb42kg2yy996eGwY,6415
+kostyl/ml/lightning/extensions/pretrained_model.py,sha256=hVIsIUu6Iryrz6S7GQTqog9vNq8LQyjJd2aoJ5Ws6KU,5253
 kostyl/ml/lightning/loggers/__init__.py,sha256=e51dszaoJbuzwBkbdugmuDsPldoSO4yaRgmZUg1Bdy0,71
-kostyl/ml/lightning/loggers/tb_logger.py,sha256=j02HK5ue8yzXXV8FWKmmXyHkFlIxgHx-ahHWk_rFCZs,893
-kostyl/ml/lightning/utils.py,sha256=imvMbgOKRtCUiiRGEcVtN-hxw-aEFKHdCWc0J_CIZp4,1980
+kostyl/ml/lightning/loggers/tb_logger.py,sha256=CpjlcEIT187cJXJgRYafqfzvcnwPgPaVZ0vLUflIr7k,899
+kostyl/ml/lightning/utils.py,sha256=DhLy_3JA5VyMQkB1v6xxRxDNHfisjXFYVjuIKPpO81M,1967
 kostyl/ml/metrics_formatting.py,sha256=U6vdNENZLvp2dT1L3HqFKtXrHwGKoDXN93hvamPGHjM,1341
 kostyl/ml/params_groups.py,sha256=nUyw5d06Pvy9QPiYtZzLYR87xwXqJLxbHthgQH8oSCM,3583
 kostyl/ml/registry_uploader.py,sha256=BbyLXvF8AL145k7g6MRkJ7gf_3Um53p3Pn5280vVD9U,4384
@@ -32,7 +32,7 @@ kostyl/ml/schedulers/linear.py,sha256=RnnnblRuRXP3LT03QVIHUaK2kNsiMP1AedrMoeyh3q
 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=LSbyQFLAIa89xPb4tcobE2BwVIHHUSaDXqOIKVzLoWs,5801
-kostyl_toolkit-0.1.36.dist-info/WHEEL,sha256=eycQt0QpYmJMLKpE3X9iDk8R04v2ZF0x82ogq-zP6bQ,79
-kostyl_toolkit-0.1.36.dist-info/METADATA,sha256=Lfyx6u3LKZ6co4s7GZgJp31zoy-NViriSGqwjIzOQFA,4269
-kostyl_toolkit-0.1.36.dist-info/RECORD,,
+kostyl/utils/logging.py,sha256=CgNFNogcK0hoZmygvBWlTcq5A3m2Pfv9eOAP_gwx0pM,6633
+kostyl_toolkit-0.1.37.dist-info/WHEEL,sha256=eycQt0QpYmJMLKpE3X9iDk8R04v2ZF0x82ogq-zP6bQ,79
+kostyl_toolkit-0.1.37.dist-info/METADATA,sha256=yHPgSAhPnm5tDQjvDIfs213-bsVX6vMfVsUbX9GboGU,4269
+kostyl_toolkit-0.1.37.dist-info/RECORD,,