nshtrainer 0.44.0__py3-none-any.whl → 1.0.0b9__py3-none-any.whl

nshtrainer/config/model/config/__init__.py

@@ -1,37 +0,0 @@
-from __future__ import annotations
-
-__codegen__ = True
-
-from typing import TYPE_CHECKING
-
-# Config/alias imports
-
-if TYPE_CHECKING:
-    from nshtrainer.model.config import BaseConfig as BaseConfig
-    from nshtrainer.model.config import CallbackConfigBase as CallbackConfigBase
-    from nshtrainer.model.config import DirectoryConfig as DirectoryConfig
-    from nshtrainer.model.config import EnvironmentConfig as EnvironmentConfig
-    from nshtrainer.model.config import MetricConfig as MetricConfig
-    from nshtrainer.model.config import TrainerConfig as TrainerConfig
-else:
-
-    def __getattr__(name):
-        import importlib
-
-        if name in globals():
-            return globals()[name]
-        if name == "MetricConfig":
-            return importlib.import_module("nshtrainer.model.config").MetricConfig
-        if name == "TrainerConfig":
-            return importlib.import_module("nshtrainer.model.config").TrainerConfig
-        if name == "BaseConfig":
-            return importlib.import_module("nshtrainer.model.config").BaseConfig
-        if name == "EnvironmentConfig":
-            return importlib.import_module("nshtrainer.model.config").EnvironmentConfig
-        if name == "DirectoryConfig":
-            return importlib.import_module("nshtrainer.model.config").DirectoryConfig
-        if name == "CallbackConfigBase":
-            return importlib.import_module("nshtrainer.model.config").CallbackConfigBase
-        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
-
-# Submodule exports

nshtrainer/config/model/mixins/logger/__init__.py

@@ -1,22 +0,0 @@
-from __future__ import annotations
-
-__codegen__ = True
-
-from typing import TYPE_CHECKING
-
-# Config/alias imports
-
-if TYPE_CHECKING:
-    from nshtrainer.model.mixins.logger import BaseConfig as BaseConfig
-else:
-
-    def __getattr__(name):
-        import importlib
-
-        if name in globals():
-            return globals()[name]
-        if name == "BaseConfig":
-            return importlib.import_module("nshtrainer.model.mixins.logger").BaseConfig
-        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
-
-# Submodule exports

nshtrainer/config/runner/__init__.py

@@ -1,22 +0,0 @@
-from __future__ import annotations
-
-__codegen__ = True
-
-from typing import TYPE_CHECKING
-
-# Config/alias imports
-
-if TYPE_CHECKING:
-    from nshtrainer.runner import BaseConfig as BaseConfig
-else:
-
-    def __getattr__(name):
-        import importlib
-
-        if name in globals():
-            return globals()[name]
-        if name == "BaseConfig":
-            return importlib.import_module("nshtrainer.runner").BaseConfig
-        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
-
-# Submodule exports

nshtrainer/ll/__init__.py

