nshtrainer 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

nshtrainer/ll/__init__.py

@@ -0,0 +1,42 @@
+import importlib
+import sys
+from types import ModuleType
+
+# The name of your new package
+NEW_PACKAGE = "nshtrainer"
+
+# Import the new package
+new_package = importlib.import_module(NEW_PACKAGE)
+
+
+# Create a custom module class that inherits from ModuleType
+class ProxyModule(ModuleType):
+    def __getattr__(self, name):
+        return getattr(new_package, name)
+
+    def __dir__(self):
+        return dir(new_package)
+
+
+# Create a new module instance
+old_module = ProxyModule(__name__)
+
+# Copy attributes from new_package to old_module
+for attr in dir(new_package):
+    if not attr.startswith("__"):
+        setattr(old_module, attr, getattr(new_package, attr))
+
+# Replace the module in sys.modules
+sys.modules[__name__] = old_module
+
+
+# Handle submodule imports
+class SubmoduleProxy:
+    def __getattr__(self, name):
+        return importlib.import_module(f"{NEW_PACKAGE}.{name}")
+
+
+# Add submodule handling to the proxy module
+old_module.__class__ = type(
+    "ProxyModuleWithSubmodules", (ProxyModule, SubmoduleProxy), {}
+)

nshtrainer/runner.py

@@ -1,7 +1,11 @@
+import copy
+import functools
+from collections.abc import Callable, Mapping, Sequence
 from typing import Generic
 
 from nshrunner import RunInfo
 from nshrunner import Runner as _Runner
+from nshrunner._runner import SnapshotArgType
 from typing_extensions import TypeVar, TypeVarTuple, Unpack, override
 
 from .model.config import BaseConfig
@@ -29,3 +33,65 @@ class Runner(
             "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: Sequence[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_session(
+        self,
+        runs: Sequence[tuple[TConfig, Unpack[TArguments]]],
+        n_batches: int = 1,
+        *,
+        snapshot: SnapshotArgType,
+        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,
+        session_name: str = "nshrunner",
+        attach: bool = True,
+        print_command: bool = True,
+        pause_before_exit: bool = False,
+    ):
+        transforms = transforms or []
+        transforms.append(
+            functools.partial(self._fast_dev_run_transform, n_batches=n_batches)
+        )
+        return self.session(
+            runs,
+            snapshot=snapshot,
+            setup_commands=setup_commands,
+            env=env,
+            transforms=transforms,
+            activate_venv=activate_venv,
+            session_name=session_name,
+            attach=attach,
+            print_command=print_command,
+            pause_before_exit=pause_before_exit,
+        )

{nshtrainer-0.4.0.dist-info → nshtrainer-0.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.4.0
+Version: 0.4.2
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -13,7 +13,7 @@ Requires-Dist: lightning
 Requires-Dist: lovely-numpy (>=0.2.13,<0.3.0)
 Requires-Dist: lovely-tensors (>=0.1.16,<0.2.0)
 Requires-Dist: nshconfig (>=0.2.0,<0.3.0)
-Requires-Dist: nshrunner (>=0.5.4,<0.6.0)
+Requires-Dist: nshrunner (>=0.5.5,<0.6.0)
 Requires-Dist: nshutils (>=0.3.0,<0.4.0)
 Requires-Dist: numpy
 Requires-Dist: pytorch-lightning

{nshtrainer-0.4.0.dist-info → nshtrainer-0.4.2.dist-info}/RECORD

@@ -26,6 +26,7 @@ nshtrainer/config.py,sha256=IXOAl_JWFNX9kPTo_iw4Nc3qXqkKrbA6-ZrvTAjqu6A,104
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=bcJBcQjh1hB1yKF_xSlT9AtEWv0BJjYc1CuH2BF-ea8,4392
 nshtrainer/data/transform.py,sha256=JeGxvytQly8hougrsdMmKG8gJ6qvFPDglJCO4Tp6STk,1795
+nshtrainer/ll/__init__.py,sha256=1cccXGP085GSR5AYM4zVg4CgTso9rgzrPJ_mZz0RRJ0,1027
 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=mn6cyizyI_stkXtg6zxIEGF9btIxMRWigUHUTlUYCSw,5221
@@ -46,7 +47,7 @@ nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,
 nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
 nshtrainer/nn/nonlinearity.py,sha256=owtU4kh4G98psD0axOJWVfBhm-OtJVgFM-TXSHmbNPU,3625
 nshtrainer/optimizer.py,sha256=kuJEA1pvB3y1FcsfhAoOJujVqEZqFHlmYO8GW6JeA1g,1527
-nshtrainer/runner.py,sha256=vyHr0EZ0PBOWZh09BtOOxio-FRQZFbVoL4cdBlI97vY,991
+nshtrainer/runner.py,sha256=zhNiDmv9R-oVmutQLrNUTDz9RLLAlHG6K0uYz77HkRM,3090
 nshtrainer/scripts/check_env.py,sha256=IMl6dSqsLYppI0XuCsVq8lK4bYqXwY9KHJkzsShz4Kg,806
 nshtrainer/scripts/find_packages.py,sha256=FbdlfmAefttFSMfaT0A46a-oHLP_ioaQKihwBfBeWeA,1467
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
@@ -58,6 +59,6 @@ nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
 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.4.0.dist-info/METADATA,sha256=V4fr_C3pnSpaXT1KHAmnDFLcDznTo7AZZDAPFPm3AUk,812
-nshtrainer-0.4.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.4.0.dist-info/RECORD,,
+nshtrainer-0.4.2.dist-info/METADATA,sha256=cUl7TywscNNfOwtMvNtwEvagUiyqirLfW0i0FJqIIB8,812
+nshtrainer-0.4.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.4.2.dist-info/RECORD,,