nshtrainer 0.3.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

nshtrainer/actsave/__init__.py

@@ -1,7 +1,3 @@
+from nshutils.actsave import *  # type: ignore # noqa: F403
+
 from ._callback import ActSaveCallback as ActSaveCallback
-from ._loader import ActivationLoader as ActivationLoader
-from ._loader import ActLoad as ActLoad
-from ._saver import Activation as Activation
-from ._saver import ActivationSaver as ActivationSaver
-from ._saver import ActSave as ActSave
-from ._saver import Transform as Transform

nshtrainer/actsave/_callback.py

@@ -3,10 +3,9 @@ from typing import TYPE_CHECKING, Literal, cast
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks.callback import Callback
+from nshutils.actsave import ActSave
 from typing_extensions import TypeAlias, override
 
-from ._saver import ActSave
-
 if TYPE_CHECKING:
     from ..model.config import BaseConfig
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.3.0
+Version: 0.4.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -13,8 +13,8 @@ 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: nshutils (>=0.2.0,<0.3.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
 Requires-Dist: rich

{nshtrainer-0.3.0.dist-info → nshtrainer-0.4.1.dist-info}/RECORD

@@ -4,10 +4,8 @@ nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGL
 nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
 nshtrainer/_experimental/flops/module_tracker.py,sha256=bUL-IRTd0aF_DwmXkZjHZAA31p4ZEhyqhc26XWKQUUY,4922
 nshtrainer/_snoop.py,sha256=2rEemPyMP3aIo2QgPzo_-AlT1oXGWYQipId4RQskMls,58
-nshtrainer/actsave/__init__.py,sha256=G1T-fELuGWkVqdhdyoePtj2dTOUtcIOW4VgsXv9JNTA,338
-nshtrainer/actsave/_callback.py,sha256=QoTa60F70f1RxB41VKixN9l5_htfFQxXDPHHSNFreuk,2770
-nshtrainer/actsave/_loader.py,sha256=fAhD32DrJa4onkYfcwc21YIeGEYzOSXCK_HVo9SZLgQ,4604
-nshtrainer/actsave/_saver.py,sha256=0EHmQDhqVxQWRWWSyt03eP1K9ETiACMQYmsZkDMt6HY,9451
+nshtrainer/actsave/__init__.py,sha256=_ZuwgRtF1-ekouXNvtZCAS1g_IDYGB4NX8BFSGNGBT8,119
+nshtrainer/actsave/_callback.py,sha256=mnHOtuG9vtHEzz9q4vCvDNC6VvjZsgb4MSSuOoUDh3M,2778
 nshtrainer/callbacks/__init__.py,sha256=I6W33ityL9Ko8jjqHh3WH_8miV59SAe9LxInhoqX5XE,1665
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/base.py,sha256=LrcRUV02bZEKXRIRvhHT9qsvw_kwoWiAdQkVMyKc5NU,3542
@@ -48,7 +46,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
@@ -60,6 +58,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.3.0.dist-info/METADATA,sha256=oME_P_Y7U4bavZunb-rF2m1R1w31vOS0Qh9l0Nfua68,812
-nshtrainer-0.3.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.3.0.dist-info/RECORD,,
+nshtrainer-0.4.1.dist-info/METADATA,sha256=HJFw7NgykuNISf63Lwa7mtnJf8Cpws2x4a_LirXSCnw,812
+nshtrainer-0.4.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.4.1.dist-info/RECORD,,

nshtrainer/actsave/_loader.py

@@ -1,144 +0,0 @@
-import pprint
-from dataclasses import dataclass, field
-from functools import cached_property
-from logging import getLogger
-from pathlib import Path
-from typing import cast, overload
-
-import numpy as np
-from typing_extensions import TypeVar, override
-
-log = getLogger(__name__)
-
-T = TypeVar("T", infer_variance=True)
-
-
-@dataclass
-class LoadedActivation:
-    base_dir: Path = field(repr=False)
-    name: str
-    num_activations: int = field(init=False)
-    activation_files: list[Path] = field(init=False, repr=False)
-
-    def __post_init__(self):
-        if not self.activation_dir.exists():
-            raise ValueError(f"Activation dir {self.activation_dir} does not exist")
-
-        # The number of activations = the * of .npy files in the activation dir
-        self.activation_files = list(self.activation_dir.glob("*.npy"))
-        # Sort the activation files by the numerical index in the filename
-        self.activation_files.sort(key=lambda p: int(p.stem))
-        self.num_activations = len(self.activation_files)
-
-    @property
-    def activation_dir(self) -> Path:
-        return self.base_dir / self.name
-
-    def _load_activation(self, item: int):
-        activation_path = self.activation_files[item]
-        if not activation_path.exists():
-            raise ValueError(f"Activation {activation_path} does not exist")
-        return cast(np.ndarray, np.load(activation_path, allow_pickle=True))
-
-    @overload
-    def __getitem__(self, item: int) -> np.ndarray: ...
-
-    @overload
-    def __getitem__(self, item: slice | list[int]) -> list[np.ndarray]: ...
-
-    def __getitem__(
-        self, item: int | slice | list[int]
-    ) -> np.ndarray | list[np.ndarray]:
-        if isinstance(item, int):
-            return self._load_activation(item)
-        elif isinstance(item, slice):
-            return [
-                self._load_activation(i)
-                for i in range(*item.indices(self.num_activations))
-            ]
-        elif isinstance(item, list):
-            return [self._load_activation(i) for i in item]
-        else:
-            raise TypeError(f"Invalid type {type(item)} for item {item}")
-
-    def __iter__(self):
-        return iter(self[i] for i in range(self.num_activations))
-
-    def __len__(self):
-        return self.num_activations
-
-    def all_activations(self):
-        return [self[i] for i in range(self.num_activations)]
-
-    @override
-    def __repr__(self):
-        return f"<LoadedActivation {self.name} ({self.num_activations} activations)>"
-
-
-class ActLoad:
-    @classmethod
-    def all_versions(cls, dir: str | Path):
-        dir = Path(dir)
-
-        # If the dir is not an activation base directory, we return None
-        if not (dir / ".activationbase").exists():
-            return None
-
-        # The contents of `dir` should be directories, each of which is a version.
-        return [
-            (subdir, int(subdir.name)) for subdir in dir.iterdir() if subdir.is_dir()
-        ]
-
-    @classmethod
-    def is_valid_activation_base(cls, dir: str | Path):
-        return cls.all_versions(dir) is not None
-
-    @classmethod
-    def from_latest_version(cls, dir: str | Path):
-        # The contents of `dir` should be directories, each of which is a version
-        # We need to find the latest version
-        if (all_versions := cls.all_versions(dir)) is None:
-            raise ValueError(f"{dir} is not an activation base directory")
-
-        path, _ = max(all_versions, key=lambda p: p[1])
-        return cls(path)
-
-    def __init__(self, dir: Path):
-        self._dir = dir
-
-    def activation(self, name: str):
-        return LoadedActivation(self._dir, name)
-
-    @cached_property
-    def activations(self):
-        dirs = list(self._dir.iterdir())
-        # Sort the dirs by the last modified time
-        dirs.sort(key=lambda p: p.stat().st_mtime)
-
-        return {p.name: LoadedActivation(self._dir, p.name) for p in dirs}
-
-    def __iter__(self):
-        return iter(self.activations.values())
-
-    def __getitem__(self, item: str):
-        return self.activations[item]
-
-    def __len__(self):
-        return len(self.activations)
-
-    @override
-    def __repr__(self):
-        acts_str = pprint.pformat(
-            {
-                name: f"<{activation.num_activations} activations>"
-                for name, activation in self.activations.items()
-            }
-        )
-        acts_str = acts_str.replace("'<", "<").replace(">'", ">")
-        return f"ActLoad({acts_str})"
-
-    def get(self, name: str, /, default: T) -> LoadedActivation | T:
-        return self.activations.get(name, default)
-
-
-ActivationLoader = ActLoad

nshtrainer/actsave/_saver.py

@@ -1,337 +0,0 @@
-import contextlib
-import fnmatch
-import tempfile
-import uuid
-import weakref
-from collections.abc import Callable, Mapping
-from dataclasses import dataclass
-from functools import wraps
-from logging import getLogger
-from pathlib import Path
-from typing import Generic, TypeAlias, cast, overload
-
-import numpy as np
-import torch
-from lightning_utilities.core.apply_func import apply_to_collection
-from typing_extensions import ParamSpec, TypeVar, override
-
-log = getLogger(__name__)
-
-Value: TypeAlias = int | float | complex | bool | str | np.ndarray | torch.Tensor | None
-ValueOrLambda = Value | Callable[..., Value]
-
-
-def _to_numpy(activation: Value) -> np.ndarray:
-    # Make sure it's not `None`
-    if activation is None:
-        raise ValueError("Activation should not be `None`")
-
-    if isinstance(activation, np.ndarray):
-        return activation
-    if isinstance(activation, torch.Tensor):
-        activation = activation.detach()
-        if activation.is_floating_point():
-            # NOTE: We need to convert to float32 because [b]float16 is not supported by numpy
-            activation = activation.float()
-        return activation.cpu().numpy()
-    if isinstance(activation, (int, float, complex, str, bool)):
-        return np.array(activation)
-
-    return activation
-
-
-T = TypeVar("T", infer_variance=True)
-
-
-# A wrapper around weakref.ref that allows for primitive types
-# To get around errors like:
-# TypeError: cannot create weak reference to 'int' object
-class WeakRef(Generic[T]):
-    _ref: Callable[[], T] | None
-
-    def __init__(self, obj: T):
-        try:
-            self._ref = cast(Callable[[], T], weakref.ref(obj))
-        except TypeError as e:
-            if "cannot create weak reference" not in str(e):
-                raise
-            self._ref = lambda: obj
-
-    def __call__(self) -> T:
-        if self._ref is None:
-            raise RuntimeError("WeakRef is deleted")
-        return self._ref()
-
-    def delete(self):
-        del self._ref
-        self._ref = None
-
-
-@dataclass
-class Activation:
-    name: str
-    ref: WeakRef[ValueOrLambda] | None
-    transformed: np.ndarray | None = None
-
-    def __post_init__(self):
-        # Update the `name` to replace `/` with `.`
-        self.name = self.name.replace("/", ".")
-
-    def __call__(self) -> np.ndarray | None:
-        # If we have a transformed value, we return it
-        if self.transformed is not None:
-            return self.transformed
-
-        if self.ref is None:
-            raise RuntimeError("Activation is deleted")
-
-        # If we have a lambda, we need to call it
-        unrwapped_ref = self.ref()
-        activation = unrwapped_ref
-        if callable(unrwapped_ref):
-            activation = unrwapped_ref()
-
-        # If we have a `None`, we return early
-        if activation is None:
-            return None
-
-        activation = apply_to_collection(activation, torch.Tensor, _to_numpy)
-        activation = _to_numpy(activation)
-
-        # Set the transformed value
-        self.transformed = activation
-
-        # Delete the reference
-        self.ref.delete()
-        del self.ref
-        self.ref = None
-
-        return self.transformed
-
-    @classmethod
-    def from_value_or_lambda(cls, name: str, value_or_lambda: ValueOrLambda):
-        return cls(name, WeakRef(value_or_lambda))
-
-    @classmethod
-    def from_dict(cls, d: Mapping[str, ValueOrLambda]):
-        return [cls.from_value_or_lambda(k, v) for k, v in d.items()]
-
-
-Transform = Callable[[Activation], Mapping[str, ValueOrLambda]]
-
-
-def _ensure_supported():
-    try:
-        import torch.distributed as dist
-
-        if dist.is_initialized() and dist.get_world_size() > 1:
-            raise RuntimeError("Only single GPU is supported at the moment")
-    except ImportError:
-        pass
-
-
-P = ParamSpec("P")
-
-
-def _ignore_if_scripting(fn: Callable[P, None]) -> Callable[P, None]:
-    @wraps(fn)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> None:
-        if torch.jit.is_scripting():
-            return
-
-        _ensure_supported()
-        fn(*args, **kwargs)
-
-    return wrapper
-
-
-class _Saver:
-    def __init__(
-        self,
-        save_dir: Path,
-        prefixes_fn: Callable[[], list[str]],
-        *,
-        filters: list[str] | None = None,
-    ):
-        # Create a directory under `save_dir` by autoincrementing
-        # (i.e., every activation save context, we create a new directory)
-        # The id = the number of activation subdirectories
-        self._id = sum(1 for subdir in save_dir.glob("*") if subdir.is_dir())
-        save_dir.mkdir(parents=True, exist_ok=True)
-
-        # Add a .activationbase file to the save_dir to indicate that this is an activation base
-        (save_dir / ".activationbase").touch(exist_ok=True)
-
-        self._save_dir = save_dir / f"{self._id:04d}"
-        # Make sure `self._save_dir` does not exist and create it
-        self._save_dir.mkdir(exist_ok=False)
-
-        self._prefixes_fn = prefixes_fn
-        self._filters = filters
-
-    def _save_activation(self, activation: Activation):
-        # If the activation value is `None`, we skip it.
-        if (activation_value := activation()) is None:
-            return
-
-        # Save the activation to self._save_dir / name / {id}.npz, where id is an auto-incrementing integer
-        file_name = ".".join(self._prefixes_fn() + [activation.name])
-        path = self._save_dir / file_name
-        path.mkdir(exist_ok=True, parents=True)
-
-        # Get the next id and save the activation
-        id = len(list(path.glob("*.npy")))
-        np.save(path / f"{id:04d}.npy", activation_value)
-
-    @_ignore_if_scripting
-    def save(
-        self,
-        acts: dict[str, ValueOrLambda] | None = None,
-        /,
-        **kwargs: ValueOrLambda,
-    ):
-        kwargs.update(acts or {})
-
-        # Build activations
-        activations = Activation.from_dict(kwargs)
-
-        for activation in activations:
-            # Make sure name matches at least one filter if filters are specified
-            if self._filters is not None and all(
-                not fnmatch.fnmatch(activation.name, f) for f in self._filters
-            ):
-                continue
-
-            # Save the current activation
-            self._save_activation(activation)
-
-        del activations
-
-
-class ActSaveProvider:
-    _saver: _Saver | None = None
-    _prefixes: list[str] = []
-
-    def initialize(self, save_dir: Path | None = None):
-        """
-        Initializes the saver with the given configuration and save directory.
-
-        Args:
-            save_dir (Path): The directory where the saved files will be stored.
-        """
-        if self._saver is None:
-            if save_dir is None:
-                save_dir = Path(tempfile.gettempdir()) / f"actsave-{uuid.uuid4()}"
-                log.critical(f"No save_dir specified, using {save_dir=}")
-            self._saver = _Saver(
-                save_dir,
-                lambda: self._prefixes,
-            )
-
-    @contextlib.contextmanager
-    def enabled(self, save_dir: Path | None = None):
-        """
-        Context manager that enables the actsave functionality with the specified configuration.
-
-        Args:
-            save_dir (Path): The directory where the saved files will be stored.
-        """
-        prev = self._saver
-        self.initialize(save_dir)
-        try:
-            yield
-        finally:
-            self._saver = prev
-
-    @override
-    def __init__(self):
-        super().__init__()
-
-        self._saver = None
-        self._prefixes = []
-
-    @contextlib.contextmanager
-    def context(self, label: str):
-        """
-        A context manager that adds a label to the current context.
-
-        Args:
-            label (str): The label for the context.
-        """
-        if torch.jit.is_scripting():
-            yield
-            return
-
-        if self._saver is None:
-            yield
-            return
-
-        _ensure_supported()
-
-        log.debug(f"Entering ActSave context {label}")
-        self._prefixes.append(label)
-        try:
-            yield
-        finally:
-            _ = self._prefixes.pop()
-
-    prefix = context
-
-    @overload
-    def __call__(
-        self,
-        acts: dict[str, ValueOrLambda] | None = None,
-        /,
-        **kwargs: ValueOrLambda,
-    ):
-        """
-        Saves the activations to disk.
-
-        Args:
-            acts (dict[str, ValueOrLambda] | None, optional): A dictionary of acts. Defaults to None.
-            **kwargs (ValueOrLambda): Additional keyword arguments.
-
-        Returns:
-            None
-
-        """
-        ...
-
-    @overload
-    def __call__(self, acts: Callable[[], dict[str, ValueOrLambda]], /):
-        """
-        Saves the activations to disk.
-
-        Args:
-            acts (Callable[[], dict[str, ValueOrLambda]]): A callable that returns a dictionary of acts.
-            **kwargs (ValueOrLambda): Additional keyword arguments.
-
-        Returns:
-            None
-
-        """
-        ...
-
-    def __call__(
-        self,
-        acts: (
-            dict[str, ValueOrLambda] | Callable[[], dict[str, ValueOrLambda]] | None
-        ) = None,
-        /,
-        **kwargs: ValueOrLambda,
-    ):
-        if torch.jit.is_scripting():
-            return
-
-        if self._saver is None:
-            return
-
-        if acts is not None and callable(acts):
-            acts = acts()
-        self._saver.save(acts, **kwargs)
-
-    save = __call__
-
-
-ActSave = ActSaveProvider()
-ActivationSaver = ActSave