checkpointer 2.13.0__tar.gz → 2.14.0__tar.gz

{checkpointer-2.13.0 → checkpointer-2.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.13.0
+Version: 2.14.0
 Summary: checkpointer adds code-aware caching to Python functions, maintaining correctness and speeding up execution as your code changes.
 Project-URL: Repository, https://github.com/Reddan/checkpointer.git
 Author: Hampus Hallman
@@ -16,7 +16,7 @@ Description-Content-Type: text/markdown
 
 # checkpointer · [![License](https://img.shields.io/badge/license-MIT-blue)](https://github.com/Reddan/checkpointer/blob/master/LICENSE) [![pypi](https://img.shields.io/pypi/v/checkpointer)](https://pypi.org/project/checkpointer/) [![pypi](https://img.shields.io/pypi/pyversions/checkpointer)](https://pypi.org/project/checkpointer/)
 
-`checkpointer` is a Python library offering a decorator-based API for memoizing (caching) function results with code-aware cache invalidation. It works with sync and async functions, supports multiple storage backends, and refreshes caches automatically when your code or dependencies change - helping you maintain correctness, speed up execution, and smooth out your workflows by skipping redundant, costly operations.
+`checkpointer` is a Python library offering a decorator-based API for memoizing (caching) function results with code-aware cache invalidation. It works with sync and async functions, supports multiple storage backends, and invalidates caches automatically when your code or dependencies change - helping you maintain correctness, speed up execution, and smooth out your workflows by skipping redundant, costly operations.
 
 ## 📦 Installation
 
@@ -42,41 +42,51 @@ result = expensive_function(4)  # Loads from the cache
 
 ## 🧠 How It Works
 
-When a `@checkpoint`-decorated function is called, `checkpointer` computes a unique identifier for the call. This identifier derives from the function's source code, its dependencies, captured variables, and the arguments passed.
+When you decorate a function with `@checkpoint` and call it, `checkpointer` computes a unique identifier that represents that specific call. This identifier is based on:
 
-It then tries to retrieve a cached result using this identifier. If a valid cached result is found, it's returned immediately. Otherwise, the original function executes, its result is stored, and then returned.
+* The function's source code and all its user-defined dependencies,
+* Global variables used by the function (if capturing is enabled or explicitly annotated),
+* The actual arguments passed to the function.
 
-### 🚨 What Triggers Cache Invalidation?
+`checkpointer` then looks up this identifier in its cache. If a valid cached result exists, it returns that immediately. Otherwise, it runs the original function, stores the result, and returns it.
 
-`checkpointer` maintains cache correctness using two types of hashes:
+`checkpointer` is designed to be flexible through features like:
 
-#### 1. Function Identity Hash (One-Time per Function)
+* **Support for decorated methods**, correctly caching results bound to instances.
+* **Support for decorated async functions**, compatible with any async runtime.
+* **Robust hashing**, covering complex Python objects and large **NumPy**/**PyTorch** arrays via its internal `ObjectHash`.
+* **Targeted hashing**, allowing you to optimize how arguments and captured variables are hashed.
+* **Multi-layered caching**, letting you stack decorators for layered caching strategies without losing cache consistency.
+
+### 🚨 What Causes Cache Invalidation?
+
+To ensure cache correctness, `checkpointer` tracks two types of hashes:
+
+#### 1. Function Identity Hash (Computed Once Per Function)
 
 This hash represents the decorated function itself and is computed once (usually on first invocation). It covers:
 
-* **Decorated Function's Code:**\
-    The function's logic and signature (excluding parameter type annotations) are hashed. Formatting changes like whitespace, newlines, comments, or trailing commas do **not** cause invalidation.
+* **Function Code and Signature:**\
+    The actual logic and parameters of the function are hashed - but *not* parameter type annotations or formatting details like whitespace, newlines, comments, or trailing commas, which do **not** trigger invalidation.
 
 * **Dependencies:**\
-    All user-defined functions and methods the function calls or uses are included recursively. Dependencies are detected by:
-    * Inspecting the function's global scope for referenced functions/objects.
-    * Inferring from argument type annotations.
+    All user-defined functions and methods that the decorated function calls or relies on, including indirect dependencies, are included recursively. Dependencies are identified by:
+    * Inspecting the function's global scope for referenced functions and objects.
+    * Inferring from the function's argument type annotations.
     * Analyzing object constructions and method calls to identify classes and methods used.
 
-* **Top-Level Module Code:**\
-    Changes unrelated to the function or its dependencies in the module do **not** trigger invalidation.
+* **Exclusions:**\
+    Changes elsewhere in the module unrelated to the function or its dependencies do **not** cause invalidation.
 
 #### 2. Call Hash (Computed on Every Function Call)
 
-Each function call's cache key (the **call hash**) combines:
+Every function call produces a call hash, combining:
 
 * **Passed Arguments:**\
     Includes positional and keyword arguments, combined with default values. Changing defaults alone doesn't necessarily trigger invalidation unless it affects actual call values.
 
 * **Captured Global Variables:**\
-    When `capture=True` or explicit capture annotations are used, `checkpointer` hashes global variables referenced by the function:
-    * `CaptureMe` variables are hashed on every call, so changes trigger invalidation.
-    * `CaptureMeOnce` variables are hashed once per session for performance optimization.
+    When `capture=True` or explicit capture annotations are used, `checkpointer` includes referenced global variables in the call hash. Variables annotated with `CaptureMe` are hashed on every call, causing immediate cache invalidation if they change. Variables annotated with `CaptureMeOnce` are hashed only once per Python session, improving performance by avoiding repeated hashing.
 
 * **Custom Argument Hashing:**\
     Using `HashBy` annotations, arguments or captured variables can be transformed before hashing (e.g., sorting lists to ignore order), allowing more precise or efficient call hashes.
@@ -103,7 +113,7 @@ Once a function is decorated with `@checkpoint`, you can interact with its cachi
 * **`expensive_function.delete(...)`**:\
     Remove the cached entry for given arguments.
 
-* **`expensive_function.reinit(recursive: bool = False)`**:\
+* **`expensive_function.reinit(recursive: bool = True)`**:\
     Recalculate the function identity hash and recapture `CaptureMeOnce` variables, updating the cached function state within the same Python session.
 
 ## ⚙️ Configuration & Customization
@@ -116,18 +126,18 @@ The `@checkpoint` decorator accepts the following parameters:
 * **`directory`** (Type: `str` or `pathlib.Path` or `None`, Default: `~/.cache/checkpoints`)\
     Base directory for disk-based checkpoints (only for `"pickle"` storage).
 
-* **`when`** (Type: `bool`, Default: `True`)\
-    Enable or disable checkpointing dynamically, useful for environment-based toggling.
-
 * **`capture`** (Type: `bool`, Default: `False`)\
     If `True`, includes global variables referenced by the function in call hashes (except those excluded via `NoHash`).
 
-* **`should_expire`** (Type: `Callable[[datetime.datetime], bool]`, Default: `None`)\
+* **`expiry`** (Type: `Callable[[datetime.datetime], bool]` or `datetime.timedelta`, Default: `None`)\
     A custom callable that receives the `datetime` timestamp of a cached result. It should return `True` if the cached result is considered expired and needs recomputation, or `False` otherwise.
 
 * **`fn_hash_from`** (Type: `Any`, Default: `None`)\
     Override the computed function identity hash with any hashable object you provide (e.g., version strings, config IDs). This gives you explicit control over the function's version and when its cache should be invalidated.
 
+* **`when`** (Type: `bool`, Default: `True`)\
+    Enable or disable checkpointing dynamically, useful for environment-based toggling.
+
 * **`verbosity`** (Type: `int` (`0`, `1`, or `2`), Default: `1`)\
     Controls the level of logging output from `checkpointer`.
     * `0`: No output.
@@ -231,30 +241,3 @@ class MyCustomStorage(Storage):
 def custom_cached_function(x: int):
     return x ** 2
 ```
-
-## ⚡ Async Support
-
-`checkpointer` works with Python's `asyncio` and other async runtimes.
-
-```python
-import asyncio
-from checkpointer import checkpoint
-
-@checkpoint
-async def async_compute_sum(a: int, b: int) -> int:
-    print(f"Asynchronously computing {a} + {b}...")
-    await asyncio.sleep(1)
-    return a + b
-
-async def main():
-    result1 = await async_compute_sum(3, 7)
-    print(f"Result 1: {result1}")
-
-    result2 = await async_compute_sum(3, 7)
-    print(f"Result 2: {result2}")
-
-    result3 = async_compute_sum.get(3, 7)
-    print(f"Result 3 (from cache): {result3}")
-
-asyncio.run(main())
-```

{checkpointer-2.13.0 → checkpointer-2.14.0}/README.md

@@ -1,6 +1,6 @@
 # checkpointer · [![License](https://img.shields.io/badge/license-MIT-blue)](https://github.com/Reddan/checkpointer/blob/master/LICENSE) [![pypi](https://img.shields.io/pypi/v/checkpointer)](https://pypi.org/project/checkpointer/) [![pypi](https://img.shields.io/pypi/pyversions/checkpointer)](https://pypi.org/project/checkpointer/)
 
-`checkpointer` is a Python library offering a decorator-based API for memoizing (caching) function results with code-aware cache invalidation. It works with sync and async functions, supports multiple storage backends, and refreshes caches automatically when your code or dependencies change - helping you maintain correctness, speed up execution, and smooth out your workflows by skipping redundant, costly operations.
+`checkpointer` is a Python library offering a decorator-based API for memoizing (caching) function results with code-aware cache invalidation. It works with sync and async functions, supports multiple storage backends, and invalidates caches automatically when your code or dependencies change - helping you maintain correctness, speed up execution, and smooth out your workflows by skipping redundant, costly operations.
 
 ## 📦 Installation
 
@@ -26,41 +26,51 @@ result = expensive_function(4)  # Loads from the cache
 
 ## 🧠 How It Works
 
-When a `@checkpoint`-decorated function is called, `checkpointer` computes a unique identifier for the call. This identifier derives from the function's source code, its dependencies, captured variables, and the arguments passed.
+When you decorate a function with `@checkpoint` and call it, `checkpointer` computes a unique identifier that represents that specific call. This identifier is based on:
 
-It then tries to retrieve a cached result using this identifier. If a valid cached result is found, it's returned immediately. Otherwise, the original function executes, its result is stored, and then returned.
+* The function's source code and all its user-defined dependencies,
+* Global variables used by the function (if capturing is enabled or explicitly annotated),
+* The actual arguments passed to the function.
 
-### 🚨 What Triggers Cache Invalidation?
+`checkpointer` then looks up this identifier in its cache. If a valid cached result exists, it returns that immediately. Otherwise, it runs the original function, stores the result, and returns it.
 
-`checkpointer` maintains cache correctness using two types of hashes:
+`checkpointer` is designed to be flexible through features like:
 
-#### 1. Function Identity Hash (One-Time per Function)
+* **Support for decorated methods**, correctly caching results bound to instances.
+* **Support for decorated async functions**, compatible with any async runtime.
+* **Robust hashing**, covering complex Python objects and large **NumPy**/**PyTorch** arrays via its internal `ObjectHash`.
+* **Targeted hashing**, allowing you to optimize how arguments and captured variables are hashed.
+* **Multi-layered caching**, letting you stack decorators for layered caching strategies without losing cache consistency.
+
+### 🚨 What Causes Cache Invalidation?
+
+To ensure cache correctness, `checkpointer` tracks two types of hashes:
+
+#### 1. Function Identity Hash (Computed Once Per Function)
 
 This hash represents the decorated function itself and is computed once (usually on first invocation). It covers:
 
-* **Decorated Function's Code:**\
-    The function's logic and signature (excluding parameter type annotations) are hashed. Formatting changes like whitespace, newlines, comments, or trailing commas do **not** cause invalidation.
+* **Function Code and Signature:**\
+    The actual logic and parameters of the function are hashed - but *not* parameter type annotations or formatting details like whitespace, newlines, comments, or trailing commas, which do **not** trigger invalidation.
 
 * **Dependencies:**\
-    All user-defined functions and methods the function calls or uses are included recursively. Dependencies are detected by:
-    * Inspecting the function's global scope for referenced functions/objects.
-    * Inferring from argument type annotations.
+    All user-defined functions and methods that the decorated function calls or relies on, including indirect dependencies, are included recursively. Dependencies are identified by:
+    * Inspecting the function's global scope for referenced functions and objects.
+    * Inferring from the function's argument type annotations.
     * Analyzing object constructions and method calls to identify classes and methods used.
 
-* **Top-Level Module Code:**\
-    Changes unrelated to the function or its dependencies in the module do **not** trigger invalidation.
+* **Exclusions:**\
+    Changes elsewhere in the module unrelated to the function or its dependencies do **not** cause invalidation.
 
 #### 2. Call Hash (Computed on Every Function Call)
 
-Each function call's cache key (the **call hash**) combines:
+Every function call produces a call hash, combining:
 
 * **Passed Arguments:**\
     Includes positional and keyword arguments, combined with default values. Changing defaults alone doesn't necessarily trigger invalidation unless it affects actual call values.
 
 * **Captured Global Variables:**\
-    When `capture=True` or explicit capture annotations are used, `checkpointer` hashes global variables referenced by the function:
-    * `CaptureMe` variables are hashed on every call, so changes trigger invalidation.
-    * `CaptureMeOnce` variables are hashed once per session for performance optimization.
+    When `capture=True` or explicit capture annotations are used, `checkpointer` includes referenced global variables in the call hash. Variables annotated with `CaptureMe` are hashed on every call, causing immediate cache invalidation if they change. Variables annotated with `CaptureMeOnce` are hashed only once per Python session, improving performance by avoiding repeated hashing.
 
 * **Custom Argument Hashing:**\
     Using `HashBy` annotations, arguments or captured variables can be transformed before hashing (e.g., sorting lists to ignore order), allowing more precise or efficient call hashes.
@@ -87,7 +97,7 @@ Once a function is decorated with `@checkpoint`, you can interact with its cachi
 * **`expensive_function.delete(...)`**:\
     Remove the cached entry for given arguments.
 
-* **`expensive_function.reinit(recursive: bool = False)`**:\
+* **`expensive_function.reinit(recursive: bool = True)`**:\
     Recalculate the function identity hash and recapture `CaptureMeOnce` variables, updating the cached function state within the same Python session.
 
 ## ⚙️ Configuration & Customization
@@ -100,18 +110,18 @@ The `@checkpoint` decorator accepts the following parameters:
 * **`directory`** (Type: `str` or `pathlib.Path` or `None`, Default: `~/.cache/checkpoints`)\
     Base directory for disk-based checkpoints (only for `"pickle"` storage).
 
-* **`when`** (Type: `bool`, Default: `True`)\
-    Enable or disable checkpointing dynamically, useful for environment-based toggling.
-
 * **`capture`** (Type: `bool`, Default: `False`)\
     If `True`, includes global variables referenced by the function in call hashes (except those excluded via `NoHash`).
 
-* **`should_expire`** (Type: `Callable[[datetime.datetime], bool]`, Default: `None`)\
+* **`expiry`** (Type: `Callable[[datetime.datetime], bool]` or `datetime.timedelta`, Default: `None`)\
     A custom callable that receives the `datetime` timestamp of a cached result. It should return `True` if the cached result is considered expired and needs recomputation, or `False` otherwise.
 
 * **`fn_hash_from`** (Type: `Any`, Default: `None`)\
     Override the computed function identity hash with any hashable object you provide (e.g., version strings, config IDs). This gives you explicit control over the function's version and when its cache should be invalidated.
 
+* **`when`** (Type: `bool`, Default: `True`)\
+    Enable or disable checkpointing dynamically, useful for environment-based toggling.
+
 * **`verbosity`** (Type: `int` (`0`, `1`, or `2`), Default: `1`)\
     Controls the level of logging output from `checkpointer`.
     * `0`: No output.
@@ -215,30 +225,3 @@ class MyCustomStorage(Storage):
 def custom_cached_function(x: int):
     return x ** 2
 ```
-
-## ⚡ Async Support
-
-`checkpointer` works with Python's `asyncio` and other async runtimes.
-
-```python
-import asyncio
-from checkpointer import checkpoint
-
-@checkpoint
-async def async_compute_sum(a: int, b: int) -> int:
-    print(f"Asynchronously computing {a} + {b}...")
-    await asyncio.sleep(1)
-    return a + b
-
-async def main():
-    result1 = await async_compute_sum(3, 7)
-    print(f"Result 1: {result1}")
-
-    result2 = await async_compute_sum(3, 7)
-    print(f"Result 2: {result2}")
-
-    result3 = async_compute_sum.get(3, 7)
-    print(f"Result 3 (from cache): {result3}")
-
-asyncio.run(main())
-```

{checkpointer-2.13.0 → checkpointer-2.14.0}/checkpointer/__init__.py

@@ -1,5 +1,6 @@
 import gc
 import tempfile
+from datetime import timedelta
 from typing import Callable
 from .checkpoint import CachedFunction, Checkpointer, CheckpointError, FunctionIdent
 from .object_hash import ObjectHash

{checkpointer-2.13.0 → checkpointer-2.14.0}/checkpointer/checkpoint.py

@@ -1,19 +1,19 @@
 from __future__ import annotations
 import re
-from datetime import datetime
+from datetime import datetime, timedelta
 from functools import cached_property, update_wrapper
 from inspect import Parameter, iscoroutine, signature, unwrap
-from itertools import chain
 from pathlib import Path
 from typing import (
   Callable, Concatenate, Coroutine, Generic, Iterable,
-  Literal, Self, Type, TypedDict, Unpack, cast, overload,
+  Literal, Self, Type, TypedDict, Unpack, overload,
 )
 from .fn_ident import Capturable, RawFunctionIdent, get_fn_ident
 from .object_hash import ObjectHash
 from .print_checkpoint import print_checkpoint
 from .storages import STORAGE_MAP, Storage, StorageType
 from .types import AwaitableValue, C, Coro, Fn, P, R, hash_by_from_annotation
+from .utils import flatten
 
 DEFAULT_DIR = Path.home() / ".cache/checkpoints"
 
@@ -25,7 +25,7 @@ class CheckpointerOpts(TypedDict, total=False):
   directory: Path | str | None
   when: bool
   verbosity: Literal[0, 1, 2]
-  should_expire: Callable[[datetime], bool] | None
+  expiry: Callable[[datetime], bool] | timedelta | None
   capture: bool
   fn_hash_from: object
 
@@ -35,7 +35,7 @@ class Checkpointer:
     self.directory = Path(opts.get("directory", DEFAULT_DIR) or ".")
     self.when = opts.get("when", True)
     self.verbosity = opts.get("verbosity", 1)
-    self.should_expire = opts.get("should_expire")
+    self.expiry = opts.get("expiry")
     self.capture = opts.get("capture", False)
     self.fn_hash_from = opts.get("fn_hash_from")
 
@@ -129,7 +129,7 @@ class CachedFunction(Generic[Fn]):
   def __init__(self, checkpointer: Checkpointer, fn: Fn):
     store_format = checkpointer.storage
     Storage = STORAGE_MAP[store_format] if isinstance(store_format, str) else store_format
-    update_wrapper(cast(Callable, self), unwrap(fn))
+    update_wrapper(self, unwrap(fn))  # type: ignore
     self.ident = FunctionIdent(self, checkpointer, fn)
     self.storage = Storage(self)
     self.bound = ()
@@ -152,13 +152,13 @@ class CachedFunction(Generic[Fn]):
 
   @property
   def fn(self) -> Fn:
-    return cast(Fn, self.ident.fn)
+    return self.ident.fn  # type: ignore
 
   @property
   def cleanup(self):
     return self.storage.cleanup
 
-  def reinit(self, recursive=False) -> CachedFunction[Fn]:
+  def reinit(self, recursive=True) -> CachedFunction[Fn]:
     depend_idents = list(self.ident.deep_idents()) if recursive else [self.ident]
     for ident in depend_idents: ident.reset()
     for ident in depend_idents: ident.fn_hash
@@ -178,12 +178,10 @@ class CachedFunction(Generic[Fn]):
       elif key == b"**":
         for key in kw.keys() - ident.arg_names:
           named_args[key] = hash_by(named_args[key])
-    named_args_iter = chain.from_iterable(sorted(named_args.items()))
-    captured = chain.from_iterable(capturable.capture() for capturable in ident.capturables)
     call_hash = ObjectHash(digest_size=16) \
-      .update(iter=named_args_iter, header="NAMED") \
-      .update(iter=pos_args, header="POS") \
-      .update(iter=captured, header="CAPTURED")
+      .update(header="NAMED", iter=flatten(sorted(named_args.items()))) \
+      .update(header="POS", iter=pos_args) \
+      .update(header="CAPTURED", iter=flatten(c.capture() for c in ident.capturables))
     return str(call_hash)
 
   def get_call_hash(self: CachedFunction[Callable[P, R]], *args: P.args, **kw: P.kwargs) -> str:
@@ -201,9 +199,7 @@ class CachedFunction(Generic[Fn]):
 
     call_hash = self._get_call_hash(args, kw)
     call_id = f"{storage.fn_id()}/{call_hash}"
-    refresh = rerun \
-      or not storage.exists(call_hash) \
-      or (params.should_expire and params.should_expire(storage.checkpoint_date(call_hash)))
+    refresh = rerun or not storage.exists(call_hash) or storage.expired(call_hash)
 
     if refresh:
       print_checkpoint(params.verbosity >= 1, "MEMORIZING", call_id, "blue")

{checkpointer-2.13.0 → checkpointer-2.14.0}/checkpointer/fn_ident.py

@@ -86,14 +86,14 @@ def extract_scope_values(code: CodeType, scope_vars: dict) -> Iterable[tuple[Att
       next_scope_vars = {**scope_vars, "LOAD_FAST": {}, "LOAD_DEREF": next_deref}
       yield from extract_scope_values(const, next_scope_vars)
 
-def resolve_class_annotations(anno: object) -> Type | None:
+def class_from_annotation(anno: object) -> Type | None:
   if anno in (None, Annotated):
     return None
-  elif is_class(anno):
+  if is_class(anno):
     return anno
-  elif get_origin(anno) is Annotated:
-    return resolve_class_annotations(next(iter(get_args(anno)), None))
-  return resolve_class_annotations(get_origin(anno))
+  if get_origin(anno) is Annotated:
+    return class_from_annotation(next(iter(get_args(anno)), None))
+  return class_from_annotation(get_origin(anno))
 
 def get_self_value(fn: Callable) -> Type | object | None:
   if isinstance(fn, MethodType):
@@ -107,9 +107,9 @@ def get_capturables(fn: Callable, capture: bool, captured_vars: dict[AttrPath, o
   module = getmodule(fn)
   if not module or not is_user_fn(fn):
     return
-  for (instruct, *attr_path), obj in captured_vars.items():
+  for (instruct_type, *attr_path), obj in captured_vars.items():
     attr_path = AttrPath(attr_path)
-    if instruct == "LOAD_GLOBAL" and not callable(obj) and not isinstance(obj, ModuleType):
+    if instruct_type == "LOAD_GLOBAL" and not callable(obj) and not isinstance(obj, ModuleType):
       anno = resolve_annotation(module, ".".join(attr_path))
       if capture or is_capture_me(anno) or is_capture_me_once(anno):
         hash_by = hash_by_from_annotation(anno)
@@ -122,7 +122,7 @@ def get_fn_captures(fn: Callable, capture: bool) -> tuple[list[Callable], list[C
     for param in signature(fn).parameters.values()
     if param.annotation is not Parameter.empty
     if param.kind not in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD)
-    if (class_anno := resolve_class_annotations(param.annotation))
+    if (class_anno := class_from_annotation(param.annotation))
   }
   if self_obj := get_self_value(fn):
     scope_vars_signature["self"] = self_obj
@@ -152,10 +152,10 @@ def get_depend_fns(fn: Callable, capture: bool, capturable_by_fn: CapturableByFn
 def get_fn_ident(fn: Callable, capture: bool) -> RawFunctionIdent:
   from .checkpoint import CachedFunction
   capturable_by_fn = get_depend_fns(fn, capture)
-  capturables = {capt for capts in capturable_by_fn.values() for capt in capts}
   depends = capturable_by_fn.keys()
   depends = distinct(fn.__func__ if isinstance(fn, MethodType) else fn for fn in depends)
   depend_callables = [fn for fn in depends if not isinstance(fn, CachedFunction)]
-  assert fn == depend_callables[0]
   fn_hash = str(ObjectHash(iter=map(get_fn_aststr, depend_callables)))
+  capturables = {capt for capts in capturable_by_fn.values() for capt in capts}
+  assert fn == depend_callables[0]
   return RawFunctionIdent(fn_hash, depends, capturables)

{checkpointer-2.13.0 → checkpointer-2.14.0}/checkpointer/fn_string.py

@@ -15,7 +15,7 @@ def get_decorator_path(node: ast.AST) -> tuple[str, ...]:
   else:
     return ()
 
-def is_empty_expression(node: ast.AST) -> bool:
+def is_lone_expression(node: ast.AST) -> bool:
   # Filter out docstrings
   return isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant)
 
@@ -26,7 +26,8 @@ class CleanFunctionTransform(ast.NodeTransformer):
 
   def is_checkpointer(self, node: ast.AST) -> bool:
     from .checkpoint import Checkpointer
-    return isinstance(get_at(self.fn_globals, *get_decorator_path(node)), Checkpointer)
+    decorator = get_at(self.fn_globals, *get_decorator_path(node))
+    return isinstance(decorator, Checkpointer) or decorator is Checkpointer
 
   def visit_FunctionDef(self, node: ast.FunctionDef | ast.AsyncFunctionDef):
     fn_type = type(node).__name__
@@ -45,7 +46,7 @@ class CleanFunctionTransform(ast.NodeTransformer):
     return ast.List([
       ast.Constant(header),
       ast.List([child for child in node.decorator_list if not self.is_checkpointer(child)], ast.Load()),
-      ast.List([self.visit(child) for child in node.body if not is_empty_expression(child)], ast.Load()),
+      ast.List([self.visit(child) for child in node.body if not is_lone_expression(child)], ast.Load()),
     ], ast.Load())
 
   def visit_AsyncFunctionDef(self, node):
@@ -66,10 +67,7 @@ def get_fn_aststr(fn: Callable) -> str:
   if fn.__name__ != "<lambda>":
     tree = CleanFunctionTransform(fn.__globals__).visit(tree)
   else:
-    for node in ast.walk(tree):
-      if isinstance(node, ast.Lambda):
-        tree = node
-        break
+    tree = ast.List([node for node in ast.walk(tree) if isinstance(node, ast.Lambda)], ast.Load())
 
   if sys.version_info >= (3, 13):
     return ast.dump(tree, annotate_fields=False, show_empty=True)

{checkpointer-2.13.0 → checkpointer-2.14.0}/checkpointer/object_hash.py

@@ -5,16 +5,19 @@ import io
 import re
 import sys
 import tokenize
+import sysconfig
 from collections import OrderedDict
 from collections.abc import Iterable
 from contextlib import nullcontext, suppress
 from decimal import Decimal
 from io import StringIO
+from inspect import getfile
 from itertools import chain
+from pathlib import Path
 from pickle import HIGHEST_PROTOCOL as PICKLE_PROTOCOL
 from types import BuiltinFunctionType, FunctionType, GeneratorType, MappingProxyType, MethodType, ModuleType, UnionType
 from typing import Callable, Self, TypeVar
-from .utils import ContextVar
+from .utils import ContextVar, flatten
 
 np, torch = None, None
 
@@ -31,8 +34,8 @@ if sys.version_info >= (3, 12):
 else:
   TypeAliasType = _Never
 
-flatten = chain.from_iterable
 nc = nullcontext()
+stdlib = Path(sysconfig.get_paths()["stdlib"]).resolve()
 
 def encode_type(t: type | FunctionType) -> str:
   return f"{t.__module__}:{t.__qualname__}"
@@ -128,7 +131,11 @@ class ObjectHash:
         self.header("builtin", obj.__qualname__)
 
       case FunctionType():
-        self.header("function", encode_type(obj)).update(get_fn_body(obj), obj.__defaults__, obj.__kwdefaults__, obj.__annotations__)
+        fn_file = Path(getfile(obj)).resolve()
+        if fn_file.is_relative_to(stdlib):
+          self.header("function-std", obj.__qualname__)
+        else:
+          self.header("function", encode_type(obj)).update(get_fn_body(obj), obj.__defaults__, obj.__kwdefaults__, obj.__annotations__)
 
       case MethodType():
         self.header("method").update(obj.__func__, obj.__self__.__class__)

{checkpointer-2.13.0 → checkpointer-2.14.0}/checkpointer/storages/memory_storage.py

@@ -31,7 +31,7 @@ class MemoryStorage(Storage):
       if key.parent == curr_key.parent:
         if invalidated and key != curr_key:
           del item_map[key]
-        elif expired and self.checkpointer.should_expire:
+        elif expired and self.checkpointer.expiry:
           for call_hash, (date, _) in list(calldict.items()):
-            if self.checkpointer.should_expire(date):
+            if self.expired_dt(date):
               del calldict[call_hash]

{checkpointer-2.13.0 → checkpointer-2.14.0}/checkpointer/storages/pickle_storage.py

@@ -9,7 +9,7 @@ def filedate(path: Path) -> datetime:
 
 class PickleStorage(Storage):
   def get_path(self, call_hash: str):
-    return self.fn_dir() / f"{call_hash}.pkl"
+    return self.fn_dir() / f"{call_hash[:2]}/{call_hash[2:]}.pkl"
 
   def store(self, call_hash, data):
     path = self.get_path(call_hash)
@@ -40,10 +40,10 @@ class PickleStorage(Storage):
       for path in old_dirs:
         shutil.rmtree(path)
       print(f"Removed {len(old_dirs)} invalidated directories for {self.cached_fn.__qualname__}")
-    if expired and self.checkpointer.should_expire:
+    if expired and self.checkpointer.expiry:
       count = 0
       for pkl_path in fn_path.glob("**/*.pkl"):
-        if self.checkpointer.should_expire(filedate(pkl_path)):
+        if self.expired_dt(filedate(pkl_path)):
           count += 1
           pkl_path.unlink(missing_ok=True)
       print(f"Removed {count} expired checkpoints for {self.cached_fn.__qualname__}")

{checkpointer-2.13.0 → checkpointer-2.14.0}/checkpointer/storages/storage.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from typing import Any, TYPE_CHECKING
 from pathlib import Path
-from datetime import datetime
+from datetime import datetime, timedelta
 
 if TYPE_CHECKING:
   from ..checkpoint import Checkpointer, CachedFunction
@@ -21,6 +21,19 @@ class Storage:
   def fn_dir(self) -> Path:
     return self.checkpointer.directory / self.fn_id()
 
+  def expired(self, call_hash: str) -> bool:
+    if not self.checkpointer.expiry:
+      return False
+    return self.expired_dt(self.checkpoint_date(call_hash))
+
+  def expired_dt(self, dt: datetime) -> bool:
+    expiry = self.checkpointer.expiry
+    if isinstance(expiry, timedelta):
+      return dt < datetime.now() - expiry
+    else:
+      if TYPE_CHECKING: assert expiry
+      return expiry(dt)
+
   def store(self, call_hash: str, data: Any) -> Any: ...
 
   def exists(self, call_hash: str) -> bool: ...

{checkpointer-2.13.0 → checkpointer-2.14.0}/checkpointer/utils.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 import inspect
 from contextlib import contextmanager, suppress
-from itertools import islice
+from itertools import chain, islice
 from pathlib import Path
 from types import FunctionType, MethodType, ModuleType
 from typing import Callable, Generic, Iterable, Self, Type, TypeGuard
 from .types import T
 
 cwd = Path.cwd().resolve()
+flatten = chain.from_iterable
 
 def is_class(obj) -> TypeGuard[Type]:
   return isinstance(obj, type)
@@ -22,7 +23,7 @@ def is_user_fn(obj) -> TypeGuard[Callable]:
   return isinstance(obj, (FunctionType, MethodType)) and is_user_file(get_file(obj))
 
 def get_cell_contents(fn: Callable) -> Iterable[tuple[str, object]]:
-  for key, cell in zip(fn.__code__.co_freevars, fn.__closure__ or []):
+  for key, cell in zip(fn.__code__.co_freevars, fn.__closure__ or ()):
     with suppress(ValueError):
       yield (key, cell.cell_contents)
 

{checkpointer-2.13.0 → checkpointer-2.14.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "checkpointer"
-version = "2.13.0"
+version = "2.14.0"
 requires-python = ">=3.11"
 dependencies = []
 authors = [

{checkpointer-2.13.0 → checkpointer-2.14.0}/uv.lock

@@ -8,7 +8,7 @@ resolution-markers = [
 
 [[package]]
 name = "checkpointer"
-version = "2.13.0"
+version = "2.14.0"
 source = { editable = "." }
 
 [package.dev-dependencies]