nshtrainer 0.35.0__py3-none-any.whl → 0.36.0__py3-none-any.whl

nshtrainer/callbacks/__init__.py

@@ -39,6 +39,8 @@ from .shared_parameters import SharedParametersConfig as SharedParametersConfig
 from .throughput_monitor import ThroughputMonitorConfig as ThroughputMonitorConfig
 from .timer import EpochTimer as EpochTimer
 from .timer import EpochTimerConfig as EpochTimerConfig
+from .wandb_upload_code import WandbUploadCodeCallback as WandbUploadCodeCallback
+from .wandb_upload_code import WandbUploadCodeConfig as WandbUploadCodeConfig
 from .wandb_watch import WandbWatchCallback as WandbWatchCallback
 from .wandb_watch import WandbWatchConfig as WandbWatchConfig
 
@@ -57,6 +59,7 @@ CallbackConfig = Annotated[
     | OnExceptionCheckpointCallbackConfig
     | SharedParametersConfig
     | RLPSanityChecksConfig
-    | WandbWatchConfig,
+    | WandbWatchConfig
+    | WandbUploadCodeConfig,
     C.Field(discriminator="name"),
 ]

nshtrainer/callbacks/wandb_upload_code.py

@@ -0,0 +1,79 @@
+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 override
+
+from .base import CallbackConfigBase
+
+log = logging.getLogger(__name__)
+
+
+class WandbUploadCodeConfig(CallbackConfigBase):
+    name: Literal["wandb_upload_code"] = "wandb_upload_code"
+
+    enabled: bool = True
+    """Enable uploading the code to wandb."""
+
+    def __bool__(self):
+        return self.enabled
+
+    @override
+    def create_callbacks(self, root_config):
+        if not self:
+            return
+
+        yield WandbUploadCodeCallback(self)
+
+
+class WandbUploadCodeCallback(Callback):
+    def __init__(self, config: WandbUploadCodeConfig):
+        super().__init__()
+
+        self.config = config
+
+    @override
+    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str):
+        if not self.config:
+            return
+
+        if not trainer.is_global_zero:
+            return
+
+        if (
+            logger := next(
+                (
+                    logger
+                    for logger in trainer.loggers
+                    if isinstance(logger, WandbLogger)
+                ),
+                None,
+            )
+        ) is None:
+            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.")
+            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
+
+        log.info(f"Uploading code from snapshot directory '{snapshot_dir}'")
+        run.log_code(str(snapshot_dir.absolute()))

nshtrainer/config.py

@@ -61,14 +61,17 @@ from nshtrainer.callbacks.throughput_monitor import (
     ThroughputMonitorConfig as ThroughputMonitorConfig,
 )
 from nshtrainer.callbacks.timer import EpochTimerConfig as EpochTimerConfig
+from nshtrainer.callbacks.wandb_upload_code import (
+    WandbUploadCodeConfig as WandbUploadCodeConfig,
+)
 from nshtrainer.callbacks.wandb_watch import WandbWatchConfig as WandbWatchConfig
+from nshtrainer.config import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.loggers._base import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers.csv import CSVLoggerConfig as CSVLoggerConfig
 from nshtrainer.loggers.tensorboard import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
 from nshtrainer.loggers.wandb import WandbLoggerConfig as WandbLoggerConfig
-from nshtrainer.lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.lr_scheduler._base import LRSchedulerConfigBase as LRSchedulerConfigBase
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     DurationConfig as DurationConfig,

nshtrainer/loggers/wandb.py

@@ -5,9 +5,10 @@ from typing import TYPE_CHECKING, Literal
 import nshconfig as C
 from lightning.pytorch import Callback, LightningModule, Trainer
 from packaging import version
-from typing_extensions import override
+from typing_extensions import assert_never, override
 
 from ..callbacks.base import CallbackConfigBase
+from ..callbacks.wandb_upload_code import WandbUploadCodeConfig
 from ..callbacks.wandb_watch import WandbWatchConfig
 from ._base import BaseLoggerConfig
 
@@ -82,15 +83,18 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
     project: str | None = None
     """WandB project name to use for the logger. If None, will use the root config's project name."""
 
-    log_model: bool | Literal["all"] = False
+    log_model: Literal["all", "latest", "none"] | bool = False
     """
     Whether to log the model checkpoints to wandb.
     Valid values are:
-        - False: Do not log the model checkpoints.
-        - True: Log the latest model checkpoint.
-        - "all": Log all model checkpoints.
+    - "all": Log all checkpoints.
+    - "latest" or True: Log only the latest checkpoint.
+    - "none" or False: Do not log any checkpoints
     """
 
