nshutils 0.31.0__tar.gz → 0.31.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {nshutils-0.31.0 → nshutils-0.31.1}/PKG-INFO +1 -1
- {nshutils-0.31.0 → nshutils-0.31.1}/pyproject.toml +1 -1
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/actsave/_saver.py +18 -1
- {nshutils-0.31.0 → nshutils-0.31.1}/README.md +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/__init__.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/__init__.pyi +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/actsave/__init__.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/actsave/_loader.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/collections.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/display.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/logging.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/__init__.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/_base.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/_monkey_patch_all.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/config.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/jax_.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/numpy_.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/torch_.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/utils.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/snoop.py +0 -0
- {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/typecheck.py +0 -0
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|
2
2
|
|
3
3
|
import contextlib
|
4
4
|
import fnmatch
|
5
|
+
import os
|
5
6
|
import tempfile
|
6
7
|
import weakref
|
7
8
|
from collections.abc import Callable, Mapping
|
@@ -266,7 +267,11 @@ class ActSaveProvider:
|
|
266
267
|
|
267
268
|
if save_dir is None:
|
268
269
|
save_dir = Path(tempfile.gettempdir()) / f"actsave-{uuid7str()}"
|
269
|
-
log.
|
270
|
+
log.warning(
|
271
|
+
f"ActSave: Using temporary directory {save_dir} for activations."
|
272
|
+
)
|
273
|
+
else:
|
274
|
+
log.info(f"ActSave enabled. Saving to {save_dir}")
|
270
275
|
self._saver = _Saver(save_dir, lambda: self._prefixes)
|
271
276
|
|
272
277
|
def disable(self):
|
@@ -307,6 +312,18 @@ class ActSaveProvider:
|
|
307
312
|
self._prefixes = []
|
308
313
|
self._disable_count = 0
|
309
314
|
|
315
|
+
# Check for environment variable `ACTSAVE` to automatically enable saving.
|
316
|
+
# If set to "1" or "true" (case-insensitive), activations are saved to a temporary directory.
|
317
|
+
# If set to a path, activations are saved to that path.
|
318
|
+
if env_var := os.environ.get("ACTSAVE"):
|
319
|
+
log.info(
|
320
|
+
f"`ACTSAVE={env_var}` detected, attempting to auto-enable activation saving."
|
321
|
+
)
|
322
|
+
if env_var.lower() in ("1", "true"):
|
323
|
+
self.enable()
|
324
|
+
else:
|
325
|
+
self.enable(Path(env_var))
|
326
|
+
|
310
327
|
@contextlib.contextmanager
|
311
328
|
def disabled(self, condition: bool | Callable[[], bool] = True):
|
312
329
|
"""
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|