nshtrainer 0.38.0__tar.gz → 0.40.0__tar.gz

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/PKG-INFO

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

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/pyproject.toml

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

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/src/nshtrainer/__init__.py

@@ -8,6 +8,7 @@ from . import model as model
 from . import nn as nn
 from . import optimizer as optimizer
 from . import profiler as profiler
+from .data import LightningDataModuleBase as LightningDataModuleBase
 from .metrics import MetricConfig as MetricConfig
 from .model import BaseConfig as BaseConfig
 from .model import LightningModuleBase as LightningModuleBase

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/src/nshtrainer/callbacks/checkpoint/_base.py

@@ -41,7 +41,7 @@ class BaseCheckpointCallbackConfig(CallbackConfigBase, ABC):
         self,
         root_config: "BaseConfig",
         dirpath: Path,
-    ) -> "CheckpointBase": ...
+    ) -> "CheckpointBase | None": ...
 
     @override
     def create_callbacks(self, root_config):
@@ -50,7 +50,8 @@ class BaseCheckpointCallbackConfig(CallbackConfigBase, ABC):
             or root_config.directory.resolve_subdirectory(root_config.id, "checkpoint")
         )
 
-        yield self.create_checkpoint(root_config, dirpath)
+        if (callback := self.create_checkpoint(root_config, dirpath)) is not None:
+            yield callback
 
 
 TConfig = TypeVar("TConfig", bound=BaseCheckpointCallbackConfig, infer_variance=True)

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -20,15 +20,26 @@ class BestCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
     metric: MetricConfig | None = None
     """Metric to monitor, or `None` to use the default metric."""
 
+    throw_on_no_metric: bool = True
+    """
+    Whether to throw an error if no metric is provided and no primary metric is found in the root config.
+    """
+
     @override
     def create_checkpoint(self, root_config, dirpath):
         # Resolve metric
         if (metric := self.metric) is None and (
             metric := root_config.primary_metric
         ) is None:
-            raise ValueError(
-                "No metric provided and no primary metric found in the root config"
+            error_msg = (
+                "No metric provided and no primary metric found in the root config. "
+                "Cannot create BestCheckpointCallback."
             )
+            if self.throw_on_no_metric:
+                raise ValueError(error_msg)
+            else:
+                log.warning(error_msg)
+                return None
 
         return BestCheckpoint(self, dirpath, metric)
 

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/src/nshtrainer/callbacks/early_stopping.py

@@ -51,7 +51,7 @@ class EarlyStoppingConfig(CallbackConfigBase):
             metric := root_config.primary_metric
         ) is None:
             raise ValueError(
-                "Either `metric` or `root_config.primary_metric` must be set."
+                "Either `metric` or `root_config.primary_metric` must be set to use EarlyStopping."
             )
 
         yield EarlyStopping(self, metric)

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/src/nshtrainer/data/__init__.py

@@ -1,4 +1,5 @@
 from . import transform as dataset_transform
 from .balanced_batch_sampler import BalancedBatchSampler as BalancedBatchSampler
+from .datamodule import LightningDataModuleBase as LightningDataModuleBase
 
 _ = dataset_transform

nshtrainer-0.40.0/src/nshtrainer/data/datamodule.py

@@ -0,0 +1,5 @@
+from lightning.pytorch import LightningDataModule
+
+
+class LightningDataModuleBase(LightningDataModule):
+    pass

nshtrainer-0.40.0/src/nshtrainer/runner.py

@@ -0,0 +1,99 @@
+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,
+        )

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/src/nshtrainer/trainer/_config.py

@@ -263,7 +263,7 @@ class CheckpointSavingConfig(CallbackConfigBase):
     """Enable checkpoint saving."""
 
     checkpoint_callbacks: Sequence[CheckpointCallbackConfig] = [
-        BestCheckpointCallbackConfig(),
+        BestCheckpointCallbackConfig(throw_on_no_metric=False),
         LastCheckpointCallbackConfig(),
         OnExceptionCheckpointCallbackConfig(),
     ]

nshtrainer-0.38.0/src/nshtrainer/runner.py

