nshtrainer 0.10.18__py3-none-any.whl → 0.11.0__py3-none-any.whl

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/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/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()

{nshtrainer-0.10.18.dist-info → nshtrainer-0.11.0.dist-info}/METADATA

@@ -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.dist-info → nshtrainer-0.11.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
 nshtrainer/_checkpoint/loader.py,sha256=48flPr1XgQHOgIPaCrRqOEvRuG0SZuV3cQ1vgHLqFqI,11025
-nshtrainer/_checkpoint/metadata.py,sha256=GlhlAyJh5gcp3R8l2Y3eAUQtQzBnitFlB0xdx-khEUQ,5579
+nshtrainer/_checkpoint/metadata.py,sha256=3yxGxHLIVwKh5K4L8LYOEK3GQ6HQXy89CGcy9zarApo,5583
 nshtrainer/_checkpoint/saver.py,sha256=z_c7a91O4Bh4lZZjqJgxT3w25qFlJsOopV3cpJtkHk8,1655
 nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ4b-PoWM,120
 nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
@@ -8,7 +8,7 @@ nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPUL
 nshtrainer/_experimental/flops/module_tracker.py,sha256=bUL-IRTd0aF_DwmXkZjHZAA31p4ZEhyqhc26XWKQUUY,4922
 nshtrainer/callbacks/__init__.py,sha256=ifXQRwtccznl4lMKwKLSuuAQC4bKFBgfzQ4rx9gOqjE,2345
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
-nshtrainer/callbacks/actsave.py,sha256=aY6T_NAzaFAVU8WMHOXnWL5wd2bi8eVxeU2S0iAs70c,4446
+nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,3555
 nshtrainer/callbacks/early_stopping.py,sha256=LGn3rdbvkFfUo9kwMzK4eMGlPAqD9uFdowDx6VdfozQ,3761
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
@@ -56,7 +56,7 @@ nshtrainer/model/config.py,sha256=4vze6tpMFAgpk532T33jmSH1lHfolHK1vAJEaa2-Vxs,54
 nshtrainer/model/modules/callback.py,sha256=K0-cyEtBcQhI7Q2e-AGTE8T-GghUPY9DYmneU6ULV6g,6401
 nshtrainer/model/modules/debug.py,sha256=Yy7XEdPou9BkCsD5hJchwJGmCVGrfUru5g9VjPM4uAw,1120
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
-nshtrainer/model/modules/logger.py,sha256=YYhehQysqTjuVFcd_EREYDh57CIlezidFBS2Ohp_xKo,5661
+nshtrainer/model/modules/logger.py,sha256=CJWSmNT8SV5GLtfml-qGYenqRPXcNOMsJRGEavAd8Hw,5464
 nshtrainer/model/modules/profiler.py,sha256=rQ_jRMcM1Z2AIROZlRnBRHM5rkTpq67afZPD6CIRfXs,825
 nshtrainer/model/modules/rlp_sanity_checks.py,sha256=I_ralr2ThQ-D_FkVQTwbdXLLlgHJEr7-s01I5wSDjps,8893
 nshtrainer/model/modules/shared_parameters.py,sha256=ZiRKkZXr6RwdwLCdZCJPl3dXe7bnT8Z9yTeRK5bXBGk,2687
@@ -79,6 +79,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.10.18.dist-info/METADATA,sha256=-r5sz2eulZvKCATa6NFmUGkwT4FFrjetLGJawsakdyM,696
-nshtrainer-0.10.18.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.10.18.dist-info/RECORD,,
+nshtrainer-0.11.0.dist-info/METADATA,sha256=IE_faJS_HOMLLu67UPzCAO1aBPZvjWZT3YBCaQ_YpC0,860
+nshtrainer-0.11.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.11.0.dist-info/RECORD,,