+    log_code: WandbUploadCodeConfig | None = None
+    """WandB code upload configuration. Used to upload code to WandB."""
+
     watch: WandbWatchConfig | None = WandbWatchConfig()
     """WandB model watch configuration. Used to log model architecture, gradients, and parameters."""
 
@@ -110,6 +114,18 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
         self.use_wandb_core = value
         return self
 
+    @property
+    def _lightning_log_model(self) -> Literal["all"] | bool:
+        match self.log_model:
+            case "all":
+                return "all"
+            case "latest" | True:
+                return True
+            case "none" | False:
+                return False
+            case _:
+                assert_never(self.log_model)
+
     @override
     def create_logger(self, root_config):
         if not self.enabled:
@@ -128,11 +144,28 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
                         f"(expected version >= 0.17.5, found version {wandb.__version__}). "
                         "Please either upgrade to a newer version of WandB or disable the `use_wandb_core` option."
                     )
-                else:
+                # W&B versions 0.18.0 use wandb-core by default
+                elif wandb_version < version.parse("0.18.0"):
                     wandb.require("core")  # type: ignore
                     log.critical("Using the `wandb-core` backend for WandB.")
             except ImportError:
                 pass
+        else:
+            # W&B versions 0.18.0 use wandb-core by default,
+            #   so if `use_wandb_core` is False, we should use the old backend
+            #   explicitly.
+            wandb_version = version.parse(importlib.metadata.version("wandb"))
+            if wandb_version >= version.parse("0.18.0"):
+                log.warning(
+                    "Explicitly using the old backend for WandB. "
+                    "If you want to use the new `wandb-core` backend, set `use_wandb_core=True`."
+                )
+                try:
+                    import wandb  # type: ignore
+
+                    wandb.require("legacy-service")  # type: ignore
+                except ImportError:
+                    pass
 
         from lightning.pytorch.loggers.wandb import WandbLogger
 
@@ -145,7 +178,7 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
             project=self.project or _project_name(root_config),
             name=root_config.run_name,
             version=root_config.id,
