nshtrainer 1.3.3__tar.gz → 1.3.5__tar.gz

{nshtrainer-1.3.3 → nshtrainer-1.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.3.3
+Version: 1.3.5
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -15,7 +15,7 @@ Requires-Dist: GitPython ; extra == "extra"
 Requires-Dist: huggingface-hub ; extra == "extra"
 Requires-Dist: lightning
 Requires-Dist: nshconfig (>0.39)
-Requires-Dist: nshrunner
+Requires-Dist: nshrunner ; extra == "extra"
 Requires-Dist: nshutils ; extra == "extra"
 Requires-Dist: numpy
 Requires-Dist: packaging

{nshtrainer-1.3.3 → nshtrainer-1.3.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nshtrainer"
-version = "1.3.3"
+version = "1.3.5"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.10,<4.0"
@@ -8,7 +8,6 @@ readme = "README.md"
 
 dependencies = [
     "nshconfig>0.39",
-    "nshrunner",
     "psutil",
     "numpy",
     "torch",
@@ -20,6 +19,7 @@ dependencies = [
 
 [project.optional-dependencies]
 extra = [
+    "nshrunner",
     "wrapt",
     "GitPython",
     "wandb",

{nshtrainer-1.3.3 → nshtrainer-1.3.5}/src/nshtrainer/_hf_hub.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import contextlib
 import logging
-import os
 import re
 from dataclasses import dataclass
 from functools import cached_property
@@ -10,7 +9,6 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
 
 import nshconfig as C
-from nshrunner._env import SNAPSHOT_DIR
 from typing_extensions import assert_never, override
 
 from ._callback import NTCallbackBase
@@ -19,6 +17,7 @@ from .callbacks.base import (
     CallbackMetadataConfig,
     callback_registry,
 )
+from .util.code_upload import get_code_dir
 
 if TYPE_CHECKING:
     from huggingface_hub import HfApi  # noqa: F401
@@ -319,20 +318,13 @@ class HFHubCallback(NTCallbackBase):
     def _save_code(self):
         # If a snapshot has been taken (which can be detected using the SNAPSHOT_DIR env),
         # then upload all contents within the snapshot directory to the repository.
-        if not (snapshot_dir := os.environ.get(SNAPSHOT_DIR)):
+        if (snapshot_dir := get_code_dir()) is None:
             log.debug("No snapshot directory found. Skipping upload.")
             return
 
         with self._with_error_handling("save code"):
-            snapshot_dir = Path(snapshot_dir)
-            if not snapshot_dir.exists() or not snapshot_dir.is_dir():
-                log.warning(
-                    f"Snapshot directory '{snapshot_dir}' does not exist or is not a directory."
-                )
-                return
-
             self.api.upload_folder(
-                folder_path=str(snapshot_dir),
+                folder_path=str(snapshot_dir.absolute()),
                 repo_id=self.repo_id,
                 repo_type="model",
                 path_in_repo="code",  # Prefix with "code" folder

{nshtrainer-1.3.3 → nshtrainer-1.3.5}/src/nshtrainer/callbacks/early_stopping.py

@@ -43,6 +43,13 @@ class EarlyStoppingCallbackConfig(CallbackConfigBase):
     the training will be stopped.
     """
 
+    skip_first_n_epochs: int = 0
+    """
+    Number of initial epochs to skip before starting to monitor for early stopping.
+    This helps avoid false early stopping when the model might temporarily perform worse
+    during early training phases.
+    """
+
     strict: bool = True
     """
     Whether to enforce that the monitored quantity must improve by at least `min_delta`
@@ -94,6 +101,16 @@ class EarlyStoppingCallback(_EarlyStopping):
         if getattr(trainer, "fast_dev_run", False):
             return
 
+        # Skip early stopping check for the first n epochs
+        if trainer.current_epoch < self.config.skip_first_n_epochs:
+            if self.verbose and trainer.current_epoch == 0:
+                self._log_info(
+                    trainer,
+                    f"Early stopping checks are disabled for the first {self.config.skip_first_n_epochs} epochs",
+                    self.log_rank_zero_only,
+                )
+            return
+
         should_stop, reason = False, None
 
         if not should_stop:

{nshtrainer-1.3.3 → nshtrainer-1.3.5}/src/nshtrainer/callbacks/wandb_upload_code.py

@@ -1,16 +1,14 @@
 from __future__ import annotations
 
 import logging
-import os
-from pathlib import Path
 from typing import Literal, cast
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks.callback import Callback
 from lightning.pytorch.loggers import WandbLogger
-from nshrunner._env import SNAPSHOT_DIR
 from typing_extensions import final, override
 
+from ..util.code_upload import get_code_dir
 from .base import CallbackConfigBase, callback_registry
 
 log = logging.getLogger(__name__)
@@ -62,22 +60,12 @@ class WandbUploadCodeCallback(Callback):
             log.warning("Wandb logger not found. Skipping code upload.")
             return
 
-        from wandb.wandb_run import Run
-
-        run = cast(Run, logger.experiment)
-
-        # If a snapshot has been taken (which can be detected using the SNAPSHOT_DIR env),
-        # then upload all contents within the snapshot directory to the repository.
-        if not (snapshot_dir := os.environ.get(SNAPSHOT_DIR)):
-            log.debug("No snapshot directory found. Skipping upload.")
+        if (snapshot_dir := get_code_dir()) is None:
+            log.info("No nshrunner snapshot found. Skipping code upload.")
             return
 
-        snapshot_dir = Path(snapshot_dir)
-        if not snapshot_dir.exists() or not snapshot_dir.is_dir():
-            log.warning(
-                f"Snapshot directory '{snapshot_dir}' does not exist or is not a directory."
-            )
-            return
+        from wandb.wandb_run import Run
 
+        run = cast(Run, logger.experiment)
         log.info(f"Uploading code from snapshot directory '{snapshot_dir}'")
         run.log_code(str(snapshot_dir.absolute()))

{nshtrainer-1.3.3 → nshtrainer-1.3.5}/src/nshtrainer/trainer/signal_connector.py

@@ -14,7 +14,6 @@ from pathlib import Path
 from types import FrameType
 from typing import Any
 
-import nshrunner as nr
 import torch.utils.data
 from lightning.fabric.plugins.environments.lsf import LSFEnvironment
 from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
@@ -34,6 +33,12 @@ _IS_WINDOWS = platform.system() == "Windows"
 
 
 def _resolve_requeue_signals():
+    try:
+        import nshrunner as nr
+    except ImportError:
+        log.debug("nshrunner not found. Skipping signal requeueing.")
+        return None
+
     if (session := nr.Session.from_current_session()) is None:
         return None
 
@@ -52,9 +57,9 @@ class _SignalConnector(_LightningSignalConnector):
 
         signals_set = set(signals)
         valid_signals: set[signal.Signals] = signal.valid_signals()
-        assert signals_set.issubset(
-            valid_signals
-        ), f"Invalid signal(s) found: {signals_set - valid_signals}"
+        assert signals_set.issubset(valid_signals), (
+            f"Invalid signal(s) found: {signals_set - valid_signals}"
+        )
         return signals
 
     def _compose_and_register(
@@ -241,9 +246,9 @@ class _SignalConnector(_LightningSignalConnector):
                 "Writing requeue script to exit script directory."
             )
             exit_script_dir = Path(exit_script_dir)
-            assert (
-                exit_script_dir.is_dir()
-            ), f"Exit script directory {exit_script_dir} does not exist"
+            assert exit_script_dir.is_dir(), (
+                f"Exit script directory {exit_script_dir} does not exist"
+            )
 
             exit_script_path = exit_script_dir / f"requeue_{job_id}.sh"
             log.info(f"Writing requeue script to {exit_script_path}")

{nshtrainer-1.3.3 → nshtrainer-1.3.5}/src/nshtrainer/util/_environment_info.py

@@ -356,12 +356,20 @@ class EnvironmentSnapshotConfig(C.Config):
 
     @classmethod
     def from_current_environment(cls):
-        draft = cls.draft()
-        if snapshot_dir := os.environ.get("NSHRUNNER_SNAPSHOT_DIR"):
-            draft.snapshot_dir = Path(snapshot_dir)
-        if modules := os.environ.get("NSHRUNNER_SNAPSHOT_MODULES"):
-            draft.modules = modules.split(",")
-        return draft.finalize()
+        try:
+            import nshrunner as nr
+
+            if (session := nr.Session.from_current_session()) is None:
+                log.warning("No active session found, skipping snapshot information")
+                return cls.empty()
+
+            draft = cls.draft()
+            draft.snapshot_dir = session.snapshot_dir
+            draft.modules = session.snapshot_modules
+            return draft.finalize()
+        except ImportError:
+            log.warning("nshrunner not found, skipping snapshot information")
+            return cls.empty()
 
 
 class EnvironmentPackageConfig(C.Config):

nshtrainer-1.3.5/src/nshtrainer/util/code_upload.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+log = logging.getLogger(__name__)
+
+
+def get_code_dir() -> Path | None:
+    try:
+        import nshrunner as nr
+
+        if (session := nr.Session.from_current_session()) is None:
+            log.debug("No active session found. Skipping code upload.")
+            return None
+
+        # New versions of nshrunner will have the code_dir attribute
+        #   in the session object. We should use that. Otherwise, use snapshot_dir.
+        try:
+            code_dir = session.code_dir  # type: ignore
+        except AttributeError:
+            code_dir = session.snapshot_dir
+
+        if code_dir is None:
+            log.debug("No code directory found. Skipping code upload.")
+            return None
+
+        assert isinstance(code_dir, Path), (
+            f"Code directory should be a Path object. Got {type(code_dir)} instead."
+        )
+        if not code_dir.exists() or not code_dir.is_dir():
+            log.warning(
+                f"Code directory '{code_dir}' does not exist or is not a directory."
+            )
+            return None
+
+        return code_dir
+    except ImportError:
+        log.debug("nshrunner not found. Skipping code upload.")
+        return None