checkpointer 2.11.2__py3-none-any.whl → 2.13.0__py3-none-any.whl

checkpointer/fn_string.py

@@ -0,0 +1,77 @@
+import ast
+import sys
+from inspect import getsource
+from textwrap import dedent
+from typing import Callable
+from .utils import drop_none, get_at
+
+def get_decorator_path(node: ast.AST) -> tuple[str, ...]:
+  if isinstance(node, ast.Call):
+    return get_decorator_path(node.func)
+  elif isinstance(node, ast.Attribute):
+    return get_decorator_path(node.value) + (node.attr,)
+  elif isinstance(node, ast.Name):
+    return (node.id,)
+  else:
+    return ()
+
+def is_empty_expression(node: ast.AST) -> bool:
+  # Filter out docstrings
+  return isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant)
+
+class CleanFunctionTransform(ast.NodeTransformer):
+  def __init__(self, fn_globals: dict):
+    self.is_root = True
+    self.fn_globals = fn_globals
+
+  def is_checkpointer(self, node: ast.AST) -> bool:
+    from .checkpoint import Checkpointer
+    return isinstance(get_at(self.fn_globals, *get_decorator_path(node)), Checkpointer)
+
+  def visit_FunctionDef(self, node: ast.FunctionDef | ast.AsyncFunctionDef):
+    fn_type = type(node).__name__
+    fn_name = None if self.is_root else node.name
+    args_by_type = [
+      node.args.posonlyargs + node.args.args,
+      drop_none([node.args.vararg]),
+      sorted(node.args.kwonlyargs, key=lambda x: x.arg),
+      drop_none([node.args.kwarg]),
+    ]
+    arg_kind_names = ",".join(f"{i}:{arg.arg}" for i, args in enumerate(args_by_type) for arg in args)
+    header = " ".join(drop_none((fn_type, fn_name, arg_kind_names or None)))
+
+    self.is_root = False
+
+    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.Load())
+
+  def visit_AsyncFunctionDef(self, node):
+    return self.visit_FunctionDef(node)
+
+def get_fn_aststr(fn: Callable) -> str:
+  try:
+    source = getsource(fn)
+  except OSError:
+    return ""
+  try:
+    tree = ast.parse(dedent(source), mode="exec")
+    tree = tree.body[0]
+  except SyntaxError:
+    # lambda functions can cause SyntaxError in ast.parse
+    return source.strip()
+
+  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
+
+  if sys.version_info >= (3, 13):
+    return ast.dump(tree, annotate_fields=False, show_empty=True)
+  else:
+    return ast.dump(tree, annotate_fields=False)

checkpointer/import_mappings.py

@@ -0,0 +1,47 @@
+import ast
+import inspect
+import sys
+from types import ModuleType
+from typing import Iterable, Type
+from .utils import cwd, get_file, is_user_file
+
+ImportTarget = tuple[str, str | None]
+
+cache: dict[tuple[str, int], dict[str, ImportTarget]] = {}
+
+def generate_import_mappings(module: ModuleType) -> Iterable[tuple[str, ImportTarget]]:
+  mod_path = get_file(module)
+  if not is_user_file(mod_path):
+    return
+  mod_parts = list(mod_path.with_suffix("").relative_to(cwd).parts)
+  source = inspect.getsource(module)
+  tree = ast.parse(source)
+  for node in ast.walk(tree):
+    if isinstance(node, ast.Import):
+      for alias in node.names:
+        yield (alias.asname or alias.name, (alias.name, None))
+    elif isinstance(node, ast.ImportFrom):
+      target_mod = node.module or ""
+      if node.level > 0:
+        target_mod_parts = target_mod.split(".") * bool(target_mod)
+        target_mod_parts = mod_parts[:-node.level] + target_mod_parts
+        target_mod = ".".join(target_mod_parts)
+      for alias in node.names:
+        yield (alias.asname or alias.name, (target_mod, alias.name))
+
+def get_import_mappings(module: ModuleType) -> dict[str, ImportTarget]:
+  cache_key = (module.__name__, id(module))
+  if cached := cache.get(cache_key):
+    return cached
+  import_mappings = dict(generate_import_mappings(module))
+  return cache.setdefault(cache_key, import_mappings)
+
+def resolve_annotation(module: ModuleType, attr_name: str | None) -> Type | None:
+  if not attr_name:
+    return None
+  if anno := module.__annotations__.get(attr_name):
+    return anno
+  if next_pair := get_import_mappings(module).get(attr_name):
+    next_module_name, next_attr_name = next_pair
+    if next_module := sys.modules.get(next_module_name):
+      return resolve_annotation(next_module, next_attr_name)

