checkpointer 1.2.0__tar.gz

checkpointer-1.2.0/.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+/dist/
+.DS_Store

checkpointer-1.2.0/LICENSE

@@ -0,0 +1,7 @@
+Copyright 2024 Hampus Hallman
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

checkpointer-1.2.0/PKG-INFO

@@ -0,0 +1,16 @@
+Metadata-Version: 2.3
+Name: checkpointer
+Version: 1.2.0
+Project-URL: Repository, https://github.com/Reddan/checkpointer.git
+Author: Hampus Hallman
+License: Copyright 2024 Hampus Hallman
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+License-File: LICENSE
+Requires-Python: >=3.8
+Requires-Dist: relib
+Requires-Dist: termcolor

checkpointer-1.2.0/checkpointer/__init__.py

@@ -0,0 +1,10 @@
+import os
+
+from .checkpoint import create_checkpointer, read_only, default_dir
+from .storage import store_on_demand, read_from_store
+from .function_body import get_function_hash
+
+storage_dir = os.environ.get('CHECKPOINTS_DIR', default_dir)
+verbosity = int(os.environ.get('CHECKPOINTS_VERBOSITY', '1'))
+
+checkpoint = create_checkpointer(root_path=storage_dir, verbosity=verbosity)

checkpointer-1.2.0/checkpointer/checkpoint.py

@@ -0,0 +1,68 @@
+import inspect
+from collections import namedtuple
+from pathlib import Path
+from functools import wraps
+import relib.hashing as hashing
+from . import storage
+from .function_body import get_function_hash
+from .utils import unwrap_func, sync_resolve_coroutine
+
+default_dir = Path.home() / '.checkpoints'
+
+def get_invoke_path(func, function_hash, args, kwargs, path):
+  if type(path) == str:
+    return path
+  elif callable(path):
+    return path(*args, **kwargs)
+  else:
+    hash = hashing.hash([function_hash, args, kwargs or 0])
+    file_name = Path(func.__code__.co_filename).name
+    name = func.__name__
+    return file_name + '/' + name + '/' + hash
+
+def create_checkpointer_from_config(config):
+  def checkpoint(opt_func=None, format=config.format, path=None, should_expire=None, when=True):
+    def receive_func(func):
+      if not (config.when and when):
+        return func
+
+      config_ = config._replace(format=format)
+      is_async = inspect.iscoroutinefunction(func)
+      unwrapped_func = unwrap_func(func)
+      function_hash = get_function_hash(unwrapped_func)
+
+      @wraps(unwrapped_func)
+      def wrapper(*args, **kwargs):
+        compute = lambda: func(*args, **kwargs)
+        recheck = kwargs.pop('recheck', False)
+        invoke_path = get_invoke_path(unwrapped_func, function_hash, args, kwargs, path)
+        coroutine = storage.store_on_demand(compute, invoke_path, config_, recheck, should_expire)
+        if is_async:
+          return coroutine
+        else:
+          return sync_resolve_coroutine(coroutine)
+
+      wrapper.checkpoint_config = config_
+
+      return wrapper
+
+    return receive_func(opt_func) if callable(opt_func) else receive_func
+
+  return checkpoint
+
+def create_checkpointer(format='pickle', root_path=default_dir, when=True, verbosity=1):
+  root_path = None if root_path is None else Path(root_path)
+  opts = locals()
+  CheckpointerConfig = namedtuple('CheckpointerConfig', sorted(opts))
+  config = CheckpointerConfig(**opts)
+  return create_checkpointer_from_config(config)
+
+def read_only(wrapper_func, config, format='pickle', path=None):
+  func = unwrap_func(wrapper_func)
+  function_hash = get_function_hash(func)
+
+  def wrapper(*args, **kwargs):
+    invoke_path = get_invoke_path(func, function_hash, args, kwargs, path)
+    return storage.read_from_store(invoke_path, config, storage=format)
+
+  return wrapper

checkpointer-1.2.0/checkpointer/function_body.py

