nshtrainer 1.3.4__py3-none-any.whl → 1.3.6__py3-none-any.whl

nshtrainer/__init__.py

@@ -19,3 +19,17 @@ try:
     from . import configs as configs
 except BaseException:
     pass
+
+try:
+    from importlib.metadata import PackageNotFoundError, version
+except ImportError:
+    # For Python <3.8
+    from importlib_metadata import (  # pyright: ignore[reportMissingImports]
+        PackageNotFoundError,
+        version,
+    )
+
+try:
+    __version__ = version(__name__)
+except PackageNotFoundError:
+    __version__ = "unknown"

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/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/configs/__init__.py

@@ -5,7 +5,6 @@ __codegen__ = True
 from nshtrainer import MetricConfig as MetricConfig
 from nshtrainer import TrainerConfig as TrainerConfig
 from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
-from nshtrainer._directory import DirectoryConfig as DirectoryConfig
 from nshtrainer._hf_hub import CallbackConfigBase as CallbackConfigBase
 from nshtrainer._hf_hub import (
     HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig,
@@ -126,9 +125,9 @@ from nshtrainer.trainer._config import (
     CheckpointCallbackConfig as CheckpointCallbackConfig,
 )
 from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
+from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
 from nshtrainer.trainer._config import EnvironmentConfig as EnvironmentConfig
 from nshtrainer.trainer._config import GradientClippingConfig as GradientClippingConfig
-from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
 from nshtrainer.trainer._config import StrategyConfig as StrategyConfig
 from nshtrainer.trainer.accelerator import CPUAcceleratorConfig as CPUAcceleratorConfig
 from nshtrainer.trainer.accelerator import (
@@ -227,7 +226,6 @@ from nshtrainer.util.config import EpochsConfig as EpochsConfig
 from nshtrainer.util.config import StepsConfig as StepsConfig
 
 from . import _checkpoint as _checkpoint
-from . import _directory as _directory
 from . import _hf_hub as _hf_hub
 from . import callbacks as callbacks
 from . import loggers as loggers
@@ -338,7 +336,6 @@ __all__ = [
     "RpropConfig",
     "SGDConfig",
     "SLURMEnvironmentPlugin",
-    "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "SiLUNonlinearityConfig",
     "SigmoidNonlinearityConfig",
@@ -367,7 +364,6 @@ __all__ = [
     "XLAEnvironmentPlugin",
     "XLAPluginConfig",
     "_checkpoint",
-    "_directory",
     "_hf_hub",
     "accelerator_registry",
     "callback_registry",

nshtrainer/configs/trainer/__init__.py

@@ -22,6 +22,9 @@ from nshtrainer.trainer._config import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
 from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
+from nshtrainer.trainer._config import (
+    DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
+)
 from nshtrainer.trainer._config import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -51,7 +54,6 @@ from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
-from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
@@ -152,6 +154,7 @@ __all__ = [
     "DebugFlagCallbackConfig",
     "DeepSpeedPluginConfig",
     "DirectoryConfig",
+    "DirectorySetupCallbackConfig",
     "DistributedPredictionWriterConfig",
     "DoublePrecisionPluginConfig",
     "EarlyStoppingCallbackConfig",
@@ -180,7 +183,6 @@ __all__ = [
     "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
     "SLURMEnvironmentPlugin",
-    "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "StrategyConfig",
     "StrategyConfigBase",

nshtrainer/configs/trainer/_config/__init__.py

@@ -18,6 +18,9 @@ from nshtrainer.trainer._config import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
 from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
+from nshtrainer.trainer._config import (
+    DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
+)
 from nshtrainer.trainer._config import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -48,7 +51,6 @@ from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import (
     RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
 )
-from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
 from nshtrainer.trainer._config import (
     SharedParametersCallbackConfig as SharedParametersCallbackConfig,
 )
@@ -70,6 +72,7 @@ __all__ = [
     "CheckpointSavingConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
+    "DirectorySetupCallbackConfig",
     "EarlyStoppingCallbackConfig",
     "EnvironmentConfig",
     "GradientClippingConfig",
@@ -86,7 +89,6 @@ __all__ = [
     "PluginConfig",
     "ProfilerConfig",
     "RLPSanityChecksCallbackConfig",
-    "SanityCheckingConfig",
     "SharedParametersCallbackConfig",
     "StrategyConfig",
     "TensorboardLoggerConfig",

nshtrainer/trainer/_config.py

@@ -26,7 +26,6 @@ from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.strategies.strategy import Strategy
 from typing_extensions import TypeAliasType, TypedDict, override
 
-from .._directory import DirectoryConfig
 from .._hf_hub import HuggingFaceHubConfig
 from ..callbacks import (
     BestCheckpointCallbackConfig,
@@ -38,6 +37,7 @@ from ..callbacks import (
 )
 from ..callbacks.base import CallbackConfigBase
 from ..callbacks.debug_flag import DebugFlagCallbackConfig
+from ..callbacks.directory_setup import DirectorySetupCallbackConfig
 from ..callbacks.log_epoch import LogEpochCallbackConfig
 from ..callbacks.lr_monitor import LearningRateMonitorConfig
 from ..callbacks.metric_validation import MetricValidationCallbackConfig
@@ -352,19 +352,74 @@ class LightningTrainerKwargs(TypedDict, total=False):
     """
 
 
-class SanityCheckingConfig(C.Config):
-    reduce_lr_on_plateau: Literal["disable", "error", "warn"] = "error"
+DEFAULT_LOGDIR_BASENAME = "nshtrainer_logs"
+"""Default base name for the log directory."""
+
+
+class DirectoryConfig(C.Config):
+    project_root: Path | None = None
     """
-    If enabled, will do some sanity checks if the `ReduceLROnPlateau` scheduler is used:
-        - If the `interval` is step, it makes sure that validation is called every `frequency` steps.
-        - If the `interval` is epoch, it makes sure that validation is called every `frequency` epochs.
-    Valid values are: "disable", "warn", "error".
+    Root directory for this project.
+
+    This isn't specific to the current run; it is the parent directory of all runs.
     """
 
+    logdir_basename: str = DEFAULT_LOGDIR_BASENAME
+    """Base name for the log directory."""
+
+    setup_callback: DirectorySetupCallbackConfig = DirectorySetupCallbackConfig()
+    """Configuration for the directory setup PyTorch Lightning callback."""
+
+    def resolve_run_root_directory(self, run_id: str) -> Path:
+        if (project_root_dir := self.project_root) is None:
+            project_root_dir = Path.cwd()
+
+        # The default base dir is $CWD/{logdir_basename}/{id}/
+        base_dir = project_root_dir / self.logdir_basename
+        base_dir.mkdir(exist_ok=True)
+
+        # Add a .gitignore file to the {logdir_basename} directory
+        #   which will ignore all files except for the .gitignore file itself
+        gitignore_path = base_dir / ".gitignore"
+        if not gitignore_path.exists():
+            gitignore_path.touch()
+            gitignore_path.write_text("*\n")
+
+        base_dir = base_dir / run_id
+        base_dir.mkdir(exist_ok=True)
+
+        return base_dir
+
+    def resolve_subdirectory(self, run_id: str, subdirectory: str) -> Path:
+        # The subdir will be $CWD/{logdir_basename}/{id}/{log, stdio, checkpoint, activation}/
+        if (subdir := getattr(self, subdirectory, None)) is not None:
+            assert isinstance(subdir, Path), (
+                f"Expected a Path for {subdirectory}, got {type(subdir)}"
+            )
+            return subdir
+
+        dir = self.resolve_run_root_directory(run_id)
+        dir = dir / subdirectory
+        dir.mkdir(exist_ok=True)
+        return dir
+
+    def _resolve_log_directory_for_logger(self, run_id: str, logger: LoggerConfig):
+        if (log_dir := logger.log_dir) is not None:
+            return log_dir
+
+        # Save to {logdir_basename}/{id}/log/{logger name}
+        log_dir = self.resolve_subdirectory(run_id, "log")
+        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)
+
+        return log_dir
+
 
 class TrainerConfig(C.Config):
     # region Active Run Configuration
-    id: str = C.Field(default_factory=lambda: TrainerConfig.generate_id())
+    id: Annotated[str, C.AllowMissing()] = C.MISSING
     """ID of the run."""
     name: list[str] = []
     """Run name in parts. Full name is constructed by joining the parts with spaces."""
@@ -393,39 +448,6 @@ class TrainerConfig(C.Config):
 
     directory: DirectoryConfig = DirectoryConfig()
     """Directory configuration options."""
-
-    _rng: ClassVar[np.random.Generator | None] = None
-
-    @classmethod
-    def generate_id(cls, *, length: int = 8) -> str:
-        """
-        Generate a random ID of specified length.
-
-        """
-        if (rng := cls._rng) is None:
-            rng = np.random.default_rng()
-
-        alphabet = list(string.ascii_lowercase + string.digits)
-
-        id = "".join(rng.choice(alphabet) for _ in range(length))
-        return id
-
-    @classmethod
-    def set_seed(cls, seed: int | None = None) -> None:
-        """
-        Set the seed for the random number generator.
-
-        Args:
-            seed (int | None, optional): The seed value to set. If None, a seed based on the current time will be used. Defaults to None.
-
-        Returns:
-            None
-        """
-        if seed is None:
-            seed = int(time.time() * 1000)
-        log.critical(f"Seeding {cls.__name__} with seed {seed}")
-        cls._rng = np.random.default_rng(seed)
-
     # endregion
 
     primary_metric: MetricConfig | None = None
@@ -755,40 +777,40 @@ class TrainerConfig(C.Config):
             None,
         )
 
-    def _nshtrainer_all_callback_configs(self) -> Iterable[CallbackConfigBase | None]:
-        yield self.directory.setup_callback
-        yield self.early_stopping
-        yield self.checkpoint_saving
-        yield self.lr_monitor
-        yield from (
-            logger_config
-            for logger_config in self.enabled_loggers()
-            if logger_config is not None
-            and isinstance(logger_config, CallbackConfigBase)
-        )
-        yield self.log_epoch
-        yield self.log_norms
-        yield self.hf_hub
-        yield self.shared_parameters
-        yield self.reduce_lr_on_plateau_sanity_checking
-        yield self.auto_set_debug_flag
-        yield self.auto_validate_metrics
-        yield from self.callbacks
+    # region Helper Methods
+    def id_(self, value: str):
+        """
+        Set the id for the trainer configuration in-place.
 
-    def _nshtrainer_all_logger_configs(self) -> Iterable[LoggerConfigBase | None]:
-        # Disable all loggers if barebones mode is enabled
-        if self.barebones:
-            return
+        Parameters
+        ----------
+        value : str
+            The id value to set
 
-        yield from self.enabled_loggers()
-        yield self.actsave_logger
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.id = value
+        return self
 
-    def _nshtrainer_validate_before_run(self):
-        # shared_parameters is not supported under barebones mode
-        if self.barebones and self.shared_parameters:
-            raise ValueError("shared_parameters is not supported under barebones mode")
+    def with_id(self, value: str):
+        """
+        Create a copy of the current configuration with an updated id.
+
+        Parameters
+        ----------
+        value : str
+            The id value to set
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the updated id
+        """
+        return copy.deepcopy(self).id_(value)
 
-    # region Helper Methods
     def fast_dev_run_(self, value: int | bool = True, /):
         """
         Enables fast_dev_run mode for the trainer.
@@ -831,6 +853,349 @@ class TrainerConfig(C.Config):
         """
         return copy.deepcopy(self).project_root_(project_root)
 
+    def name_(self, *parts: str):
+        """
+        Set the name for the trainer configuration in-place.
+
+        Parameters
+        ----------
+        *parts : str
+            The parts of the name to set. Will be joined with spaces.
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.name = list(parts)
+        return self
+
+    def with_name(self, *parts: str):
+        """
+        Create a copy of the current configuration with an updated name.
+
+        Parameters
+        ----------
+        *parts : str
+            The parts of the name to set. Will be joined with spaces.
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the updated name
+        """
+        return copy.deepcopy(self).name_(*parts)
+
+    def project_(self, project: str | None):
+        """
+        Set the project name for the trainer configuration in-place.
+
+        Parameters
+        ----------
+        project : str | None
+            The project name to set
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.project = project
+        return self
+
+    def with_project(self, project: str | None):
+        """
+        Create a copy of the current configuration with an updated project name.
+
+        Parameters
+        ----------
+        project : str | None
+            The project name to set
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the updated project name
+        """
+        return copy.deepcopy(self).project_(project)
+
+    def tags_(self, *tags: str):
+        """
+        Set the tags for the trainer configuration in-place.
+
+        Parameters
+        ----------
+        *tags : str
+            The tags to set
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.tags = list(tags)
+        return self
+
+    def with_tags(self, *tags: str):
+        """
+        Create a copy of the current configuration with updated tags.
+
+        Parameters
+        ----------
+        *tags : str
+            The tags to set
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the updated tags
+        """
+        return copy.deepcopy(self).tags_(*tags)
+
+    def add_tags_(self, *tags: str):
+        """
+        Add tags to the trainer configuration in-place.
+
+        Parameters
+        ----------
+        *tags : str
+            The tags to add
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.tags.extend(tags)
+        return self
+
+    def with_added_tags(self, *tags: str):
+        """
+        Create a copy of the current configuration with additional tags.
+
+        Parameters
+        ----------
+        *tags : str
+            The tags to add
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the additional tags
+        """
+        return copy.deepcopy(self).add_tags_(*tags)
+
+    def notes_(self, *notes: str):
+        """
+        Set the notes for the trainer configuration in-place.
+
+        Parameters
+        ----------
+        *notes : str
+            The notes to set
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.notes = list(notes)
+        return self
+
+    def with_notes(self, *notes: str):
+        """
+        Create a copy of the current configuration with updated notes.
+
+        Parameters
+        ----------
+        *notes : str
+            The notes to set
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the updated notes
+        """
+        return copy.deepcopy(self).notes_(*notes)
+
+    def add_notes_(self, *notes: str):
+        """
+        Add notes to the trainer configuration in-place.
+
+        Parameters
+        ----------
+        *notes : str
+            The notes to add
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.notes.extend(notes)
+        return self
+
+    def with_added_notes(self, *notes: str):
+        """
+        Create a copy of the current configuration with additional notes.
+
+        Parameters
+        ----------
+        *notes : str
+            The notes to add
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the additional notes
+        """
+        return copy.deepcopy(self).add_notes_(*notes)
+
+    def meta_(self, meta: dict[str, Any] | None = None, /, **kwargs: Any):
+        """
+        Update the `meta` dictionary in-place with the provided key-value pairs.
+
+        This method allows updating the meta information associated with the trainer
+        configuration by either passing a dictionary or keyword arguments.
+
+        Parameters
+        ----------
+        meta : dict[str, Any] | None, optional
+            A dictionary containing meta information to be added, by default None
+        **kwargs : Any
+            Additional key-value pairs to be added to the meta dictionary
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        if meta is not None:
+            self.meta.update(meta)
+        self.meta.update(kwargs)
+        return self
+
+    def with_meta(self, meta: dict[str, Any] | None = None, /, **kwargs: Any):
+        """
+        Create a copy of the current configuration with updated meta information.
+
+        This method is similar to `meta_`, but it returns a new instance of the configuration
+        with the updated meta information instead of modifying the current instance.
+
+        Parameters
+        ----------
+        meta : dict[str, Any] | None, optional
+            A dictionary containing meta information to be added, by default None
+        **kwargs : Any
+            Additional key-value pairs to be added to the meta dictionary
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with updated meta information
+        """
+
+        return self.model_copy(deep=True).meta_(meta, **kwargs)
+
+    def debug_(self, value: bool = True):
+        """
+        Set the debug flag for the trainer configuration in-place.
+
+        Parameters
+        ----------
+        value : bool, optional
+            The debug flag value to set, by default True
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.debug = value
+        return self
+
+    def with_debug(self, value: bool = True):
+        """
+        Create a copy of the current configuration with an updated debug flag.
+
+        Parameters
+        ----------
+        value : bool, optional
+            The debug flag value to set, by default True
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the updated debug flag
+        """
+        return copy.deepcopy(self).debug_(value)
+
+    def ckpt_path_(self, path: Literal["none"] | str | Path | None):
+        """
+        Set the checkpoint path for the trainer configuration in-place.
+
+        Parameters
+        ----------
+        path : Literal["none"] | str | Path | None
+            The checkpoint path to set
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.ckpt_path = path
+        return self
+
+    def with_ckpt_path(self, path: Literal["none"] | str | Path | None):
+        """
+        Create a copy of the current configuration with an updated checkpoint path.
+
+        Parameters
+        ----------
+        path : Literal["none"] | str | Path | None
+            The checkpoint path to set
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the updated checkpoint path
+        """
+        return copy.deepcopy(self).ckpt_path_(path)
+
+    def barebones_(self, value: bool = True):
+        """
+        Set the barebones flag for the trainer configuration in-place.
+
+        Parameters
+        ----------
+        value : bool, optional
+            The barebones flag value to set, by default True
+
+        Returns
+        -------
+        self
+            Returns self for method chaining
+        """
+        self.barebones = value
+        return self
+
+    def with_barebones(self, value: bool = True):
+        """
+        Create a copy of the current configuration with an updated barebones flag.
+
+        Parameters
+        ----------
+        value : bool, optional
+            The barebones flag value to set, by default True
+
+        Returns
+        -------
+        TrainerConfig
+            A new instance of the configuration with the updated barebones flag
+        """
+        return copy.deepcopy(self).barebones_(value)
+
     def reset_run(
         self,
         *,
@@ -873,3 +1238,84 @@ class TrainerConfig(C.Config):
         return config
 
     # endregion
+
+    # region Random ID Generation
+    _rng: ClassVar[np.random.Generator | None] = None
+
+    @classmethod
+    def generate_id(cls, *, length: int = 8) -> str:
+        """
+        Generate a random ID of specified length.
+
+        """
+        if (rng := cls._rng) is None:
+            rng = np.random.default_rng()
+
+        alphabet = list(string.ascii_lowercase + string.digits)
+
+        id = "".join(rng.choice(alphabet) for _ in range(length))
+        return id
+
+    @classmethod
+    def set_seed(cls, seed: int | None = None) -> None:
+        """
+        Set the seed for the random number generator.
+
+        Args:
+            seed (int | None, optional): The seed value to set. If None, a seed based on the current time will be used. Defaults to None.
+
+        Returns:
+            None
+        """
+        if seed is None:
+            seed = int(time.time() * 1000)
+        log.critical(f"Seeding {cls.__name__} with seed {seed}")
+        cls._rng = np.random.default_rng(seed)
+
+    # endregion
+
+    # region Internal Methods
+    def _nshtrainer_all_callback_configs(self) -> Iterable[CallbackConfigBase | None]:
+        yield self.directory.setup_callback
+        yield self.early_stopping
+        yield self.checkpoint_saving
+        yield self.lr_monitor
+        yield from (
+            logger_config
+            for logger_config in self.enabled_loggers()
+            if logger_config is not None
+            and isinstance(logger_config, CallbackConfigBase)
+        )
+        yield self.log_epoch
+        yield self.log_norms
+        yield self.hf_hub
+        yield self.shared_parameters
+        yield self.reduce_lr_on_plateau_sanity_checking
+        yield self.auto_set_debug_flag
+        yield self.auto_validate_metrics
+        yield from self.callbacks
+
+    def _nshtrainer_all_logger_configs(self) -> Iterable[LoggerConfigBase | None]:
+        # Disable all loggers if barebones mode is enabled
+        if self.barebones:
+            return
+
+        yield from self.enabled_loggers()
+        yield self.actsave_logger
+
+    def _nshtrainer_validate_before_run(self):
+        # shared_parameters is not supported under barebones mode
+        if self.barebones and self.shared_parameters:
+            raise ValueError("shared_parameters is not supported under barebones mode")
+
+    def _nshtrainer_set_id_if_missing(self):
+        """
+        Set the ID for the configuration object if it is missing.
+        """
+        if self.id is C.MISSING:
+            self.id = self.generate_id()
+            log.info(f"TrainerConfig's run ID is missing, setting to {self.id}.")
+        else:
+            log.debug(f"TrainerConfig's run ID is already set to {self.id}.")
+
+    # endregion

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/trainer/trainer.py

@@ -316,6 +316,7 @@ class Trainer(LightningTrainer):
                 f"Trainer hparams must either be an instance of {hparams_cls} or a mapping. "
                 f"Got {type(hparams)=} instead."
             )
+        hparams._nshtrainer_set_id_if_missing()
         hparams = hparams.model_deep_validate()
         hparams._nshtrainer_validate_before_run()
 

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.4.dist-info → nshtrainer-1.3.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.3.4
+Version: 1.3.6
 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.4.dist-info → nshtrainer-1.3.6.dist-info}/RECORD

@@ -1,11 +1,10 @@
 nshtrainer/.nshconfig.generated.json,sha256=yZd6cn1RhvNNJUgiUTRYut8ofZYvbulnpPG-rZIRhi4,106
-nshtrainer/__init__.py,sha256=VcqBfL8RgCcZDaY645nxeDmOspqerx4x46wggCMnS0E,692
+nshtrainer/__init__.py,sha256=RI_2B_IUWa10B6H5TAuWtE5FWX1X4ue-J4dTDaF2-lQ,1035
 nshtrainer/_callback.py,sha256=ZDppiJ4d65tRXTEWYPZLH_F1xFizdz1pkWJe_sQ5uII,12564
 nshtrainer/_checkpoint/metadata.py,sha256=Hh5a7OkdknUEbkEwX6vS88-XLEeuVDoR6a3en2uLzQE,5597
 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
@@ -30,13 +29,12 @@ 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
+nshtrainer/configs/__init__.py,sha256=-yJ5Uk9VkANqfk-QnX2aynL0jSf7cJQuQNzT1GAE1x8,15684
 nshtrainer/configs/_checkpoint/__init__.py,sha256=6s7Y68StboqscY2G4P_QG443jz5aiym5SjOogIljWLg,342
 nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=oOPfYkXTjKgm6pluGsG6V1TPyCEGjsQpHVL-LffSUFQ,290
-nshtrainer/configs/_directory/__init__.py,sha256=_oO7vM9DhzHSxtZcv86sTi7hZIptnK1gr-AP9mqQ370,386
 nshtrainer/configs/_hf_hub/__init__.py,sha256=ciFLbV-JV8SVzqo2SyythEuDMnk7gGfdIacB18QYnkY,511
 nshtrainer/configs/callbacks/__init__.py,sha256=tP9urR73NIanyxpbi4EERsxOnGNiptbQpmsj-v53a38,4774
 nshtrainer/configs/callbacks/actsave/__init__.py,sha256=JvjSZtEoA28FC4u-QT3skQzBDVbN9eq07rn4u2ydW-E,377
@@ -85,8 +83,8 @@ nshtrainer/configs/profiler/_base/__init__.py,sha256=ekYfPg-VDhCAFM5nJka2TxUYdRD
 nshtrainer/configs/profiler/advanced/__init__.py,sha256=-ThpUat16Ij_0avkMUVVA8wCWDG_q_tM7KQofnWQCtg,308
 nshtrainer/configs/profiler/pytorch/__init__.py,sha256=soAU1s2_Pa1na4gW8CK-iysJBO5M_7YeZC2_x40iEdg,294
 nshtrainer/configs/profiler/simple/__init__.py,sha256=3Wb11lPuFuyasq8xS1CZ4WLuBCLS_nVSQGVllvOOi0Y,289
-nshtrainer/configs/trainer/__init__.py,sha256=YLlDOUYDp_qURHhcmhCxTcY6K5AbmoTxdzBPB9SEZII,8040
-nshtrainer/configs/trainer/_config/__init__.py,sha256=6DXdtP-uH11TopQ7kzId9fco-wVkD7ZfevbBqDpN6TE,3817
+nshtrainer/configs/trainer/__init__.py,sha256=DM2PlB4WRDZ_dqEeW91LbKRFa4sIF_pETU0T9GYJ5-g,8073
+nshtrainer/configs/trainer/_config/__init__.py,sha256=z5UpuXktBanLOYNkkbgbbHE06iQtcSuAKTpnx2TLmCo,3850
 nshtrainer/configs/trainer/accelerator/__init__.py,sha256=3H6R3wlwbKL1TzDqGCChZk78-BcE2czLouo7Djiq3nA,898
 nshtrainer/configs/trainer/plugin/__init__.py,sha256=NkHQxMPkrtTtdIAO4dQUE9SWEcHRDB0yUXLkTjnl4dA,3332
 nshtrainer/configs/trainer/plugin/base/__init__.py,sha256=slW5z1FZw2qICXO9l9DnLIDB1Yl7KOcxPEZkyYIHrp4,276
@@ -135,7 +133,7 @@ nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5N
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
 nshtrainer/trainer/__init__.py,sha256=jRaHdaFK8wxNrN1bleT9cf29iZahL_-XkWo5TWz2CmA,550
-nshtrainer/trainer/_config.py,sha256=SohR7uxANnP3xrrcW_mAjk6TuDamsW5Qdk3dlnPinDw,33457
+nshtrainer/trainer/_config.py,sha256=x3YjP_0IykqSRh8YIilCxq3nPt_fZXoVcxfR13ulmV0,45578
 nshtrainer/trainer/_distributed_prediction_result.py,sha256=bQw8Z6PT694UUf-zQPkech6CxyUSy8bAIexfSfPej0U,2507
 nshtrainer/trainer/_log_hparams.py,sha256=XH2lZ4U_3AZBhOt91ocsEhdL_NRz35oWvqLCUFDohUs,2389
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
@@ -146,11 +144,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/trainer/trainer.py,sha256=iQWu0KfwLY-1q9EEsg0xPlyUN1fsJ9iXfSQbPmiMlac,24177
+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 +159,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.4.dist-info/METADATA,sha256=Dm6wgfQh8ZC42IeftejuUZ-KZ2YBWBjnpHa_pYNi7Kc,960
-nshtrainer-1.3.4.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.3.4.dist-info/RECORD,,
+nshtrainer-1.3.6.dist-info/METADATA,sha256=KidLM7J5P7mALTfYQsveSDjAOJZ-Gcq4_e1-Xgrms68,979
+nshtrainer-1.3.6.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.3.6.dist-info/RECORD,,

nshtrainer/_directory.py

@@ -1,72 +0,0 @@
-from __future__ import annotations
-
-import logging
-from pathlib import Path
-
-import nshconfig as C
-
-from .callbacks.directory_setup import DirectorySetupCallbackConfig
-from .loggers import LoggerConfig
-
-log = logging.getLogger(__name__)
-
-
-class DirectoryConfig(C.Config):
-    project_root: Path | None = None
-    """
-    Root directory for this project.
-
-    This isn't specific to the run; it is the parent directory of all runs.
-    """
-
-    logdir_basename: str = "nshtrainer"
-    """Base name for the log directory."""
-
-    setup_callback: DirectorySetupCallbackConfig = DirectorySetupCallbackConfig()
-    """Configuration for the directory setup PyTorch Lightning callback."""
-
-    def resolve_run_root_directory(self, run_id: str) -> Path:
-        if (project_root_dir := self.project_root) is None:
-            project_root_dir = Path.cwd()
-
-        # The default base dir is $CWD/{logdir_basename}/{id}/
-        base_dir = project_root_dir / self.logdir_basename
-        base_dir.mkdir(exist_ok=True)
-
-        # Add a .gitignore file to the {logdir_basename} directory
-        #   which will ignore all files except for the .gitignore file itself
-        gitignore_path = base_dir / ".gitignore"
-        if not gitignore_path.exists():
-            gitignore_path.touch()
-            gitignore_path.write_text("*\n")
-
-        base_dir = base_dir / run_id
-        base_dir.mkdir(exist_ok=True)
-
-        return base_dir
-
-    def resolve_subdirectory(self, run_id: str, subdirectory: str) -> Path:
-        # The subdir will be $CWD/{logdir_basename}/{id}/{log, stdio, checkpoint, activation}/
-        if (subdir := getattr(self, subdirectory, None)) is not None:
-            assert isinstance(subdir, Path), (
-                f"Expected a Path for {subdirectory}, got {type(subdir)}"
-            )
-            return subdir
-
-        dir = self.resolve_run_root_directory(run_id)
-        dir = dir / subdirectory
-        dir.mkdir(exist_ok=True)
-        return dir
-
-    def _resolve_log_directory_for_logger(self, run_id: str, logger: LoggerConfig):
-        if (log_dir := logger.log_dir) is not None:
-            return log_dir
-
-        # Save to {logdir_basename}/{id}/log/{logger name}
-        log_dir = self.resolve_subdirectory(run_id, "log")
-        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)
-
-        return log_dir

nshtrainer/configs/_directory/__init__.py

@@ -1,15 +0,0 @@
-from __future__ import annotations
-
-__codegen__ = True
-
-from nshtrainer._directory import DirectoryConfig as DirectoryConfig
-from nshtrainer._directory import (
-    DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
-)
-from nshtrainer._directory import LoggerConfig as LoggerConfig
-
-__all__ = [
-    "DirectoryConfig",
-    "DirectorySetupCallbackConfig",
-    "LoggerConfig",
-]