checkpointer/object_hash.py

@@ -5,32 +5,35 @@ import io
 import re
 import sys
 import tokenize
+from collections import OrderedDict
 from collections.abc import Iterable
 from contextlib import nullcontext, suppress
 from decimal import Decimal
 from io import StringIO
 from itertools import chain
 from pickle import HIGHEST_PROTOCOL as PICKLE_PROTOCOL
-from types import BuiltinFunctionType, FunctionType, GeneratorType, MethodType, ModuleType, UnionType
-from typing import Callable, TypeVar
+from types import BuiltinFunctionType, FunctionType, GeneratorType, MappingProxyType, MethodType, ModuleType, UnionType
+from typing import Callable, Self, TypeVar
 from .utils import ContextVar
 
 np, torch = None, None
 
-with suppress(Exception):
-  import numpy as np
-with suppress(Exception):
-  import torch
-
 class _Never:
   def __getattribute__(self, _: str):
     pass
 
+with suppress(Exception):
+  import numpy as np
+with suppress(Exception):
+  import torch
 if sys.version_info >= (3, 12):
   from typing import TypeAliasType
 else:
   TypeAliasType = _Never
 
+flatten = chain.from_iterable
+nc = nullcontext()
+
 def encode_type(t: type | FunctionType) -> str:
   return f"{t.__module__}:{t.__qualname__}"
 
@@ -43,14 +46,14 @@ class ObjectHashError(Exception):
     self.obj = obj
 
 class ObjectHash:
-  def __init__(self, *objs: object, iter: Iterable[object] = (), digest_size=64, tolerate_errors=False) -> None:
+  def __init__(self, *objs: object, iter: Iterable[object] = (), digest_size=64, tolerable=False) -> None:
     self.hash = hashlib.blake2b(digest_size=digest_size)
     self.current: dict[int, int] = {}
-    self.tolerate_errors = ContextVar(tolerate_errors)
+    self.tolerable = ContextVar(tolerable)
     self.update(iter=chain(objs, iter))
 
   def copy(self) -> "ObjectHash":
-    new = ObjectHash(tolerate_errors=self.tolerate_errors.value)
+    new = ObjectHash(tolerable=self.tolerable.value)
     new.hash = self.hash.copy()
     return new
 
@@ -63,26 +66,29 @@ class ObjectHash:
     return isinstance(value, ObjectHash) and str(self) == str(value)
 
   def nested_hash(self, *objs: object) -> str:
-    return ObjectHash(iter=objs, tolerate_errors=self.tolerate_errors.value).hexdigest()
+    return ObjectHash(iter=objs, tolerable=self.tolerable.value).hexdigest()
 
-  def write_bytes(self, *data: bytes, iter: Iterable[bytes] = ()) -> "ObjectHash":
+  def write_bytes(self, *data: bytes, iter: Iterable[bytes] = ()) -> Self:
     for d in chain(data, iter):
       self.hash.update(d)
     return self
 
-  def write_text(self, *data: str, iter: Iterable[str] = ()) -> "ObjectHash":
+  def write_text(self, *data: str, iter: Iterable[str] = ()) -> Self:
     return self.write_bytes(iter=(d.encode() for d in chain(data, iter)))
 
-  def header(self, *args: object) -> "ObjectHash":
+  def header(self, *args: object) -> Self:
     return self.write_bytes(":".join(map(str, args)).encode())
 
-  def update(self, *objs: object, iter: Iterable[object] = (), tolerate_errors: bool | None=None) -> "ObjectHash":
-    with nullcontext() if tolerate_errors is None else self.tolerate_errors.set(tolerate_errors):
+  def update(self, *objs: object, iter: Iterable[object] = (), tolerable: bool | None=None, header: str | None = None) -> Self:
+    with nc if tolerable is None else self.tolerable.set(tolerable):
       for obj in chain(objs, iter):
+        if header is not None:
+          self.write_bytes(header.encode())
+          header = None
         try:
           self._update_one(obj)
         except Exception as ex:
-          if self.tolerate_errors.value:
+          if self.tolerable.value:
             self.header("error").update(type(ex))
           else:
             raise ObjectHashError(obj, ex) from ex
@@ -102,11 +108,11 @@ class ObjectHash:
 
       case set() | frozenset():
         try:
-          items = sorted(obj)
           header = "set"
+          items = sorted(obj)
         except:
-          items = sorted(map(self.nested_hash, obj))
           header = "set-unsortable"
+          items = sorted(map(self.nested_hash, obj))
         self.header(header, encode_type_of(obj), len(obj)).update(iter=items)
 
       case TypeVar():
