decorator memoize to file using cPickle (mostly use with caching vectorizers and featurizers so can focus on optimizing model [hyper]params)

disk_memoize.py

"""
from here: https://github.com/petered/plato/pull/56/files
adapted by Motoki
"""

import hashlib
from collections import OrderedDict
import numpy as np
import cPickle as pickle
import os

__author__ = 'peter'

__all__ = ["memoize_to_disk"]

PACKAGE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))


def get_local_path(relative_path=''):
    return os.path.join(PACKAGE_ROOT, relative_path)


def make_file_dir(full_file_path):
    full_local_dir, _ = os.path.split(full_file_path)
    try:
        os.makedirs(full_local_dir)
    except OSError:
        pass
    return full_file_path


MEMO_WRITE_ENABLED = True
MEMO_READ_ENABLED = True
MEMO_DIR = get_local_path('memoize_to_disk')


def memoize_to_disk(fcn, force=False):
    """
    Save (memoize) computed results to disk, so that the same function, called with the
    same arguments, does not need to be recomputed.  This is useful if you have a long-running
    function that is often being given the same arguments.  Note: this does NOT check for the state
    of Global variables/time/whatever else the function may use, so you need to make sure your
    function is truly a function in that outputs only depend on inputs.  Otherwise, this will
    give you misleading results.
    e.g.
        @memoize_to_disk
        def fcn(a, b, c = None):
            results = ...
            return results

    You can also use this without the decorator.
    e.g.
        result = memoize_to_disk(fcn)(a, b, c=3)

    This is useful if:
    a) The decorator can/should not be visible from where the function is defined.
    b) You only want to memoize the function in one use-case, but not all.

    :param fcn: The function you're decorating
    :return: A wrapper around the function that checks for memos and loads old results if they exist.
    """

    def check_memos(*args, **kwargs):

        if MEMO_READ_ENABLED:
            filepath = get_function_hash_filename(fcn, args, kwargs)
            file_found = os.path.exists(filepath)
            if file_found and not force:
                with open(filepath) as f:
                    result = pickle.load(f)
                    print "load from", filepath
            else:
                result = fcn(*args, **kwargs)
        else:
            result = fcn(*args, **kwargs)

        if MEMO_WRITE_ENABLED and (not file_found or force):
            filepath = get_function_hash_filename(fcn, args, kwargs)
            make_file_dir(filepath)
            with open(filepath, 'w') as f:
                pickle.dump(result, f)
                print "saved to", filepath
        return result

    check_memos.wrapped_fcn = fcn

    return check_memos


def get_function_hash_filename(fcn, args, kwargs):
    args_code = compute_fixed_hash((args, kwargs))
    return os.path.join(MEMO_DIR, '%s-%s.pkl' % (fcn.__name__, args_code))


def get_memo_files_for_function(fcn):
    all_memos = os.listdir(MEMO_DIR) if os.path.exists(MEMO_DIR) else []
    matching_memos = [os.path.join(MEMO_DIR, m) for m in all_memos if m.startswith(fcn.wrapped_fcn.__name__)]
    return matching_memos


def clear_memo_files_for_function(fcn):
    memos = get_memo_files_for_function(fcn)
    for m in memos:
        os.remove(m)


def compute_fixed_hash(obj, hasher = None):
    """
    Given an object, return a hash that will always be the same (not just for the lifetime of the
    object, but for all future runs of the program too).
    :param obj: Some nested container of primitives
    :param hasher: (for internal use)
    :return:
    """

    if hasher is None:
        hasher = hashlib.md5()

    hasher.update(obj.__class__.__name__)

    if isinstance(obj, (int, str, float, bool)):
        hasher.update(pickle.dumps(obj))
    elif isinstance(obj, (list, tuple)):
        hasher.update(str(len(obj)))  # Necessary to distinguish ([a, b], c) from ([a, b, c])
        for el in obj:
            compute_fixed_hash(el, hasher=hasher)
    elif isinstance(obj, np.ndarray):
        hasher.update(pickle.dumps(obj.dtype))
        hasher.update(pickle.dumps(obj.shape))
        hasher.update(obj.tostring())
    elif isinstance(obj, dict):
        hasher.update(str(len(obj)))  # Necessary to distinguish ([a, b], c) from ([a, b, c])
        keys = obj.keys() if isinstance(obj, OrderedDict) else sorted(obj.keys())
        for k in keys:
            compute_fixed_hash(k, hasher=hasher)
            compute_fixed_hash(obj[k], hasher=hasher)
    elif obj is None:
        hasher.update(pickle.dumps(obj))
    else:
        raise NotImplementedError("Don't have a method for hashing this %s" % (obj, ))

    return hasher.hexdigest()


if __name__ == '__main__':
    # test with scipy sparse arrays
    import functools
    from scipy.sparse import csr_matrix

    @functools.partial(memoize_to_disk, force=False)
    def dummy_function():
        a = csr_matrix((3, 4), dtype=np.int8)
        return a.toarray()

    result_from_cache = dummy_function()
    print type(result_from_cache), result_from_cache

    @functools.partial(memoize_to_disk, force=True)
    def dummy_function():
        a = csr_matrix((3, 4), dtype=np.int8)
        return a.toarray()

    result = dummy_function()
    print type(result), result

    assert np.array_equal(result, result_from_cache), "result and result_from_cache different"

2 limitations:

• doesn't include optional arguments for the file name hashing
• doesn't include scope of the functions so functions with same name will be treated the same.