nshtrainer 1.0.0b37__py3-none-any.whl → 1.0.0b39__py3-none-any.whl

nshtrainer/_directory.py

@@ -81,7 +81,7 @@ class DirectoryConfig(C.Config):
 
         # Save to nshtrainer/{id}/log/{logger name}
         log_dir = self.resolve_subdirectory(run_id, "log")
-        log_dir = log_dir / getattr(logger, "name")
+        log_dir = log_dir / logger.resolve_logger_dirname()
         # ^ NOTE: Logger must have a `name` attribute, as this is
         # the discriminator for the logger registry
         log_dir.mkdir(exist_ok=True)

nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -21,11 +21,8 @@ log = logging.getLogger(__name__)
 class LastCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
     name: Literal["last_checkpoint"] = "last_checkpoint"
 
-    save_on_time_interval: bool = True
-    """Whether to save checkpoints based on time interval."""
-
-    interval: timedelta = timedelta(hours=12)
-    """Time interval between checkpoints when save_on_time_interval is True."""
+    save_on_time_interval: timedelta | None = None
+    """Save a checkpoint every `save_on_time_interval` seconds. If `None`, this feature is disabled."""
 
     @override
     def create_checkpoint(self, trainer_config, dirpath):
@@ -38,8 +35,6 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
         super().__init__(config, dirpath)
         self.start_time = time.time()
         self.last_checkpoint_time = self.start_time
-        self.interval_seconds = config.interval.total_seconds()
-        self.save_on_time_interval = config.save_on_time_interval
 
     @override
     def name(self):
@@ -57,12 +52,18 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
     def topk_sort_reverse(self):
         return True
 
-    def _should_checkpoint(self) -> bool:
-        if not self.save_on_time_interval:
+    def _local_should_checkpoint(self) -> bool:
+        if (interval := self.config.save_on_time_interval) is None:
             return False
+
         current_time = time.time()
         elapsed_time = current_time - self.last_checkpoint_time
-        return elapsed_time >= self.interval_seconds
+        return elapsed_time >= interval.total_seconds()
+
+    def _should_checkpoint(self, trainer: Trainer):
+        if self.config.save_on_time_interval is None:
+            return False
+        return trainer.strategy.broadcast(self._local_should_checkpoint(), src=0)
 
     def _format_duration(self, seconds: float) -> str:
         """Format duration in seconds to a human-readable string."""
@@ -98,7 +99,7 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
         *args,
         **kwargs,
     ):
-        if not self._should_checkpoint():
+        if not self._should_checkpoint(trainer):
             return
         self.save_checkpoints(trainer)
 
@@ -110,5 +111,5 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
     def save_checkpoints(self, trainer):
         super().save_checkpoints(trainer)
 
-        if self.save_on_time_interval:
+        if self.config.save_on_time_interval is not None:
             self.last_checkpoint_time = time.time()

nshtrainer/loggers/base.py

@@ -30,5 +30,14 @@ class LoggerConfigBase(C.Config, ABC):
     def __bool__(self):
         return self.enabled
 
+    def resolve_logger_dirname(self) -> str:
+        if not (name := getattr(self, "name", None)):
+            raise ValueError(
+                "Logger must have a name attribute to resolve the directory name.\n"
+                "Otherwise, you must override `resolve_logger_dirname`."
+            )
+
+        return name
+
 
 logger_registry = C.Registry(LoggerConfigBase, discriminator="name")

{nshtrainer-1.0.0b37.dist-info → nshtrainer-1.0.0b39.dist-info}/METADATA

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

{nshtrainer-1.0.0b37.dist-info → nshtrainer-1.0.0b39.dist-info}/RECORD

@@ -3,7 +3,7 @@ nshtrainer/__init__.py,sha256=g_moPnfQxSxFZX5NB9ILQQOJrt4RTRuiFt9N0STIpxM,874
 nshtrainer/_callback.py,sha256=tXQCDzS6CvMTuTY5lQSH5qZs1pXUi-gt9bQdpXMVdEs,12715
 nshtrainer/_checkpoint/metadata.py,sha256=PHy-54Cg-o3OtCffAqrVv6ZVMU7zhRo_-sZiSEEno1Y,5019
 nshtrainer/_checkpoint/saver.py,sha256=LOP8jjKF0Dw9x9H-BKrLMWlEp1XTan2DUK0zQUCWw5U,1360
