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

kostyl/ml/base_uploader.py

@@ -0,0 +1,17 @@
+from abc import ABC
+from abc import abstractmethod
+from pathlib import Path
+
+from kostyl.utils.logging import setup_logger
+
+
+logger = setup_logger()
+
+
+class ModelCheckpointUploader(ABC):
+    """Abstract base class for uploading model checkpoints to a registry backend."""
+
+    @abstractmethod
+    def upload_checkpoint(self, path: str | Path) -> None:
+        """Upload the checkpoint located at the given path to the configured registry backend."""
+        raise NotImplementedError

kostyl/ml/configs/__init__.py

@@ -1,8 +1,8 @@
-from .base_model import KostylBaseModel
 from .hyperparams import HyperparamsConfig
 from .hyperparams import Lr
 from .hyperparams import Optimizer
 from .hyperparams import WeightDecay
+from .mixins import ConfigLoadingMixin
 from .training_settings import CheckpointConfig
 from .training_settings import DataConfig
 from .training_settings import DDPStrategyConfig
@@ -15,12 +15,12 @@ from .training_settings import TrainingSettings
 
 __all__ = [
     "CheckpointConfig",
+    "ConfigLoadingMixin",
     "DDPStrategyConfig",
     "DataConfig",
     "EarlyStoppingConfig",
     "FSDP1StrategyConfig",
     "HyperparamsConfig",
-    "KostylBaseModel",
     "LightningTrainerParameters",
     "Lr",
     "Optimizer",

kostyl/ml/configs/mixins.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+from pydantic import BaseModel as PydanticBaseModel
+
+from kostyl.utils.fs import load_config
+
+
+class ConfigLoadingMixin[TConfig: PydanticBaseModel]:
+    """Mixin providing configuration loading functionality for Pydantic models."""
+
+    @classmethod
+    def from_file(
+        cls: type[TConfig],  # pyright: ignore
+        path: str | Path,
+    ) -> TConfig:
+        """
+        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[TConfig],  # pyright: ignore
+        state_dict: dict,
+    ) -> TConfig:
+        """
+        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

kostyl/ml/{data_processing_utils.py → data_collator.py}

@@ -36,6 +36,7 @@ class BatchCollatorWithKeyAlignment:
         keys_mapping: A dictionary mapping original keys to new keys.
         keys_to_keep: A set of keys to retain as-is from the original items.
         max_length: If provided, truncates "input_ids" and "attention_mask" to this length.
+            Only 1D tensors/lists are supported.
 
         Raises:
             ValueError: If both `keys_mapping` and `keys_to_keep` are None.
@@ -59,14 +60,16 @@ class BatchCollatorWithKeyAlignment:
     def _truncate_data(self, key: str, value: Any) -> Any:
         match value:
             case torch.Tensor():
-                if value.dim() > 2:
+                if value.dim() >= 2:
                     raise ValueError(
-                        f"Expected value with dim <= 2 for key {key}, got {value.dim()}"
+                        f"Expected tensor with dim < 2 for key {key}, got {value.dim()}. "
+                        "Check your data or disable truncation with `max_length=None`."
                     )
             case list():
                 if isinstance(value[0], list):
                     raise ValueError(
-                        f"Expected value with dim <= 2 for key {key}, got nested lists"
+                        f"Expected value with dim <= 2 for key {key}, got nested lists. "
+                        "Check your data or disable truncation with `max_length=None`."
                     )
         value = value[: self.max_length]
         return value

kostyl/ml/dist_utils.py

@@ -4,47 +4,69 @@ 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 the current process only."
+        )
+        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
 
 
 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,
