nshtrainer 0.30.0__py3-none-any.whl → 0.31.0__py3-none-any.whl

nshtrainer/model/{modules → mixins}/callback.py

@@ -165,7 +165,7 @@ class CallbackModuleMixin(
     CallbackRegistrarModuleMixin,
     mixin_base_type(LightningModule),
 ):
-    def _gather_all_callbacks(self):
+    def _nshtrainer_gather_all_callbacks(self):
         modules: list[Any] = []
         if isinstance(self, CallbackRegistrarModuleMixin):
             modules.append(self)
@@ -189,7 +189,7 @@ class CallbackModuleMixin(
             callbacks = [callbacks]
 
         callbacks = list(callbacks)
-        for callback_fn in self._gather_all_callbacks():
+        for callback_fn in self._nshtrainer_gather_all_callbacks():
             callback_result = callback_fn()
             if callback_result is None:
                 continue

nshtrainer/model/{modules → mixins}/logger.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
 import torchmetrics
-from lightning.pytorch import LightningDataModule, LightningModule
+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
@@ -23,15 +23,13 @@ class _LogContext:
     kwargs: dict[str, Any] = field(default_factory=dict)
 
 
-class LoggerModuleMixin:
+class LoggerModuleMixin(mixin_base_type(LightningModule)):
     @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:
+        """
+        The directory where logs are saved.
+        """
+        if (trainer := self._trainer) is None:
             raise RuntimeError("trainer is not defined")
 
         if (logger := trainer.logger) is None:
@@ -44,16 +42,15 @@ class LoggerModuleMixin:
 
     @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"
+        """
+        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"
             )
 
-        trainer = self._trainer if isinstance(self, LightningModule) else self.trainer
-        if trainer is None:
-            return True
-
-        return trainer._logger_connector.should_update_logs
+        return self._trainer._logger_connector.should_update_logs
 
 
 class LoggerLightningModuleMixin(LoggerModuleMixin, mixin_base_type(LightningModule)):

nshtrainer/profiler/__init__.py

@@ -0,0 +1,13 @@
+from typing import Annotated, TypeAlias
+
+import nshconfig as C
+
+from ._base import BaseProfilerConfig as BaseProfilerConfig
+from .advanced import AdvancedProfilerConfig as AdvancedProfilerConfig
+from .pytorch import PyTorchProfilerConfig as PyTorchProfilerConfig
+from .simple import SimpleProfilerConfig as SimpleProfilerConfig
+
+ProfilerConfig: TypeAlias = Annotated[
+    SimpleProfilerConfig | AdvancedProfilerConfig | PyTorchProfilerConfig,
+    C.Discriminator("name"),
+]

nshtrainer/profiler/_base.py

@@ -0,0 +1,29 @@
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import nshconfig as C
+from lightning.pytorch.profilers import Profiler
+
+if TYPE_CHECKING:
+    from ..model import BaseConfig
+
+log = logging.getLogger(__name__)
+
+
+class BaseProfilerConfig(C.Config, ABC):
+    dirpath: str | Path | None = None
+    """
+    Directory path for the ``filename``. If ``dirpath`` is ``None`` but ``filename`` is present, the
+        ``trainer.log_dir`` (from :class:`~lightning.pytorch.loggers.tensorboard.TensorBoardLogger`)
+        will be used.
+    """
+    filename: str | None = None
+    """
+    If present, filename where the profiler results will be saved instead of printing to stdout.
+        The ``.txt`` extension will be used automatically.
+    """
+
+    @abstractmethod
+    def create_profiler(self, root_config: "BaseConfig") -> Profiler | None: ...

nshtrainer/profiler/advanced.py

@@ -0,0 +1,37 @@
+import logging
+from typing import Literal
+
+from typing_extensions import override
+
+from ._base import BaseProfilerConfig
+
+log = logging.getLogger(__name__)
+
+
+class AdvancedProfilerConfig(BaseProfilerConfig):
+    name: Literal["advanced"] = "advanced"
+
+    line_count_restriction: float = 1.0
+    """
+    This can be used to limit the number of functions
+        reported for each action. either an integer (to select a count of lines),
+        or a decimal fraction between 0.0 and 1.0 inclusive (to select a percentage of lines)
+    """
+
+    @override
+    def create_profiler(self, root_config):
+        from lightning.pytorch.profilers.advanced import AdvancedProfiler
+
+        if (dirpath := self.dirpath) is None:
+            dirpath = root_config.directory.resolve_subdirectory(
+                root_config.id, "profile"
+            )
+
+        if (filename := self.filename) is None:
+            filename = f"{root_config.id}_profile.txt"
+
+        return AdvancedProfiler(
+            line_count_restriction=self.line_count_restriction,
+            dirpath=dirpath,
+            filename=filename,
+        )

nshtrainer/profiler/pytorch.py

@@ -0,0 +1,83 @@
+import logging
+from typing import Any, Literal
+
+from typing_extensions import override
+
+from ._base import BaseProfilerConfig
+
+log = logging.getLogger(__name__)
+
+
+class PyTorchProfilerConfig(BaseProfilerConfig):
+    name: Literal["pytorch"] = "pytorch"
+
+    group_by_input_shapes: bool = False
+    """Include operator input shapes and group calls by shape."""
+
+    emit_nvtx: bool = False
+    """
+    Context manager that makes every autograd operation emit an NVTX range
+        Run::
+
+            nvprof --profile-from-start off -o trace_name.prof -- <regular command here>
+
+        To visualize, you can either use::
+
+            nvvp trace_name.prof
+            torch.autograd.profiler.load_nvprof(path)
+    """
+
+    export_to_chrome: bool = True
+    """
+    Whether to export the sequence of profiled operators for Chrome.
+        It will generate a ``.json`` file which can be read by Chrome.
+    """
+
+    row_limit: int = 20
+    """
+    Limit the number of rows in a table, ``-1`` is a special value that
+        removes the limit completely.
+    """
+
+    sort_by_key: str | None = None
+    """
+    Attribute used to sort entries. By default
+        they are printed in the same order as they were registered.
+        Valid keys include: ``cpu_time``, ``cuda_time``, ``cpu_time_total``,
+        ``cuda_time_total``, ``cpu_memory_usage``, ``cuda_memory_usage``,
+        ``self_cpu_memory_usage``, ``self_cuda_memory_usage``, ``count``.
+    """
+
+    record_module_names: bool = True
+    """Whether to add module names while recording autograd operation."""
+
+    table_kwargs: dict[str, Any] | None = None
+    """Dictionary with keyword arguments for the summary table."""
+
+    additional_profiler_kwargs: dict[str, Any] = {}
+    """Keyword arguments for the PyTorch profiler. This depends on your PyTorch version"""
+
+    @override
+    def create_profiler(self, root_config):
+        from lightning.pytorch.profilers.pytorch import PyTorchProfiler
+
+        if (dirpath := self.dirpath) is None:
+            dirpath = root_config.directory.resolve_subdirectory(
+                root_config.id, "profile"
+            )
+
+        if (filename := self.filename) is None:
+            filename = f"{root_config.id}_profile.txt"
+
+        return PyTorchProfiler(
+            group_by_input_shapes=self.group_by_input_shapes,
+            emit_nvtx=self.emit_nvtx,
+            export_to_chrome=self.export_to_chrome,
+            row_limit=self.row_limit,
+            sort_by_key=self.sort_by_key,
+            record_module_names=self.record_module_names,
+            table_kwargs=self.table_kwargs,
+            dirpath=dirpath,
+            filename=filename,
+            **self.additional_profiler_kwargs,
+        )

nshtrainer/profiler/simple.py

@@ -0,0 +1,36 @@
+import logging
+from typing import Literal
+
+from typing_extensions import override
+
+from ._base import BaseProfilerConfig
+
+log = logging.getLogger(__name__)
+
+
+class SimpleProfilerConfig(BaseProfilerConfig):
+    name: Literal["simple"] = "simple"
+
+    extended: bool = True
+    """
+    If ``True``, adds extra columns representing number of calls and percentage of
+        total time spent onrespective action.
+    """
+
+    @override
+    def create_profiler(self, root_config):
+        from lightning.pytorch.profilers.simple import SimpleProfiler
+
+        if (dirpath := self.dirpath) is None:
+            dirpath = root_config.directory.resolve_subdirectory(
+                root_config.id, "profile"
+            )
+
+        if (filename := self.filename) is None:
+            filename = f"{root_config.id}_profile.txt"
+
+        return SimpleProfiler(
+            extended=self.extended,
+            dirpath=dirpath,
+            filename=filename,
+        )