Repository: joblib/joblib
Pull request #1566 (2024-04-08 08:23:30)

DOC add gotcha on non-reproducible pickling
Fixes https://github.com/joblib/joblib/issues/1534


Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 5bf9c95..e191018 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -4,11 +4,14 @@ Latest changes
 In development
 --------------
 
+- Allow caching co-routines with `Memory.cache`.
+  https://github.com/joblib/joblib/pull/894
+  
 - Try to cast ``n_jobs`` to int in parallel and raise an error if
   it fails. This means that ``n_jobs=2.3`` will now result in
   ``effective_n_jobs=2`` instead of failing.
   https://github.com/joblib/joblib/pull/1539
-
+  
 - Ensure that errors in the task generator given to Parallel's call
   are raised in the results consumming thread.
   https://github.com/joblib/joblib/pull/1491
diff --git a/continuous_integration/install.sh b/continuous_integration/install.sh
index c6fed28..9668b74 100755
--- a/continuous_integration/install.sh
+++ b/continuous_integration/install.sh
@@ -34,7 +34,7 @@ else
 fi
 
 # Install pytest timeout to fasten failure in deadlocking tests
-PIP_INSTALL_PACKAGES="pytest-timeout threadpoolctl"
+PIP_INSTALL_PACKAGES="pytest-timeout pytest-asyncio==0.21.1 threadpoolctl"
 
 if [ -n "$NUMPY_VERSION" ]; then
     # We want to ensure no memory copies are performed only when numpy is
diff --git a/doc/memory.rst b/doc/memory.rst
index 37a889e..9bf1444 100644
--- a/doc/memory.rst
+++ b/doc/memory.rst
@@ -383,6 +383,26 @@ Gotchas
   that are constructs (or contain references to constructs) defined in the
   upgraded library should potentially be invalidated after the upgrade.
 
+* **Cache-miss with objects that have non-reproducible pickle representations**.
+  The identifier of the cache entry is based on the pickle's representation of
+  the input arguments. Therefor, for objects that don't have a deterministic
+  pickle representation, or objects whose representation depends on the way they
+  are constructed, the cache will not work. In particular, ``pytorch.Tensor``
+  are known to have non-deterministic pickle representation (see this
+  `issue <https://github.com/pytorch/pytorch/issues/32165>`_). A good way
+  to debug this is to check that two calls to the following script with ``args``
+  and ``kwargs`` being the cached function's inputs give the same output::
+
+    from joblib import hash
+
+    for x in args:
+      print(f"{hash(x)}")
+    for k, x in kwargs.items():
+      print(f"hash({k})={hash(x)}")
+
+  To avoid this issue, a good practice is to use ``Memory.cache`` with
+  functions that take simple input arguments when possible.
+
 
 Ignoring some arguments
 -----------------------
diff --git a/joblib/_store_backends.py b/joblib/_store_backends.py
index ce91460..68e207c 100644
--- a/joblib/_store_backends.py
+++ b/joblib/_store_backends.py
@@ -7,6 +7,7 @@ import os.path
 import datetime
 import json
 import shutil
+import time
 import warnings
 import collections
 import operator
@@ -15,6 +16,7 @@ from abc import ABCMeta, abstractmethod
 
 from .backports import concurrency_safe_rename
 from .disk import mkdirp, memstr_to_bytes, rm_subdirs
