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:
@sklam is checking on this.  Serialized functions should carry a UUID that allows us to detect when identical functions are being deserialized.
Weird, I don't get the same dispatcher object if I `cloudpickle.loads` with `%time`.

```

In [37]: id(cloudpickle.loads(s))
Out[37]: 4480579288

In [38]: id(cloudpickle.loads(s))
Out[38]: 4480579288

In [39]: id(cloudpickle.loads(s))
Out[39]: 4480579288

In [40]: id(cloudpickle.loads(s))
Out[40]: 4480579288

In [41]: id(cloudpickle.loads(s))
Out[41]: 4480579288

In [42]: %time id(cloudpickle.loads(s))
CPU times: user 288 µs, sys: 1 µs, total: 289 µs
Wall time: 293 µs
Out[42]: 4480579768

In [43]: %time id(cloudpickle.loads(s))
CPU times: user 281 µs, sys: 1e+03 ns, total: 282 µs
Wall time: 287 µs
Out[43]: 4480580008

In [44]: %time id(cloudpickle.loads(s))
CPU times: user 287 µs, sys: 1e+03 ns, total: 288 µs
Wall time: 293 µs
Out[44]: 4480580248
```
Ok, it's because our deserialization memo is a WeakValueDictionary.  The repeated compilation is due to the function going away immediately after `pickle.loads(b)(x)` is done.  If something else keeps a reference to the deserialized function, then future deserialization will give the previous still alive object.
So, the question is: Do we want to make the deserialized function immortal, or cleared on a time delay?
A modest time delay would suit my needs well.

On Fri, Jun 8, 2018, 3:25 PM Siu Kwan Lam <notifications@github.com> wrote:

> So, the question is: Do we want to make the deserialized function
> immortable, or cleared on a time delay?
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/numba/numba/issues/3026#issuecomment-395864116>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszO5wkkQFGH8nuDbuyQFJPWqSwSJmks5t6s-cgaJpZM4Ugn0->
> .
>

Or alternatively a rolling buffer of the last 100 deserialized functions or so.  Given that we have bytes we could also limit this by size of the serialized result.
So something like 1MB of serialized functions
More generically -- could you specify an LRU cache that has configurable limits (either TTL or size).
1MB like to be too small for many.
Are you suggesting that I do this on my end?  (that would be fine, though less pleasant than Numba handling it)

> 1MB like to be too small for many.

To be clear here we're talking about the serialized form of the function *pre-compilation*.  1MB seems like a decent number because it's small relative to the size of the CPython runtime but likely pretty large relative to function size.  That being said, I agree that extra configurability would be nice.

```python
In [1]: import numba

In [2]: @numba.njit
   ...: def f(x):
   ...:     total = 0
   ...:     for i in range(len(x)):
   ...:         total += x[i]
   ...:     return total
   ...: 

In [3]: import cloudpickle

In [4]: len(cloudpickle.dumps(f))
Out[4]: 435
```
I don't know if you could.

I'm interested in this because we serialize callable objects to perform distributed RPC - I'm fairly certain that (in total) our serialized python objects and LLVM objects would be >> 1MB. Currently we avoid this by only serializing pure Python functions. It would be nice to be able to serialize the JITed functions too --  re-compilation at every deserialization would be a killer.
> It would be nice to be able to serialize the JITed functions too -- re-compilation at every deserialization would be a killer

In the application I care about re-compilation once per process would be fine, so serialization of just Python source suffices.
This remains a pain point for some applications that use Numba and Dask together.  I'm curious if there is anything I can do to push on this.  I tried working around it in Dask but caching at the Dask level but ran into issues. I suspect that Dask isn't caching all of the right things, especially when there are closures and such.

Resolving patch:
diff --git a/numba/config.py b/numba/config.py
--- a/numba/config.py
+++ b/numba/config.py
@@ -152,6 +152,10 @@ def optional_str(x):
         # 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
--- a/numba/dispatcher.py
+++ b/numba/dispatcher.py
@@ -514,6 +514,9 @@ class attribute.
         }
     # 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 @@ def _set_uuid(self, u):
         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
--- 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 @@ def _set_uuid(self, u):
         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
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -9,6 +9,8 @@
 import threading
 import warnings
 import inspect
+import pickle
+import weakref
 
 try:
     import jinja2
@@ -26,6 +28,7 @@
                       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 @@ def bar():
         [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
--- a/numba/tests/test_extending.py
+++ b/numba/tests/test_extending.py
@@ -676,7 +676,13 @@ def codegen(context, builder, signature, args):
         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))