-            log_model=self.log_model,
+            log_model=self._lightning_log_model,
             notes=(
                 "\n".join(f"- {note}" for note in root_config.notes)
                 if root_config.notes
@@ -161,3 +194,6 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
 
         if self.watch:
             yield from self.watch.create_callbacks(root_config)
+
+        if self.log_code:
+            yield from self.log_code.create_callbacks(root_config)

nshtrainer/model/mixins/callback.py

@@ -13,11 +13,11 @@ CallbackFn: TypeAlias = Callable[[], Callback | Iterable[Callback] | None]
 
 
 class CallbackRegistrarModuleMixin:
-    @override
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self._nshtrainer_callbacks: list[CallbackFn] = []
+    @property
+    def _nshtrainer_callbacks(self) -> list[CallbackFn]:
+        if not hasattr(self, "_private_nshtrainer_callbacks_list"):
+            self._private_nshtrainer_callbacks_list = []
+        return self._private_nshtrainer_callbacks_list
 
     def register_callback(
         self,

{nshtrainer-0.35.0.dist-info → nshtrainer-0.36.0.dist-info}/METADATA

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

{nshtrainer-0.35.0.dist-info → nshtrainer-0.36.0.dist-info}/RECORD

@@ -6,7 +6,7 @@ nshtrainer/_checkpoint/saver.py,sha256=MbX_WjkDtHHAf9Ms-KXDlknkjiPXVoGIe2ciO28Ad
 nshtrainer/_directory.py,sha256=RjnW6vKTeKlz2vQWT3cG0Jje5BkFXA7HpUubDhcSiq4,2993
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
 nshtrainer/_hf_hub.py,sha256=0bkXkqhve5D1onMW-fCfuvVKlTn0i6jv_6uMNgZ7OHQ,12974
-nshtrainer/callbacks/__init__.py,sha256=1SBLpMsx7BzgimO35MwQViYBcbgxlkyvTMz1JKUKK-0,3060
+nshtrainer/callbacks/__init__.py,sha256=y4QKDiI4IykcHncDOE-OemFFQqegxyeoD_3v9i2OwFw,3248
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=NpjeKmonJ1Kaz5_39XSn3LlDwvbGjk6WV8BpHSNCvI4,3508
@@ -29,8 +29,9 @@ nshtrainer/callbacks/rlp_sanity_checks.py,sha256=c30G9jAu42QLLIS5LnusdSnI3wqyIHg
 nshtrainer/callbacks/shared_parameters.py,sha256=fqlDweFDXPV_bfcAWpRgaJIad9i5AehYDtuJjDtUum4,2922
 nshtrainer/callbacks/throughput_monitor.py,sha256=H_ocXErZxUO3dxFk8Tx_VQdpI9E_Ztvqof5WtFevLyQ,1838
 nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
+nshtrainer/callbacks/wandb_upload_code.py,sha256=OWG4UkL2SfW6oj6AGRXeBJsZmgsqeHLW2Fj8Jm4ga3I,2298
 nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
-nshtrainer/config.py,sha256=6U7B-kCIMrfEnF_y92RuBm1WfASW7k05Zsm2uHBzRrk,8205
+nshtrainer/config.py,sha256=6jApGtO9DVFoXKr9_Z7-MFG5R4WXjbpzZ6jkNI3yD-Y,8306
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=ybMJF-CguaZ17fLEweZ5suaGOiHOMEm3Bn8rQfGTzGQ,5445
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
@@ -55,7 +56,7 @@ nshtrainer/loggers/__init__.py,sha256=C_xk0A3_qKbNdTmzK85AgjRHFD3w-jPRS2ig-iPhfE
 nshtrainer/loggers/_base.py,sha256=xiZKEK0ALJkcqf4OpVNRY0QbZsamR_WR7x7m_68YHXQ,705
 nshtrainer/loggers/csv.py,sha256=D_lYyd94bZ8jAgnRo-ARtFgVcInaD9zktxtsUD9RWCI,1052
 nshtrainer/loggers/tensorboard.py,sha256=wL2amRSdP68zbslZvBeM0ZQBnjF3hIKsz-_lBbdomaM,2216
-nshtrainer/loggers/wandb.py,sha256=C-yGX9e2FUSfbUxur7-meNUjpB3D8hIdVCOgPzGm3QM,5140
+nshtrainer/loggers/wandb.py,sha256=td8J2v8T1nvGQI7OYQ1El6k8FGsXZxbnuY97s8KzCiY,6643
 nshtrainer/lr_scheduler/__init__.py,sha256=uEvgaFAs-4s_bAEMaildy0GT6OvgpgOEKTuzqutESHE,736
 nshtrainer/lr_scheduler/_base.py,sha256=7xOIuxQ86YHbFWG5a3gX46emQj1WN_LaY4-i0Q1TDBg,3659
 nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=E8LW78uuby7bIsoLPpcF1bmNK4lSko-r3qPL-vuHWXQ,5370
@@ -65,7 +66,7 @@ nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ
 nshtrainer/model/__init__.py,sha256=2i_VEy6u_Y1LUGKljHXWeekvhnUcanZM2QyaaBM1Bmw,261
 nshtrainer/model/base.py,sha256=NasbYZJBuEly6Hm9t9HVZk-CUHmy4T7p1v-Ye981XA4,18609
 nshtrainer/model/config.py,sha256=Q4Wong6w3cp_Sq7s8iZdABKF-LZBbSCFn_TQPYkhkrI,6572
-nshtrainer/model/mixins/callback.py,sha256=lvX9Q2ErETXmGFd79CscSAOJAlTWq-mwMKVC0d0uH1c,2324
+nshtrainer/model/mixins/callback.py,sha256=rbe8P22iEjPkH1df6rfEo3Txw7EwSz6Dkm0TWO_AysM,2419
 nshtrainer/model/mixins/logger.py,sha256=xOymSTofukEYZGkGojXsMEO__ZlBI5lIPZVmlotMEX8,5291
 nshtrainer/nn/__init__.py,sha256=0QPFl02a71WZQjLMGOlFNMmsYP5aa1q3eABHmnWH58Q,1427
 nshtrainer/nn/mlp.py,sha256=V0FrScpIUdg_IgIO8GMtIsGEtmHjwF14i2IWxmZrsqg,5952
@@ -98,6 +99,6 @@ 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.35.0.dist-info/METADATA,sha256=NBZegh-RUfnkVt_ERUPdH7fdCFZriQZXoMskq_8HB60,916
-nshtrainer-0.35.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.35.0.dist-info/RECORD,,
+nshtrainer-0.36.0.dist-info/METADATA,sha256=eWSmYvpViGa536HEJ1zX2IUoRpt-wlgiNmfv8mPm7Yg,916
+nshtrainer-0.36.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.36.0.dist-info/RECORD,,