nshtrainer 0.39.0__tar.gz → 0.40.0__tar.gz

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

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

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

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

{nshtrainer-0.39.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.39.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.39.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.39.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]":