nshtrainer 0.1.0__py3-none-any.whl

nshtrainer/model/modules/callback.py

@@ -0,0 +1,157 @@
+from collections import abc
+from collections.abc import Callable, Iterable
+from logging import getLogger
+from typing import Any, TypeAlias, cast, final
+
+from lightning.pytorch import Callback, LightningModule
+from lightning.pytorch.callbacks import LambdaCallback
+from typing_extensions import override
+
+from ...util.typing_utils import mixin_base_type
+
+log = getLogger(__name__)
+
+CallbackFn: TypeAlias = Callable[[], Callback | Iterable[Callback] | None]
+
+
+class CallbackRegistrarModuleMixin:
+    @override
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._ll_callbacks: list[CallbackFn] = []
+
+    def register_callback(
+        self,
+        callback: Callback | Iterable[Callback] | CallbackFn | None = None,
+        *,
+        setup: Callable | None = None,
+        teardown: Callable | None = None,
+        on_fit_start: Callable | None = None,
+        on_fit_end: Callable | None = None,
+        on_sanity_check_start: Callable | None = None,
+        on_sanity_check_end: Callable | None = None,
+        on_train_batch_start: Callable | None = None,
+        on_train_batch_end: Callable | None = None,
+        on_train_epoch_start: Callable | None = None,
+        on_train_epoch_end: Callable | None = None,
+        on_validation_epoch_start: Callable | None = None,
+        on_validation_epoch_end: Callable | None = None,
+        on_test_epoch_start: Callable | None = None,
+        on_test_epoch_end: Callable | None = None,
+        on_validation_batch_start: Callable | None = None,
+        on_validation_batch_end: Callable | None = None,
+        on_test_batch_start: Callable | None = None,
+        on_test_batch_end: Callable | None = None,
+        on_train_start: Callable | None = None,
+        on_train_end: Callable | None = None,
+        on_validation_start: Callable | None = None,
+        on_validation_end: Callable | None = None,
+        on_test_start: Callable | None = None,
+        on_test_end: Callable | None = None,
+        on_exception: Callable | None = None,
+        on_save_checkpoint: Callable | None = None,
+        on_load_checkpoint: Callable | None = None,
+        on_before_backward: Callable | None = None,
+        on_after_backward: Callable | None = None,
+        on_before_optimizer_step: Callable | None = None,
+        on_before_zero_grad: Callable | None = None,
+        on_predict_start: Callable | None = None,
+        on_predict_end: Callable | None = None,
+        on_predict_batch_start: Callable | None = None,
+        on_predict_batch_end: Callable | None = None,
+        on_predict_epoch_start: Callable | None = None,
+        on_predict_epoch_end: Callable | None = None,
+    ):
+        if callback is None:
+            callback = LambdaCallback(
+                setup=setup,
+                teardown=teardown,
+                on_fit_start=on_fit_start,
+                on_fit_end=on_fit_end,
+                on_sanity_check_start=on_sanity_check_start,
+                on_sanity_check_end=on_sanity_check_end,
+                on_train_batch_start=on_train_batch_start,
+                on_train_batch_end=on_train_batch_end,
+                on_train_epoch_start=on_train_epoch_start,
+                on_train_epoch_end=on_train_epoch_end,
+                on_validation_epoch_start=on_validation_epoch_start,
+                on_validation_epoch_end=on_validation_epoch_end,
+                on_test_epoch_start=on_test_epoch_start,
+                on_test_epoch_end=on_test_epoch_end,
+                on_validation_batch_start=on_validation_batch_start,
+                on_validation_batch_end=on_validation_batch_end,
+                on_test_batch_start=on_test_batch_start,
+                on_test_batch_end=on_test_batch_end,
+                on_train_start=on_train_start,
+                on_train_end=on_train_end,
+                on_validation_start=on_validation_start,
+                on_validation_end=on_validation_end,
+                on_test_start=on_test_start,
+                on_test_end=on_test_end,
+                on_exception=on_exception,
+                on_save_checkpoint=on_save_checkpoint,
+                on_load_checkpoint=on_load_checkpoint,
+                on_before_backward=on_before_backward,
+                on_after_backward=on_after_backward,
+                on_before_optimizer_step=on_before_optimizer_step,
+                on_before_zero_grad=on_before_zero_grad,
+                on_predict_start=on_predict_start,
+                on_predict_end=on_predict_end,
+                on_predict_batch_start=on_predict_batch_start,
+                on_predict_batch_end=on_predict_batch_end,
+                on_predict_epoch_start=on_predict_epoch_start,
+                on_predict_epoch_end=on_predict_epoch_end,
+            )
+
+        if not callable(callback):
+            callback_ = cast(CallbackFn, lambda: callback)
+        else:
+            callback_ = callback
+
+        self._ll_callbacks.append(callback_)
+
+
+class CallbackModuleMixin(
+    CallbackRegistrarModuleMixin,
+    mixin_base_type(LightningModule),
+):
+    def _gather_all_callbacks(self):
+        modules: list[Any] = []
+        if isinstance(self, CallbackRegistrarModuleMixin):
+            modules.append(self)
+        if (
+            datamodule := getattr(self.trainer, "datamodule", None)
+        ) is not None and isinstance(datamodule, CallbackRegistrarModuleMixin):
+            modules.append(datamodule)
+        modules.extend(
+            module
+            for module in self.children()
+            if isinstance(module, CallbackRegistrarModuleMixin)
+        )
+        for module in modules:
+            yield from module._ll_callbacks
+
+    @final
+    @override
+    def configure_callbacks(self):
+        callbacks = super().configure_callbacks()
+        if not isinstance(callbacks, abc.Sequence):
+            callbacks = [callbacks]
+
+        callbacks = list(callbacks)
+        for callback_fn in self._gather_all_callbacks():
+            callback_result = callback_fn()
+            if callback_result is None:
+                continue
+
+            if not isinstance(callback_result, abc.Iterable):
+                callback_result = [callback_result]
+
+            for callback in callback_result:
+                log.info(
+                    f"Registering {callback.__class__.__qualname__} callback {callback}"
+                )
+                callbacks.append(callback)
+
+        return callbacks

