Caching: functions captured in closure

Author: sk1pCreated Sep 21, 2020Updated Sep 15, 2026

Labelsfeature_requestcaching

Feature request

I would like to capture functions in closures and be able to cache the compilation. Currently, running the following code results in re-compilation on each run (that is, when starting in a new Python interpreter):

python

import numpy as np
import numba


@numba.njit
def f1(in_arr):
    return in_arr * 2


def compose(fn):
    @numba.njit(cache=True)
    def _inner_fn(arr):
        return fn(arr)
    return _inner_fn


composed = compose(f1)
composed(np.ones((128, 128)))

Each time, a new version of composed is saved in __pycache__. This is caused by the UUID that is part of the Dispatcher - when looking at cvarbytes, only the UUID differs between runs.

As a proof of concept, the following hack for Cache._index_key makes the compose consistently cache-able:

diff

diff --git a/numba/core/caching.py b/numba/core/caching.py
index 45f82f314..f4f766790 100644
--- a/numba/core/caching.py
+++ b/numba/core/caching.py
@@ -701,11 +701,17 @@ class Cache(_Cache):
         codebytes = self._py_func.__code__.co_code
         if self._py_func.__closure__ is not None:
             cvars = tuple([x.cell_contents for x in self._py_func.__closure__])
+            cvars = tuple(var.__code__.co_code if hasattr(var, '__code__') else var
+                          for var in cvars)
             cvarbytes = dumps(cvars)
         else:
             cvarbytes = b''
 
         hasher = lambda x: hashlib.sha256(x).hexdigest()
         return (sig, codegen.magic_tuple(), (hasher(codebytes),
                                              hasher(cvarbytes),))

Implementing this as-is doesn't invalidate the cache in all cases, for example if there are multiple levels of closures. Maybe there is an easier and/or better way? Thanks!

Source: numba/numba

View original on GitHub View discussion on GitHub