+    verbose_level: Literal["only-zero-rank", "world"] | None = None,
 ) -> dict[str, float]:
     """
     Scale learning-rate configuration values to match the active distributed world size.
@@ -56,9 +78,8 @@ def scale_lrs_by_world_size(
         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.
+        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 +98,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,
+                f"lr {name.upper()}: {new_value}; OLD: {old_value}",
+                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/integrations/clearml/__init__.py

@@ -0,0 +1,7 @@
+try:
+    import clearml  # noqa: F401
+except ImportError as e:
+    raise ImportError(
+        "ClearML integration requires the 'clearml' package. "
+        "Please install it via 'pip install clearml'."
+    ) from e

kostyl/ml/{registry_uploader.py → integrations/clearml/checkpoint_uploader.py}

@@ -1,5 +1,3 @@
-from abc import ABC
-from abc import abstractmethod
 from collections.abc import Callable
 from functools import partial
 from pathlib import Path
@@ -7,22 +5,14 @@ from typing import override
 
 from clearml import OutputModel
 
+from kostyl.ml.base_uploader import ModelCheckpointUploader
 from kostyl.utils.logging import setup_logger
 
 
 logger = setup_logger()
 
 
-class RegistryUploaderCallback(ABC):
-    """Abstract Lightning callback responsible for tracking and uploading the best-performing model checkpoint."""
-
-    @abstractmethod
-    def upload_checkpoint(self, path: str | Path) -> None:
-        """Upload the checkpoint located at the given path to the configured registry backend."""
-        raise NotImplementedError
-
-
-class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
+class ClearMLCheckpointUploader(ModelCheckpointUploader):
     """PyTorch Lightning callback to upload the best model checkpoint to ClearML."""
 
     def __init__(
@@ -38,7 +28,7 @@ class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
         verbose: bool = True,
     ) -> None:
         """
-        Initializes the ClearMLRegistryUploaderCallback.
+        Initializes the ClearMLRegistryUploader.
 
         Args:
             model_name: The name for the newly created model.

kostyl/ml/{configs/base_model.py → integrations/clearml/config_mixin.py}

@@ -1,75 +1,25 @@
 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.ml.configs import ConfigLoadingMixin
 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."""
+class BaseModelWithClearmlSyncing[TConfig: ConfigLoadingMixin]:
+    """Mixin providing ClearML task configuration syncing functionality for Pydantic models."""
 
     @classmethod
     def connect_as_file(
-        cls: type[Self],  # pyright: ignore
+        cls: type[TConfig],  # pyright: ignore
         task: Task,
         path: str | Path,
         alias: str | None = None,
-    ) -> Self:
+    ) -> TConfig:
         """
         Connects the configuration file to a ClearML task and creates an instance of the class from it.
 
