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.
Files changed (21) hide show
  1. {nshutils-0.31.0 → nshutils-0.31.1}/PKG-INFO +1 -1
  2. {nshutils-0.31.0 → nshutils-0.31.1}/pyproject.toml +1 -1
  3. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/actsave/_saver.py +18 -1
  4. {nshutils-0.31.0 → nshutils-0.31.1}/README.md +0 -0
  5. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/__init__.py +0 -0
  6. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/__init__.pyi +0 -0
  7. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/actsave/__init__.py +0 -0
  8. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/actsave/_loader.py +0 -0
  9. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/collections.py +0 -0
  10. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/display.py +0 -0
  11. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/logging.py +0 -0
  12. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/__init__.py +0 -0
  13. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/_base.py +0 -0
  14. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/_monkey_patch_all.py +0 -0
  15. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/config.py +0 -0
  16. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/jax_.py +0 -0
  17. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/numpy_.py +0 -0
  18. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/torch_.py +0 -0
  19. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/lovely/utils.py +0 -0
  20. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/snoop.py +0 -0
  21. {nshutils-0.31.0 → nshutils-0.31.1}/src/nshutils/typecheck.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: nshutils
3
- Version: 0.31.0
3
+ Version: 0.31.1
4
4
  Summary:
5
5
  Author: Nima Shoghi
6
6
  Author-email: nimashoghi@gmail.com
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "nshutils"
3
- version = "0.31.0"
3
+ version = "0.31.1"
4
4
  description = ""
5
5
  authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
6
6
  requires-python = ">=3.9,<4.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.critical(f"No save_dir specified, using {save_dir=}")
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