nshtrainer/model/modules/debug.py

@@ -0,0 +1,42 @@
+from logging import getLogger
+
+import torch
+import torch.distributed
+
+log = getLogger(__name__)
+
+
+class DebugModuleMixin:
+    @torch.jit.unused
+    def breakpoint(self, rank_zero_only: bool = True):
+        if (
+            not rank_zero_only
+            or not torch.distributed.is_initialized()
+            or torch.distributed.get_rank() == 0
+        ):
+            breakpoint()
+
+        if rank_zero_only and torch.distributed.is_initialized():
+            _ = torch.distributed.barrier()
+
+    @torch.jit.unused
+    def ensure_finite(
+        self,
+        tensor: torch.Tensor,
+        name: str | None = None,
+        throw: bool = False,
+    ):
+        name_parts: list[str] = ["Tensor"]
+        if name is not None:
+            name_parts.append(name)
+        name = " ".join(name_parts)
+
+        not_finite = ~torch.isfinite(tensor)
+        if not_finite.any():
+            msg = f"{name} has {not_finite.sum().item()}/{not_finite.numel()} non-finite values."
+            if throw:
+                raise RuntimeError(msg)
+            else:
+                log.warning(msg)
+            return False
+        return True

nshtrainer/model/modules/distributed.py

@@ -0,0 +1,70 @@
+from typing import Any, Literal, cast
+
+import torch.distributed
+from lightning.pytorch import LightningModule
+from torch.distributed import ReduceOp
+from typing_extensions import TypeVar
+
+from ...util.typing_utils import mixin_base_type
+
+T = TypeVar("T", infer_variance=True)
+
+ReduceOpStr = Literal[
+    "avg",
+    "mean",
+    "band",
+    "bor",
+    "bxor",
+    "max",
+    "min",
+    "premul_sum",
+    "product",
+    "sum",
+]
+VALID_REDUCE_OPS = (
+    "avg",
+    "mean",
+    "band",
+    "bor",
+    "bxor",
+    "max",
+    "min",
+    "premul_sum",
+    "product",
+    "sum",
+)
+
+
+class DistributedMixin(mixin_base_type(LightningModule)):
+    def all_gather_object(
+        self,
+        object: T,
+        group: torch.distributed.ProcessGroup | None = None,
+    ) -> list[T]:
+        if (
+            not torch.distributed.is_available()
+            or not torch.distributed.is_initialized()
+        ):
+            return [object]
+
+        object_list = [cast(T, None) for _ in range(self.trainer.world_size)]
+        torch.distributed.all_gather_object(object_list, object, group=group)
+        return object_list
+
+    def barrier(self, name: str | None = None):
+        self.trainer.strategy.barrier(name=name)
+
+    def reduce(
+        self,
+        tensor: torch.Tensor,
+        reduce_op: ReduceOp.RedOpType | ReduceOpStr,
+        group: Any | None = None,
+    ) -> torch.Tensor:
+        if isinstance(reduce_op, str):
+            # validate reduce_op
+            if reduce_op not in VALID_REDUCE_OPS:
+                raise ValueError(
+                    f"reduce_op must be one of {VALID_REDUCE_OPS}, got {reduce_op}"
+                )
+
+        return self.trainer.strategy.reduce(tensor, group=group, reduce_op=reduce_op)