@@ -104,11 +54,11 @@ class BaseModelWithClearmlSyncing(BaseModelWithConfigLoading):
 
     @classmethod
     def connect_as_dict(
-        cls: type[Self],  # pyright: ignore
+        cls: type[TConfig],  # pyright: ignore
         task: Task,
         path: str | Path,
         alias: str | None = None,
-    ) -> Self:
+    ) -> TConfig:
         """
         Connects configuration from a file as a dictionary to a ClearML task and creates an instance of the class.
 
@@ -135,9 +85,3 @@ class BaseModelWithClearmlSyncing(BaseModelWithConfigLoading):
 
         model = cls.from_dict(state_dict=config)
         return model
-
-
-class KostylBaseModel(BaseModelWithClearmlSyncing):
-    """A Pydantic model class with basic configuration loading functionality."""
-
-    pass

kostyl/ml/{clearml/pulling_utils.py → integrations/clearml/loading_utils.py}

@@ -9,9 +9,26 @@ from transformers import AutoTokenizer
 from transformers import PreTrainedModel
 from transformers import PreTrainedTokenizerBase
 
-from kostyl.ml.lightning.extensions.pretrained_model import (
-    LightningCheckpointLoaderMixin,
-)
+
+try:
+    from kostyl.ml.integrations.lightning import (
+        LightningCheckpointLoaderMixin,  # pyright: ignore[reportAssignmentType]
+    )
+
+    LIGHTING_MIXIN_AVAILABLE = True
+except ImportError:
+
+    class LightningCheckpointLoaderMixin(PreTrainedModel):  # noqa: D101
+        pass  # type: ignore
+
+    @classmethod
+    def from_lightning_checkpoint(cls, *args: Any, **kwargs: Any) -> Any:  # noqa: D103
+        raise ImportError(
+            "Loading from Lightning checkpoints requires lightning integration. "
+            "Please package install via 'pip install lightning' to enable this functionality."
+        )
+
+    LIGHTING_MIXIN_AVAILABLE = False
 
 
 def get_tokenizer_from_clearml(
@@ -89,13 +106,23 @@ def get_model_from_clearml[
     local_path = Path(input_model.get_local_copy(raise_on_error=True))
 
     if local_path.is_dir() and input_model._is_package():
+        if not issubclass(model, (PreTrainedModel, AutoModel)):
+            raise ValueError(
+                f"Model class {model.__name__} must be a subclass of PreTrainedModel or AutoModel for directory loads."
+            )
         model_instance = model.from_pretrained(local_path, **kwargs)
     elif local_path.suffix == ".ckpt":
+        if not LIGHTING_MIXIN_AVAILABLE:
+            raise ImportError(
+                "Loading from Lightning checkpoints requires lightning integration. "
+                "Please package install via 'pip install lightning' to enable this functionality."
+            )
         if not issubclass(model, LightningCheckpointLoaderMixin):
             raise ValueError(
-                f"Model class {model.__name__} is not compatible with Lightning checkpoints."
+                f"Model class {model.__name__} is not compatible with Lightning checkpoints "
+                "(must inherit from LightningCheckpointLoaderMixin)."
             )
-        model_instance = model.from_lightning_checkpoint(local_path, **kwargs)
+        model_instance = model.from_lightning_checkpoint(local_path, **kwargs)  # type: ignore
     else:
         raise ValueError(
             f"Unsupported model format for path: {local_path}. "

kostyl/ml/integrations/lightning/__init__.py

@@ -0,0 +1,14 @@
+try:
+    import lightning  # noqa: F401
+except ImportError as e:
+    raise ImportError(
+        "Lightning integration requires the 'lightning' package. "
+        "Please install it via 'pip install lightning'."
+    ) from e
+
+
+from .mixins import LightningCheckpointLoaderMixin
+from .module import KostylLightningModule
+
+
+__all__ = ["KostylLightningModule", "LightningCheckpointLoaderMixin"]

kostyl/ml/{lightning → integrations/lightning}/callbacks/checkpoint.py

@@ -9,17 +9,16 @@ import torch.distributed as dist
 from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.callbacks import ModelCheckpoint
 
+from kostyl.ml.base_uploader import ModelCheckpointUploader
 from kostyl.ml.configs import CheckpointConfig
-from kostyl.ml.dist_utils import is_main_process
-from kostyl.ml.lightning import KostylLightningModule
-from kostyl.ml.registry_uploader import RegistryUploaderCallback
+from kostyl.ml.dist_utils import is_local_zero_rank
 from kostyl.utils import setup_logger
 
 
 logger = setup_logger("callbacks/checkpoint.py")
 
 
-class ModelCheckpointWithRegistryUploader(ModelCheckpoint):
+class ModelCheckpointWithCheckpointUploader(ModelCheckpoint):
     r"""
     Save the model after every epoch by monitoring a quantity. Every logged metrics are passed to the
     :class:`~lightning.pytorch.loggers.logger.Logger` for the version it gets saved in the same directory as the
@@ -229,8 +228,8 @@ class ModelCheckpointWithRegistryUploader(ModelCheckpoint):
 
     def __init__(  # noqa: D107
         self,
-        registry_uploader_callback: RegistryUploaderCallback,
-        uploading_mode: Literal["only-best", "every-checkpoint"] = "only-best",
+        checkpoint_uploader: ModelCheckpointUploader,
+        upload_strategy: Literal["only-best", "every-checkpoint"] = "only-best",
         dirpath: _PATH | None = None,
         filename: str | None = None,
         monitor: str | None = None,
@@ -247,9 +246,9 @@ class ModelCheckpointWithRegistryUploader(ModelCheckpoint):
         save_on_train_epoch_end: bool | None = None,
         enable_version_counter: bool = True,
     ) -> None:
-        self.registry_uploader_callback = registry_uploader_callback
+        self.registry_uploader = checkpoint_uploader
         self.process_group: dist.ProcessGroup | None = None
-        self.uploading_mode = uploading_mode
+        self.upload_strategy = upload_strategy
         super().__init__(
             dirpath=dirpath,
             filename=filename,
@@ -269,40 +268,26 @@ class ModelCheckpointWithRegistryUploader(ModelCheckpoint):
         )
         return
 
-    @override
-    def setup(
-        self,
-        trainer: pl.Trainer,
-        pl_module: pl.LightningModule | KostylLightningModule,
-        stage: str,
-    ) -> None:
-        super().setup(trainer, pl_module, stage)
-        if isinstance(pl_module, KostylLightningModule):
-            self.process_group = pl_module.get_process_group()
-        return
-
     @override
     def _save_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
         super()._save_checkpoint(trainer, filepath)
