nshtrainer 1.0.0b37__tar.gz → 1.0.0b39__tar.gz

{nshtrainer-1.0.0b37 → nshtrainer-1.0.0b39}/PKG-INFO

@@ -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 → nshtrainer-1.0.0b39}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "1.0.0-beta37"
+version = "1.0.0-beta39"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-1.0.0b37 → nshtrainer-1.0.0b39}/src/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-1.0.0b37 → nshtrainer-1.0.0b39}/src/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-1.0.0b37 → nshtrainer-1.0.0b39}/src/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")