@@ -1,118 +0,0 @@
-import copy
-import functools
-from collections.abc import Callable, Iterable, Mapping, Sequence
-from typing import Generic
-
-from nshrunner import RunInfo, Snapshot
-from nshrunner import Runner as _Runner
-from nshrunner._submit import screen
-from typing_extensions import TypeVar, TypeVarTuple, Unpack, 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)
-
-
-class Runner(
-    _Runner[TReturn, TConfig, Unpack[TArguments]],
-    Generic[TReturn, TConfig, Unpack[TArguments]],
-):
-    @override
-    @classmethod
-    def default_validate_fn(cls, config: TConfig, *args: Unpack[TArguments]) -> None:
-        super().default_validate_fn(config, *args)
-
-    @override
-    @classmethod
-    def default_info_fn(cls, config: TConfig, *args: Unpack[TArguments]) -> RunInfo:
-        run_info = super().default_info_fn(config, *args)
-        return {
-            **run_info,
-            "id": config.id,
-            "base_dir": config.directory.project_root,
-        }
-
-    def _fast_dev_run_transform(
-        self,
-        config: TConfig,
-        *args: Unpack[TArguments],
-        n_batches: int,
-    ):
-        config = copy.deepcopy(config)
-        config.trainer.fast_dev_run = n_batches
-        return (config, *args)
-
-    def fast_dev_run(
-        self,
-        runs: Iterable[tuple[TConfig, Unpack[TArguments]]],
-        n_batches: int = 1,
-        *,
-        env: Mapping[str, str] | None = None,
-        transforms: list[
-            Callable[[TConfig, Unpack[TArguments]], tuple[TConfig, Unpack[TArguments]]]
-        ]
-        | None = None,
-    ):
-        transforms = transforms or []
-        transforms.append(
-            functools.partial(self._fast_dev_run_transform, n_batches=n_batches)
-        )
-        return self.local(
-            runs,
-            env=env,
-            transforms=transforms,
-        )
-
-    def fast_dev_run_generator(
-        self,
-        runs: Iterable[tuple[TConfig, Unpack[TArguments]]],
-        n_batches: int = 1,
-        *,
-        env: Mapping[str, str] | None = None,
-        transforms: list[
-            Callable[[TConfig, Unpack[TArguments]], tuple[TConfig, Unpack[TArguments]]]
-        ]
-        | None = None,
-    ):
-        transforms = transforms or []
-        transforms.append(
-            functools.partial(self._fast_dev_run_transform, n_batches=n_batches)
-        )
-        return self.local_generator(
-            runs,
-            env=env,
-            transforms=transforms,
-        )
-
-    def fast_dev_run_session(
-        self,
-        runs: Iterable[tuple[TConfig, Unpack[TArguments]]],
-        options: screen.ScreenJobKwargs = {},
-        n_batches: int = 1,
-        *,
-        snapshot: Snapshot,
-        setup_commands: Sequence[str] | None = None,
-        env: Mapping[str, str] | None = None,
-        transforms: list[
-            Callable[[TConfig, Unpack[TArguments]], tuple[TConfig, Unpack[TArguments]]]
-        ]
-        | None = None,
-        activate_venv: bool = True,
-        print_command: bool = True,
-    ):
-        transforms = transforms or []
-        transforms.append(
-            functools.partial(self._fast_dev_run_transform, n_batches=n_batches)
-        )
-        return self.session(
-            runs,
-            options,
-            snapshot=snapshot,
-            setup_commands=setup_commands,
-            env=env,
-            transforms=transforms,
-            activate_venv=activate_venv,
-            print_command=print_command,
-        )

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/src/nshtrainer/config.py

@@ -95,10 +95,10 @@ from nshtrainer.nn.nonlinearity import MishNonlinearityConfig as MishNonlinearit
 from nshtrainer.nn.nonlinearity import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn.nonlinearity import PReLUConfig as PReLUConfig
 from nshtrainer.nn.nonlinearity import ReLUNonlinearityConfig as ReLUNonlinearityConfig
-from nshtrainer.nn.nonlinearity import SiLUNonlinearityConfig as SiLUNonlinearityConfig
 from nshtrainer.nn.nonlinearity import (
     SigmoidNonlinearityConfig as SigmoidNonlinearityConfig,
 )
+from nshtrainer.nn.nonlinearity import SiLUNonlinearityConfig as SiLUNonlinearityConfig
 from nshtrainer.nn.nonlinearity import (
     SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig,
 )
@@ -137,13 +137,13 @@ from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
 from nshtrainer.trainer._config import ReproducibilityConfig as ReproducibilityConfig
 from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
 from nshtrainer.trainer._config import TrainerConfig as TrainerConfig
-from nshtrainer.util._environment_info import (
-    EnvironmentCUDAConfig as EnvironmentCUDAConfig,
-)
 from nshtrainer.util._environment_info import (
     EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
 )
 from nshtrainer.util._environment_info import EnvironmentConfig as EnvironmentConfig
+from nshtrainer.util._environment_info import (
+    EnvironmentCUDAConfig as EnvironmentCUDAConfig,
+)
 from nshtrainer.util._environment_info import (
     EnvironmentGPUConfig as EnvironmentGPUConfig,
 )
@@ -151,10 +151,10 @@ from nshtrainer.util._environment_info import (
     EnvironmentHardwareConfig as EnvironmentHardwareConfig,
 )
 from nshtrainer.util._environment_info import (
-    EnvironmentLSFInformationConfig as EnvironmentLSFInformationConfig,
+    EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
 )
 from nshtrainer.util._environment_info import (
-    EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
+    EnvironmentLSFInformationConfig as EnvironmentLSFInformationConfig,
 )
 from nshtrainer.util._environment_info import (
     EnvironmentPackageConfig as EnvironmentPackageConfig,

{nshtrainer-0.38.0 → nshtrainer-0.40.0}/src/nshtrainer/nn/module_list.py

@@ -12,10 +12,10 @@ class TypedModuleList(nn.ModuleList, Generic[TModule]):
         super().__init__(modules)
 
     @overload
-    def __getitem__(self, idx: int) -> TModule: ...
+    def __getitem__(self, idx: slice) -> "TypedModuleList[TModule]": ...
 
     @overload
-    def __getitem__(self, idx: slice) -> "TypedModuleList[TModule]": ...
+    def __getitem__(self, idx: int) -> TModule: ...
 
     @override
     def __getitem__(self, idx: int | slice) -> TModule | "TypedModuleList[TModule]":