kostyl-toolkit 0.1.36__tar.gz → 0.1.37__tar.gz

{kostyl_toolkit-0.1.36 → kostyl_toolkit-0.1.37}/PKG-INFO

@@ -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.37/kostyl/ml/dist_utils.py

@@ -0,0 +1,129 @@
+import math
+import os
+from typing import Literal
+
+import torch.distributed as dist
+
+from kostyl.utils.logging import KostylLogger
+from kostyl.utils.logging import setup_logger
+
+
+module_logger = setup_logger()
+
+
+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.
+        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.
+
+    """
+    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 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":
+            log_attr(msg)
+        case _:
+            raise ValueError(f"Invalid logging verbosity level: {log_scope}")
+    return
+
+
+def scale_lrs_by_world_size(
+    lrs: dict[str, float],
+    group: dist.ProcessGroup | None = None,
+    config_name: str = "",
+    inv_scale: bool = False,
+    verbose_level: Literal["only-zero-rank", "world"] | None = None,
+) -> dict[str, float]:
+    """
+    Scale learning-rate configuration values to match the active distributed world size.
+
+    Note:
+        The value in the `lrs` will be modified in place.
+
+    Args:
+        lrs (dict[str, float]): A dictionary of learning rate names and their corresponding values to be scaled.
+        group (dist.ProcessGroup | None): Optional process group used to determine
+            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_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.
+
+    Returns:
+        dict[str, float]: The learning-rate configuration with scaled values.
+
+    """
+    world_size = dist.get_world_size(group=group)
+
+    if inv_scale:
+        scale = 1 / math.sqrt(world_size)
+    else:
+        scale = math.sqrt(world_size)
+
+    for name, value in lrs.items():
+        old_value = value
+        new_value = value * scale
+        if verbose_level is not None:
+            log_dist(
+                f"New {config_name} lr {name.upper()}: {new_value}; OLD: {old_value}",
+                log_scope=verbose_level,
+                group=group,
+            )
+        lrs[name] = new_value
+    return lrs
+
+
+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_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_toolkit-0.1.36 → kostyl_toolkit-0.1.37}/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_toolkit-0.1.36 → kostyl_toolkit-0.1.37}/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_toolkit-0.1.36 → kostyl_toolkit-0.1.37}/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_toolkit-0.1.36 → kostyl_toolkit-0.1.37}/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_toolkit-0.1.36 → kostyl_toolkit-0.1.37}/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_toolkit-0.1.36 → kostyl_toolkit-0.1.37}/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 → kostyl_toolkit-0.1.37}/pyproject.toml

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

kostyl_toolkit-0.1.36/kostyl/ml/dist_utils.py

@@ -1,107 +0,0 @@
-import math
-import os
-from typing import Literal
-
-import torch.distributed as dist
-
-from kostyl.utils.logging import setup_logger
-
-
-logger = setup_logger(add_rank=True)
-
-
-def log_dist(msg: str, how: Literal["only-zero-rank", "world"]) -> 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.
-            - "only-zero-rank": Log only from the main process (rank 0).
-            - "world": Log from all processes in the distributed environment.
-
-    """
-    match how:
-        case _ if not dist.is_initialized():
-            logger.warning_once(
-                "Distributed logging requested but torch.distributed is not initialized."
-            )
-            logger.info(msg)
-        case "only-zero-rank":
-            if is_main_process():
-                logger.info(msg)
-        case "world":
-            logger.info(msg)
-        case _:
-            logger.warning_once(
-                f"Invalid logging verbosity level requested: {how}. Message not logged."
-            )
-    return
-
-
-def scale_lrs_by_world_size(
-    lrs: dict[str, float],
-    group: dist.ProcessGroup | None = None,
-    config_name: str = "",
-    inv_scale: bool = False,
-    verbose: Literal["only-zero-rank", "world"] | None = None,
-) -> dict[str, float]:
-    """
-    Scale learning-rate configuration values to match the active distributed world size.
-
-    Note:
-        The value in the `lrs` will be modified in place.
-
-    Args:
-        lrs (dict[str, float]): A dictionary of learning rate names and their corresponding values to be scaled.
-        group (dist.ProcessGroup | None): Optional process group used to determine
-            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.
-            - "only-zero-rank": Log only from the main process (rank 0).
-            - "world": Log from all processes in the distributed environment.
-            -  None: No logging.
-
-    Returns:
-        dict[str, float]: The learning-rate configuration with scaled values.
-
-    """
-    world_size = dist.get_world_size(group=group)
-
-    if inv_scale:
-        scale = 1 / math.sqrt(world_size)
-    else:
-        scale = math.sqrt(world_size)
-
-    for name, value in lrs.items():
-        old_value = value
-        new_value = value * scale
-        if verbose is not None:
-            log_dist(
-                f"New {config_name} lr {name.upper()}: {new_value}; OLD: {old_value}",
-                verbose,
-            )
-        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"])
-    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()
-    if rank != 0:
-        return False
-    return True