+from .logger import format_time
 from . import numpy_pickle
 
 CacheItemInfo = collections.namedtuple('CacheItemInfo',
@@ -151,12 +153,19 @@ class StoreBackendMixin(object):
     file-like object.
     """
 
-    def load_item(self, path, verbose=1, msg=None):
-        """Load an item from the store given its path as a list of
-           strings."""
-        full_path = os.path.join(self.location, *path)
+    def load_item(self, call_id, verbose=1, timestamp=None, metadata=None):
+        """Load an item from the store given its id as a list of str."""
+        full_path = os.path.join(self.location, *call_id)
 
         if verbose > 1:
+            ts_string = ('{: <16}'.format(format_time(time.time() - timestamp))
+                         if timestamp is not None else '')
+            signature = os.path.basename(call_id[0])
+            if metadata is not None and 'input_args' in metadata:
+                kwargs = ', '.join('{}={}'.format(*item)
+                                   for item in metadata['input_args'].items())
+                signature += '({})'.format(kwargs)
+            msg = '[Memory]{}: Loading {}'.format(ts_string, signature)
             if verbose < 10:
                 print('{0}...'.format(msg))
             else:
@@ -178,11 +187,10 @@ class StoreBackendMixin(object):
             item = numpy_pickle.load(filename, mmap_mode=mmap_mode)
         return item
 
-    def dump_item(self, path, item, verbose=1):
-        """Dump an item in the store at the path given as a list of
-           strings."""
+    def dump_item(self, call_id, item, verbose=1):
+        """Dump an item in the store at the id given as a list of str."""
         try:
-            item_path = os.path.join(self.location, *path)
+            item_path = os.path.join(self.location, *call_id)
             if not self._item_exists(item_path):
                 self.create_location(item_path)
             filename = os.path.join(item_path, 'output.pkl')
@@ -210,39 +218,37 @@ class StoreBackendMixin(object):
                 CacheWarning
             )
 
-    def clear_item(self, path):
-        """Clear the item at the path, given as a list of strings."""
-        item_path = os.path.join(self.location, *path)
+    def clear_item(self, call_id):
+        """Clear the item at the id, given as a list of str."""
+        item_path = os.path.join(self.location, *call_id)
         if self._item_exists(item_path):
             self.clear_location(item_path)
 
-    def contains_item(self, path):
-        """Check if there is an item at the path, given as a list of
-           strings"""
-        item_path = os.path.join(self.location, *path)
+    def contains_item(self, call_id):
+        """Check if there is an item at the id, given as a list of str."""
+        item_path = os.path.join(self.location, *call_id)
         filename = os.path.join(item_path, 'output.pkl')
 
         return self._item_exists(filename)
 
-    def get_item_info(self, path):
+    def get_item_info(self, call_id):
         """Return information about item."""
-        return {'location': os.path.join(self.location,
-                                         *path)}
+        return {'location': os.path.join(self.location, *call_id)}
 
-    def get_metadata(self, path):
+    def get_metadata(self, call_id):
         """Return actual metadata of an item."""
         try:
-            item_path = os.path.join(self.location, *path)
+            item_path = os.path.join(self.location, *call_id)
             filename = os.path.join(item_path, 'metadata.json')
             with self._open_item(filename, 'rb') as f:
                 return json.loads(f.read().decode('utf-8'))
         except:  # noqa: E722
             return {}
 
-    def store_metadata(self, path, metadata):
+    def store_metadata(self, call_id, metadata):
         """Store metadata of a computation."""
         try:
-            item_path = os.path.join(self.location, *path)
+            item_path = os.path.join(self.location, *call_id)
             self.create_location(item_path)
             filename = os.path.join(item_path, 'metadata.json')
 
@@ -254,20 +260,20 @@ class StoreBackendMixin(object):
         except:  # noqa: E722
             pass
 
-    def contains_path(self, path):
+    def contains_path(self, call_id):
         """Check cached function is available in store."""
-        func_path = os.path.join(self.location, *path)
+        func_path = os.path.join(self.location, *call_id)
         return self.object_exists(func_path)
 
-    def clear_path(self, path):
+    def clear_path(self, call_id):
         """Clear all items with a common path in the store."""
-        func_path = os.path.join(self.location, *path)
+        func_path = os.path.join(self.location, *call_id)
         if self._item_exists(func_path):
             self.clear_location(func_path)
 
-    def store_cached_func_code(self, path, func_code=None):
+    def store_cached_func_code(self, call_id, func_code=None):
         """Store the code of the cached function."""
-        func_path = os.path.join(self.location, *path)
+        func_path = os.path.join(self.location, *call_id)
         if not self._item_exists(func_path):
             self.create_location(func_path)
 
@@ -276,19 +282,18 @@ class StoreBackendMixin(object):
             with self._open_item(filename, 'wb') as f:
                 f.write(func_code.encode('utf-8'))
 
-    def get_cached_func_code(self, path):
+    def get_cached_func_code(self, call_id):
         """Store the code of the cached function."""
-        path += ['func_code.py', ]
-        filename = os.path.join(self.location, *path)
+        filename = os.path.join(self.location, *call_id, 'func_code.py')
         try:
             with self._open_item(filename, 'rb') as f:
                 return f.read().decode('utf-8')
         except:  # noqa: E722
             raise
 
-    def get_cached_func_info(self, path):
+    def get_cached_func_info(self, call_id):
         """Return information related to the cached function if it exists."""
-        return {'location': os.path.join(self.location, *path)}
+        return {'location': os.path.join(self.location, *call_id)}
 
     def clear(self):
         """Clear the whole store content."""
diff --git a/joblib/memory.py b/joblib/memory.py
index 8171d24..6f87f80 100644
--- a/joblib/memory.py
+++ b/joblib/memory.py
@@ -9,32 +9,28 @@ is called with the same input arguments.
 # License: BSD Style, 3 clauses.
 
 
-from __future__ import with_statement
+import asyncio
+import datetime
+import functools
+import inspect
 import logging
 import os
-from textwrap import dedent
-import time
 import pathlib
 import pydoc
 import re
-import functools
+import textwrap
+import time
+import tokenize
 import traceback
 import warnings
-import inspect
 import weakref
-from datetime import timedelta
-
-from tokenize import open as open_py_source
 
-# Local imports
 from . import hashing
-from .func_inspect import get_func_code, get_func_name, filter_args
-from .func_inspect import format_call
-from .func_inspect import format_signature
-from .logger import Logger, format_time, pformat
-from ._store_backends import StoreBackendBase, FileSystemStoreBackend
 from ._store_backends import CacheWarning  # noqa
-
+from ._store_backends import FileSystemStoreBackend, StoreBackendBase
+from .func_inspect import (filter_args, format_call, format_signature,
+                           get_func_code, get_func_name)
+from .logger import Logger, format_time, pformat
 
 FIRST_LINE_TEXT = "# first line:"
 
@@ -141,45 +137,11 @@ def _store_backend_factory(backend, location, verbose=0, backend_options=None):
     return None
 
 
-def _get_func_fullname(func):
-    """Compute the part of part associated with a function."""
-    modules, funcname = get_func_name(func)
-    modules.append(funcname)
-    return os.path.join(*modules)
-
-
 def _build_func_identifier(func):
     """Build a roughly unique identifier for the cached function."""
-    parts = []
-    if isinstance(func, str):
-        parts.append(func)
-    else:
-        parts.append(_get_func_fullname(func))
-
+    modules, funcname = get_func_name(func)
     # We reuse historical fs-like way of building a function identifier
-    return os.path.join(*parts)
-
-
-def _format_load_msg(func_id, args_id, timestamp=None, metadata=None):
-    """ Helper function to format the message when loading the results.
-    """
-    signature = ""
-    try:
-        if metadata is not None:
-            args = ", ".join(['%s=%s' % (name, value)
-                              for name, value
-                              in metadata['input_args'].items()])
-            signature = "%s(%s)" % (os.path.basename(func_id), args)
-        else:
-            signature = os.path.basename(func_id)
-    except KeyError:
-        pass
-
-    if timestamp is not None:
-        ts_string = "{0: <16}".format(format_time(time.time() - timestamp))
-    else:
-        ts_string = ""
-    return '[Memory]{0}: Loading {1}'.format(ts_string, str(signature))
+    return os.path.join(*modules, funcname)
 
 
 # An in-memory store to avoid looking at the disk-based function
@@ -220,15 +182,10 @@ class MemorizedResult(Logger):
     timestamp, metadata: string
         for internal use only.
     """
-    def __init__(self, location, func, args_id, backend='local',
-                 mmap_mode=None, verbose=0, timestamp=None, metadata=None):
+    def __init__(self, location, call_id, backend='local', mmap_mode=None,
+                 verbose=0, timestamp=None, metadata=None):
         Logger.__init__(self)
-        self.func_id = _build_func_identifier(func)
-        if isinstance(func, str):
-            self.func = func
-        else:
-            self.func = self.func_id
-        self.args_id = args_id
+        self._call_id = call_id
         self.store_backend = _store_backend_factory(backend, location,
                                                     verbose=verbose)
         self.mmap_mode = mmap_mode
@@ -236,13 +193,24 @@ class MemorizedResult(Logger):
         if metadata is not None:
             self.metadata = metadata
         else:
-            self.metadata = self.store_backend.get_metadata(
-                [self.func_id, self.args_id])
+            self.metadata = self.store_backend.get_metadata(self._call_id)
 
         self.duration = self.metadata.get('duration', None)
         self.verbose = verbose
         self.timestamp = timestamp
 
+    @property
+    def func(self):
+        return self.func_id
+
+    @property
+    def func_id(self):
+        return self._call_id[0]
+
+    @property
+    def args_id(self):
+        return self._call_id[1]
+
     @property
     def argument_hash(self):
         warnings.warn(
@@ -254,38 +222,29 @@ class MemorizedResult(Logger):
 
     def get(self):
         """Read value from cache and return it."""
-        if self.verbose:
-            msg = _format_load_msg(self.func_id, self.args_id,
-                                   timestamp=self.timestamp,
-                                   metadata=self.metadata)
-        else:
-            msg = None
-
         try:
             return self.store_backend.load_item(
-                [self.func_id, self.args_id], msg=msg, verbose=self.verbose)
+                self._call_id,
+                timestamp=self.timestamp,
+                metadata=self.metadata,
+                verbose=self.verbose
+            )
         except ValueError as exc:
             new_exc = KeyError(
                 "Error while trying to load a MemorizedResult's value. "
                 "It seems that this folder is corrupted : {}".format(
-                    os.path.join(
-                        self.store_backend.location, self.func_id,
-                        self.args_id)
-                ))
+                    os.path.join(self.store_backend.location, *self._call_id)))
             raise new_exc from exc
 
     def clear(self):
         """Clear value from cache"""
