nshtrainer 0.44.1__py3-none-any.whl → 1.0.0b10__py3-none-any.whl

nshtrainer/trainer/_config.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
+import copy
 import logging
+import os
+import string
+import time
 from collections.abc import Iterable, Sequence
 from datetime import timedelta
 from pathlib import Path
@@ -8,6 +12,7 @@ from typing import (
     TYPE_CHECKING,
     Annotated,
     Any,
+    ClassVar,
     Literal,
     Protocol,
     TypeAlias,
@@ -15,6 +20,7 @@ from typing import (
 )
 
 import nshconfig as C
+import numpy as np
 from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
 from lightning.pytorch.accelerators import Accelerator
@@ -28,6 +34,7 @@ from lightning.pytorch.strategies.strategy import Strategy
 from typing_extensions import TypedDict, TypeVar, override
 
 from .._checkpoint.loader import CheckpointLoadingConfig
+from .._directory import DirectoryConfig
 from .._hf_hub import HuggingFaceHubConfig
 from ..callbacks import (
     BestCheckpointCallbackConfig,
@@ -47,10 +54,10 @@ from ..loggers import (
     TensorboardLoggerConfig,
     WandbLoggerConfig,
 )
+from ..loggers.actsave import ActSaveLoggerConfig
+from ..metrics._config import MetricConfig
 from ..profiler import ProfilerConfig
-
-if TYPE_CHECKING:
-    from ..model.config import BaseConfig
+from ..util._environment_info import EnvironmentConfig
 
 log = logging.getLogger(__name__)
 
@@ -71,7 +78,7 @@ class LoggingConfig(CallbackConfigBase):
     log_epoch: LogEpochCallbackConfig | None = LogEpochCallbackConfig()
     """If enabled, will log the fractional epoch number to the logger."""
 
-    actsave_logged_metrics: bool = False
+    actsave_logger: ActSaveLoggerConfig | None = None
     """If enabled, will automatically save logged metrics using ActSave (if nshutils is installed)."""
 
     @property
@@ -103,12 +110,12 @@ class LoggingConfig(CallbackConfigBase):
             None,
         )
 
-    def create_loggers(self, root_config: "BaseConfig"):
+    def create_loggers(self, trainer_config: TrainerConfig):
         """
         Constructs and returns a list of loggers based on the provided root configuration.
 
         Args:
-            root_config (BaseConfig): The root configuration object.
+            trainer_config (TrainerConfig): The root configuration object.
 
         Returns:
             list[Logger]: A list of constructed loggers.
@@ -123,12 +130,16 @@ class LoggingConfig(CallbackConfigBase):
         ):
             if not logger_config.enabled:
                 continue
-            if (logger := logger_config.create_logger(root_config)) is None:
+            if (logger := logger_config.create_logger(trainer_config)) is None:
                 continue
             yield logger
 
+        # If the actsave_metrics is enabled, add the ActSave logger
+        if self.actsave_logger:
+            yield self.actsave_logger.create_logger(trainer_config)
+
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         if self.log_lr:
             from lightning.pytorch.callbacks import LearningRateMonitor
 
@@ -139,13 +150,13 @@ class LoggingConfig(CallbackConfigBase):
             yield LearningRateMonitor(logging_interval=logging_interval)
 
         if self.log_epoch:
-            yield from self.log_epoch.create_callbacks(root_config)
+            yield from self.log_epoch.create_callbacks(trainer_config)
 
         for logger in self.loggers:
             if not logger or not isinstance(logger, CallbackConfigBase):
                 continue
 
-            yield from logger.create_callbacks(root_config)
+            yield from logger.create_callbacks(trainer_config)
 
 
 class GradientClippingConfig(C.Config):
@@ -172,7 +183,7 @@ class OptimizationConfig(CallbackConfigBase):
     """Gradient clipping configuration, or None to disable gradient clipping."""
 
     @override
-    def create_callbacks(self, root_config):
+    def create_callbacks(self, trainer_config):
         from ..callbacks.norm_logging import NormLoggingCallbackConfig
 
         yield from NormLoggingCallbackConfig(
@@ -180,7 +191,7 @@ class OptimizationConfig(CallbackConfigBase):
             log_grad_norm_per_param=self.log_grad_norm_per_param,
             log_param_norm=self.log_param_norm,
             log_param_norm_per_param=self.log_param_norm_per_param,
-        ).create_callbacks(root_config)
+        ).create_callbacks(trainer_config)
 
 
 TPlugin = TypeVar(
@@ -274,22 +285,22 @@ class CheckpointSavingConfig(CallbackConfigBase):
         self.enabled = False
         return self
 
-    def should_save_checkpoints(self, root_config: "BaseConfig"):
+    def should_save_checkpoints(self, trainer_config: TrainerConfig):
         if not self.enabled:
             return False
 
-        if root_config.trainer.fast_dev_run:
+        if trainer_config.fast_dev_run:
             return False
 
         return True
 
     @override
-    def create_callbacks(self, root_config: "BaseConfig"):
-        if not self.should_save_checkpoints(root_config):
+    def create_callbacks(self, trainer_config: TrainerConfig):
+        if not self.should_save_checkpoints(trainer_config):
             return
 
         for callback_config in self.checkpoint_callbacks:
-            yield from callback_config.create_callbacks(root_config)
+            yield from callback_config.create_callbacks(trainer_config)
 
 
 class LightningTrainerKwargs(TypedDict, total=False):
@@ -541,6 +552,74 @@ class SanityCheckingConfig(C.Config):
 
 
 class TrainerConfig(C.Config):
+    # region Active Run Configuration
+    id: str = C.Field(default_factory=lambda: TrainerConfig.generate_id())
+    """ID of the run."""
+    name: list[str] = []
+    """Run name in parts. Full name is constructed by joining the parts with spaces."""
+    project: str | None = None
+    """Project name."""
+    tags: list[str] = []
+    """Tags for the run."""
+    notes: list[str] = []
+    """Human readable notes for the run."""
+
+    @property
+    def full_name(self):
+        return " ".join(self.name)
+
+    debug: bool = False
+    """Whether to run in debug mode. This will enable debug logging and enable debug code paths."""
+
+    environment: Annotated[EnvironmentConfig, C.Field(repr=False)] = (
+        EnvironmentConfig.empty()
+    )
+    """A snapshot of the current environment information (e.g. python version, slurm info, etc.). This is automatically populated by the run script."""
+
+    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
+    """Primary metric configuration options. This is used in the following ways:
+    - To determine the best model checkpoint to save with the ModelCheckpoint callback.
+    - To monitor the primary metric during training and stop training based on the `early_stopping` configuration.
+    - For the ReduceLROnPlateau scheduler.
+    """
+
     ckpt_path: Literal["none"] | str | Path | None = None
     """Path to a checkpoint to load and resume training from. If ``"none"``, will not load a checkpoint."""
 
@@ -788,3 +867,71 @@ class TrainerConfig(C.Config):
         yield self.reduce_lr_on_plateau_sanity_checking
         yield self.auto_set_debug_flag
         yield from self.callbacks
+
+    # region Helper Methods
+    def with_fast_dev_run(self, value: int | bool = True, /):
+        """
+        Enables fast_dev_run mode for the trainer.
+        This will run the training loop for a specified number of batches,
+        if an integer is provided, or for a single batch if True is provided.
+        """
+        config = copy.deepcopy(self)
+        config.fast_dev_run = value
+        return config
+
+    def with_project_root(self, project_root: str | Path | os.PathLike):
+        """
+        Set the project root directory for the trainer.
+
+        Args:
+            project_root (Path): The base directory to use.
+
+        Returns:
+            self: The current instance of the class.
+        """
+        config = copy.deepcopy(self)
+        config.directory.project_root = Path(project_root)
+        return config
+
+    def reset_run(
+        self,
+        *,
+        id: bool = True,
+        basic: bool = True,
+        project_root: bool = True,
+        environment: bool = True,
+    ):
+        """
+        Reset the configuration object to its initial state.
+
+        Parameters:
+        - id (bool): If True, generate a new ID for the configuration object.
+        - basic (bool): If True, reset basic attributes like name, project, tags, and notes.
+        - project_root (bool): If True, reset the directory configuration to its initial state.
+        - environment (bool): If True, reset the environment configuration to its initial state.
+        - meta (bool): If True, reset the meta dictionary to an empty dictionary.
+
+        Returns:
+        - self: The updated configuration object.
+
+        """
+        config = copy.deepcopy(self)
+
+        if id:
+            config.id = config.generate_id()
+
+        if basic:
+            config.name = []
+            config.project = None
+            config.tags = []
+            config.notes = []
+
+        if project_root:
+            config.directory = DirectoryConfig()
+
+        if environment:
+            config.environment = EnvironmentConfig.empty()
+
+        return config
+
+    # endregion

nshtrainer/trainer/checkpoint_connector.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, cast
 
 from lightning.pytorch.trainer.connectors.checkpoint_connector import (
     _CheckpointConnector as _LightningCheckpointConnector,
@@ -12,8 +11,6 @@ from typing_extensions import override
 
 from .._checkpoint.loader import CheckpointLoadingConfig, _resolve_checkpoint
 
-if TYPE_CHECKING:
-    from ..model.config import BaseConfig
 log = logging.getLogger(__name__)
 
 
@@ -32,8 +29,7 @@ class _CheckpointConnector(_LightningCheckpointConnector):
             return None
 
         # Now, resolve the checkpoint loader config.
-        root_config = cast("BaseConfig", trainer._base_module.config)
-        ckpt_loader_config = root_config.trainer.checkpoint_loading
+        ckpt_loader_config = trainer.hparams.checkpoint_loading
         match ckpt_loader_config:
             case "auto":
                 ckpt_loader_config = CheckpointLoadingConfig.auto(ckpt_path, state_fn)
@@ -44,9 +40,7 @@ class _CheckpointConnector(_LightningCheckpointConnector):
         log.debug(f"Checkpoint loader config: {ckpt_loader_config}")
 
         # Use the config to resolve the checkpoint.
-        if (
-            ckpt_path := _resolve_checkpoint(ckpt_loader_config, root_config, trainer)
-        ) is None:
+        if (ckpt_path := _resolve_checkpoint(ckpt_loader_config, trainer)) is None:
             log.info(
                 "No checkpoint found for the current trainer state. "
                 "Training will start from scratch."
@@ -69,3 +63,24 @@ class _CheckpointConnector(_LightningCheckpointConnector):
         return super()._parse_ckpt_path(
             state_fn, ckpt_path, model_provided, model_connected
         )
+
+    @override
+    def dump_checkpoint(self, weights_only: bool = False):
+        checkpoint = super().dump_checkpoint(weights_only)
+
+        # Save the trainer's config.
+        _add_trainer_config_to_checkpoint_(checkpoint, self.trainer)
+
+        return checkpoint
+
+
+def _add_trainer_config_to_checkpoint_(checkpoint: dict, trainer):
+    from .trainer import Trainer
+
+    # If this isn't an `nshtrainer` trainer (which I don't know why it wouldn't be),
+    # then we just return.
+    if isinstance(trainer, Trainer):
+        return None
+
+    # Save the trainer's config.
+    checkpoint[trainer.CHECKPOINT_HYPER_PARAMS_KEY] = dict(trainer.hparams)