nshtrainer 0.35.1__tar.gz → 0.37.0__tar.gz

{nshtrainer-0.35.1 → nshtrainer-0.37.0}/PKG-INFO

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

{nshtrainer-0.35.1 → nshtrainer-0.37.0}/pyproject.toml

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

{nshtrainer-0.35.1 → nshtrainer-0.37.0}/src/nshtrainer/_checkpoint/saver.py

@@ -30,6 +30,8 @@ def _link_checkpoint(
                     linkpath.unlink(missing_ok=True)
         except Exception:
             log.warning(f"Failed to remove {linkpath}", exc_info=True)
+        else:
+            log.debug(f"Removed {linkpath=}")
 
         if metadata:
             _remove_checkpoint_metadata(linkpath)

{nshtrainer-0.35.1 → nshtrainer-0.37.0}/src/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-0.37.0/src/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-0.35.1 → nshtrainer-0.37.0}/src/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-0.35.1 → nshtrainer-0.37.0}/src/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-0.35.1 → nshtrainer-0.37.0}/src/nshtrainer/util/path.py

@@ -80,11 +80,25 @@ def try_symlink_or_copy(
     link_path: Path,
     target_is_directory: bool = False,
     relative: bool = True,
+    remove_existing: bool = True,
 ):
     """
     Symlinks on Unix, copies on Windows.
     """
 
+    # If the link already exists, remove it
+    if remove_existing:
+        try:
+            if link_path.exists():
+                if link_path.is_dir():
+                    shutil.rmtree(link_path)
+                else:
+                    link_path.unlink(missing_ok=True)
+        except Exception:
+            log.warning(f"Failed to remove {link_path}", exc_info=True)
+        else:
+            log.debug(f"Removed {link_path=}")
+
     symlink_target = get_relative_path(link_path, file_path) if relative else file_path
     try:
         if platform.system() == "Windows":