-        self.store_backend.clear_item([self.func_id, self.args_id])
+        self.store_backend.clear_item(self._call_id)
 
     def __repr__(self):
-        return ('{class_name}(location="{location}", func="{func}", '
-                'args_id="{args_id}")'
-                .format(class_name=self.__class__.__name__,
-                        location=self.store_backend.location,
-                        func=self.func,
-                        args_id=self.args_id
-                        ))
+        return '{}(location="{}", func="{}", args_id="{}")'.format(
+            self.__class__.__name__, self.store_backend.location,
+            *self._call_id
+        )
 
     def __getstate__(self):
         state = self.__dict__.copy()
@@ -369,6 +328,14 @@ class NotMemorizedFunc(object):
         return False
 
 
+###############################################################################
+# class `AsyncNotMemorizedFunc`
+###############################################################################
+class AsyncNotMemorizedFunc(NotMemorizedFunc):
+    async def call_and_shelve(self, *args, **kwargs):
+        return NotMemorizedResult(await self.func(*args, **kwargs))
+
+
 ###############################################################################
 # class `MemorizedFunc`
 ###############################################################################
@@ -429,10 +396,8 @@ class MemorizedFunc(Logger):
         self.compress = compress
         self.func = func
         self.cache_validation_callback = cache_validation_callback
-
-        if ignore is None:
-            ignore = []
-        self.ignore = ignore
+        self.func_id = _build_func_identifier(func)
+        self.ignore = ignore if ignore is not None else []
         self._verbose = verbose
 
         # retrieve store object from backend type and location.
@@ -444,17 +409,13 @@ class MemorizedFunc(Logger):
                                                     )
         if self.store_backend is not None:
             # Create func directory on demand.
-            self.store_backend.store_cached_func_code([
-                _build_func_identifier(self.func)
-            ])
+            self.store_backend.store_cached_func_code([self.func_id])
 
-        if timestamp is None:
-            timestamp = time.time()
-        self.timestamp = timestamp
+        self.timestamp = timestamp if timestamp is not None else time.time()
         try:
             functools.update_wrapper(self, func)
         except Exception:
-            " Objects like ufunc don't like that "
+            pass  # Objects like ufunc don't like that
         if inspect.isfunction(func):
             doc = pydoc.TextDoc().document(func)
             # Remove blank line
@@ -469,7 +430,7 @@ class MemorizedFunc(Logger):
         self._func_code_info = None
         self._func_code_id = None
 
-    def _is_in_cache_and_valid(self, path):
+    def _is_in_cache_and_valid(self, call_id):
         """Check if the function call is cached and valid for given arguments.
 
         - Compare the function code with the one from the cached function,
@@ -485,22 +446,23 @@ class MemorizedFunc(Logger):
             return False
 
         # Check if this specific call is in the cache
-        if not self.store_backend.contains_item(path):
+        if not self.store_backend.contains_item(call_id):
             return False
 
         # Call the user defined cache validation callback
-        metadata = self.store_backend.get_metadata(path)
+        metadata = self.store_backend.get_metadata(call_id)
         if (self.cache_validation_callback is not None and
                 not self.cache_validation_callback(metadata)):
-            self.store_backend.clear_item(path)
+            self.store_backend.clear_item(call_id)
             return False
 
         return True
 
-    def _cached_call(self, args, kwargs, shelving=False):
+    def _cached_call(self, args, kwargs, shelving):
         """Call wrapped function and cache result, or read cache if available.
 
