nshtrainer 0.38.0__py3-none-any.whl → 0.40.0__py3-none-any.whl

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/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/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/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/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/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/data/datamodule.py

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

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]":

nshtrainer/runner.py

@@ -1,12 +1,12 @@
 import copy
-import functools
+import logging
 from collections.abc import Callable, Iterable, Mapping, Sequence
+from pathlib import Path
 from typing import Generic
 
-from nshrunner import RunInfo, Snapshot
-from nshrunner import Runner as _Runner
+import nshrunner as nr
 from nshrunner._submit import screen
-from typing_extensions import TypeVar, TypeVarTuple, Unpack, override
+from typing_extensions import TypeVar, TypeVarTuple, Unpack, deprecated, override
 
 from .model.config import BaseConfig
 
@@ -15,34 +15,27 @@ TArguments = TypeVarTuple("TArguments", default=Unpack[tuple[()]])
 TReturn = TypeVar("TReturn", infer_variance=True)
 
 
+@deprecated("Use nshrunner.Runner instead.")
 class Runner(
-    _Runner[TReturn, TConfig, Unpack[TArguments]],
+    nr.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(
+    def __init__(
         self,
-        config: TConfig,
-        *args: Unpack[TArguments],
-        n_batches: int,
+        run_fn: Callable[[TConfig, Unpack[TArguments]], TReturn],
+        config: nr.RunnerConfig | None = None,
     ):
-        config = copy.deepcopy(config)
-        config.trainer.fast_dev_run = n_batches
-        return (config, *args)
+        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,
@@ -50,20 +43,15 @@ class Runner(
         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,
-        )
+        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,
@@ -71,20 +59,15 @@ class Runner(
         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,
-        )
+        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,
@@ -92,27 +75,25 @@ class Runner(
         options: screen.ScreenJobKwargs = {},
         n_batches: int = 1,
         *,
-        snapshot: Snapshot,
+        snapshot: nr.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)
-        )
+        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,
+            runs_updated,
             options,
             snapshot=snapshot,
             setup_commands=setup_commands,
             env=env,
-            transforms=transforms,
             activate_venv=activate_venv,
             print_command=print_command,
         )

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

@@ -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.dist-info → nshtrainer-0.40.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-nshtrainer/__init__.py,sha256=flMI50Hj1Ie8c1YMSUQ759AqtNBQLT_zHaV2J9EUmOs,573
+nshtrainer/__init__.py,sha256=8hx3uqzroMYOQM7S4U3Eznw7MHF3YDqUkO1MDQeC4cM,642
 nshtrainer/_callback.py,sha256=A1zLsTy4b_wOYnInLLXGSRdHzT2yNa6mPEql-ozm0u0,1013
 nshtrainer/_checkpoint/loader.py,sha256=5vjg-OFChXJjgiOVv8vnV8nwTscfdDtEdxQRz6uPfDE,14158
 nshtrainer/_checkpoint/metadata.py,sha256=5D4PgKodzhLsmQvuF3xxkH49epKaegxi4wh_ImDTtns,4737
@@ -11,13 +11,13 @@ nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHm
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=NpjeKmonJ1Kaz5_39XSn3LlDwvbGjk6WV8BpHSNCvI4,3508
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=g-3zIthupERKqWZQw-A_busQPaPRkto6iHBV-M7nK1Y,527
-nshtrainer/callbacks/checkpoint/_base.py,sha256=vvlwuD-20NozYVIolGGShmUdkkNYeuwN6xCoFnK4GiU,6157
-nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=8BHgLAd3Tuzf5sup0guEAKF1jJiAwYsjdKBFYZw98ac,2171
+nshtrainer/callbacks/checkpoint/_base.py,sha256=X6844ceZ-EMz9UCVLJnA1d-ej2AkkRH90v2hy2zPDkc,6215
+nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=MOC5jLlQnnc61yuiVc2-O4NeYclnkXroOPDV9-2zh5w,2553
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CWWv0cSwQ1VAX26N7hAyMxbNCk26Keh39oQguBEK5To,1102
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
 nshtrainer/callbacks/debug_flag.py,sha256=T7rkY9hYQ_-PsPo2XiQ4eVZ9bBsTd2knpZWctCbjxXc,2011
 nshtrainer/callbacks/directory_setup.py,sha256=c0uY0oTqLcQ3egInHO7G6BeQQgk_xvOLoHH8FR-9U0U,2629
-nshtrainer/callbacks/early_stopping.py,sha256=VWuJz0oN87b6SwBeVc32YNpeJr1wts8K45k8JJJmG9I,4617
+nshtrainer/callbacks/early_stopping.py,sha256=upCNtjkXmwGPwLJmDZacXdBzySG2i_sGIK9QjnaE6tU,4638
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
 nshtrainer/callbacks/finite_checks.py,sha256=gJC_RUr3ais3FJI0uB6wUZnDdE3WRwCix3ppA3PwQXA,2077
 nshtrainer/callbacks/gradient_skipping.py,sha256=EBNkANDnD3BTszWjnG-jwY8FEj-iRqhE3e1x5LQF6M8,3393
