nshtrainer 0.10.18__tar.gz → 0.11.0__tar.gz

{nshtrainer-0.10.18 → nshtrainer-0.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.10.18
+Version: 0.11.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -9,7 +9,8 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: GitPython
+Provides-Extra: extra
+Requires-Dist: GitPython ; extra == "extra"
 Requires-Dist: lightning
 Requires-Dist: nshconfig
 Requires-Dist: nshrunner
@@ -17,10 +18,12 @@ Requires-Dist: nshutils
 Requires-Dist: numpy
 Requires-Dist: psutil
 Requires-Dist: pytorch-lightning
+Requires-Dist: tensorboard ; extra == "extra"
 Requires-Dist: torch
-Requires-Dist: torchmetrics
+Requires-Dist: torchmetrics ; extra == "extra"
 Requires-Dist: typing-extensions
-Requires-Dist: wrapt
+Requires-Dist: wandb ; extra == "extra"
+Requires-Dist: wrapt ; extra == "extra"
 Description-Content-Type: text/markdown
 
 

{nshtrainer-0.10.18 → nshtrainer-0.11.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.10.18"
+version = "0.11.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -19,6 +19,8 @@ pytorch-lightning = "*"
 torchmetrics = { version = "*", optional = true }
 wrapt = { version = "*", optional = true }
 GitPython = { version = "*", optional = true }
+wandb = { version = "*", optional = true }
+tensorboard = { version = "*", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 pyright = "^1.1.372"
@@ -40,3 +42,6 @@ reportPrivateImportUsage = false
 
 [tool.ruff.lint]
 ignore = ["F722", "F821", "E731", "E741"]
+
+[tool.poetry.extras]
+extra = ["torchmetrics", "wrapt", "GitPython", "wandb", "tensorboard"]

{nshtrainer-0.10.18 → nshtrainer-0.11.0}/src/nshtrainer/_checkpoint/metadata.py

@@ -92,7 +92,7 @@ def _write_checkpoint_metadata(
     except Exception as e:
         log.warning(f"Failed to write metadata to {checkpoint_path}: {e}")
     else:
-        log.info(f"Checkpoint metadata written to {checkpoint_path}")
+        log.debug(f"Checkpoint metadata written to {checkpoint_path}")
 
     # Write the hparams to the checkpoint directory
     try:
@@ -101,7 +101,7 @@ def _write_checkpoint_metadata(
     except Exception as e:
         log.warning(f"Failed to write hparams to {checkpoint_path}: {e}")
     else:
-        log.info(f"Checkpoint metadata written to {checkpoint_path}")
+        log.debug(f"Checkpoint metadata written to {checkpoint_path}")
 
 
 def _remove_checkpoint_metadata(checkpoint_path: Path):
@@ -112,7 +112,7 @@ def _remove_checkpoint_metadata(checkpoint_path: Path):
         except Exception as e:
             log.warning(f"Failed to remove {path}: {e}")
         else:
-            log.info(f"Removed {path}")
+            log.debug(f"Removed {path}")
 
 
 def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Path):
@@ -133,7 +133,7 @@ def _link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Pat
         except Exception as e:
             log.warning(f"Failed to link {path} to {linked_path}: {e}")
         else:
-            log.info(f"Linked {path} to {linked_path}")
+            log.debug(f"Linked {path} to {linked_path}")
 
 
 def _checkpoint_sort_key_fn(key: Callable[[CheckpointMetadata, Path], Any]):

{nshtrainer-0.10.18 → nshtrainer-0.11.0}/src/nshtrainer/callbacks/actsave.py

@@ -4,15 +4,11 @@ from typing import Literal
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks.callback import Callback
+from nshutils import ActSave
 from typing_extensions import TypeAlias, override
 
 from .base import CallbackConfigBase
 
-try:
-    from nshutils import ActSave  # type: ignore
-except ImportError:
-    ActSave = None
-
 Stage: TypeAlias = Literal["train", "validation", "test", "predict"]
 
 
@@ -51,11 +47,6 @@ class ActSaveCallback(Callback):
         if not self.config:
             return
 
-        if ActSave is None:
-            raise ImportError(
-                "ActSave is not installed. Please install nshutils to use the ActSaveCallback."
-            )
-
         context = ActSave.enabled(self.save_dir)
         context.__enter__()
         self._enabled_context = context
@@ -77,11 +68,6 @@ class ActSaveCallback(Callback):
         if not self.config:
             return
 
-        if ActSave is None:
-            raise ImportError(
-                "ActSave is not installed. Please install nshutils to use the ActSaveCallback."
-            )
-
         # If we have an active context manager for this stage, exit it
         if active_contexts := self._active_contexts.get(stage):
             active_contexts.__exit__(None, None, None)

{nshtrainer-0.10.18 → nshtrainer-0.11.0}/src/nshtrainer/model/modules/logger.py

@@ -9,16 +9,12 @@ 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 nshutils import ActSave
 from typing_extensions import override
 
 from ...util.typing_utils import mixin_base_type
 from ..config import BaseConfig
 
-try:
-    from nshutils import ActSave  # type: ignore
-except ImportError:
-    ActSave = None
-
 
 @dataclass(frozen=True, kw_only=True)
 class _LogContext:
@@ -162,10 +158,6 @@ class LoggerLightningModuleMixin(LoggerModuleMixin, mixin_base_type(LightningMod
         if not hparams.trainer.logging.actsave_logged_metrics:
             return
 
-        if ActSave is None:
-            rank_zero_warn("ActSave is not available, skipping logging of metrics")
-            return
-
         ActSave.save(
             lambda: {
                 f"logger.{name}": lambda: value.compute()