nshtrainer 1.3.3__py3-none-any.whl → 1.3.5__py3-none-any.whl

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

{nshtrainer-1.3.3.dist-info → nshtrainer-1.3.5.dist-info}/METADATA

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

@@ -5,7 +5,7 @@ nshtrainer/_checkpoint/metadata.py,sha256=Hh5a7OkdknUEbkEwX6vS88-XLEeuVDoR6a3en2
 nshtrainer/_checkpoint/saver.py,sha256=utcrYKSosd04N9m2GIylufO5DO05D90qVU3mvadfApU,1658
 nshtrainer/_directory.py,sha256=RAG8e0y3VZwGIyy_D-GXgDMK5OvitQU6qEWxHTpWEeY,2490
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
-nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
+nshtrainer/_hf_hub.py,sha256=kfN0wDxK5JWKKGZnX_706i0KXGhaS19p581LDTPxlRE,13996
 nshtrainer/callbacks/__init__.py,sha256=m6eJuprZfBELuKpngKXre33B9yPXkG7jlKVmI-0yXRQ,4000
 nshtrainer/callbacks/actsave.py,sha256=NSXIIu62MNYe5gz479SMW33bdoKYoYtWtd_iTWFpKpc,3881
 nshtrainer/callbacks/base.py,sha256=K9aom1WVVRYxl-tHWgtmDUQZ1o63NgznvLsjauTKcCc,4225
@@ -17,7 +17,7 @@ nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g
 nshtrainer/callbacks/debug_flag.py,sha256=96fuP0C7C6dSs1GiMeUYzzs0X3Q4Pjt9JVWg3b75fU4,1748
 nshtrainer/callbacks/directory_setup.py,sha256=Ln6f0tCgoBscHeigIAWtCCoAmuWB-kPyaf7SylU7MYo,2773
 nshtrainer/callbacks/distributed_prediction_writer.py,sha256=PvxV9E9lHT-NQ-h1ld7WugajqiFyFXECsreUt3e7pxk,5440
-nshtrainer/callbacks/early_stopping.py,sha256=rC_qYKCQWjRQJFo0ky46uG0aDJdYP8vsSlKunk0bUVI,4765
+nshtrainer/callbacks/early_stopping.py,sha256=LTwOME4-_Zld08UjOeeoNxPOg-hCN7o9MAUWVdzGDdk,5467
 nshtrainer/callbacks/ema.py,sha256=dBFiUXG0xmyCw8-ayuSzJMKqSbepl6Ii5VIbhFlT5ug,12255
 nshtrainer/callbacks/finite_checks.py,sha256=3lZ3kEIjmYQfqTF0DcrgZ9_98ZLQhQj8usH7SgWst3o,2185
 nshtrainer/callbacks/gradient_skipping.py,sha256=8g7oC7PF0LTAEzwiNoaS5tWOnkjk_EB0QG3JdHkQ8ek,3523
@@ -30,7 +30,7 @@ nshtrainer/callbacks/print_table.py,sha256=VaS4JgI963do79laXK4lUkFQx8v6aRSy22W0z
 nshtrainer/callbacks/rlp_sanity_checks.py,sha256=Df9Prq2QKXnaeMBIvMQBhDhJTDeru5UbiuXJOJR16Gk,10050
 nshtrainer/callbacks/shared_parameters.py,sha256=s94jJTAIbDGukYJu6l247QonVOCudGClU4t5kLt8XrY,3076
 nshtrainer/callbacks/timer.py,sha256=gDcw_K_ikf0bkVgxQ0cDhvvNvz6GLZVLcatuKfh0ORU,4731
-nshtrainer/callbacks/wandb_upload_code.py,sha256=shV7UtnXgY2bUlXdVrXiaDs0PNLlIt7TzNJkJPkzvzI,2414
+nshtrainer/callbacks/wandb_upload_code.py,sha256=4X-mpiX5ghj9vnEreK2i8Xyvimqt0K-PNWA2HtT-B6I,1940
 nshtrainer/callbacks/wandb_watch.py,sha256=VB14Dy5ZRXQ3di0fPv0K_DFJurLhroLPytnuwQBiJFg,3037
 nshtrainer/configs/.gitattributes,sha256=VeZmarvNEqiRBOHGcllpKm90nL6C8u4tBu7SEm7fj-E,26
 nshtrainer/configs/__init__.py,sha256=KD3uClMwnA4LfQ7rY5phDdUbp3j8NoZfaGbGPbpaJVs,15848
@@ -146,11 +146,12 @@ nshtrainer/trainer/plugin/environment.py,sha256=SSXRWHjyFUA6oFx3duD_ZwhM59pWUjR1
 nshtrainer/trainer/plugin/io.py,sha256=OmFSKLloMypletjaUr_Ptg6LS0ljqTVIp2o4Hm3eZoE,1926
 nshtrainer/trainer/plugin/layer_sync.py,sha256=-BbEyWZ063O7tZme7Gdu1lVxK6p1NeuLcOfPXnvH29s,663
 nshtrainer/trainer/plugin/precision.py,sha256=7lf7KZd_yFyPmhLApjEIv0pkoDB5zdxi-7in0wRj3z8,5436
-nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
+nshtrainer/trainer/signal_connector.py,sha256=ZgbSkbthoe8MYN6rBoFf-7UDpQtc9fs9pG_FNvTYSfs,10962
 nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTus,1368
 nshtrainer/trainer/trainer.py,sha256=6oky6E8cjGqUNzJGyyTO551pE9A6YueOv5oxg1fZVR0,24129
-nshtrainer/util/_environment_info.py,sha256=MT8mBe6ZolRfKiwU-les1P-lPNPqXpHQcfADrh_A3uY,24629
+nshtrainer/util/_environment_info.py,sha256=j-wyEHKirsu3rIXTtqC2kLmIIkRe6obWjxPVWaqg2ow,24887
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
+nshtrainer/util/code_upload.py,sha256=CpbZEBbA8EcBElUVoCPbP5zdwtNzJhS20RLaOB-q-2k,1257
 nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDWlnZZM,272
 nshtrainer/util/config/dtype.py,sha256=Fn_MhhQoHPyFAnFPSwvcvLiGR3yWFIszMba02CJiC4g,2213
 nshtrainer/util/config/duration.py,sha256=mM-UfU_HvhXwW33TYEDg0x58n80tnle2e6VaWtxZTjk,764
@@ -160,6 +161,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.3.3.dist-info/METADATA,sha256=K_xd3BrF1Yz7gGbNQgywkjysCFuwXi3GCBoQ5EaFVKY,960
-nshtrainer-1.3.3.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.3.3.dist-info/RECORD,,
+nshtrainer-1.3.5.dist-info/METADATA,sha256=GUU8QgA8rxeCX1Z9FfwSvZQ46f0xsMvtm4p1Uz8uEwE,979
+nshtrainer-1.3.5.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.3.5.dist-info/RECORD,,