checkpoint.py

import cPickle as pickle
import os
import numpy as np
import __main__ as main

# TODO: maintain cache for more than one set of arguments
# TODO: allow arbitrary serializers/deserializers

def meta_path(key):
    return ".%s.meta" % key

def data_path(key):
    return ".%s.data" % key

def np_data_path(key):
    return ".%s.npy" % key

def _compute_key(fn):
    return "%s.%s" % (os.path.splitext(os.path.basename(main.__file__))[0], fn.__name__)

def cached(function):
    key = _compute_key(function)

    def wrapped(*args, **kwargs):
        # save args and kwargs in meta file
        # if args match and output file exists, just return pickled value
        # otherwise call function and pickle value then return it
        if os.path.exists(meta_path(key)) and os.path.exists(data_path(key)):
            with open(meta_path(key), 'r') as meta_file:
                cached_args, cached_kwargs = pickle.load(meta_file)
            if (args, kwargs) == (cached_args, cached_kwargs):
                with open(data_path(key), 'r') as data_file:
                    return pickle.load(data_file)

        data = function(*args, **kwargs)
        with open(meta_path(key), 'w') as meta_file:
            pickle.dump((args, kwargs), meta_file)
        with open(data_path(key), 'w') as data_file:
            pickle.dump(data, data_file)
        return data

    return wrapped

def np_cached(key, function, *args, **kwargs):
    # save args and kwargs in meta file
    # if args match and output file exists, just return pickled value
    # otherwise call function and pickle value then return it
    if os.path.exists(meta_path(key)) and os.path.exists(data_path(key)):
        with open(meta_path(key), 'r') as meta_file:
            cached_args, cached_kwargs = pickle.load(meta_file)
        if (args, kwargs) == (cached_args, cached_kwargs):
            with open(np_data_path(key), 'r') as data_file:
                return np.load(data_file)

    data = function(*args, **kwargs)
    with open(meta_path(key), 'w') as meta_file:
        pickle.dump((args, kwargs), meta_file)
    with open(np_data_path(key), 'w') as data_file:
        np.save(data_file, data)

    return data