@@ -167,23 +173,25 @@ class ObjectHash:
           match obj:
             case list() | tuple():
               self.header("list", encode_type_of(obj), len(obj)).update(iter=obj)
-            case dict():
-              try:
-                items = sorted(obj.items())
-                header = "dict"
-              except:
-                items = sorted((self.nested_hash(key), val) for key, val in obj.items())
-                header = "dict-unsortable"
-              self.header(header, encode_type_of(obj), len(obj)).update(iter=chain.from_iterable(items))
+            case dict() | MappingProxyType():
+              header = "dict"
+              items = obj.items()
+              if not isinstance(obj, OrderedDict):
+                try:
+                  items = sorted(items)
+                except:
+                  header = "dict-unsortable"
+                  items = sorted((self.nested_hash(key), val) for key, val in items)
+              self.header(header, encode_type_of(obj), len(obj)).update(iter=flatten(items))
             case _:
               self._update_object(obj)
         finally:
           del self.current[id(obj)]
 
-  def _update_iterator(self, obj: Iterable) -> "ObjectHash":
+  def _update_iterator(self, obj: Iterable) -> Self:
     return self.header("iterator", encode_type_of(obj)).update(iter=obj).header("iterator-end")
 
-  def _update_object(self, obj: object) -> "ObjectHash":
+  def _update_object(self, obj: object) -> Self:
     self.header("instance", encode_type_of(obj))
     get_hash = hasattr(obj, "__objecthash__") and getattr(obj, "__objecthash__")
     if callable(get_hash):

checkpointer/storages/__init__.py

@@ -1,9 +1,11 @@
-from typing import Type
-from .storage import Storage
-from .pickle_storage import PickleStorage
+from typing import Literal, Type
 from .memory_storage import MemoryStorage
+from .pickle_storage import PickleStorage
+from .storage import Storage
+
+StorageType = Literal["pickle", "memory"]
 
-STORAGE_MAP: dict[str, Type[Storage]] = {
+STORAGE_MAP: dict[StorageType, Type[Storage]] = {
   "pickle": PickleStorage,
   "memory": MemoryStorage,
 }

checkpointer/storages/storage.py

@@ -8,17 +8,18 @@ if TYPE_CHECKING:
 
 class Storage:
   checkpointer: Checkpointer
-  cached_fn: CachedFunction
+  ident: CachedFunction
 
   def __init__(self, cached_fn: CachedFunction):
-    self.checkpointer = cached_fn.checkpointer
+    self.checkpointer = cached_fn.ident.checkpointer
     self.cached_fn = cached_fn
 
   def fn_id(self) -> str:
-    return f"{self.cached_fn.fn_dir}/{self.cached_fn.ident.fn_hash}"
+    ident = self.cached_fn.ident
+    return f"{ident.fn_dir}/{ident.fn_hash}"
 
   def fn_dir(self) -> Path:
-    return self.checkpointer.root_path / self.fn_id()
+    return self.checkpointer.directory / self.fn_id()
 
   def store(self, call_hash: str, data: Any) -> Any: ...
 

checkpointer/types.py

@@ -1,4 +1,7 @@
-from typing import Annotated, Callable, Coroutine, Generic, ParamSpec, TypeVar
+from typing import (
+  Annotated, Callable, Coroutine, Generic,
+  ParamSpec, TypeVar, get_args, get_origin,
+)
 
 Fn = TypeVar("Fn", bound=Callable)
 P = ParamSpec("P")
@@ -9,7 +12,32 @@ T = TypeVar("T")
 class HashBy(Generic[Fn]):
   pass
 
-NoHash = Annotated[T, HashBy[lambda _: None]]
+class Captured:
+  pass
+
+class CapturedOnce:
+  pass
+
+def to_none(_):
+  return None
+
+def get_annotated_args(anno: object) -> tuple[object, ...]:
+  return get_args(anno) if get_origin(anno) is Annotated else ()
+
+def hash_by_from_annotation(anno: object) -> Callable[[object], object] | None:
+  for arg in get_annotated_args(anno):
+    if get_origin(arg) is HashBy:
+      return get_args(arg)[0]
+
+def is_capture_me(anno: object) -> bool:
+  return Captured in get_annotated_args(anno)
+
+def is_capture_me_once(anno: object) -> bool:
+  return CapturedOnce in get_annotated_args(anno)
+
+NoHash = Annotated[T, HashBy[to_none]]
+CaptureMe = Annotated[T, Captured]
+CaptureMeOnce = Annotated[T, CapturedOnce]
 Coro = Coroutine[object, object, R]
 
 class AwaitableValue(Generic[T]):

checkpointer/utils.py

@@ -1,23 +1,93 @@
-from contextlib import contextmanager
-from typing import Callable, Generic, Iterable, cast
-from .types import Fn, T
+from __future__ import annotations
+import inspect
+from contextlib import contextmanager, suppress
+from itertools import islice
+from pathlib import Path
+from types import FunctionType, MethodType, ModuleType
+from typing import Callable, Generic, Iterable, Self, Type, TypeGuard
+from .types import T
 
-def distinct(seq: Iterable[T]) -> list[T]:
-  return list(dict.fromkeys(seq))
+cwd = Path.cwd().resolve()
+
+def is_class(obj) -> TypeGuard[Type]:
+  return isinstance(obj, type)
+
+def get_file(obj: Callable | ModuleType) -> Path:
+  return Path(inspect.getfile(obj)).resolve()
+
+def is_user_file(path: Path) -> bool:
+  return cwd in path.parents and ".venv" not in path.parts
+
+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 []):
-    try:
+    with suppress(ValueError):
       yield (key, cell.cell_contents)
