Repository: numba/numba (Python)

Issue:
@guvectorized functions over dask distributed
### UPDATE 2024
Statically defined gufuncs don't work with distributed anymore, either. Read: https://github.com/numba/numba/issues/4314#issuecomment-1992103833

### Original post

Dynamically defined functions - namely, anything from a Jupyter notebook - can be executed over a dask distributed grid thanks to cloudpickle.
Functions decorated with @guvectorize can be pickled, but only as long as they are defined in a static module.

Type in a jupyter notebook:

```python
import distributed
from numba import guvectorize, i8

def f(x):
    return x * 2


@guvectorize([(i8, i8[:])], '()->()')
def g(x, out):
    out[0] = x * 2

    
client = distributed.Client()

print(f(2))
print(g(2))
print(client.submit(f, 2).result())
print(client.submit(g, 2).result())
```
Output:
```
4
4
4

AttributeError: module '__main__' has no attribute 'g'
```

Note that the above is not the same as just writing ``cloudpickle.loads(cloudpickle.dumps(g))``, which works (but isn't terribly useful) because loads and dumps are running in the same interpreter.


Workaround:
```python
def cp_guvectorize(*args, **kwargs):
    """Same as :func:`numba.guvectorize`, but can be used to decorate dynamically
    defined function and then pickle them with
    `cloudpickle <https://pypi.org/project/cloudpickle/>`_.
    On the other hand, it can't be called from another jit-compiled function.
    """

    def decorator(func):
        return _PickleableGUVectorized(func, args, kwargs)

    return decorator


class _PickleableGUVectorized:
    def __init__(self, func, guvectorize_args, guvectorize_kwargs):
        self.args = func, guvectorize_args, guvectorize_kwargs
        decorator = numba.guvectorize(*guvectorize_args, **guvectorize_kwargs)
        self.ufunc = decorator(func)

    def __reduce__(self):
        return _PickleableGUVectorized, self.args

    def __call__(self, *args, **kwargs):
        return self.ufunc(*args, **kwargs)

    
@cp_guvectorize([(i8, i8[:])], '()->()')
def h(x, out):
    out[0] = x * 2

print(client.submit(h, 2).result())
```
Output:
```
4
```


Discussion:
I would also love to see this.  This might also be fixed by adding `__setstate__` and `__getstate__` methods.
Oh, I see that this returns a proper `numpy.ufunc` typed object.  Perhaps some wrapper solution like what @crusaderky suggests is necessary.
In the issue cited just above, I experienced a very similar serialization problem with a statically defined numba function (inside a module inside an installed package). Is there any documentation of successfully using dask.distributed + numba functions?
@rabernat last time I checked, ``@guvectorize``'d functions are pickleable, but ``@vectorize``'d ones aren't. ``@jit``'ed functions are pickleable only when you explicitly define a signature. It's been a while since I last tested this so I may be wrong.
I can replicate the issue with `multiprocessing`:

```python
from numba import guvectorize, f8
import pickle
import multiprocessing as mproc
import numpy as np



def make_g():

    @guvectorize([(f8, f8[:])], '()->()')
    def g(x, out):
        out[0] = x * 2

    return g


if __name__ == '__main__':
    spawner = mproc.get_context("spawn")
    x = np.ones(10)
    with spawner.Pool(1) as pool:
        g = make_g()
        res = pool.apply(g, [x])
        print(res)
```

This shows the error:

```pytb
  File "/Users/siu/miniconda3/envs/dev-numba-py37/lib/python3.7/site-packages/numpy/core/__init__.py", line 125, in _ufunc_reconstruct
    return getattr(mod, name)
AttributeError: module '__mp_main__' has no attribute 'g'
```

It looks like this is a numba gufunc problem. Numba `@vectorize`d and `@njit` functions are fine with and without explicit signatures.

This issue is caused by https://github.com/numpy/numpy/blob/f71d9937d1e8a1e709f325f689f1e971e64c26a7/numpy/core/__init__.py#L124-L135

`@vectorize` is unaffected because we subclass from ufunc to make a `DUfunc` class and that bypasses the numpy registered reduce/reconstruct pair. We will need to do the same for gufunc.
The numpy ufunc serialization code is similar to the builtin pickle procedure for Python functions, which uses the `__name__` and `__module__` to find the object. This will work for statically defined ufuncs, because the deserializer can import the module and find the function. Like picking regular functions, this breaks when the function isn't in a static location, for example a closure. Cloudpickle solves this problem for closures and dynamically defined functions by serializing the bytecode and closure cells directly so that the function object can be reconstructed by the deserializer. This works because, given the same version of Python, the bytecode is a uniform representation for function we want to transfer. Arbitrary numpy ufunc objects do not have a uniform representation. Instead, ufunc objects contain a bunch of C function pointers which help select more C functions for actually performing the operations. There is no general way to represent native code, and any higher level information needed to do so will have been erased by the time the ufunc object is created. For that reason, the numpy serialization method is likely as good as they can do for ufuncs. Numba still has the higher level information which could be used to serialize arbitrary code, so if `guvectorize` was made to return a box like proposed, then the ufunc-like objects could be serialized just like `vectorize` created ufuncs.
Is there any movement on this issue? 

I'm trying to compute a `@guvectorize` function over a distributed cluster (created using [dask-cloudprovider](https://github.com/dask/dask-cloudprovider)), and I'm running into the same error described [above](https://github.com/numba/numba/issues/4314#issuecomment-583444435).

I've found a couple of different issues regarding `ufuncs` and `distributed`. Some of them offer some workarounds, like importing the function differently. This works for me on a local cluster, but not for the one running on ECS instances created through cloudprovider. 
This came up again in a Pangeo call.  Given that Numpy moves slowly, I wonder if there is a workaround.  It looks like @sklam proposed a solution here:

> @vectorize is unaffected because we subclass from ufunc to make a DUfunc class and that bypasses the numpy registered reduce/reconstruct pair. We will need to do the same for gufunc.

Do we have a sense for how long this would take?  Are there people or resources around that would make sense to allocate here?
Thanks for the push on this @mrocklin. 

First, I want to note that I think this problem does also apply to `@vectorize`. Here is an example of a `@vectorize` function defined in a package:
https://github.com/xgcm/fastjmd95/blob/c828391890d18c82ff5e44484409055605111c1e/fastjmd95/jmd95numba.py#L150-L154
And here is an example of that function failing in distributed mode:
https://nbviewer.jupyter.org/github/xgcm/fastjmd95/blob/master/doc/fastjmd95_tutorial.ipynb

I would love if someone with more numba experience could verify this by pip installing `fastjmd95` and calling `jmd95numba.rho` in distributed mode. Perhaps I've done something wrong in packaging my functions.

More generally, I would love to support the overall goal of making all numba functions serializable and compatible with dask.distributed. In Pangeo, we fell like this is a great path to scaling out performant code. But the challenges described in this issue require annoying workarounds.
xref https://github.com/numba/numba/pull/6234
Resurrecting this as the situation has somewhat gotten worse.
As of numba 0.59, the code in my initial post doesn't work anymore even when the function is defined in a static module.
`@vectorize` still works.

The issue lays in `GUFunc.__call__`:
https://github.com/numba/numba/blob/54433810bfa60458209ed2c24f62e793b12aca6b/numba/np/ufunc/gufunc.py#L171-L172

These two lines undo the serialization hack implemented by GUFunc:
https://github.com/numba/numba/blob/54433810bfa60458209ed2c24f62e793b12aca6b/numba/np/ufunc/gufunc.py#L33-L46

By the time control reaches `dask.array.Array.__array_ufunc__`, it's too late - we've lost the reference to the GUFunc object as well as the wrapped pure-python object. I can't find a way to hack around it in our serialization layer.

The workaround is to bypass `GUFunc.__call__`:
```python
@numba.guvectorize(["f8,f8[:]"], "()->()")
def double(x, out):
    out[:] = x * 2

x = da.random.randint(...)
y = x.__array_ufunc__(double, "__call__", x)
client = distributed.Client()
y.compute()  # Works
```
the above calls `GUFunc.__call__` for the first time on the remote worker, on the wrapped numpy data.

# Proposed design
`GUFunc.__call__` should special-case dask.
```python
try:
    from dask.base import is_dask_collection
except ModuleNotFoundError:
    def is_dask_collection(x):
        return False

class GUFunc:
    def call(self, *args, **kwargs):
        if is_dask_collection(args[0]):
            return args[0].__array_ufunc__(self, "__call__", *args, **kwargs)
        else:
            return self.ufunc(*args, **kwargs) 
```
The above will work transparently with dask wrappers too, such as xarray and pint.



Patch:
diff --git a/docs/upcoming_changes/9495.new_feature.rst b/docs/upcoming_changes/9495.new_feature.rst
new file mode 100644
index 00000000000..e65e6e97c75
--- /dev/null
+++ b/docs/upcoming_changes/9495.new_feature.rst
@@ -0,0 +1,5 @@
+Schedule guvectorize'd functions over dask.distributed
+------------------------------------------------------
+
+Functions decorated with `@guvectorize` can now be scheduled over distributed
+`Dask <https://www.dask.org/>`_ clusters.
diff --git a/numba/np/ufunc/gufunc.py b/numba/np/ufunc/gufunc.py
index b133e9b1ad5..103ef7b5d95 100644
--- a/numba/np/ufunc/gufunc.py
+++ b/numba/np/ufunc/gufunc.py
@@ -249,7 +249,16 @@ def __call__(self, *args, **kwargs):
         # If compilation is disabled OR it is NOT a dynamic gufunc
         # call the underlying gufunc
         if self._frozen or not self.is_dynamic:
-            return self.ufunc(*args, **kwargs)
+            # Do not unwrap the ufunc if the argument is a wrapper that will
+            # potentially pickle the ufunc after it receives it in
+            # __array_ufunc__. The same logic in theory should be replicated
+            # for reduce(), outer(), etc., but they're not implemented in dask.
+            if args and _is_array_wrapper(args[0]):
+                return args[0].__array_ufunc__(
+                    self, "__call__", *args, **kwargs
+                )
+            else:
+                return self.ufunc(*args, **kwargs)
         elif "out" in kwargs:
             # If "out" argument is supplied
             args += (kwargs.pop("out"),)
@@ -277,3 +286,38 @@ def __call__(self, *args, **kwargs):
             self.build_ufunc()
 
         return self.ufunc(*args, **kwargs)
+
+
+def _is_array_wrapper(obj):
+    """Return True if obj wraps around numpy or another numpy-like library
+    and is likely going to apply the ufunc to the wrapped array; False
+    otherwise.
+
+    At the moment, this returns True for
+
+    - dask.array.Array
+    - dask.dataframe.DataFrame
+    - dask.dataframe.Series
+    - xarray.DataArray
+    - xarray.Dataset
+    - xarray.Variable
+    - pint.Quantity
+    - other potential wrappers around dask array or dask dataframe
+
+    We may need to add other libraries that pickle ufuncs from their
+    __array_ufunc__ method in the future.
+
+    Note that the below test is a lot more naive than
+    `dask.base.is_dask_collection`
+    (https://github.com/dask/dask/blob/5949e54bc04158d215814586a44d51e0eb4a964d/dask/base.py#L209-L249),  # noqa: E501
+    because it doesn't need to find out if we're actually dealing with
+    a dask collection, only that we're dealing with a wrapper.
+    Namely, it will return True for a pint.Quantity wrapping around a plain float, a
+    numpy.ndarray, or a dask.array.Array, and it's OK because in all cases
+    Quantity.__array_ufunc__ is going to forward the ufunc call inwards.
+    """
+    return (
+        not isinstance(obj, type)
+        and hasattr(obj, "__dask_graph__")
+        and hasattr(obj, "__array_ufunc__")
+    )


Test patch:
diff --git a/numba/tests/npyufunc/test_ufuncbuilding.py b/numba/tests/npyufunc/test_ufuncbuilding.py
index 556e6659f6d..df6e37dcd5c 100644
--- a/numba/tests/npyufunc/test_ufuncbuilding.py
+++ b/numba/tests/npyufunc/test_ufuncbuilding.py
@@ -1,4 +1,8 @@
+import pickle
+import unittest
+
 import numpy as np
+from numpy.testing import assert_array_equal
 
 from numba.np.ufunc.ufuncbuilder import GUFuncBuilder
 from numba import vectorize, guvectorize
@@ -6,7 +10,6 @@
 from numba.np.ufunc.dufunc import DUFunc as UFuncBuilder
 from numba.tests.support import tag, TestCase
 from numba.core import config
-import unittest
 
 
 class TestUfuncBuilding(TestCase):
@@ -317,6 +320,58 @@ def test_guvectorize_invalid_layout(self):
         # (error message depends on Numpy version)
 
 
+class NEP13Array:
+    """https://numpy.org/neps/nep-0013-ufunc-overrides.html"""
+    def __init__(self, array):
+        self.array = array
+
+    def __array__(self):
+        return self.array
+
+    def tolist(self):
+        return self.array.tolist()
+
+    def __array_ufunc__(self, ufunc, method, *args, **kwargs):
+        if method != "__call__":
+            return NotImplemented
+
+        return NEP13Array(ufunc(*[np.asarray(x) for x in args], **kwargs))
+
+
+class FakeDaskArray:
+    """This class defines both the NEP13 protocol and the dask collection protocol
+    (https://docs.dask.org/en/stable/custom-collections.html). This is a stand-in for
+    dask array, dask dataframe, and for any wrapper around them (e.g. xarray or pint).
+    """
+
+    def __init__(self, array):
+        self.array = array
+
+    def __array_ufunc__(self, ufunc, method, *args, **kwargs):
+        if method != "__call__":
+            return NotImplemented
+
+        # Simulate sending the ufunc over the network and applying it on a remote worker
+        ufunc = pickle.loads(pickle.dumps(ufunc))
+        args = [x.array if isinstance(x, FakeDaskArray) else x for x in args]
+        return FakeDaskArray(ufunc(*args, **kwargs))
+
+    def _dask_method(self, *args, **kwargs):
+        raise AssertionError("called potentially expensive method")
+
+    __array__ = _dask_method
+    __dask_graph__ = _dask_method
+    __dask_keys__ = _dask_method
+    __dask_optimize__ = _dask_method
+    __dask_postcompute__ = _dask_method
+    __dask_postpersist__ = _dask_method
+    __dask_scheduler__ = _dask_method
+    __dask_tokenize__ = _dask_method
+    compute = _dask_method
+    persist = _dask_method
+    visualize = _dask_method
+
+
 class TestNEP13WithoutSignature(TestCase):
 
     def test_all(self):
@@ -326,23 +381,6 @@ def test_all(self):
         def new_ufunc(hundreds, tens, ones):
             return 100*hundreds + 10*tens + ones
 
-        # https://numpy.org/neps/nep-0013-ufunc-overrides.html
-        class NEP13Array:
-            def __init__(self, array):
-                self.array = array
-
-            def __array__(self):
-                return self.array
-
-            def tolist(self):
-                return self.array.tolist()
-
-            def __array_ufunc__(self, ufunc, method, *args, **kwargs):
-                if method != "__call__":
-                    return NotImplemented
-
-                return NEP13Array(ufunc(*[np.asarray(x) for x in args], **kwargs))
-
         # give it integers
         a = np.array([1, 2, 3], dtype=np.int64)
         b = np.array([4, 5, 6], dtype=np.int64)
@@ -386,6 +424,39 @@ def __array_ufunc__(self, ufunc, method, *args, **kwargs):
         self.assertEqual(nep13_3.tolist(), [161.7, 283.8, 405.9])
 
 
+class TestDask(unittest.TestCase):
+    """Test that numba ufuncs are compatible with dask collections and wrappers around
+    dask (e.g. xarray or pint) and that they can be serialized, sent over the network,
+    deserialized on a different host and applied remotely.
+    """
+
+    def test_dask_array(self):
+        a = FakeDaskArray(np.arange(4, dtype=np.float64))
+        expect = np.arange(4, dtype=np.float64) * 2
+
+        @vectorize(["f8(f8)"])
+        def double_static_vectorize(x):
+            return x * 2
+
+        @vectorize()
+        def double_dynamic_vectorize(x):
+            return x * 2
+
+        @guvectorize(["f8,f8[:]"], "()->()")
+        def double_guvectorize(x, out):
+            out[:] = x * 2
+
+        for func in (
+            double_static_vectorize,
+            double_dynamic_vectorize,
+            double_guvectorize,
+        ):
+            with self.subTest(func):
+                b = func(a)
+                assert isinstance(b, FakeDaskArray)
+                assert_array_equal(b.array, expect)
+
+
 class TestVectorizeDecorJitDisabled(TestVectorizeDecor):
 
     def setUp(self):