Repository: numba/numba

Issue:
Slow performance on repeated deserialization
I'm running into performance problems when I deserialize a numba function many times within the same process.  I have a small example here:

### Process 1

```python
import numba
@numba.njit
def f(x):
    total = 0
    for i in range(len(x)):
        total += x[i]
    return total
import cloudpickle
cloudpickle.dumps(f)
```

```
b'\x80\x04\x95\xa8\x01\x00\x00\x00\x00\x00\x00\x8c\x0fnumba.serialize\x94\x8c\x12_rebuild_reduction\x94\x93\x94(\x8c\x16numba.targets.registry\x94\x8c\rCPUDispatcher\x94\x93\x94\x8c$64ebbc90-6b3a-11e8-963f-49110aa33909\x94(K\x04C\x043\r\r\n\x94C\xa4\xe3\x01\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x04\x00\x00\x00C\x00\x00\x00s*\x00\x00\x00d\x01}\x01x t\x00t\x01|\x00\x83\x01\x83\x01D\x00]\x10}\x02|\x01|\x00|\x02\x19\x007\x00}\x01q\x12W\x00|\x01S\x00)\x02N\xe9\x00\x00\x00\x00)\x02\xda\x05range\xda\x03len)\x03\xda\x01x\xda\x05total\xda\x01i\xa9\x00r\x07\x00\x00\x00\xfa\x1e<ipython-input-2-359c566f9f92>\xda\x01f\x01\x00\x00\x00s\x08\x00\x00\x00\x00\x02\x04\x01\x12\x01\x10\x01\x94\x87\x94}\x94(\x8c\x05range\x94\x8c\x08builtins\x94\x8c\x05range\x94\x93\x94\x8c\x03len\x94\x8c\x08builtins\x94\x8c\x03len\x94\x93\x94\x8c\x08__name__\x94\x8c\x08__main__\x94u\x8c\x01f\x94Nt\x94}\x94}\x94\x8c\x08nopython\x94\x88s\x8c\x06direct\x94\x88]\x94t\x94R\x94.'
```

### Process 2

```python
import pickle
import numpy as np
x = np.random.random(1000)
b =  b'\x80\x04\x95\xa8\x01\x00\x00\x00\x00\x00\x00\x8c\x0fnumba.serialize\x94\x8c\x12_rebuild_reduction\x94\x93\x94(\x8c\x16numba.targets.registry\x94\x8c\rCPUDispatcher\x94\x93\x94\x8c$64ebbc90-6b3a-11e8-963f-49110aa33909\x94(K\x04C\x043\r\r\n\x94C\xa4\xe3\x01\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x04\x00\x00\x00C\x00\x00\x00s*\x00\x00\x00d\x01}\x01x t\x00t\x01|\x00\x83\x01\x83\x01D\x00]\x10}\x02|\x01|\x00|\x02\x19\x007\x00}\x01q\x12W\x00|\x01S\x00)\x02N\xe9\x00\x00\x00\x00)\x02\xda\x05range\xda\x03len)\x03\xda\x01x\xda\x05total\xda\x01i\xa9\x00r\x07\x00\x00\x00\xfa\x1e<ipython-input-2-359c566f9f92>\xda\x01f\x01\x00\x00\x00s\x08\x00\x00\x00\x00\x02\x04\x01\x12\x01\x10\x01\x94\x87\x94}\x94(\x8c\x05range\x94\x8c\x08builtins\x94\x8c\x05range\x94\x93\x94\x8c\x03len\x94\x8c\x08builtins\x94\x8c\x03len\x94\x93\x94\x8c\x08__name__\x94\x8c\x08__main__\x94u\x8c\x01f\x94Nt\x94}\x94}\x94\x8c\x08nopython\x94\x88s\x8c\x06direct\x94\x88]\x94t\x94R\x94.'
```

We see that repeated calls of deserializing and then calling the function remain somewhat slow at 60ms per call.