-        This function returns the wrapped function output and some metadata.
+        This function returns the wrapped function output or a reference to
+        the cached result.
 
         Arguments:
         ----------
@@ -514,35 +476,22 @@ class MemorizedFunc(Logger):
 
         Returns
         -------
-        output: value or tuple or None
-            Output of the wrapped function.
-            If shelving is True and the call has been already cached,
-            output is None.
-
-        argument_hash: string
-            Hash of function arguments.
-
-        metadata: dict
-            Some metadata about wrapped function call (see _persist_input()).
+        Output of the wrapped function if shelving is false, or a
+        MemorizedResult reference to the value if shelving is true.
         """
-        func_id, args_id = self._get_output_identifiers(*args, **kwargs)
-        metadata = None
-        msg = None
-
-        # Whether or not the memorized function must be called
-        must_call = False
+        args_id = self._get_args_id(*args, **kwargs)
+        call_id = (self.func_id, args_id)
+        _, func_name = get_func_name(self.func)
+        func_info = self.store_backend.get_cached_func_info([self.func_id])
+        location = func_info['location']
 
         if self._verbose >= 20:
             logging.basicConfig(level=logging.INFO)
-            _, name = get_func_name(self.func)
-            location = self.store_backend.get_cached_func_info([func_id])[
-                'location']
             _, signature = format_signature(self.func, *args, **kwargs)
-
             self.info(
-                dedent(
+                textwrap.dedent(
                     f"""
-                        Querying {name} with signature
+                        Querying {func_name} with signature
                         {signature}.
 
                         (argument hash {args_id})
@@ -555,58 +504,30 @@ class MemorizedFunc(Logger):
         # Compare the function code with the previous to see if the
         # function code has changed and check if the results are present in
         # the cache.
-        if self._is_in_cache_and_valid([func_id, args_id]):
-            try:
-                t0 = time.time()
-                if self._verbose:
-                    msg = _format_load_msg(func_id, args_id,
-                                           timestamp=self.timestamp,
-                                           metadata=metadata)
-
-                if not shelving:
-                    # When shelving, we do not need to load the output
-                    out = self.store_backend.load_item(
-                        [func_id, args_id],
-                        msg=msg,
-                        verbose=self._verbose)
-                else:
-                    out = None
+        if self._is_in_cache_and_valid(call_id):
+            if shelving:
+                return self._get_memorized_result(call_id)
 
+            try:
+                start_time = time.time()
+                output = self._load_item(call_id)
                 if self._verbose > 4:
-                    t = time.time() - t0
-                    _, name = get_func_name(self.func)
-                    msg = '%s cache loaded - %s' % (name, format_time(t))
-                    print(max(0, (80 - len(msg))) * '_' + msg)
+                    self._print_duration(time.time() - start_time,
+                                         context='cache loaded ')
+                return output
             except Exception:
                 # XXX: Should use an exception logger
                 _, signature = format_signature(self.func, *args, **kwargs)
                 self.warn('Exception while loading results for '
                           '{}\n {}'.format(signature, traceback.format_exc()))
 
-                must_call = True
-        else:
-            if self._verbose > 10:
-                _, name = get_func_name(self.func)
-                self.warn('Computing func {0}, argument hash {1} '
-                          'in location {2}'
-                          .format(name, args_id,
-                                  self.store_backend.
-                                  get_cached_func_info([func_id])['location']))
-            must_call = True
-
-        if must_call:
-            out, metadata = self.call(*args, **kwargs)
-            if self.mmap_mode is not None:
-                # Memmap the output at the first call to be consistent with
-                # later calls
-                if self._verbose:
-                    msg = _format_load_msg(func_id, args_id,
-                                           timestamp=self.timestamp,
-                                           metadata=metadata)
-                out = self.store_backend.load_item([func_id, args_id], msg=msg,
-                                                   verbose=self._verbose)
-
-        return (out, args_id, metadata)
+        if self._verbose > 10:
+            self.warn(
+                f"Computing func {func_name}, argument hash {args_id} "
+                f"in location {location}"
+            )
+
+        return self._call(call_id, args, kwargs, shelving)
 
     @property
     def func_code_info(self):
@@ -646,13 +567,10 @@ class MemorizedFunc(Logger):
             class "NotMemorizedResult" is used when there is no cache
             activated (e.g. location=None in Memory).
         """
-        _, args_id, metadata = self._cached_call(args, kwargs, shelving=True)
-        return MemorizedResult(self.store_backend, self.func, args_id,
-                               metadata=metadata, verbose=self._verbose - 1,
-                               timestamp=self.timestamp)
+        return self._cached_call(args, kwargs, shelving=True)
 
     def __call__(self, *args, **kwargs):
-        return self._cached_call(args, kwargs)[0]
+        return self._cached_call(args, kwargs, shelving=False)
 
     def __getstate__(self):
         # Make sure self.func's source is introspected prior to being pickled -
@@ -682,22 +600,17 @@ class MemorizedFunc(Logger):
             Whether or not the result of the function has been cached
             for the input arguments that have been passed.
         """
-        func_id, args_id = self._get_output_identifiers(*args, **kwargs)
-        return self.store_backend.contains_item((func_id, args_id))
+        call_id = (self.func_id, self._get_args_id(*args, **kwargs))
+        return self.store_backend.contains_item(call_id)
 
     # ------------------------------------------------------------------------
     # Private interface
     # ------------------------------------------------------------------------
 
-    def _get_argument_hash(self, *args, **kwargs):
+    def _get_args_id(self, *args, **kwargs):
+        """Return the input parameter hash of a result."""
         return hashing.hash(filter_args(self.func, self.ignore, args, kwargs),
-                            coerce_mmap=(self.mmap_mode is not None))
-
-    def _get_output_identifiers(self, *args, **kwargs):
-        """Return the func identifier and input parameter hash of a result."""
-        func_id = _build_func_identifier(self.func)
-        argument_hash = self._get_argument_hash(*args, **kwargs)
-        return func_id, argument_hash
+                            coerce_mmap=self.mmap_mode is not None)
 
     def _hash_func(self):
         """Hash a function to key the online cache"""
