nshtrainer 0.39.0__py3-none-any.whl → 0.40.1__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/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/loggers/wandb.py

@@ -92,7 +92,7 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
     - "none" or False: Do not log any checkpoints
     """
 
-    log_code: WandbUploadCodeConfig | None = None
+    log_code: WandbUploadCodeConfig | None = WandbUploadCodeConfig()
     """WandB code upload configuration. Used to upload code to WandB."""
 
     watch: WandbWatchConfig | None = WandbWatchConfig()

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

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

{nshtrainer-0.39.0.dist-info → nshtrainer-0.40.1.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
@@ -32,8 +32,9 @@ nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50
 nshtrainer/callbacks/wandb_upload_code.py,sha256=OWG4UkL2SfW6oj6AGRXeBJsZmgsqeHLW2Fj8Jm4ga3I,2298
 nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
 nshtrainer/config.py,sha256=pZyRZOkBRR7eBFRiHpHjQFNEFjaX9tYZIAqZtvKi6cA,8312
-nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
+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
@@ -56,7 +57,7 @@ nshtrainer/loggers/__init__.py,sha256=C_xk0A3_qKbNdTmzK85AgjRHFD3w-jPRS2ig-iPhfE
 nshtrainer/loggers/_base.py,sha256=xiZKEK0ALJkcqf4OpVNRY0QbZsamR_WR7x7m_68YHXQ,705
 nshtrainer/loggers/csv.py,sha256=D_lYyd94bZ8jAgnRo-ARtFgVcInaD9zktxtsUD9RWCI,1052
 nshtrainer/loggers/tensorboard.py,sha256=wL2amRSdP68zbslZvBeM0ZQBnjF3hIKsz-_lBbdomaM,2216
-nshtrainer/loggers/wandb.py,sha256=td8J2v8T1nvGQI7OYQ1El6k8FGsXZxbnuY97s8KzCiY,6643
+nshtrainer/loggers/wandb.py,sha256=GMlEzkFesgS3qgVgYqATPHJ3rvs9etwFlk3HbidOyRQ,6662
 nshtrainer/lr_scheduler/__init__.py,sha256=uEvgaFAs-4s_bAEMaildy0GT6OvgpgOEKTuzqutESHE,736
 nshtrainer/lr_scheduler/_base.py,sha256=7xOIuxQ86YHbFWG5a3gX46emQj1WN_LaY4-i0Q1TDBg,3659
 nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=E8LW78uuby7bIsoLPpcF1bmNK4lSko-r3qPL-vuHWXQ,5370
@@ -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,7 +80,7 @@ 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=YqpGb4RodkUg87TVE5WBSc4CQkUF0z3qDRdil1HRxoM,29198
@@ -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.39.0.dist-info/METADATA,sha256=bIWwvsGuePEZeT3Q8dYS8IK5Y6ZM4yqe9v41Ybs9OGM,916
-nshtrainer-0.39.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.39.0.dist-info/RECORD,,
+nshtrainer-0.40.1.dist-info/METADATA,sha256=m9hT1zAMnexEhlbXEvTFJBbDsWffVWw11M9NxqqTrZg,916
+nshtrainer-0.40.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.40.1.dist-info/RECORD,,