```python
In [5]: %time pickle.loads(b)(x)
CPU times: user 255 ms, sys: 23.8 ms, total: 279 ms
Wall time: 281 ms
Out[5]: 515.6796588752262

In [6]: %time pickle.loads(b)(x)
CPU times: user 63.4 ms, sys: 3.84 ms, total: 67.2 ms
Wall time: 66.2 ms
Out[6]: 515.6796588752262

In [7]: %time pickle.loads(b)(x)
CPU times: user 68.3 ms, sys: 0 ns, total: 68.3 ms
Wall time: 67.1 ms
Out[7]: 515.6796588752262
```

The cost of deserializing is relatively low

```python
In [8]: %time f = pickle.loads(b)
CPU times: user 619 µs, sys: 66 µs, total: 685 µs
Wall time: 696 µs
```

And if we call this function many times then things are ok

```python
In [9]: %time f(x)
CPU times: user 65.9 ms, sys: 3.42 ms, total: 69.3 ms
Wall time: 68.6 ms
Out[9]: 515.6796588752262

In [10]: %time f(x)
CPU times: user 25 µs, sys: 2 µs, total: 27 µs
Wall time: 34.6 µs
Out[10]: 515.6796588752262

In [11]: %time f(x)
CPU times: user 36 µs, sys: 3 µs, total: 39 µs
Wall time: 52.7 µs
Out[11]: 515.6796588752262
```

So my guess is that we should be deduplicating things in some way.  Is this in scope for Numba to resolve or is this something that I should be handling on my end?


Discussion:

Patch:
diff --git a/docs/source/reference/envvars.rst b/docs/source/reference/envvars.rst
index 3041af0f2..45e7cdf70 100644
--- a/docs/source/reference/envvars.rst
+++ b/docs/source/reference/envvars.rst
@@ -235,6 +235,20 @@ Compilation options
         portable code (portable within the same architecture and OS),
         simply set ``NUMBA_CPU_NAME=generic``.
 
+.. envvar:: NUMBA_FUNCTION_CACHE_SIZE
+
+    Override the size of the function cache for retaining recently
+    deserialized functions in memory.  In systems like
+    `Dask <http://dask.pydata.org>`_, it is common for functions to be deserialized
+    multiple times.  Numba will cache functions as long as there is a
+    reference somewhere in the interpreter.  This cache size variable controls
+    how many functions that are no longer referenced will also be retained,
+    just in case they show up in the future.  The implementation of this is
+    not a true LRU, but the large size of the cache should be sufficient for
+    most situations.
+
+    *Default value:* 128
+
 
 GPU support
 -----------
diff --git a/numba/config.py b/numba/config.py
index defaaba98..b7f3a4e8d 100644
--- a/numba/config.py
+++ b/numba/config.py
@@ -152,6 +152,10 @@ class _EnvReloader(object):
         # Enable debugging of front-end operation (up to and including IR generation)
         DEBUG_FRONTEND = _readenv("NUMBA_DEBUG_FRONTEND", int, 0)
 
+        # How many recently deserialized functions to retain regardless
+        # of external references
+        FUNCTION_CACHE_SIZE = _readenv("NUMBA_FUNCTION_CACHE_SIZE", int, 128)
+
         # Enable logging of cache operation
         DEBUG_CACHE = _readenv("NUMBA_DEBUG_CACHE", int, DEBUG)
 
diff --git a/numba/dispatcher.py b/numba/dispatcher.py
index 471ea592b..58f52e4c4 100644
--- a/numba/dispatcher.py
+++ b/numba/dispatcher.py
@@ -514,6 +514,9 @@ class Dispatcher(_DispatcherBase):
         }
     # A {uuid -> instance} mapping, for deserialization
     _memo = weakref.WeakValueDictionary()
+    # hold refs to last N functions deserialized, retaining them in _memo
+    # regardless of whether there is another reference
+    _recent = collections.deque(maxlen=config.FUNCTION_CACHE_SIZE)
     __uuid = None
     __numba__ = 'py_func'
 
@@ -624,6 +627,7 @@ class Dispatcher(_DispatcherBase):
         assert self.__uuid is None
         self.__uuid = u
         self._memo[u] = self
+        self._recent.append(self)
 
     def compile(self, sig):
         if not self._can_compile:
diff --git a/numba/extending.py b/numba/extending.py
index 623dc8138..72451f20c 100644
--- a/numba/extending.py
+++ b/numba/extending.py
@@ -2,8 +2,9 @@
 import inspect
 import uuid
 import weakref
+import collections
 
-from numba import types
+from numba import types, config
 
 # Exported symbols
 from .typing.typeof import typeof_impl
@@ -216,6 +217,10 @@ class _Intrinsic(object):
     Dummy callable for intrinsic
     """
     _memo = weakref.WeakValueDictionary()
+    # hold refs to last N functions deserialized, retaining them in _memo
+    # regardless of whether there is another reference
+    _recent = collections.deque(maxlen=config.FUNCTION_CACHE_SIZE)
+
     __uuid = None
 
     def __init__(self, name, defn, support_literals=False):
@@ -241,6 +246,7 @@ class _Intrinsic(object):
         assert self.__uuid is None
         self.__uuid = u
         self._memo[u] = self
+        self._recent.append(self)
 
     def _register(self):
         from .typing.templates import make_intrinsic_template, infer_global


Test patch:
diff --git a/numba/tests/test_dispatcher.py b/numba/tests/test_dispatcher.py
index 38b8c862d..38d657a72 100644
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -9,6 +9,8 @@ import sys
 import threading
 import warnings
 import inspect
+import pickle
+import weakref
 
 try:
     import jinja2
@@ -26,6 +28,7 @@ from .support import (TestCase, tag, temp_directory, import_dynamic,
                       override_env_config, capture_cache_log, captured_stdout)
 from numba.targets import codegen
 from numba.caching import _UserWideCacheLocator
+from numba.dispatcher import Dispatcher
 
 import llvmlite.binding as ll
 
@@ -337,6 +340,50 @@ class TestDispatcher(BaseTest):
         [cr] = bar.overloads.values()
         self.assertEqual(len(cr.lifted), 1)
 
+    def test_serialization(self):
+        """
+        Test serialization of Dispatcher objects
+        """
+        @jit(nopython=True)
+        def foo(x):
+            return x + 1
+
+        self.assertEqual(foo(1), 2)
+
+        # get serialization memo
+        memo = Dispatcher._memo
+        Dispatcher._recent.clear()
+        memo_size = len(memo)
+
+        # pickle foo and check memo size
+        serialized_foo = pickle.dumps(foo)
+        # increases the memo size
+        self.assertEqual(memo_size + 1, len(memo))
+
+        # unpickle
+        foo_rebuilt = pickle.loads(serialized_foo)
+        self.assertEqual(memo_size + 1, len(memo))
+
+        self.assertIs(foo, foo_rebuilt)
+
+        # do we get the same object even if we delete all the explict references?
+        id_orig = id(foo_rebuilt)
+        del foo
+        del foo_rebuilt
+        self.assertEqual(memo_size + 1, len(memo))
+        new_foo = pickle.loads(serialized_foo)
+        self.assertEqual(id_orig, id(new_foo))
+
+        # now clear the recent cache
+        ref = weakref.ref(new_foo)
+        del new_foo
+        Dispatcher._recent.clear()
+        self.assertEqual(memo_size, len(memo))
+
+        # show that deserializing creates a new object
+        newer_foo = pickle.loads(serialized_foo)
+        self.assertIs(ref(), None)
+
 
 class TestSignatureHandling(BaseTest):
     """
diff --git a/numba/tests/test_extending.py b/numba/tests/test_extending.py
index 3b377ce1b..565726ccb 100644
--- a/numba/tests/test_extending.py
+++ b/numba/tests/test_extending.py
@@ -676,7 +676,13 @@ class TestIntrinsic(TestCase):
         memo_size += 1
         self.assertEqual(memo_size, len(memo))
         del original  # remove original before unpickling
-        # by deleting, the memo entry is removed
+
+        # by deleting, the memo entry is NOT removed due to recent
+        # function queue
+        self.assertEqual(memo_size, len(memo))
+
+        # Manually force clear of _recent queue
+        _Intrinsic._recent.clear()
         memo_size -= 1
         self.assertEqual(memo_size, len(memo))