-        if dist.is_initialized():
-            dist.barrier(group=self.process_group)
-        if trainer.is_global_zero and self.registry_uploader_callback is not None:
-            match self.uploading_mode:
+        if trainer.is_global_zero and self.registry_uploader is not None:
+            match self.upload_strategy:
                 case "every-checkpoint":
-                    self.registry_uploader_callback.upload_checkpoint(filepath)
+                    self.registry_uploader.upload_checkpoint(filepath)
                 case "only-best":
                     if filepath == self.best_model_path:
-                        self.registry_uploader_callback.upload_checkpoint(filepath)
+                        self.registry_uploader.upload_checkpoint(filepath)
         return
 
 
 def setup_checkpoint_callback(
     dirpath: Path,
     ckpt_cfg: CheckpointConfig,
-    registry_uploader_callback: RegistryUploaderCallback | None = None,
-    uploading_strategy: Literal["only-best", "every-checkpoint"] | None = None,
+    checkpoint_uploader: ModelCheckpointUploader | None = None,
+    upload_strategy: Literal["only-best", "every-checkpoint"] | None = None,
     remove_folder_if_exists: bool = True,
-) -> ModelCheckpointWithRegistryUploader | ModelCheckpoint:
+) -> ModelCheckpointWithCheckpointUploader | ModelCheckpoint:
     """
     Create and configure a checkpoint callback for model saving.
 
@@ -313,33 +298,33 @@ def setup_checkpoint_callback(
     Args:
         dirpath: Path to the directory for saving checkpoints.
         ckpt_cfg: Checkpoint configuration (filename, monitor, mode, save_top_k).
-        registry_uploader_callback: Optional callback for uploading checkpoints to a remote registry.
-            Must be specified together with uploading_strategy.
-        uploading_strategy: Checkpoint upload mode:
+        checkpoint_uploader: Optional checkpoint uploader instance. If provided, enables
+            uploading of checkpoints to a remote registry.
+        upload_strategy: Checkpoint upload mode:
             - "only-best": only the best checkpoint is uploaded
             - "every-checkpoint": every saved checkpoint is uploaded
-            Must be specified together with registry_uploader_callback.
+            Must be specified together with checkpoint_uploader.
         remove_folder_if_exists: If True, removes existing checkpoint directory before creating a new one.
 
     Returns:
-        ModelCheckpointWithRegistryUploader if registry_uploader_callback is provided,
+        ModelCheckpointWithCheckpointUploader if checkpoint_uploader is provided,
         otherwise standard ModelCheckpoint.
 
     Raises:
-        ValueError: If only one of registry_uploader_callback or uploading_mode is None.
+        ValueError: If only one of checkpoint_uploader or uploading_mode is None.
 
     Note:
         If the dirpath directory already exists, it will be removed and recreated
         (only on the main process in distributed training) if remove_folder_if_exists is True.
 
     """
-    if (registry_uploader_callback is None) != (uploading_strategy is None):
+    if (checkpoint_uploader is None) != (upload_strategy is None):
         raise ValueError(
-            "Both registry_uploader_callback and uploading_mode must be provided or neither."
+            "Both checkpoint_uploader and upload_strategy must be provided or neither."
         )
 
     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)
@@ -348,8 +333,8 @@ def setup_checkpoint_callback(
         logger.info(f"Creating checkpoint directory {dirpath}.")
         dirpath.mkdir(parents=True, exist_ok=True)
 
-    if (registry_uploader_callback is not None) and (uploading_strategy is not None):
-        checkpoint_callback = ModelCheckpointWithRegistryUploader(
+    if (checkpoint_uploader is not None) and (upload_strategy is not None):
+        checkpoint_callback = ModelCheckpointWithCheckpointUploader(
             dirpath=dirpath,
             filename=ckpt_cfg.filename,
             save_top_k=ckpt_cfg.save_top_k,
@@ -357,8 +342,8 @@ def setup_checkpoint_callback(
             mode=ckpt_cfg.mode,
             verbose=True,
             save_weights_only=ckpt_cfg.save_weights_only,
-            registry_uploader_callback=registry_uploader_callback,
-            uploading_mode=uploading_strategy,
+            checkpoint_uploader=checkpoint_uploader,
+            upload_strategy=upload_strategy,
         )
     else:
         checkpoint_callback = ModelCheckpoint(

kostyl/ml/{lightning → integrations/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}.")