@@ -31,9 +31,10 @@ nshtrainer/callbacks/throughput_monitor.py,sha256=H_ocXErZxUO3dxFk8Tx_VQdpI9E_Zt
 nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
 nshtrainer/callbacks/wandb_upload_code.py,sha256=OWG4UkL2SfW6oj6AGRXeBJsZmgsqeHLW2Fj8Jm4ga3I,2298
 nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
-nshtrainer/config.py,sha256=HHzDTE-DYDLn7KuX7UHd7JBPVOW_EDiNAUblHPhjpbI,8312
-nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
+nshtrainer/config.py,sha256=pZyRZOkBRR7eBFRiHpHjQFNEFjaX9tYZIAqZtvKi6cA,8312
+nshtrainer/data/__init__.py,sha256=fnigdEtNcB8DJqz4MHvcI6UwOZICFzWItskY78lLTMA,224
 nshtrainer/data/balanced_batch_sampler.py,sha256=ybMJF-CguaZ17fLEweZ5suaGOiHOMEm3Bn8rQfGTzGQ,5445
+nshtrainer/data/datamodule.py,sha256=Dk06LdcbK13sCLQWacdCoXU_CHWZxoKEZ_xpAg7nZQg,113
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
 nshtrainer/ll/__init__.py,sha256=L-aTi1V1bbvnZjOro8NvI393zbHQSFR9movWSRK9Mds,2477
 nshtrainer/ll/_experimental.py,sha256=oBQCKOEVYoxuUU9eLb-Fg2B2mzZD7SA0zfAO6lmWZ88,53
@@ -71,7 +72,7 @@ nshtrainer/model/mixins/logger.py,sha256=xOymSTofukEYZGkGojXsMEO__ZlBI5lIPZVmlot
 nshtrainer/nn/__init__.py,sha256=0QPFl02a71WZQjLMGOlFNMmsYP5aa1q3eABHmnWH58Q,1427
 nshtrainer/nn/mlp.py,sha256=V0FrScpIUdg_IgIO8GMtIsGEtmHjwF14i2IWxmZrsqg,5952
 nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,2329
-nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
+nshtrainer/nn/module_list.py,sha256=bbR62bnQM4No6HtmBllijOhLwDtWwmcPcO2DklPbsbc,1719
 nshtrainer/nn/nonlinearity.py,sha256=4sYE4MN5zojc-go1k0PYtqssVRuXrM7D4tbpIXp5K-E,6078
 nshtrainer/optimizer.py,sha256=kuJEA1pvB3y1FcsfhAoOJujVqEZqFHlmYO8GW6JeA1g,1527
 nshtrainer/profiler/__init__.py,sha256=cJ_wAm8j3Bz6cKgNQ_9gQLZz9nddKW53VE81UKg8l8g,480
@@ -79,10 +80,10 @@ nshtrainer/profiler/_base.py,sha256=YF5lsJBIl9qts9GLW5Z62JuYeo4SnIArhlFwTGkfTb4,
 nshtrainer/profiler/advanced.py,sha256=44asloha0aGUW8YwjQt3lm3ve8H-N6mM4QgseUSLT30,1112
 nshtrainer/profiler/pytorch.py,sha256=tGeRvoPP5ulWX2RkfXrQvMBoki1T95dpz5p8mwyon1I,2709
 nshtrainer/profiler/simple.py,sha256=MbMfsJvligd0mtGiltxJ0T8MQVDP9T9BzQZFwswl66Y,957
-nshtrainer/runner.py,sha256=USAjrExHkN5oVNVunsoPnLxfQrEHSaa54S3RipOe544,3605
+nshtrainer/runner.py,sha256=llW5PThXfzlu8ym1waFVzED1RGI91OHqN3j02sepoZo,3153
 nshtrainer/scripts/find_packages.py,sha256=ixYivZobumyyGsf2B9oYMLyLTRcBzY_vUv-u3bNW-hs,1424
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
-nshtrainer/trainer/_config.py,sha256=ZIodM5Ek1lpkWFhQ_VfmKR7q1mZFFwtjfx8FH72H8WM,29174
+nshtrainer/trainer/_config.py,sha256=YqpGb4RodkUg87TVE5WBSc4CQkUF0z3qDRdil1HRxoM,29198
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=r0ir4xYSdf_jebM0x09qaO6nJsvsiRQDyM0fs80ppOQ,2347
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
@@ -99,6 +100,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.38.0.dist-info/METADATA,sha256=0LdRFR5fEfmvO704YWcxFHUunc1A8mVhKTfi_WLjgPI,916
-nshtrainer-0.38.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.38.0.dist-info/RECORD,,
+nshtrainer-0.40.0.dist-info/METADATA,sha256=E8dGtOSZ6jdCxugNpIutAeUyd7CcYSRH9cBchBDY7UE,916
+nshtrainer-0.40.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.40.0.dist-info/RECORD,,