@@ -712,12 +625,10 @@ class MemorizedFunc(Logger):
         # sometimes have several functions named the same way in a
         # file. This is bad practice, but joblib should be robust to bad
         # practice.
-        func_id = _build_func_identifier(self.func)
         func_code = u'%s %i\n%s' % (FIRST_LINE_TEXT, first_line, func_code)
-        self.store_backend.store_cached_func_code([func_id], func_code)
+        self.store_backend.store_cached_func_code([self.func_id], func_code)
 
         # Also store in the in-memory store of function hashes
-        is_named_callable = False
         is_named_callable = (hasattr(self.func, '__name__') and
                              self.func.__name__ != '<lambda>')
         if is_named_callable:
@@ -755,12 +666,9 @@ class MemorizedFunc(Logger):
         # changing code and collision. We cannot inspect.getsource
         # because it is not reliable when using IPython's magic "%run".
         func_code, source_file, first_line = self.func_code_info
-        func_id = _build_func_identifier(self.func)
-
         try:
-            old_func_code, old_first_line =\
-                extract_first_line(
-                    self.store_backend.get_cached_func_code([func_id]))
+            old_func_code, old_first_line = extract_first_line(
+                self.store_backend.get_cached_func_code([self.func_id]))
         except (IOError, OSError):  # some backend can also raise OSError
             self._write_func_code(func_code, first_line)
             return False
@@ -789,11 +697,10 @@ class MemorizedFunc(Logger):
         # file has not changed, but the name we have is pointing to a new
         # code block.
         if not old_first_line == first_line and source_file is not None:
-            possible_collision = False
             if os.path.exists(source_file):
                 _, func_name = get_func_name(self.func, resolv_alias=False)
                 num_lines = len(func_code.split('\n'))
-                with open_py_source(source_file) as f:
+                with tokenize.open(source_file) as f:
                     on_disk_func_code = f.readlines()[
                         old_first_line - 1:old_first_line - 1 + num_lines - 1]
                 on_disk_func_code = ''.join(on_disk_func_code)
@@ -814,14 +721,13 @@ class MemorizedFunc(Logger):
         if self._verbose > 10:
             _, func_name = get_func_name(self.func, resolv_alias=False)
             self.warn("Function {0} (identified by {1}) has changed"
-                      ".".format(func_name, func_id))
+                      ".".format(func_name, self.func_id))
         self.clear(warn=True)
         return False
 
     def clear(self, warn=True):
         """Empty the function's cache."""
-        func_id = _build_func_identifier(self.func)
-
+        func_id = self.func_id
         if self._verbose > 0 and warn:
             self.warn("Clearing function cache identified by %s" % func_id)
         self.store_backend.clear_path([func_id, ])
@@ -846,27 +752,38 @@ class MemorizedFunc(Logger):
         -------
         output : object
             The output of the function call.
-        metadata : dict
-            The metadata associated with the call.
         """
+        call_id = (self.func_id, self._get_args_id(*args, **kwargs))
+        return self._call(call_id, args, kwargs)
+
+    def _call(self, call_id, args, kwargs, shelving=False):
+        self._before_call(args, kwargs)
         start_time = time.time()
-        func_id, args_id = self._get_output_identifiers(*args, **kwargs)
+        output = self.func(*args, **kwargs)
+        return self._after_call(call_id, args, kwargs, shelving,
+                                output, start_time)
+
+    def _before_call(self, args, kwargs):
         if self._verbose > 0:
             print(format_call(self.func, args, kwargs))
-        output = self.func(*args, **kwargs)
-        self.store_backend.dump_item(
-            [func_id, args_id], output, verbose=self._verbose)
 
+    def _after_call(self, call_id, args, kwargs, shelving, output, start_time):
+        self.store_backend.dump_item(call_id, output, verbose=self._verbose)
         duration = time.time() - start_time
-        metadata = self._persist_input(duration, args, kwargs)
-
         if self._verbose > 0:
-            _, name = get_func_name(self.func)
-            msg = '%s - %s' % (name, format_time(duration))
-            print(max(0, (80 - len(msg))) * '_' + msg)
-        return output, metadata
-
-    def _persist_input(self, duration, args, kwargs, this_duration_limit=0.5):
+            self._print_duration(duration)
+        metadata = self._persist_input(duration, call_id, args, kwargs)
+        if shelving:
+            return self._get_memorized_result(call_id, metadata)
+
+        if self.mmap_mode is not None:
+            # Memmap the output at the first call to be consistent with
+            # later calls
+            output = self._load_item(call_id, metadata)
+        return output
+
+    def _persist_input(self, duration, call_id, args, kwargs,
+                       this_duration_limit=0.5):
         """ Save a small summary of the call using json format in the
             output directory.
 
@@ -894,8 +811,7 @@ class MemorizedFunc(Logger):
             "duration": duration, "input_args": input_repr, "time": start_time,
         }
 
-        func_id, args_id = self._get_output_identifiers(*args, **kwargs)
-        self.store_backend.store_metadata([func_id, args_id], metadata)
+        self.store_backend.store_metadata(call_id, metadata)
 
         this_duration = time.time() - start_time
         if this_duration > this_duration_limit:
@@ -914,6 +830,21 @@ class MemorizedFunc(Logger):
                           % this_duration, stacklevel=5)
         return metadata
 