-nshtrainer/_directory.py,sha256=xY8Z9POZJw0Uh56yqffZbnNZvdA_tnWCucT31dhwFCM,3183
+nshtrainer/_directory.py,sha256=TJR9ccyuzRlAVfVjGyeQ3E2AFAcz-XbBCxWfiXo2SlY,3191
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
 nshtrainer/callbacks/__init__.py,sha256=4giOYT8A709UOLRtQEt16QbOAFUHCjJ_aLB7ITTwXJI,3577
@@ -12,7 +12,7 @@ nshtrainer/callbacks/base.py,sha256=Alaou1IHAIlMEM7g58d_02ozY2xWlshBN7fsw5Ee21s,
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=l8tkHc83_mLiU0-wT09SWdRzwpm2ulbkLzcuCmuTwzE,620
 nshtrainer/callbacks/checkpoint/_base.py,sha256=ZVEUVl5kjCSSe69Q0rMUbKBNNUog0pxBwWkeyuxG2w0,6304
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=2CQuhPJ3Fi7lDw7z-J8kXXXuDU8-4HcU48oZxR49apk,2667
-nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=MJcNB0biOebx2si2IBFaSUiVOSLSCZTzxB-RcEgO2gY,3482
+nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM6kUSHYNwHJeYG82j2dMcc,3554
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g8mxaeFK5HxA8My0dlIPzIUscSMWyk,3525
 nshtrainer/callbacks/debug_flag.py,sha256=96fuP0C7C6dSs1GiMeUYzzs0X3Q4Pjt9JVWg3b75fU4,1748
 nshtrainer/callbacks/directory_setup.py,sha256=wPas_Ren8ANejogmIdKhqqgj4ulxz9AS_8xVIAfRXa0,2565
@@ -101,7 +101,7 @@ nshtrainer/data/datamodule.py,sha256=lSOgH32nysJWa6Y7ba1QyOdUV0DVVdO98qokP8wigjk
 nshtrainer/data/transform.py,sha256=qd0lIocO59Fk_m90xyOHgFezbymd1mRwly8nbYIfHGc,2263
 nshtrainer/loggers/__init__.py,sha256=Ddd3JJXVzew_ZpwHA9kGnGmvq4OwhItwghDL5PzNhDc,614
 nshtrainer/loggers/actsave.py,sha256=wgNrpBB6wQM7qff8iLDb_sQnbiAcYHRmH56pcEJPB3o,1409
-nshtrainer/loggers/base.py,sha256=1-HoPmOiyXevQvMLXboiKe-4GOE1V5SvjURohOHakVc,882
+nshtrainer/loggers/base.py,sha256=ON92XbwTSgadQOSyw5PiRRFzyH6uJ-xLtE0nB3cbgPc,1205
 nshtrainer/loggers/csv.py,sha256=xJ8mSmw4vJwinIfqhF6t2HWmh_1dXEYyLfGuXwL7WHo,1160
 nshtrainer/loggers/tensorboard.py,sha256=E7iO_fDt9bfH02hBL430bXPLljOo5iGgq2QyPqmx2gQ,2324
 nshtrainer/loggers/wandb.py,sha256=KZXAUWrrmdX_L8rqej77oUHaM0JxZRM8y9z6JP9PISw,6856
@@ -151,6 +151,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.0.0b37.dist-info/METADATA,sha256=ObMgpZ_qJLmBAkeRDN7ufTuRSTltiB_LYPFTphNvWks,988
-nshtrainer-1.0.0b37.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-nshtrainer-1.0.0b37.dist-info/RECORD,,
+nshtrainer-1.0.0b39.dist-info/METADATA,sha256=zzE6nHlj-clB3HJs5_-bBePCHSOrtTkZTi9z_NrSeRY,988
+nshtrainer-1.0.0b39.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+nshtrainer-1.0.0b39.dist-info/RECORD,,