Repository: numba/numba

Issue:
Cannot cache functions with callable arguments
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I have tried using ~~the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG)~~ d58a63a870587bf8732038ca5b2229a4de218378.
- [x] I have included below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->
```python
import numba

@numba.njit(cache=True)
def mul_self(x):
	return x * x

@numba.njit(cache=True)
def compose_self_apply(f, x):
	return f(f(x))

compose_self_apply(mul_self, 2)
```
gives
<details>

<summary>Traceback</summary>

```python
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
c:\users\wiese\repos\ga\numba\numba\core\caching.py in IndexDataCacheFile.save(self, key, data)
    482 try:
    483     # If key already exists, we will overwrite the file
--> 484     data_name = overloads[key]
    485 except KeyError:
    486     # Find an available name for the data file

KeyError: ((type(CPUDispatcher(<function mul_self at 0x0000025918224C10>)), int64), ('x86_64-pc-windows-msvc', 'broadwell', '+64bit,+adx,-aes,+avx,+avx2,-avx512bf16,-avx512bitalg,-avx512bw,-avx512cd,-avx512dq,-avx512er,-avx512f,-avx512ifma,-avx512pf,-avx512vbmi,-avx512vbmi2,-avx512vl,-avx512vnni,-avx512vpopcntdq,+bmi,+bmi2,-cldemote,-clflushopt,-clwb,-clzero,+cmov,+cx16,+cx8,-enqcmd,+f16c,+fma,-fma4,+fsgsbase,+fxsr,-gfni,+invpcid,-lwp,+lzcnt,+mmx,+movbe,-movdir64b,-movdiri,-mpx,-mwaitx,+pclmul,-pconfig,-pku,+popcnt,-prefetchwt1,+prfchw,-ptwrite,-rdpid,+rdrnd,+rdseed,-rtm,+sahf,-sgx,-sha,-shstk,+sse,+sse2,+sse3,+sse4.1,+sse4.2,-sse4a,+ssse3,-tbm,-vaes,-vpclmulqdq,-waitpkg,-wbnoinvd,-xop,+xsave,-xsavec,+xsaveopt,-xsaves'), ('95d9e4857cb9755b29e175e557e615d8467821eb42c66a521b4270d86f6b8c06', 'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'))

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-7-74c10e900da8> in <module>
      7 @numba.njit(cache=True)
      8 def compose_self_apply(f, x):
      9         return f(f(x))
---> 11 compose_self_apply(mul_self, 2)

c:\users\wiese\repos\ga\numba\numba\core\dispatcher.py in _DispatcherBase._compile_for_args(self, *args, **kws)
    431         e.patch_message('\n'.join((str(e).rstrip(), help_msg)))
    432 # ignore the FULL_TRACEBACKS config, this needs reporting!
--> 433 raise e

c:\users\wiese\repos\ga\numba\numba\core\dispatcher.py in _DispatcherBase._compile_for_args(self, *args, **kws)
    364         argtypes.append(self.typeof_pyval(a))
    365 try:
--> 366     return self.compile(tuple(argtypes))
    367 except errors.ForceLiteralArg as e:
    368     # Received request for compiler re-entry with the list of arguments
    369     # indicated by e.requested_args.
    370     # First, check if any of these args are already Literal-ized
    371     already_lit_pos = [i for i in e.requested_args
    372                        if isinstance(args[i], types.Literal)]

c:\users\wiese\repos\ga\numba\numba\core\compiler_lock.py in _CompilerLock.__call__.<locals>._acquire_compile_lock(*args, **kwargs)
     29 @functools.wraps(func)
     30 def _acquire_compile_lock(*args, **kwargs):
     31     with self:
---> 32         return func(*args, **kwargs)

c:\users\wiese\repos\ga\numba\numba\core\dispatcher.py in Dispatcher.compile(self, sig)
    822     raise e.bind_fold_arguments(folded)
    823 self.add_overload(cres)
--> 824 self._cache.save_overload(sig, cres)
    825 return cres.entry_point

c:\users\wiese\repos\ga\numba\numba\core\caching.py in Cache.save_overload(self, sig, data)
    664 """
    665 Save the data for the given signature in the cache.
    666 """
    667 with self._guard_against_spurious_io_errors():
--> 668     self._save_overload(sig, data)

c:\users\wiese\repos\ga\numba\numba\core\caching.py in Cache._save_overload(self, sig, data)
    676 key = self._index_key(sig, _get_codegen(data))
    677 data = self._impl.reduce(data)
--> 678 self._cache_file.save(key, data)

c:\users\wiese\repos\ga\numba\numba\core\caching.py in IndexDataCacheFile.save(self, key, data)
    491             break
    492     overloads[key] = data_name
--> 493     self._save_index(overloads)
    494 self._save_data(data_name, data)

c:\users\wiese\repos\ga\numba\numba\core\caching.py in IndexDataCacheFile._save_index(self, overloads)
    537 def _save_index(self, overloads):
    538     data = self._source_stamp, overloads
--> 539     data = self._dump(data)
    540     with self._open_for_write(self._index_path) as f:
    541         pickle.dump(self._version, f, protocol=-1)

c:\users\wiese\repos\ga\numba\numba\core\caching.py in IndexDataCacheFile._dump(self, obj)
    566 def _dump(self, obj):
--> 567     return pickle.dumps(obj, protocol=-1)

TypeError: cannot pickle 'weakref' object
```