+    def _get_memorized_result(self, call_id, metadata=None):
+        return MemorizedResult(self.store_backend, call_id,
+                               metadata=metadata, timestamp=self.timestamp,
+                               verbose=self._verbose - 1)
+
+    def _load_item(self, call_id, metadata=None):
+        return self.store_backend.load_item(call_id, metadata=metadata,
+                                            timestamp=self.timestamp,
+                                            verbose=self._verbose)
+
+    def _print_duration(self, duration, context=''):
+        _, name = get_func_name(self.func)
+        msg = f"{name} {context}- {format_time(duration)}"
+        print(max(0, (80 - len(msg))) * '_' + msg)
+
     # ------------------------------------------------------------------------
     # Private `object` interface
     # ------------------------------------------------------------------------
@@ -925,6 +856,30 @@ class MemorizedFunc(Logger):
             location=self.store_backend.location,)
 
 
+###############################################################################
+# class `AsyncMemorizedFunc`
+###############################################################################
+class AsyncMemorizedFunc(MemorizedFunc):
+    async def __call__(self, *args, **kwargs):
+        out = super().__call__(*args, **kwargs)
+        return await out if asyncio.iscoroutine(out) else out
+
+    async def call_and_shelve(self, *args, **kwargs):
+        out = super().call_and_shelve(*args, **kwargs)
+        return await out if asyncio.iscoroutine(out) else out
+
+    async def call(self, *args, **kwargs):
+        out = super().call(*args, **kwargs)
+        return await out if asyncio.iscoroutine(out) else out
+
+    async def _call(self, call_id, args, kwargs, shelving=False):
+        self._before_call(args, kwargs)
+        start_time = time.time()
+        output = await self.func(*args, **kwargs)
+        return self._after_call(call_id, args, kwargs, shelving,
+                                output, start_time)
+
+
 ###############################################################################
 # class `Memory`
 ###############################################################################
@@ -1072,14 +1027,20 @@ class Memory(Logger):
                 cache_validation_callback=cache_validation_callback
             )
         if self.store_backend is None:
-            return NotMemorizedFunc(func)
+            cls = (AsyncNotMemorizedFunc
+                   if asyncio.iscoroutinefunction(func)
+                   else NotMemorizedFunc)
+            return cls(func)
         if verbose is None:
             verbose = self._verbose
         if mmap_mode is False:
             mmap_mode = self.mmap_mode
         if isinstance(func, MemorizedFunc):
             func = func.func