@@ -1,59 +0,0 @@
-from __future__ import annotations
-
-from typing import TypeAlias
-
-from . import _experimental as _experimental
-from . import actsave as actsave
-from . import callbacks as callbacks
-from . import data as data
-from . import lr_scheduler as lr_scheduler
-from . import model as model
-from . import nn as nn
-from . import optimizer as optimizer
-from . import snapshot as snapshot
-from . import typecheck as typecheck
-from .actsave import ActLoad as ActLoad
-from .actsave import ActSave as ActSave
-from .config import MISSING as MISSING
-from .config import AllowMissing as AllowMissing
-from .config import Field as Field
-from .config import MissingField as MissingField
-from .config import PrivateAttr as PrivateAttr
-from .config import TypedConfig as TypedConfig
-from .data import dataset_transform as dataset_transform
-from .log import init_python_logging as init_python_logging
-from .log import lovely as lovely
-from .log import pretty as pretty
-from .lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
-from .model import BaseConfig as BaseConfig
-from .model import CheckpointLoadingConfig as CheckpointLoadingConfig
-from .model import CheckpointSavingConfig as CheckpointSavingConfig
-from .model import DirectoryConfig as DirectoryConfig
-from .model import (
-    EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
-)
-from .model import EnvironmentConfig as EnvironmentConfig
-from .model import (
-    EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
-)
-from .model import (
-    EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
-)
-from .model import GradientClippingConfig as GradientClippingConfig
-from .model import LightningModuleBase as LightningModuleBase
-from .model import LoggingConfig as LoggingConfig
-from .model import MetricConfig as MetricConfig
-from .model import OptimizationConfig as OptimizationConfig
-from .model import ReproducibilityConfig as ReproducibilityConfig
-from .model import SanityCheckingConfig as SanityCheckingConfig
-from .model import TrainerConfig as TrainerConfig
-from .nn import TypedModuleDict as TypedModuleDict
-from .nn import TypedModuleList as TypedModuleList
-from .optimizer import OptimizerConfig as OptimizerConfig
-from .runner import Runner as Runner
-from .runner import SnapshotConfig as SnapshotConfig
-from .snoop import snoop as snoop
-from .trainer import Trainer as Trainer
-
-PrimaryMetricConfig: TypeAlias = MetricConfig
-ConfigList: TypeAlias = list[tuple[BaseConfig, type[LightningModuleBase]]]

nshtrainer/ll/_experimental.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshtrainer._experimental import *  # noqa: F403

nshtrainer/ll/actsave.py

@@ -1,6 +0,0 @@
-from __future__ import annotations
-
-from nshutils.actsave import *  # type: ignore # noqa: F403
-
-from nshtrainer.callbacks.actsave import ActSaveCallback as ActSaveCallback
-from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig

nshtrainer/ll/callbacks.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshtrainer.callbacks import *  # noqa: F403

nshtrainer/ll/config.py

@@ -1,6 +0,0 @@
-from __future__ import annotations
-
-from nshconfig import *  # type: ignore  # noqa: F403
-from nshconfig import Config as TypedConfig  # type: ignore  # noqa: F401
-
-_ = TypedConfig

nshtrainer/ll/data.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshtrainer.data import *  # noqa: F403

nshtrainer/ll/log.py

@@ -1,5 +0,0 @@
-from __future__ import annotations
-
-from nshutils import init_python_logging as init_python_logging
-from nshutils import lovely as lovely
-from nshutils import pretty as pretty

nshtrainer/ll/lr_scheduler.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshtrainer.lr_scheduler import *  # noqa: F403

nshtrainer/ll/model.py

@@ -1,21 +0,0 @@
-from __future__ import annotations
-
-from nshtrainer.model import *  # noqa: F403
-
-from ..trainer._config import CheckpointLoadingConfig as CheckpointLoadingConfig
-from ..trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
-from ..trainer._config import GradientClippingConfig as GradientClippingConfig
-from ..trainer._config import LoggingConfig as LoggingConfig
-from ..trainer._config import OptimizationConfig as OptimizationConfig
-from ..trainer._config import ReproducibilityConfig as ReproducibilityConfig
-from ..trainer._config import SanityCheckingConfig as SanityCheckingConfig
-from ..util._environment_info import (
-    EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
-)
-from ..util._environment_info import EnvironmentConfig as EnvironmentConfig
-from ..util._environment_info import (
-    EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
-)
-from ..util._environment_info import (
-    EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
-)

nshtrainer/ll/nn.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshtrainer.nn import *  # noqa: F403

nshtrainer/ll/optimizer.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshtrainer.optimizer import *  # noqa: F403

nshtrainer/ll/runner.py

@@ -1,5 +0,0 @@
-from __future__ import annotations
-
-from nshrunner import SnapshotConfig as SnapshotConfig
-
-from nshtrainer.runner import *  # type: ignore # noqa: F403

nshtrainer/ll/snapshot.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshsnap import *  # pyright: ignore[reportWildcardImportFromLibrary] # noqa: F403

nshtrainer/ll/snoop.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshutils import snoop as snoop

nshtrainer/ll/trainer.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshtrainer.trainer import *  # noqa: F403

nshtrainer/ll/typecheck.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshutils.typecheck import *  # type: ignore # noqa: F403