@@ -0,0 +1,49 @@
+import inspect
+from types import FunctionType, CodeType
+import relib.hashing as hashing
+from pathlib import Path
+from .utils import unwrap_func
+
+cwd = Path.cwd()
+
+def get_fn_path(fn):
+  return Path(inspect.getfile(fn)).absolute()
+
+def get_function_body(fn):
+  # TODO: Strip comments
+  lines = inspect.getsourcelines(fn)[0]
+  lines = [line.rstrip() for line in lines]
+  lines = [line for line in lines if line]
+  return '\n'.join(lines)
+
+def get_code_children(__code__):
+  consts = [const for const in __code__.co_consts if isinstance(const, CodeType)]
+  children = [child for const in consts for child in get_code_children(const)]
+  return list(__code__.co_names) + children
+
+def is_user_fn(candidate_fn, cleared_fns):
+  return isinstance(candidate_fn, FunctionType) \
+    and candidate_fn not in cleared_fns \
+    and cwd in get_fn_path(candidate_fn).parents
+
+def append_fn_children(fn, cleared_fns):
+  code_children = get_code_children(fn.__code__)
+  fn_children = [unwrap_func(fn.__globals__.get(co_name, None)) for co_name in code_children]
+  fn_children = [child for child in fn_children if is_user_fn(child, cleared_fns)]
+
+  for fn in fn_children:
+    cleared_fns.add(fn)
+
+  for child_fn in fn_children:
+    append_fn_children(child_fn, cleared_fns)
+
+def get_fn_children(fn):
+  cleared_fns = set()
+  append_fn_children(fn, cleared_fns)
+  return sorted(cleared_fns, key=lambda fn: fn.__name__)
+
+def get_function_hash(fn):
+  fns = [fn] + get_fn_children(fn)
+  fn_bodies = list(map(get_function_body, fns))
+  fn_bodies_hash = hashing.hash(fn_bodies)
+  return fn_bodies_hash

checkpointer-1.2.0/checkpointer/storage.py

@@ -0,0 +1,60 @@
+import inspect
+from termcolor import colored
+from .storages import memory_storage, pickle_storage, bcolz_storage
+
+storages = {
+  'memory': memory_storage,
+  'pickle': pickle_storage,
+  'bcolz': bcolz_storage,
+}
+
+initialized_storages = set()
+
+def create_logger(should_log):
+  def log(color, title, text):
+    if should_log:
+      title_log = colored(f' {title} ', 'grey', 'on_' + color)
+      rest_log = colored(text, color)
+      print(title_log + ' ' + rest_log)
+  return log
+
+def get_storage(storage):
+  if type(storage) == str:
+    storage = storages[storage]
+  if storage not in initialized_storages:
+    if hasattr(storage, 'initialize'):
+      storage.initialize()
+    initialized_storages.add(storage)
+  return storage
+
+async def store_on_demand(get_data, name, config, force=False, should_expire=None):
+  storage = get_storage(config.format)
+  should_log = storage != memory_storage and config.verbosity != 0
+  log = create_logger(should_log)
+  refresh = force \
+    or storage.get_is_expired(config, name) \
+    or (should_expire and storage.should_expire(config, name, should_expire))
+
+  if refresh:
+    log('blue', 'MEMORIZING', name)
+    data = get_data()
+    if inspect.iscoroutine(data):
+      data = await data
+    return storage.store_data(config, name, data)
+  else:
+    try:
+      data = storage.load_data(config, name)
+      log('green', 'REMEMBERED', name)
+      return data
+    except (EOFError, FileNotFoundError):
+      log('yellow', 'CORRUPTED', name)
+      storage.delete_data(config, name)
+      result = await store_on_demand(get_data, name, config, force, should_expire)
+      return result
+
+def read_from_store(name, config, storage='pickle'):
+  storage = get_storage(storage)
+  try:
+    return storage.load_data(config, name)
+  except:
+    return None

checkpointer-1.2.0/checkpointer/storages/bcolz_storage.py

@@ -0,0 +1,92 @@
+import shutil
+from pathlib import Path
+from datetime import datetime
+
+def get_data_type_str(x):
+  if isinstance(x, tuple):
+    return 'tuple'
+  elif isinstance(x, dict):
+    return 'dict'
+  elif isinstance(x, list):
+    return 'list'
+  elif isinstance(x, str) or not hasattr(x, '__len__'):
+    return 'other'
+  else:
+    return 'ndarray'
+
+def get_paths(root_path, invoke_path):
+  full_path = Path(invoke_path) if root_path is None else root_path / invoke_path
+  meta_full_path = full_path.with_name(full_path.name + '_meta')
+  return full_path, meta_full_path
+
+def get_collection_timestamp(config, path):
+  import bcolz
+  _, meta_full_path = get_paths(config.root_path, path)
+  meta_data = bcolz.open(meta_full_path)[:][0]
+  return meta_data['created']
+
+def get_is_expired(config, path):
+  try:
+    get_collection_timestamp(config, path)
+    return False
+  except (FileNotFoundError, EOFError):
+    return True
+
+def should_expire(config, path, expire_fn):
+  return expire_fn(get_collection_timestamp(config, path))
+
+def insert_data(path, data):
+  import bcolz
+  c = bcolz.carray(data, rootdir=path, mode='w')
+  c.flush()
+
+def store_data(config, path, data, expire_in=None):
+  full_path, meta_full_path = get_paths(config.root_path, path)
+  full_path.parent.mkdir(parents=True, exist_ok=True)
+  created = datetime.now()
+  data_type_str = get_data_type_str(data)
+  if data_type_str == 'tuple':
+    fields = list(range(len(data)))
+  elif data_type_str == 'dict':
+    fields = sorted(data.keys())
+  else:
+    fields = []
+  meta_data = {'created': created, 'data_type_str': data_type_str, 'fields': fields}
+  insert_data(meta_full_path, meta_data)
+  if data_type_str in ['tuple', 'dict']:
+    for i in range(len(fields)):
+      sub_path = f"{path} ({i})"
+      store_data(config, sub_path, data[fields[i]])
+  else:
+    insert_data(full_path, data)
+  return data
+
+def load_data(config, path):
+  import bcolz
+  full_path, meta_full_path = get_paths(config.root_path, path)
+  meta_data = bcolz.open(meta_full_path)[:][0]
+  data_type_str = meta_data['data_type_str']
+  if data_type_str in ['tuple', 'dict']:
+    fields = meta_data['fields']
+    partitions = range(len(fields))
+    data = [load_data(config, f"{path} ({i})") for i in partitions]
+    if data_type_str == 'tuple':
+      return tuple(data)
+    else:
+      return dict(zip(fields, data))
+  else:
+    data = bcolz.open(full_path)
+    if data_type_str == 'list':
+      return list(data)
+    elif data_type_str == 'other':
+      return data[0]
+    else:
+      return data[:]
+
+def delete_data(config, path):
+  full_path, meta_full_path = get_paths(config.root_path, path)
+  try:
+    shutil.rmtree(meta_full_path)
+    shutil.rmtree(full_path)
+  except FileNotFoundError:
+    pass

