nshtrainer 0.44.1__py3-none-any.whl → 1.0.0b10__py3-none-any.whl

nshtrainer/model/mixins/debug.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import torch
+
+log = logging.getLogger(__name__)
+
+
+def _trainer(module: Any):
+    if torch.jit.is_scripting():
+        return None
+
+    if hasattr(module, "_trainer"):
+        trainer = module._trainer
+    else:
+        try:
+            trainer = module.trainer
+        except RuntimeError:
+            return None
+
+    from ...trainer import Trainer
+
+    if not isinstance(trainer, Trainer):
+        return None
+
+    return trainer
+
+
+class _DebugModuleMixin:
+    @property
+    def nshtrainer_or_none(self):
+        return _trainer(self)
+
+    @property
+    def nshtrainer(self):
+        if (trainer := _trainer(self)) is None:
+            raise RuntimeError("Could not resolve trainer.")
+        return trainer
+
+    @property
+    def debug(self) -> bool:
+        if (trainer := _trainer(self)) is None:
+            return False
+        return trainer.debug
+
+    @debug.setter
+    def debug(self, value: bool):
+        if (trainer := _trainer(self)) is None:
+            return
+        trainer.debug = value
+
+    @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/mixins/logger.py

@@ -1,166 +1,163 @@
 from __future__ import annotations
 
+import copy
+import dataclasses
 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
+from typing import Any, ClassVar
 
-import torchmetrics
 from lightning.pytorch import LightningModule
 from lightning.pytorch.utilities.types import _METRIC
 from lightning_utilities.core.rank_zero import rank_zero_warn
-from nshutils import ActSave
-from typing_extensions import override
+from typing_extensions import Self, override
 
 from ...util.typing_utils import mixin_base_type
-from ..config import BaseConfig
 
 
-@dataclass(frozen=True, kw_only=True)
-class _LogContext:
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class _LogContextKwargs:
+    __ignore_fields__: ClassVar[set[str]] = {"prefix", "disabled"}
+
     prefix: str | None = None
     disabled: bool | None = None