nshtrainer/ll/util.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from nshtrainer.util import *  # noqa: F403

nshtrainer/model/config.py

@@ -1,218 +0,0 @@
-from __future__ import annotations
-
-import copy
-import logging
-import os
-import string
-import time
-from collections.abc import Iterable
-from pathlib import Path
-from typing import Annotated, Any, ClassVar
-
-import nshconfig as C
-import numpy as np
-import torch
-from typing_extensions import Self
-
-from .._directory import DirectoryConfig
-from ..callbacks.base import CallbackConfigBase
-from ..metrics import MetricConfig
-from ..trainer._config import TrainerConfig
-from ..util._environment_info import EnvironmentConfig
-
-log = logging.getLogger(__name__)
-
-
-class BaseConfig(C.Config):
-    id: str = C.Field(default_factory=lambda: BaseConfig.generate_id())
-    """ID of the run."""
-    name: str | None = None
-    """Run name."""
-    name_parts: list[str] = []
-    """A list of parts used to construct the run name. This is useful for constructing the run name dynamically."""
-    project: str | None = None
-    """Project name."""
-    tags: list[str] = []
-    """Tags for the run."""
-    notes: list[str] = []
-    """Human readable notes for the run."""
-
-    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."""
-    trainer: TrainerConfig = TrainerConfig()
-    """PyTorch Lightning trainer configuration options. Check Lightning's `Trainer` documentation for more information."""
-
-    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.
-    """
-
-    meta: dict[str, Any] = {}
-    """Additional metadata for this run. This can be used to store arbitrary data that is not part of the config schema."""
-
-    @property
-    def run_name(self) -> str:
-        parts = self.name_parts.copy()
-        if self.name is not None:
-            parts = [self.name] + parts
-        name = "-".join(parts)
-        if not name:
-            name = self.id
-        return name
-
-    def clone(self, with_new_id: bool = True) -> Self:
-        c = copy.deepcopy(self)
-        if with_new_id:
-            c.id = BaseConfig.generate_id()
-        return c
-
-    def subdirectory(self, subdirectory: str) -> Path:
-        return self.directory.resolve_subdirectory(self.id, subdirectory)
-
-    # region Helper methods
-    def 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.trainer.fast_dev_run = value
-        return config
-
-    def with_project_root_(self, project_root: str | Path | os.PathLike) -> Self:
-        """
-        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.
-        """
-        self.directory.project_root = Path(project_root)
-        return self
-
-    def reset_(
-        self,
-        *,
-        id: bool = True,
-        basic: bool = True,
-        project_root: bool = True,
-        environment: bool = True,
-        meta: 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.
-
-        """
-        if id:
-            self.id = self.generate_id()
-
-        if basic:
-            self.name = None
-            self.name_parts = []
-            self.project = None
-            self.tags = []
-            self.notes = []
-
-        if project_root:
-            self.directory = DirectoryConfig()
-
-        if environment:
-            self.environment = EnvironmentConfig.empty()
-
-        if meta:
-            self.meta = {}
-
-        return self
-
-    def concise_repr(self) -> str:
-        """Get a concise representation of the configuration object."""
-
-        def _truncate(s: str, max_len: int = 50):
-            return s if len(s) <= max_len else f"{s[:max_len - 3]}..."
-
-        cls_name = self.__class__.__name__
-
-        parts: list[str] = []
-        parts.append(f"name={self.run_name}")
-        if self.project:
-            parts.append(f"project={_truncate(self.project)}")
-
-        return f"{cls_name}({', '.join(parts)})"
-
-    # endregion
-
-    # region Seeding
-
-    _rng: ClassVar[np.random.Generator | None] = None
-
-    @staticmethod
-    def generate_id(*, length: int = 8) -> str:
-        """
-        Generate a random ID of specified length.
-
-        """
-        if (rng := BaseConfig._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
-
-    @staticmethod
-    def set_seed(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 BaseConfig with seed {seed}")
-        BaseConfig._rng = np.random.default_rng(seed)
-
-    # endregion
-
-    @classmethod
-    def from_checkpoint(
-        cls,
-        path: str | Path,
-        *,
-        hparams_key: str = "hyper_parameters",
-    ):
-        ckpt = torch.load(path)
-        if (hparams := ckpt.get(hparams_key)) is None:
-            raise ValueError(
-                f"The checkpoint does not contain the `{hparams_key}` attribute. "
-                "Are you sure this is a valid Lightning checkpoint?"
-            )
-        return cls.model_validate(hparams)
-
-    def _nshtrainer_all_callback_configs(self) -> Iterable[CallbackConfigBase | None]:
-        yield from self.trainer._nshtrainer_all_callback_configs()

nshtrainer/runner.py

@@ -1,101 +0,0 @@
-from __future__ import annotations
-
-import copy
-import logging
-from collections.abc import Callable, Iterable, Mapping, Sequence
-from pathlib import Path
-from typing import Generic
-
-import nshrunner as nr
-from nshrunner._submit import screen
-from typing_extensions import TypeVar, TypeVarTuple, Unpack, deprecated, override
-
-from .model.config import BaseConfig
-
-TConfig = TypeVar("TConfig", bound=BaseConfig, infer_variance=True)
-TArguments = TypeVarTuple("TArguments", default=Unpack[tuple[()]])
-TReturn = TypeVar("TReturn", infer_variance=True)
-
-
-@deprecated("Use nshrunner.Runner instead.")
-class Runner(
-    nr.Runner[TReturn, TConfig, Unpack[TArguments]],
-    Generic[TReturn, TConfig, Unpack[TArguments]],
-):
-    @override
-    def __init__(
-        self,
-        run_fn: Callable[[TConfig, Unpack[TArguments]], TReturn],
-        config: nr.RunnerConfig | None = None,
-    ):
-        if config is None:
-            working_dir = Path.cwd() / "nshrunner"
-            working_dir.mkdir(exist_ok=True)
-
-            logging.warning(
-                f"`config` is not provided. Using default working directory of {working_dir}."
-            )
-            config = nr.RunnerConfig(working_dir=working_dir)
-
-        super().__init__(run_fn, config)
-
-    def fast_dev_run(
-        self,
-        runs: Iterable[tuple[TConfig, Unpack[TArguments]]],
-        n_batches: int = 1,
-        *,
-        env: Mapping[str, str] | None = None,
-    ):
-        runs_updated: list[tuple[TConfig, Unpack[TArguments]]] = []
-        for args in runs:
-            config = copy.deepcopy(args[0])
-            config.trainer.fast_dev_run = n_batches
-            runs_updated.append((config, *args[1:]))
-        del runs
-
-        return self.local(runs_updated, env=env)
-
-    def fast_dev_run_generator(
-        self,
-        runs: Iterable[tuple[TConfig, Unpack[TArguments]]],
-        n_batches: int = 1,
-        *,
-        env: Mapping[str, str] | None = None,
-    ):
-        runs_updated: list[tuple[TConfig, Unpack[TArguments]]] = []
-        for args in runs:
-            config = copy.deepcopy(args[0])
-            config.trainer.fast_dev_run = n_batches
-            runs_updated.append((config, *args[1:]))
-        del runs
-
-        return self.local_generator(runs_updated, env=env)
-
-    def fast_dev_run_session(
-        self,
-        runs: Iterable[tuple[TConfig, Unpack[TArguments]]],
-        options: screen.ScreenJobKwargs = {},
-        n_batches: int = 1,
-        *,
-        snapshot: nr.Snapshot,
-        setup_commands: Sequence[str] | None = None,
-        env: Mapping[str, str] | None = None,
-        activate_venv: bool = True,
-        print_command: bool = True,
-    ):
-        runs_updated: list[tuple[TConfig, Unpack[TArguments]]] = []
-        for args in runs:
-            config = copy.deepcopy(args[0])
-            config.trainer.fast_dev_run = n_batches
-            runs_updated.append((config, *args[1:]))
-        del runs
-
-        return self.session(
-            runs_updated,
-            options,
-            snapshot=snapshot,
-            setup_commands=setup_commands,
-            env=env,
-            activate_venv=activate_venv,
-            print_command=print_command,
-        )