checkpointer-1.2.0/checkpointer/storages/memory_storage.py

@@ -0,0 +1,18 @@
+from datetime import datetime
+
+store = {}
+date_stored = {}
+
+def get_is_expired(config, path):
+  return path not in store
+
+def should_expire(config, path, expire_fn):
+  return expire_fn(date_stored[path])
+
+def store_data(config, path, data):
+  store[path] = data
+  date_stored[path] = datetime.now()
+  return data
+
+def load_data(config, path):
+  return store[path]

checkpointer-1.2.0/checkpointer/storages/pickle_storage.py

@@ -0,0 +1,49 @@
+import pickle
+from pathlib import Path
+from datetime import datetime
+
+def get_paths(root_path, invoke_path):
+  p = Path(invoke_path) if root_path is None else root_path / invoke_path
+  meta_full_path = p.with_name(p.name + '_meta.pkl')
+  pkl_full_path = p.with_name(p.name + '.pkl')
+  return meta_full_path, pkl_full_path
+
+def get_collection_timestamp(config, path):
+  meta_full_path, pkl_full_path = get_paths(config.root_path, path)
+  with meta_full_path.open('rb') as file:
+    meta_data = pickle.load(file)
+    return meta_data['created']
+
+def get_is_expired(config, path):
+  try:
+    get_collection_timestamp(config, path)
+    return False
+  except (FileNotFoundError, EOFError):
+    return True
+
+def should_expire(config, path, expire_fn):
+  return expire_fn(get_collection_timestamp(config, path))
+
+def store_data(config, path, data):
+  created = datetime.now()
+  meta_data = {'created': created}
+  meta_full_path, pkl_full_path = get_paths(config.root_path, path)
+  pkl_full_path.parent.mkdir(parents=True, exist_ok=True)
+  with pkl_full_path.open('wb') as file:
+    pickle.dump(data, file, -1)
+  with meta_full_path.open('wb') as file:
+    pickle.dump(meta_data, file, -1)
+  return data
+
+def load_data(config, path):
+  _, full_path = get_paths(config.root_path, path)
+  with full_path.open('rb') as file:
+    return pickle.load(file)
+
+def delete_data(config, path):
+  meta_full_path, pkl_full_path = get_paths(config.root_path, path)
+  try:
+    meta_full_path.unlink()
+    pkl_full_path.unlink()
+  except FileNotFoundError:
+    pass

checkpointer-1.2.0/checkpointer/utils.py

@@ -0,0 +1,17 @@
+import types
+
+def unwrap_func(func):
+  while hasattr(func, '__wrapped__'):
+    func = func.__wrapped__
+  return func
+
+@types.coroutine
+def coroutine_as_generator(coroutine):
+  val = yield from coroutine
+  return val
+
+def sync_resolve_coroutine(coroutine):
+  try:
+    next(coroutine_as_generator(coroutine))
+  except StopIteration as ex:
+    return ex.value

checkpointer-1.2.0/pyproject.toml

@@ -0,0 +1,24 @@
+[project]
+name = "checkpointer"
+version = "1.2.0"
+requires-python = ">=3.8"
+dependencies = [
+  "relib",
+  "termcolor",
+]
+authors = [
+  {name = "Hampus Hallman"}
+]
+# description = ""
+readme = "README.md"
+license = {file = "LICENSE"}
+
+[project.urls]
+Repository = "https://github.com/Reddan/checkpointer.git"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["checkpointer", "checkpointer.storages"]