-    kwargs: dict[str, Any] = field(default_factory=dict)
-
-
-class LoggerModuleMixin(mixin_base_type(LightningModule)):
-    @property
-    def log_dir(self):
-        """
-        The directory where logs are saved.
-        """
-        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):
-        """
-        Whether logs should be updated. This is true once every `log_every_n_steps` steps.
-        """
-        if self._trainer is None:
-            raise RuntimeError(
-                "`should_update_logs` can only be used after the module is attached to a trainer"
-            )
-
-        return self._trainer._logger_connector.should_update_logs
-
-
-class LoggerLightningModuleMixin(LoggerModuleMixin, mixin_base_type(LightningModule)):
+    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
+
+    def copy_from(self, other: Self):
+        kwargs = copy.deepcopy(self)
+
+        # Copy over all the not-None values from the other object
+        updates = {}
+        for field in dataclasses.fields(self):
+            # Ignore disabled fields
+            if field.name in self.__ignore_fields__:
+                continue
+
+            if (value := getattr(other, field.name, None)) is None:
+                continue
+            # setattr(kwargs, field.name, value)
+            updates[field.name] = value
+
+        return dataclasses.replace(kwargs, **updates)
+
+    def to_dict(self):
+        d = dataclasses.asdict(self)
+        for field in self.__ignore_fields__:
+            d.pop(field, None)
+        return d
+
+
+class LoggerLightningModuleMixin(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,
-                )
+        self._logger_prefix_stack = deque[_LogContextKwargs]()
+
+    @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]:
+        self._logger_prefix_stack.append(
+            _LogContextKwargs(
+                prefix=prefix,
+                disabled=disabled,
+                prog_bar=prog_bar,
+                logger=logger,
+                on_step=on_step,
+                on_epoch=on_epoch,
+                reduce_fx=reduce_fx,
+                enable_graph=enable_graph,
+                sync_dist=sync_dist,
+                sync_dist_group=sync_dist_group,
+                add_dataloader_idx=add_dataloader_idx,
+                batch_size=batch_size,
+                rank_zero_only=rank_zero_only,
             )
-            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.logging.actsave_logged_metrics:
-            return
-
-        ActSave.save(
-            lambda: {
-                f"logger.{name}": lambda: value.compute()
-                if isinstance(value, torchmetrics.Metric)
-                else value
-            }
+        )
+        try:
+            yield
+        finally:
+            _ = self._logger_prefix_stack.pop()
+
+    @override
+    def log(
+        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:
+        # 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 = _LogContextKwargs()
+        for c in self._logger_prefix_stack:
+            fn_kwargs = fn_kwargs.copy_from(c)
+        fn_kwargs = fn_kwargs.copy_from(
+            _LogContextKwargs(
+                prog_bar=prog_bar,
+                logger=logger,
+                on_step=on_step,
+                on_epoch=on_epoch,
+                reduce_fx=reduce_fx,
+                enable_graph=enable_graph,
+                sync_dist=sync_dist,
+                sync_dist_group=sync_dist_group,
+                add_dataloader_idx=add_dataloader_idx,
+                batch_size=batch_size,
+                rank_zero_only=rank_zero_only,
+            )
+        )
+        return super().log(
+            name,
+            value,
+            metric_attribute=metric_attribute,
+            **fn_kwargs.to_dict(),
         )

nshtrainer/profiler/_base.py

@@ -9,7 +9,7 @@ import nshconfig as C
 from lightning.pytorch.profilers import Profiler
 
 if TYPE_CHECKING:
-    from ..model import BaseConfig
+    from ..trainer._config import TrainerConfig
 
 log = logging.getLogger(__name__)
 
@@ -28,4 +28,4 @@ class BaseProfilerConfig(C.Config, ABC):
     """
 
     @abstractmethod
-    def create_profiler(self, root_config: "BaseConfig") -> Profiler | None: ...
+    def create_profiler(self, trainer_config: TrainerConfig) -> Profiler | None: ...

nshtrainer/profiler/advanced.py

@@ -21,16 +21,16 @@ class AdvancedProfilerConfig(BaseProfilerConfig):
     """
 
     @override
-    def create_profiler(self, root_config):
+    def create_profiler(self, trainer_config):
         from lightning.pytorch.profilers.advanced import AdvancedProfiler
 
         if (dirpath := self.dirpath) is None:
-            dirpath = root_config.directory.resolve_subdirectory(
-                root_config.id, "profile"
+            dirpath = trainer_config.directory.resolve_subdirectory(
+                trainer_config.id, "profile"
             )
 
         if (filename := self.filename) is None:
-            filename = f"{root_config.id}_profile.txt"
+            filename = f"{trainer_config.id}_profile.txt"
 
         return AdvancedProfiler(
             line_count_restriction=self.line_count_restriction,

nshtrainer/profiler/pytorch.py

@@ -60,16 +60,16 @@ class PyTorchProfilerConfig(BaseProfilerConfig):
     """Keyword arguments for the PyTorch profiler. This depends on your PyTorch version"""
 
     @override
-    def create_profiler(self, root_config):
+    def create_profiler(self, trainer_config):
         from lightning.pytorch.profilers.pytorch import PyTorchProfiler
 
         if (dirpath := self.dirpath) is None:
-            dirpath = root_config.directory.resolve_subdirectory(
-                root_config.id, "profile"
+            dirpath = trainer_config.directory.resolve_subdirectory(
+                trainer_config.id, "profile"
             )
 
         if (filename := self.filename) is None:
-            filename = f"{root_config.id}_profile.txt"
+            filename = f"{trainer_config.id}_profile.txt"
 
         return PyTorchProfiler(
             group_by_input_shapes=self.group_by_input_shapes,

nshtrainer/profiler/simple.py

@@ -20,16 +20,16 @@ class SimpleProfilerConfig(BaseProfilerConfig):
     """
 
     @override
-    def create_profiler(self, root_config):
+    def create_profiler(self, trainer_config):
         from lightning.pytorch.profilers.simple import SimpleProfiler
 
         if (dirpath := self.dirpath) is None:
-            dirpath = root_config.directory.resolve_subdirectory(
-                root_config.id, "profile"
+            dirpath = trainer_config.directory.resolve_subdirectory(
+                trainer_config.id, "profile"
             )
 
         if (filename := self.filename) is None:
-            filename = f"{root_config.id}_profile.txt"
+            filename = f"{trainer_config.id}_profile.txt"
 
         return SimpleProfiler(
             extended=self.extended,

nshtrainer/trainer/__init__.py

@@ -1,3 +1,4 @@
 from __future__ import annotations
 
+from ._config import TrainerConfig as TrainerConfig
 from .trainer import Trainer as Trainer