</details>



Discussion:
Thanks for the report - I can reproduce this. I've added the "needtriage" label for now, as I'm not sure whether the error message just needs to be better (along the lines of "Caching of functions with callable arguments is unsupported") or whether it should be supported and accepted as valid, and would be classed as a failure to compile.
If I attach a debugger, I can  use the following function
```python
def find_unpickleable(obj):
    def rec(obj):
        if isinstance(obj, tuple):
            for i, obj_i in enumerate(obj):
                r = rec(obj_i)
                if r is not None:
                    return lambda s: r(f"{s}[{i}]")
            return None
        elif isinstance(obj, dict):
            for i, (k, v) in enumerate(obj.items()):
                r = rec(k)
                if r is not None:
                    return lambda s: r(f'list({s}.keys())[{i}]')
                r = rec(v)
                if r is not None:
                    return lambda s: r(f'{s}[{k!r}]')
            return None

        if hasattr(obj, '__reduce__'):
            try:
                r = rec(obj.__reduce__())
            except TypeError:
                pass
            else:
                if r is not None:
                    return lambda s: r(f'{s}.__reduce__()')

        try:
            pickle.dumps(obj)
        except TypeError:
            return lambda s: s
    return rec(obj)
```
to find the unpickleable object
```python
ipdb> find_unpickleable(obj)('obj')
"list(obj[1].keys())[0][0][0].__reduce__()[1][2]['_wr']"
ipdb> p list(obj[1].keys())[0][0][0].__reduce__()[1][2]['_wr']
<weakref at 0x000002591961A220; to 'CPUDispatcher' at 0x00000259184C9970 (mul_self)>
```

At a glance, it seems that replacing `weakref.ref` with this class in `WeakType` seems to do the trick:
```python
class _PickleableWeakRef(weakref.ref):
    """
    Allow a weakref to be pickled.

    Note that this only unpickles usefully if the object referred to
    is kept alive elsewhere in the pickle, else the weakref will immediately expire.
    """
    def __getnewargs__(self):
        obj = self()
        if obj is None:
            raise ReferenceError("underlying object has vanished")
        return (obj,)
```
Thanks for the investigation! I've labelled this as a failure to compile for now.

Resolving patch:
diff --git a/numba/core/types/functions.py b/numba/core/types/functions.py
--- a/numba/core/types/functions.py
+++ b/numba/core/types/functions.py
@@ -457,6 +457,20 @@ class MakeFunctionLiteral(Literal, Opaque):
     pass
 
 
+class _PickleableWeakRef(weakref.ref):
+    """
+    Allow a weakref to be pickled.
+
+    Note that if the object referred to is not kept alive elsewhere in the
+    pickle, the weakref will immediately expire after being constructed.
+    """
+    def __getnewargs__(self):
+        obj = self()
+        if obj is None:
+            raise ReferenceError("underlying object has vanished")
+        return (obj,)
+
+
 class WeakType(Type):
     """
     Base class for types parametered by a mortal object, to which only
@@ -464,7 +478,7 @@ class WeakType(Type):
     """
 
     def _store_object(self, obj):
-        self._wr = weakref.ref(obj)
+        self._wr = _PickleableWeakRef(obj)
 
     def _get_object(self):
         obj = self._wr()


Test patch:
diff --git a/numba/tests/cache_usecases.py b/numba/tests/cache_usecases.py
--- a/numba/tests/cache_usecases.py
+++ b/numba/tests/cache_usecases.py
@@ -163,6 +163,21 @@ def runTest(self):
         raise NotImplementedError
 
 
+@jit(cache=True)
+def first_class_function_mul(x):
+    return x * x
+
+
+@jit(cache=True)
+def first_class_function_add(x):
+    return x + x
+
+
+@jit(cache=True)
+def first_class_function_usecase(f, x):
+    return f(x)
+
+
 def self_test():
     mod = sys.modules[__name__]
     _TestModule().check_module(mod)
diff --git a/numba/tests/test_dispatcher.py b/numba/tests/test_dispatcher.py
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -1266,6 +1266,17 @@ def test_closure(self):
             self.assertPreciseEqual(f(3), 12) # 3 + 9 = 12
             self.check_pycache(5) # 1 nbi, 4 nbc
 
+    def test_first_class_function(self):
+        mod = self.import_module()
+        f = mod.first_class_function_usecase
+        self.assertEqual(f(mod.first_class_function_mul, 1), 1)
+        self.assertEqual(f(mod.first_class_function_mul, 10), 100)
+        self.assertEqual(f(mod.first_class_function_add, 1), 2)
+        self.assertEqual(f(mod.first_class_function_add, 10), 20)
+        # 1 + 1 + 1 nbi, 1 + 1 + 2 nbc - a separate cache for each call to `f`
+        # with a different callback.
+        self.check_pycache(7)
+
     def test_cache_reuse(self):
         mod = self.import_module()
         mod.add_usecase(2, 3)