-        return MemorizedFunc(
+        cls = (AsyncMemorizedFunc
+               if asyncio.iscoroutinefunction(func)
+               else MemorizedFunc)
+        return cls(
             func, location=self.store_backend, backend=self.backend,
             ignore=ignore, mmap_mode=mmap_mode, compress=self.compress,
             verbose=verbose, timestamp=self.timestamp,
@@ -1187,7 +1148,7 @@ def expires_after(days=0, seconds=0, microseconds=0, milliseconds=0, minutes=0,
     days, seconds, microseconds, milliseconds, minutes, hours, weeks: numbers
         argument passed to a timedelta.
     """
-    delta = timedelta(
+    delta = datetime.timedelta(
         days=days, seconds=seconds, microseconds=microseconds,
         milliseconds=milliseconds, minutes=minutes, hours=hours, weeks=weeks
     )
diff --git a/joblib/parallel.py b/joblib/parallel.py
index 91cc03a..e42385d 100644
--- a/joblib/parallel.py
+++ b/joblib/parallel.py
@@ -254,7 +254,7 @@ class parallel_config:
 
         Alternatively the backend can be passed directly as an instance.
 
-    n_jobs: int, default: None
+    n_jobs: int, default=None
         The maximum number of concurrently running jobs, such as the number
         of Python worker processes when ``backend="loky"`` or the size of the
         thread-pool when ``backend="threading"``.
@@ -280,7 +280,7 @@ class parallel_config:
         The frequency of the messages increases with the verbosity level.
         If it more than 10, all iterations are reported.
 
-    temp_folder: str, default=None
+    temp_folder: str or None, default=None
         Folder to be used by the pool for memmapping large arrays
         for sharing memory with worker processes. If None, this will try in
         order:
@@ -950,7 +950,7 @@ class Parallel(Logger):
 
         Parameters
         ----------
-        n_jobs: int, default: None
+        n_jobs: int, default=None
             The maximum number of concurrently running jobs, such as the number
             of Python worker processes when ``backend="loky"`` or the size of
             the thread-pool when ``backend="threading"``.
@@ -969,7 +969,7 @@ class Parallel(Logger):
             unless the call is performed under a :func:`~parallel_config`
             context manager that sets another value for ``n_jobs``.
             If n_jobs = 0 then a ValueError is raised.
-        backend: str, ParallelBackendBase instance or None, default: 'loky'
+        backend: str, ParallelBackendBase instance or None, default='loky'
             Specify the parallelization backend implementation.
             Supported backends are:
 
@@ -996,8 +996,7 @@ class Parallel(Logger):
             soft hints (prefer) or hard constraints (require) so as to make it
             possible for library users to change the backend from the outside
             using the :func:`~parallel_config` context manager.
-        return_as: str in {'list', 'generator', 'generator_unordered'},
-            default: 'list'
+        return_as: str in {'list', 'generator', 'generator_unordered'}, default='list'
             If 'list', calls to this instance will return a list, only when
             all results have been processed and retrieved.
             If 'generator', it will return a generator that yields the results
@@ -1007,32 +1006,32 @@ class Parallel(Logger):
             available results independently of the submission order. The output
             order is not deterministic in this case because it depends on the
             concurrency of the workers.
-        prefer: str in {'processes', 'threads'} or None, default: None
+        prefer: str in {'processes', 'threads'} or None, default=None
             Soft hint to choose the default backend if no specific backend
             was selected with the :func:`~parallel_config` context manager.
             The default process-based backend is 'loky' and the default
             thread-based backend is 'threading'. Ignored if the ``backend``
             parameter is specified.
-        require: 'sharedmem' or None, default None
+        require: 'sharedmem' or None, default=None
             Hard constraint to select the backend. If set to 'sharedmem',
             the selected backend will be single-host and thread-based even
             if the user asked for a non-thread based backend with
             :func:`~joblib.parallel_config`.
-        verbose: int, optional
+        verbose: int, default=0
             The verbosity level: if non zero, progress messages are
             printed. Above 50, the output is sent to stdout.
             The frequency of the messages increases with the verbosity level.
             If it more than 10, all iterations are reported.
-        timeout: float, optional
+        timeout: float or None, default=None
             Timeout limit for each task to complete.  If any task takes longer
             a TimeOutError will be raised. Only applied when n_jobs != 1
-        pre_dispatch: {'all', integer, or expression, as in '3*n_jobs'}
+        pre_dispatch: {'all', integer, or expression, as in '3*n_jobs'}, default='2*n_jobs'
             The number of batches (of tasks) to be pre-dispatched.
             Default is '2*n_jobs'. When batch_size="auto" this is reasonable
             default and the workers should never starve. Note that only basic
             arithmetics are allowed here and no modules can be used in this
             expression.
-        batch_size: int or 'auto', default: 'auto'
+        batch_size: int or 'auto', default='auto'
             The number of atomic tasks to dispatch at once to each
             worker. When individual evaluations are very fast, dispatching
             calls to workers can be slower than sequential computation because
@@ -1046,7 +1045,7 @@ class Parallel(Logger):
             batches of a single task at a time as the threading backend has
             very little overhead and using larger batch size has not proved to
             bring any gain in that case.
-        temp_folder: str, optional
+        temp_folder: str or None, default=None
             Folder to be used by the pool for memmapping large arrays
             for sharing memory with worker processes. If None, this will try in
             order:
@@ -1061,13 +1060,13 @@ class Parallel(Logger):
               variables, typically /tmp under Unix operating systems.
 
             Only active when ``backend="loky"`` or ``"multiprocessing"``.
-        max_nbytes int, str, or None, optional, 1M by default
+        max_nbytes int, str, or None, optional, default='1M'
             Threshold on the size of arrays passed to the workers that
             triggers automated memory mapping in temp_folder. Can be an int
             in Bytes, or a human-readable string, e.g., '1M' for 1 megabyte.
             Use None to disable memmapping of large arrays.
             Only active when ``backend="loky"`` or ``"multiprocessing"``.
-        mmap_mode: {None, 'r+', 'r', 'w+', 'c'}, default: 'r'
+        mmap_mode: {None, 'r+', 'r', 'w+', 'c'}, default='r'
             Memmapping mode for numpy arrays passed to workers. None will
             disable memmapping, other modes defined in the numpy.memmap doc:
             https://numpy.org/doc/stable/reference/generated/numpy.memmap.html
@@ -1194,7 +1193,7 @@ class Parallel(Logger):
         [Parallel(n_jobs=2)]: Done 6 out of 6 | elapsed:  0.0s remaining: 0.0s
         [Parallel(n_jobs=2)]: Done 6 out of 6 | elapsed:  0.0s finished
 
-    '''
+    '''  # noqa: E501
     def __init__(
         self,
         n_jobs=default_parallel_config["n_jobs"],


Test patch:
diff --git a/continuous_integration/run_tests.sh b/continuous_integration/run_tests.sh
index 396bbe2..fedfe7f 100755
--- a/continuous_integration/run_tests.sh
+++ b/continuous_integration/run_tests.sh
@@ -30,7 +30,7 @@ if [[ "$SKLEARN_TESTS" == "true" ]]; then
     # TODO: unpin pip once either https://github.com/pypa/pip/issues/10825
     # accepts invalid HTML or Anaconda is fixed.
     conda install -y -c conda-forge cython pillow numpy scipy "pip<22"
-    pip install --pre --extra-index https://pypi.anaconda.org/scipy-wheels-nightly/simple scikit-learn
+    pip install --pre --extra-index https://pypi.anaconda.org/scientific-python-nightly-wheels/simple scikit-learn
     python -c "import sklearn; print('Testing scikit-learn', sklearn.__version__)"
 
     # Move to a dedicated folder to avoid being polluted by joblib specific conftest.py
diff --git a/joblib/test/test_memory.py b/joblib/test/test_memory.py
index ec9761f..120987b 100644
--- a/joblib/test/test_memory.py
+++ b/joblib/test/test_memory.py
@@ -571,8 +571,8 @@ def test_func_dir(tmpdir):
     assert g._check_previous_func_code()
 
     # Test the robustness to failure of loading previous results.
-    func_id, args_id = g._get_output_identifiers(1)
-    output_dir = os.path.join(g.store_backend.location, func_id, args_id)
+    args_id = g._get_args_id(1)
+    output_dir = os.path.join(g.store_backend.location, g.func_id, args_id)
     a = g(1)
     assert os.path.exists(output_dir)
     os.remove(os.path.join(output_dir, 'output.pkl'))
@@ -587,10 +587,10 @@ def test_persistence(tmpdir):
 
     h = pickle.loads(pickle.dumps(g))
 
-    func_id, args_id = h._get_output_identifiers(1)
-    output_dir = os.path.join(h.store_backend.location, func_id, args_id)
+    args_id = h._get_args_id(1)
+    output_dir = os.path.join(h.store_backend.location, h.func_id, args_id)
     assert os.path.exists(output_dir)
-    assert output == h.store_backend.load_item([func_id, args_id])
+    assert output == h.store_backend.load_item([h.func_id, args_id])
     memory2 = pickle.loads(pickle.dumps(memory))
     assert memory.store_backend.location == memory2.store_backend.location
 
@@ -668,9 +668,9 @@ def test_call_and_shelve_lazily_load_stored_result(tmpdir):
 
     memory = Memory(location=tmpdir.strpath, verbose=0)
     func = memory.cache(f)
-    func_id, argument_hash = func._get_output_identifiers(2)
+    args_id = func._get_args_id(2)
     result_path = os.path.join(memory.store_backend.location,
-                               func_id, argument_hash, 'output.pkl')
+                               func.func_id, args_id, 'output.pkl')
     assert func(2) == 5
     first_access_time = os.stat(result_path).st_atime
     time.sleep(1)
@@ -875,7 +875,7 @@ def _setup_toy_cache(tmpdir, num_inputs=10):
         get_1000_bytes(arg)
 
     func_id = _build_func_identifier(get_1000_bytes)
-    hash_dirnames = [get_1000_bytes._get_output_identifiers(arg)[1]
+    hash_dirnames = [get_1000_bytes._get_args_id(arg)
                      for arg in inputs]
 
     full_hashdirs = [os.path.join(get_1000_bytes.store_backend.location,
diff --git a/joblib/test/test_memory_async.py b/joblib/test/test_memory_async.py
new file mode 100644
index 0000000..ecad0c9
--- /dev/null
+++ b/joblib/test/test_memory_async.py
@@ -0,0 +1,149 @@
+import asyncio
+import gc
+import shutil
+
+import pytest
+
+from joblib.memory import (AsyncMemorizedFunc, AsyncNotMemorizedFunc,
+                           MemorizedResult, Memory, NotMemorizedResult)
+from joblib.test.common import np, with_numpy
+from joblib.testing import raises
+
+from .test_memory import (corrupt_single_cache_item,
+                          monkeypatch_cached_func_warn)
+
+
+async def check_identity_lazy_async(func, accumulator, location):
+    """ Similar to check_identity_lazy_async for coroutine functions"""
+    memory = Memory(location=location, verbose=0)
+    func = memory.cache(func)
+    for i in range(3):
+        for _ in range(2):
+            value = await func(i)
+            assert value == i
+            assert len(accumulator) == i + 1
+
+
+@pytest.mark.asyncio
+async def test_memory_integration_async(tmpdir):
+    accumulator = list()
+
+    async def f(n):
+        await asyncio.sleep(0.1)
+        accumulator.append(1)
+        return n
+
+    await check_identity_lazy_async(f, accumulator, tmpdir.strpath)
+
+    # Now test clearing
+    for compress in (False, True):
+        for mmap_mode in ('r', None):
+            memory = Memory(location=tmpdir.strpath, verbose=10,
+                            mmap_mode=mmap_mode, compress=compress)
+            # First clear the cache directory, to check that our code can
+            # handle that
+            # NOTE: this line would raise an exception, as the database
+            # file is still open; we ignore the error since we want to
+            # test what happens if the directory disappears
+            shutil.rmtree(tmpdir.strpath, ignore_errors=True)
+            g = memory.cache(f)
+            await g(1)
+            g.clear(warn=False)
+            current_accumulator = len(accumulator)
+            out = await g(1)
+
+        assert len(accumulator) == current_accumulator + 1
+        # Also, check that Memory.eval works similarly
+        evaled = await memory.eval(f, 1)
+        assert evaled == out
+        assert len(accumulator) == current_accumulator + 1
+
+    # Now do a smoke test with a function defined in __main__, as the name
+    # mangling rules are more complex
+    f.__module__ = '__main__'
+    memory = Memory(location=tmpdir.strpath, verbose=0)
+    await memory.cache(f)(1)
+
+
+@pytest.mark.asyncio
+async def test_no_memory_async():
+    accumulator = list()
+
+    async def ff(x):
+        await asyncio.sleep(0.1)
+        accumulator.append(1)
+        return x
+
+    memory = Memory(location=None, verbose=0)
+    gg = memory.cache(ff)
+    for _ in range(4):
+        current_accumulator = len(accumulator)
+        await gg(1)
+        assert len(accumulator) == current_accumulator + 1
+
+
+@with_numpy
+@pytest.mark.asyncio
+async def test_memory_numpy_check_mmap_mode_async(tmpdir, monkeypatch):
+    """Check that mmap_mode is respected even at the first call"""
+
+    memory = Memory(location=tmpdir.strpath, mmap_mode='r', verbose=0)
+
+    @memory.cache()
+    async def twice(a):
+        return a * 2
+
+    a = np.ones(3)
+    b = await twice(a)
+    c = await twice(a)
+
+    assert isinstance(c, np.memmap)
+    assert c.mode == 'r'
+
+    assert isinstance(b, np.memmap)
+    assert b.mode == 'r'
+
+    # Corrupts the file,  Deleting b and c mmaps
+    # is necessary to be able edit the file
+    del b
+    del c
+    gc.collect()
+    corrupt_single_cache_item(memory)
+
+    # Make sure that corrupting the file causes recomputation and that
+    # a warning is issued.
+    recorded_warnings = monkeypatch_cached_func_warn(twice, monkeypatch)
+    d = await twice(a)
+    assert len(recorded_warnings) == 1
+    exception_msg = 'Exception while loading results'
+    assert exception_msg in recorded_warnings[0]
+    # Asserts that the recomputation returns a mmap
+    assert isinstance(d, np.memmap)
+    assert d.mode == 'r'
+
+
+@pytest.mark.asyncio
+async def test_call_and_shelve_async(tmpdir):
+    async def f(x, y=1):
+        await asyncio.sleep(0.1)
+        return x ** 2 + y
+
+    # Test MemorizedFunc outputting a reference to cache.
+    for func, Result in zip((AsyncMemorizedFunc(f, tmpdir.strpath),
+                             AsyncNotMemorizedFunc(f),
+                             Memory(location=tmpdir.strpath,
+                                    verbose=0).cache(f),
+                             Memory(location=None).cache(f),
+                             ),
+                            (MemorizedResult, NotMemorizedResult,
+                             MemorizedResult, NotMemorizedResult,
+                             )):
+        for _ in range(2):
+            result = await func.call_and_shelve(2)
+            assert isinstance(result, Result)
+            assert result.get() == 5
+
+        result.clear()
+        with raises(KeyError):
+            result.get()
+        result.clear()  # Do nothing if there is no cache.