nshtrainer/model/modules/logger.py

@@ -0,0 +1,170 @@
+from collections import deque
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import torchmetrics
+from lightning.pytorch import LightningDataModule, LightningModule
+from lightning.pytorch.utilities.types import _METRIC
+from lightning_utilities.core.rank_zero import rank_zero_warn
+from typing_extensions import override
+
+from ...actsave import ActSave
+from ...util.typing_utils import mixin_base_type
+from ..config import BaseConfig
+
+
+@dataclass(frozen=True, kw_only=True)
+class _LogContext:
+    prefix: str | None = None
+    disabled: bool | None = None
+    kwargs: dict[str, Any] = field(default_factory=dict)
+
+
+class LoggerModuleMixin:
+    @property
+    def log_dir(self):
+        if not isinstance(self, (LightningModule, LightningDataModule)):
+            raise TypeError(
+                "log_dir can only be used on LightningModule or LightningDataModule"
+            )
+
+        if (trainer := self.trainer) is None:
+            raise RuntimeError("trainer is not defined")
+
+        if (logger := trainer.logger) is None:
+            raise RuntimeError("trainer.logger is not defined")
+
+        if (log_dir := logger.log_dir) is None:
+            raise RuntimeError("trainer.logger.log_dir is not defined")
+
+        return Path(log_dir)
+
+    @property
+    def should_update_logs(self):
+        if not isinstance(self, (LightningModule, LightningDataModule)):
+            raise TypeError(
+                "should_update_logs can only be used on LightningModule or LightningDataModule"
+            )
+
+        trainer = self._trainer if isinstance(self, LightningModule) else self.trainer
+        if trainer is None:
+            return True
+
+        return trainer._logger_connector.should_update_logs
+
+
+class LoggerLightningModuleMixin(LoggerModuleMixin, mixin_base_type(LightningModule)):
+    @override
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._logger_prefix_stack = deque[_LogContext]()
+
+    if TYPE_CHECKING:
+
+        @contextmanager
+        def log_context(
+            self,
+            prefix: str | None = None,
+            *,
+            disabled: bool | None = None,
+            prog_bar: bool | None = None,
+            logger: bool | None = None,
+            on_step: bool | None = None,
+            on_epoch: bool | None = None,
+            reduce_fx: str | Callable | None = None,
+            enable_graph: bool | None = None,
+            sync_dist: bool | None = None,
+            sync_dist_group: Any | None = None,
+            add_dataloader_idx: bool | None = None,
+            batch_size: int | None = None,
+            rank_zero_only: bool | None = None,
+        ) -> Generator[None, None, None]: ...
+
+    else:
+
+        @contextmanager
+        def log_context(
+            self, prefix: str | None = None, *, disabled: bool | None = None, **kwargs
+        ) -> Generator[None, None, None]:
+            self._logger_prefix_stack.append(
+                _LogContext(
+                    prefix=prefix,
+                    disabled=disabled,
+                    kwargs=kwargs,
+                )
+            )
+            try:
+                yield
+            finally:
+                _ = self._logger_prefix_stack.pop()
+
+    if TYPE_CHECKING:
+
+        @override
+        def log(  # type: ignore[override]
+            self,
+            name: str,
+            value: _METRIC,
+            *,
+            prog_bar: bool = False,
+            logger: bool | None = None,
+            on_step: bool | None = None,
+            on_epoch: bool | None = None,
+            reduce_fx: str | Callable = "mean",
+            enable_graph: bool = False,
+            sync_dist: bool = False,
+            sync_dist_group: Any | None = None,
+            add_dataloader_idx: bool = True,
+            batch_size: int | None = None,
+            metric_attribute: str | None = None,
+            rank_zero_only: bool = False,
+        ) -> None: ...
+
+    else:
+
+        @override
+        def log(self, name: str, value: _METRIC, **kwargs) -> None:
+            # join all prefixes
+            prefix = "".join(c.prefix for c in self._logger_prefix_stack if c.prefix)
+            name = f"{prefix}{name}"
+
+            # check for disabled context:
+            # if the topmost non-null context is disabled, then we don't log
+            for c in reversed(self._logger_prefix_stack):
+                if c.disabled is not None:
+                    if c.disabled:
+                        rank_zero_warn(
+                            f"Skipping logging of {name} due to disabled context"
+                        )
+                        return
+                    else:
+                        break
+
+            fn_kwargs = {}
+            for c in self._logger_prefix_stack:
+                fn_kwargs.update(c.kwargs)
+            fn_kwargs.update(kwargs)
+
+            self._logger_actsave(name, value)
+
+            return super().log(name, value, **fn_kwargs)
+
+    def _logger_actsave(self, name: str, value: _METRIC) -> None:
+        hparams = cast(BaseConfig, self.hparams)
+        if (
+            not hparams.trainer.actsave
+            or not hparams.trainer.actsave.auto_save_logged_metrics
+        ):
+            return
+
+        ActSave.save(
+            {
+                f"logger.{name}": lambda: value.compute()
+                if isinstance(value, torchmetrics.Metric)
+                else value
+            }
+        )

