nshtrainer 0.10.11__py3-none-any.whl → 0.10.13__py3-none-any.whl

nshtrainer/_checkpoint/metadata.py

@@ -10,7 +10,7 @@ import nshconfig as C
 import numpy as np
 import torch
 
-from ..model._environment import EnvironmentConfig
+from ..util._environment_info import EnvironmentConfig
 
 if TYPE_CHECKING:
     from ..model import BaseConfig, LightningModuleBase

nshtrainer/_checkpoint/saver.py

@@ -36,7 +36,8 @@ def _link_checkpoint(
             # fall back to copying the file
             shutil.copy(filepath, linkpath)
 
-        _link_checkpoint_metadata(filepath, linkpath)
+        if metadata:
+            _link_checkpoint_metadata(filepath, linkpath)
     if barrier:
         trainer.strategy.barrier()
 
@@ -44,9 +45,17 @@ def _link_checkpoint(
 def _remove_checkpoint(
     trainer: Trainer,
     filepath: str | Path | os.PathLike,
-    remove_metadata: bool = True,
+    *,
+    metadata: bool,
+    barrier: bool,
 ):
     if not isinstance(filepath, Path):
         filepath = Path(filepath)
-    trainer.strategy.remove_checkpoint(filepath)
-    _remove_checkpoint_metadata(filepath)
+
+    if trainer.is_global_zero:
+        trainer.strategy.remove_checkpoint(filepath)
+        if metadata:
+            _remove_checkpoint_metadata(filepath)
+
+    if barrier:
+        trainer.strategy.barrier()

nshtrainer/callbacks/latest_epoch_checkpoint.py

@@ -69,7 +69,7 @@ class LatestEpochCheckpoint(Checkpoint):
 
     def _remove_checkpoints(self, trainer: Trainer, ckpt_paths: list[Path]):
         for ckpt_path in ckpt_paths:
-            _remove_checkpoint(trainer, ckpt_path, remove_metadata=True)
+            _remove_checkpoint(trainer, ckpt_path, metadata=True, barrier=False)
 
     def _remove_old_checkpoints(self, trainer: Trainer):
         if (latest_k := self.config.latest_k) == "all":

nshtrainer/callbacks/model_checkpoint.py

@@ -202,4 +202,4 @@ class ModelCheckpoint(_ModelCheckpoint):
 
     @override
     def _remove_checkpoint(self, trainer: Trainer, filepath: str):
-        return _remove_checkpoint(trainer, filepath, remove_metadata=True)
+        return _remove_checkpoint(trainer, filepath, metadata=True, barrier=False)

nshtrainer/ll/model.py

@@ -1 +1,12 @@
 from nshtrainer.model import *  # noqa: F403
+
+from ..util._environment_info import (
+    EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
+)
+from ..util._environment_info import EnvironmentConfig as EnvironmentConfig
+from ..util._environment_info import (
+    EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
+)
+from ..util._environment_info import (
+    EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
+)

nshtrainer/model/__init__.py

@@ -1,16 +1,5 @@
 from typing_extensions import TypeAlias
 
-from ._environment import (
-    EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
-)
-from ._environment import EnvironmentConfig as EnvironmentConfig
-from ._environment import (
-    EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
-)
-from ._environment import (
-    EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
-)
-from ._environment import EnvironmentSnapshotConfig as EnvironmentSnapshotConfig
 from .base import Base as Base
 from .base import LightningModuleBase as LightningModuleBase
 from .config import BaseConfig as BaseConfig

nshtrainer/model/base.py

@@ -11,7 +11,7 @@ from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from typing_extensions import Self, TypeVar, override
 
-from ._environment import EnvironmentConfig
+from ..util._environment_info import EnvironmentConfig
 from .config import BaseConfig
 from .modules.callback import CallbackModuleMixin
 from .modules.debug import DebugModuleMixin

nshtrainer/model/config.py

@@ -44,7 +44,7 @@ from ..callbacks import (
 )
 from ..callbacks.base import CallbackConfigBase
 from ..metrics import MetricConfig
-from ._environment import EnvironmentConfig
+from ..util._environment_info import EnvironmentConfig
 
 log = logging.getLogger(__name__)
 

nshtrainer/{model/_environment.py → util/_environment_info.py}

@@ -14,11 +14,11 @@ import psutil
 import torch
 from typing_extensions import Self
 
-from ..util.slurm import parse_slurm_node_list
+from .slurm import parse_slurm_node_list
 
 if TYPE_CHECKING:
-    from .base import LightningModuleBase
-    from .config import BaseConfig
+    from ..model.base import LightningModuleBase
+    from ..model.config import BaseConfig
 
 
 log = logging.getLogger(__name__)

{nshtrainer-0.10.11.dist-info → nshtrainer-0.10.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.10.11
+Version: 0.10.13
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.10.11.dist-info → nshtrainer-0.10.13.dist-info}/RECORD

@@ -1,7 +1,7 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
 nshtrainer/_checkpoint/loader.py,sha256=48flPr1XgQHOgIPaCrRqOEvRuG0SZuV3cQ1vgHLqFqI,11025
-nshtrainer/_checkpoint/metadata.py,sha256=B6kPmWsq2TQh0gTzBx-1pLIwTVEs_Qw5v0nHEeTBdO4,5636
-nshtrainer/_checkpoint/saver.py,sha256=KZp9ITUVHwj2Ttu81zXKdlS_h-fKkHearspwuAijDpM,1501
+nshtrainer/_checkpoint/metadata.py,sha256=soK9tXVs6EOpzhlnIxTEF51KmdkaCDUj0Rdyid3uREk,5640
+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
 nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
@@ -15,9 +15,9 @@ nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,1
 nshtrainer/callbacks/finite_checks.py,sha256=gJC_RUr3ais3FJI0uB6wUZnDdE3WRwCix3ppA3PwQXA,2077
 nshtrainer/callbacks/gradient_skipping.py,sha256=pqu5AELx4ctJxR2Y7YSSiGd5oGauVCTZFCEIIS6s88w,3665
 nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvLfaw,8080
-nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=t4vWa4PvJDO3rKXKZbuegm7iLl7xCEd17wNif0Bp-BA,4138
+nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=zUeYAGfeQby0R6IwQBJH3lng-MD0vkckdX4aIOm-VIc,4146
 nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
-nshtrainer/callbacks/model_checkpoint.py,sha256=MaDkD8Ismcj8u6l2flCFlqJR3-k1Tc4xzhxNWNux4n0,6556
+nshtrainer/callbacks/model_checkpoint.py,sha256=8D0wWLhr_KiksAA1fjfIuby42Mq6XokCvAnVUhjADd8,6564
 nshtrainer/callbacks/norm_logging.py,sha256=T2psu8mYsw9iahPKT6aUPjkGrZ4TIzm6_UUUmE09GJs,6274
 nshtrainer/callbacks/on_exception_checkpoint.py,sha256=x42BYZ2ejf2rhqPLCmT5nyWKhA9qBEosiV8ZNhhZ6lI,3355
 nshtrainer/callbacks/print_table.py,sha256=_FdAHhqylWGk4Z0c2FrLFeiMA4jhfA_beZRK_BHpzmE,2837
@@ -35,7 +35,7 @@ nshtrainer/ll/config.py,sha256=fKumJf42HY2FITX1QUM1OTXkYD6U2np2ciyd4PFRPZ8,145
 nshtrainer/ll/data.py,sha256=zRG0FRje-jtSHximVzkHIHzpwsyQxpHCoACFihNKLPM,44
 nshtrainer/ll/log.py,sha256=d4BB3TyM8imK65EXOiOeUTF0zFM1ropbe7Vq3DeB0xU,140
 nshtrainer/ll/lr_scheduler.py,sha256=7xjhN6L69BCUzFhcy33NtMtPuCzHiB611zVWFg92lQ0,52
-nshtrainer/ll/model.py,sha256=6I9gQjEFT2Veer-UmcPy05Pt3mnvaxYX1b3sOdaF96A,45
+nshtrainer/ll/model.py,sha256=cxFQfFc-2mAYBGwDpP8m5tjQBs7M47cZ6JoPXksPaoI,473
 nshtrainer/ll/nn.py,sha256=8qiRDFwojIxkB7-LtNWk4mLL2tJbaskHYofDsOIHiNg,42
 nshtrainer/ll/optimizer.py,sha256=3T-VZtT73jVvwCNJGDjgGEbzs-1LFTzMQH-SB_58mSo,49
 nshtrainer/ll/runner.py,sha256=B0m5VEhNKIjF1aFmqPkonkQxDoRL2jeHZGsV3zwhSVE,117
@@ -50,10 +50,9 @@ nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=mn6cyizyI_stkXtg6zxIEGF9b
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-WKArwxxf7ZQizhmbCw,2782
 nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
 nshtrainer/metrics/_config.py,sha256=hWWS4IXENRyH3RmJ7z1Wx1n3Lt1sNMlGOrcU6PW15o0,1104
-nshtrainer/model/__init__.py,sha256=TbexTxiE20WHYg5q3L88Hysk4LlHeKk_isv33aSBREA,1918
-nshtrainer/model/_environment.py,sha256=oTtecQeF5oY2RV7UkkSLnzDy3clz4AUkf9oocD6-e54,23115
-nshtrainer/model/base.py,sha256=WtCj0-nLWeW04Tu2TWVjIq0D-jW_kMN2hg--4VWVnvE,17505
-nshtrainer/model/config.py,sha256=orGBrp8TXnHksfAzXxNJVDdo0X_iIn_nda6BZDS9N70,53349
+nshtrainer/model/__init__.py,sha256=NpvyQHmGaHB8xdraHmm8l7kDHLmvJSgBNQKkfYqtgyI,1454
+nshtrainer/model/base.py,sha256=AXRfEsFAT0Ln7zjYVPU5NgtHS_c8FZM-M4pyLamO7OA,17516
+nshtrainer/model/config.py,sha256=65UDzt3ZZFUQaHMlK7f9wzwyGH3cDyHGtjZ2eOjHvVo,53360
 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
@@ -74,11 +73,12 @@ nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2f
 nshtrainer/trainer/checkpoint_connector.py,sha256=xoqI2dcPnlNFPPLVIU6dBOvRPC9PtfX5qu__xV1lx0Y,2124
 nshtrainer/trainer/signal_connector.py,sha256=llwc8pdKAWxREFpjdi14Bpy8rGVMEJsmJx_s2p4gI8E,10689
 nshtrainer/trainer/trainer.py,sha256=tFyzIsF8c-FABTH6wwDOR9y8kydVJqeVO7PDNFMvhSU,16950
+nshtrainer/util/_environment_info.py,sha256=yPtAbgjCY4tkvh5wp9sjNsF0Z45TYwzEAM_N2_b5BbY,23123
 nshtrainer/util/environment.py,sha256=AeW_kLl-N70wmb6L_JLz1wRj0kA70xs6RCmc9iUqczE,4159
 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.11.dist-info/METADATA,sha256=9WAsp25_csjDcchr5X22g7ocQpQ-d-ewB3gS9EAZSE8,696
-nshtrainer-0.10.11.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.10.11.dist-info/RECORD,,
+nshtrainer-0.10.13.dist-info/METADATA,sha256=HpGl8_E6q2l2nQrIzU5ibNEyUXj8adF8cxMzouUSpAg,696
+nshtrainer-0.10.13.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.10.13.dist-info/RECORD,,