-    except ValueError:
-      pass
 
-def unwrap_fn(fn: Fn, cached_fn=False) -> Fn:
-  from .checkpoint import CachedFunction
-  while True:
-    if (cached_fn and isinstance(fn, CachedFunction)) or not hasattr(fn, "__wrapped__"):
-      return cast(Fn, fn)
-    fn = getattr(fn, "__wrapped__")
+def drop_none(iterable: Iterable[T | None]) -> list[T]:
+  return [x for x in iterable if x is not None]
+
+def distinct(seq: Iterable[T]) -> list[T]:
+  return list(dict.fromkeys(seq))
+
+def takewhile(iter: Iterable[tuple[bool, T]]) -> Iterable[T]:
+  for condition, value in iter:
+    if not condition:
+      return
+    yield value
+
+class seekable(Generic[T]):
+  def __init__(self, iterable: Iterable[T]):
+    self.index = 0
+    self.source = iter(iterable)
+    self.sink: list[T] = []
+
+  def __iter__(self):
+    return self
+
+  def __next__(self) -> T:
+    if len(self.sink) > self.index:
+      item = self.sink[self.index]
+    else:
+      item = next(self.source)
+      self.sink.append(item)
+    self.index += 1
+    return item
+
+  def __bool__(self):
+    return bool(self.lookahead(1))
+
+  def seek(self, index: int) -> Self:
+    remainder = index - len(self.sink)
+    if remainder > 0:
+      next(islice(self, remainder, remainder), None)
+    self.index = max(0, min(index, len(self.sink)))
+    return self
+
+  def step(self, count: int) -> Self:
+    return self.seek(self.index + count)
+
+  @contextmanager
+  def freeze(self):
+    initial_index = self.index
+    try:
+      yield
+    finally:
+      self.seek(initial_index)
+
+  def lookahead(self, count: int) -> list[T]:
+    with self.freeze():
+      return list(islice(self, count))
+
+def get_at(obj: object, *attrs: str) -> object:
+  for attr in attrs:
+    if type(obj) is dict:
+      obj = obj.get(attr, None)
+    else:
+      obj = getattr(obj, attr, None)
+  return obj
 
 class AttrDict(dict):
   def __init__(self, *args, **kwargs):
@@ -30,23 +100,6 @@ class AttrDict(dict):
   def __setattr__(self, name: str, value: object):
     super().__setattr__(name, value)
 
-  def set(self, d: dict) -> "AttrDict":
-    if not d:
-      return self
-    return AttrDict({**self, **d})
-
-  def delete(self, *attrs: str) -> "AttrDict":
-    d = AttrDict(self)
-    for attr in attrs:
-      del d[attr]
-    return d
-
-  def get_at(self, attrs: tuple[str, ...]) -> object:
-    d = self
-    for attr in attrs:
-      d = getattr(d, attr, None)
-    return d
-
 class ContextVar(Generic[T]):
   def __init__(self, value: T):
     self.value = value
@@ -58,26 +111,3 @@ class ContextVar(Generic[T]):
       yield
     finally:
       self.value = old
-
-class iterate_and_upcoming(Generic[T]):
-  def __init__(self, it: Iterable[T]) -> None:
-    self.it = iter(it)
-    self.previous: tuple[()] | tuple[T] = ()
-    self.tracked = self._tracked_iter()
-
-  def __iter__(self):
-    return self
-
-  def __next__(self) -> tuple[T, Iterable[T]]:
-    try:
-      item = self.previous[0] if self.previous else next(self.it)
-      self.previous = ()
-      return item, self.tracked
-    except StopIteration:
-      self.tracked.close()
-      raise
-
-  def _tracked_iter(self):
-    for x in self.it:
-      self.previous = (x,)
-      yield x