nshtrainer/model/modules/profiler.py

@@ -0,0 +1,24 @@
+from lightning.pytorch import LightningDataModule, LightningModule
+from lightning.pytorch.profilers import PassThroughProfiler
+
+from ...util.typing_utils import mixin_base_type
+
+
+class ProfilerMixin(mixin_base_type(LightningModule)):
+    @property
+    def profiler(self):
+        if not isinstance(self, (LightningModule, LightningDataModule)):
+            raise TypeError(
+                "`profiler` can only be used on LightningModule or LightningDataModule"
+            )
+
+        if (trainer := self.trainer) is None:
+            raise RuntimeError("trainer is not defined")
+
+        if not hasattr(trainer, "profiler"):
+            raise RuntimeError("trainer does not have profiler")
+
+        if (profiler := getattr(trainer, "profiler")) is None:
+            profiler = PassThroughProfiler()
+
+        return profiler

nshtrainer/model/modules/rlp_sanity_checks.py

@@ -0,0 +1,202 @@
+from collections.abc import Mapping
+from logging import getLogger
+from typing import cast
+
+import torch
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.utilities.types import (
+    LRSchedulerConfigType,
+    LRSchedulerTypeUnion,
+)
+from typing_extensions import Protocol, override, runtime_checkable
+
+from ...util.typing_utils import mixin_base_type
+from ..config import BaseConfig
+from .callback import CallbackModuleMixin
+
+log = getLogger(__name__)
+
+
+def _on_train_start_callback(trainer: Trainer, pl_module: LightningModule):
+    # If we're in PL's "sanity check" mode, we don't need to run this check
+    if trainer.sanity_checking:
+        return
+
+    config = cast(BaseConfig, pl_module.hparams)
+    if config.trainer.sanity_checking.reduce_lr_on_plateau == "disable":
+        return
+
+    # if no lr schedulers, return
+    if not trainer.lr_scheduler_configs:
+        return
+
+    errors: list[str] = []
+    disable_message = (
+        "Otherwise, set `config.trainer.sanity_checking.reduce_lr_on_plateau='disable'` "
+        "to disable this sanity check."
+    )
+
+    for lr_scheduler_config in trainer.lr_scheduler_configs:
+        if not lr_scheduler_config.reduce_on_plateau:
+            continue
+
+        match lr_scheduler_config.interval:
+            case "epoch":
+                # we need to make sure that the trainer runs val every `frequency` epochs
+
+                # If `trainer.check_val_every_n_epoch` is None, then Lightning
+                # will run val every `int(trainer.val_check_interval)` steps.
+                # So, first we need to make sure that `trainer.val_check_interval` is not None first.
+                if trainer.check_val_every_n_epoch is None:
+                    errors.append(
+                        "Trainer is not running validation at epoch intervals "
+                        "(i.e., `trainer.check_val_every_n_epoch` is None) but "
+                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} is used."
+                        f"Please set `config.trainer.check_val_every_n_epoch={lr_scheduler_config.frequency}`. "
+                        + disable_message
+                    )
+
+                # Second, we make sure that the trainer runs val at least every `frequency` epochs
+                if (
+                    trainer.check_val_every_n_epoch is not None
+                    and lr_scheduler_config.frequency % trainer.check_val_every_n_epoch
+                    != 0
+                ):
+                    errors.append(
+                        f"Trainer is not running validation every {lr_scheduler_config.frequency} epochs but "
+                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} and frequency={lr_scheduler_config.frequency} is used."
+                        f"Please set `config.trainer.check_val_every_n_epoch` to a multiple of {lr_scheduler_config.frequency}. "
+                        + disable_message
+                    )
+
+            case "step":
+                # In this case, we need to make sure that the trainer runs val at step intervals
+                # that are multiples of `frequency`.
+
+                # First, we make sure that validation is run at step intervals
+                if trainer.check_val_every_n_epoch is not None:
+                    errors.append(
+                        "Trainer is running validation at epoch intervals "
+                        "(i.e., `trainer.check_val_every_n_epoch` is not None) but "
+                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} is used."
+                        "Please set `config.trainer.check_val_every_n_epoch=None` "
+                        f"and `config.trainer.val_check_interval={lr_scheduler_config.frequency}`. "
+                        + disable_message
+                    )
+
+                # Second, we make sure `trainer.val_check_interval` is an integer
+                if not isinstance(trainer.val_check_interval, int):
+                    errors.append(
+                        f"Trainer is not running validation at step intervals "
+                        f"(i.e., `trainer.val_check_interval` is not an integer) but "
+                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} is used."
+                        "Please set `config.trainer.val_check_interval=None` "
+                        f"and `config.trainer.val_check_interval={lr_scheduler_config.frequency}`. "
+                        + disable_message
+                    )
+
+                # Third, we make sure that the trainer runs val at least every `frequency` steps
+                if (
+                    isinstance(trainer.val_check_interval, int)
+                    and trainer.val_check_interval % lr_scheduler_config.frequency != 0
+                ):
+                    errors.append(
+                        f"Trainer is not running validation every {lr_scheduler_config.frequency} steps but "
+                        f"a ReduceLRPlateau scheduler with interval={lr_scheduler_config.interval} and frequency={lr_scheduler_config.frequency} is used."
+                        "Please set `config.trainer.val_check_interval` "
+                        f"to a multiple of {lr_scheduler_config.frequency}. "
+                        + disable_message
+                    )
+
+            case _:
+                pass
+
+    if not errors:
+        return
+
+    message = (
+        "ReduceLRPlateau sanity checks failed with the following errors:\n"
+        + "\n".join(errors)
+    )
+    match config.trainer.sanity_checking.reduce_lr_on_plateau:
+        case "warn":
+            log.warning(message)
+        case "error":
+            raise ValueError(message)
+        case _:
+            pass
+
+
+@runtime_checkable
+class CustomRLPImplementation(Protocol):
+    __reduce_lr_on_plateau__: bool
+
+
+class RLPSanityCheckModuleMixin(mixin_base_type(CallbackModuleMixin)):
+    @override
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        global _on_train_start_callback
+        self.register_callback(on_train_start=_on_train_start_callback)
+
+    def reduce_lr_on_plateau_config(
+        self,
+        lr_scheduler: LRSchedulerTypeUnion | LRSchedulerConfigType,
+    ) -> LRSchedulerConfigType:
+        if (trainer := self._trainer) is None:
+            raise RuntimeError(
+                "Could not determine the frequency of ReduceLRPlateau scheduler "
+                "because `self.trainer` is None."
+            )
+
+        # First, resolve the LR scheduler from the provided config.
+        lr_scheduler_config: LRSchedulerConfigType
+        match lr_scheduler:
+            case Mapping():
+                lr_scheduler_config = cast(LRSchedulerConfigType, lr_scheduler)
+            case _:
+                lr_scheduler_config = {"scheduler": lr_scheduler}
+
+        # Make sure the scheduler is a ReduceLRPlateau scheduler. Otherwise, warn the user.
+        if (
+            not isinstance(
+                lr_scheduler_config["scheduler"],
+                torch.optim.lr_scheduler.ReduceLROnPlateau,
+            )
+        ) and (
+            not isinstance(lr_scheduler_config["scheduler"], CustomRLPImplementation)
+            or not lr_scheduler_config["scheduler"].__reduce_lr_on_plateau__
+        ):
+            log.warning(
+                "`reduce_lr_on_plateau_config` should only be used with a ReduceLRPlateau scheduler. "
+                f"The provided scheduler, {lr_scheduler_config['scheduler']}, does not subclass "
+                "`torch.optim.lr_scheduler.ReduceLROnPlateau`. "
+                "Please ensure that the scheduler is a ReduceLRPlateau scheduler. "
+                "If you are using a custom ReduceLRPlateau scheduler implementation, "
+                "please either (1) make sure that it subclasses `torch.optim.lr_scheduler.ReduceLROnPlateau`, "
+                "or (2) set the scheduler's `__reduce_lr_on_plateau__` attribute to `True`."
+            )
+
+        # If trainer.check_val_every_n_epoch is an integer, then we run val at epoch intervals.
+        if trainer.check_val_every_n_epoch is not None:
+            return {
+                "reduce_on_plateau": True,
+                "interval": "epoch",
+                "frequency": trainer.check_val_every_n_epoch,
+                **lr_scheduler_config,
+            }
+
+        # Otherwise, we run val at step intervals.
+        if not isinstance(trainer.val_check_batch, int):
+            raise ValueError(
+                "Could not determine the frequency of ReduceLRPlateau scheduler "
+                f"because {trainer.val_check_batch=} is not an integer."
+            )
+
+        return {
+            "reduce_on_plateau": True,
+            "interval": "step",
+            "frequency": trainer.val_check_batch,
+            **lr_scheduler_config,
+        }