Repository: dask/dask
Pull request #11697 (2025-01-24 10:36:09)

array: Never share buffers after compute()
- ~Fail loudly~ issue DeprecationWarning when a user attempts `np.asarray(x, copy=False)`, in accordance with the Array API standard. This came up during the work to introduce dask support for scipy through Array API compliance.
- Ensure that `x.compute()` never returns a buffer that is shared with the dask graph or a process-local `Worker.data`, so that users can't write back to it
- Ensure test coverage for both numpy 2.0 and 2.2

FYI @lithomas1 @lucascolley 

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     15 files  ±  0       15 suites  ±0   4h 26m 37s ⏱️ + 6m 32s
 17 170 tests + 29   15 990 ✅ + 29   1 180 💤 ±  0  0 ❌ ±0 
211 589 runs  +377  194 450 ✅ +263  17 139 💤 +114  0 ❌ ±0 

Results for commit 80006b79. ± Comparison against base commit d8eb38e9.

[test-results]:data:application/gzip;base64,H4sIADN5k2cC/02MOw7DIBAFr2JRp2BtloVcJiIYJBR/IgxVlLsHW8SmnHlP82E+TG5j9w7w1rEth3TRmKNJYV0O1pqKKms6dgLif35s2dp6auQrvHcJ6nLehKm4S7gY11hNzMue7gFQ6cpnWguB/JQ1TTCcxyZ9cFu26zyHVIApzrl8khYGRwJlucBByd4rAi2Fk1KQcYDIvj8N1JBAGgEAAA==

crusaderky: @fjetter @phofl @hendrikmakait 
Any comments?

Patch:
diff --git a/.gitignore b/.gitignore
index 20b3da0ed..c6df6f8bb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -23,5 +23,6 @@ coverage.xml
 hdfs-initialized-indicator
 .ipynb_checkpoints
 .vscode/
+.env
 .history
 env.yaml
diff --git a/continuous_integration/environment-3.11.yaml b/continuous_integration/environment-3.11.yaml
index d99e71028..ab3536413 100644
--- a/continuous_integration/environment-3.11.yaml
+++ b/continuous_integration/environment-3.11.yaml
@@ -25,7 +25,7 @@ dependencies:
   - moto<5
   # Optional dependencies
   - mimesis
-  - numpy
+  - numpy=2.0
   - pandas=2
   - numba
   - flask
diff --git a/continuous_integration/environment-3.12.yaml b/continuous_integration/environment-3.12.yaml
index 9b6614ab1..079c5afef 100644
--- a/continuous_integration/environment-3.12.yaml
+++ b/continuous_integration/environment-3.12.yaml
@@ -24,7 +24,7 @@ dependencies:
   - moto<5
   # Optional dependencies
   - mimesis
-  - numpy >=2  # only tested here
+  - numpy>=2
   - pandas
   - numba
   - flask
diff --git a/continuous_integration/environment-3.13.yaml b/continuous_integration/environment-3.13.yaml
index ec50b14dd..2f559866f 100644
--- a/continuous_integration/environment-3.13.yaml
+++ b/continuous_integration/environment-3.13.yaml
@@ -23,7 +23,7 @@ dependencies:
   - moto<5
   # Optional dependencies
   - mimesis
-  - numpy >=2  # only tested here
+  - numpy>=2.2  # only tested here
   - pandas
   #- numba  # not available for py 13
   - flask
diff --git a/dask/array/core.py b/dask/array/core.py
index b54ec1e44..cc6c18816 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1303,7 +1303,16 @@ def finalize(results):
             return concatenate3(results)
         else:
             results2 = results2[0]
-    return unpack_singleton(results)
+
+    results = unpack_singleton(results)
+    # Single chunk. There is a risk that the result holds a buffer stored in the
+    # graph or on a process-local Worker. Deep copy to make sure that nothing can
+    # accidentally write back to it.
+    try:
+        return results.copy()  # numpy, sparse, scipy.sparse (any version)
+    except AttributeError:
+        # Not an Array API object
+        return results
 
 
 CHUNKS_NONE_ERROR_MESSAGE = """
@@ -1717,13 +1726,29 @@ class Array(DaskMethodsMixin):
 
     __array_priority__ = 11  # higher than numpy.ndarray and numpy.matrix
 
-    def __array__(self, dtype=None, **kwargs):
+    def __array__(self, dtype=None, copy=None, **kwargs):
+        if kwargs:
+            warnings.warn(
+                f"Extra keyword arguments {kwargs} are ignored and won't be "
+                "accepted in the future",
+                FutureWarning,
+            )
+        if copy is False:
+            warnings.warn(
+                "Can't acquire a memory view of a Dask array. "
+                "This will raise in the future.",
+                FutureWarning,
+            )
+
         x = self.compute()
-        if dtype and x.dtype != dtype:
-            x = x.astype(dtype)
-        if not isinstance(x, np.ndarray):
-            x = np.array(x)
-        return x
+
+        # Apply requested dtype and convert non-numpy backends to numpy.
+        # If copy is True, numpy is going to perform its own deep copy
+        # after this method returns.
+        # If copy is None, finalize() ensures that the returned object
+        # does not share memory with an object stored in the graph or on a
+        # process-local Worker.
+        return np.asarray(x, dtype=dtype)
 
     def __array_function__(self, func, types, args, kwargs):
         import dask.array as module


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 0b890456e..f6873bf41 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -52,7 +52,7 @@ from dask.array.core import (
     stack,
     store,
 )
-from dask.array.numpy_compat import NUMPY_GE_200
+from dask.array.numpy_compat import NUMPY_GE_200, NUMPY_GE_210
 from dask.array.reshape import _not_implemented_message
 from dask.array.utils import assert_eq, same_keys
 from dask.base import collections_to_dsk, compute_as_if_collection, tokenize
@@ -410,6 +410,136 @@ def test_Array_computation():
     assert float(a[0, 0]) == 1
 
 
+@pytest.mark.parametrize("asarray", [np.asarray, np.asanyarray, np.array])
+def test_numpy_asarray_dtype(asarray):
+    """Test dtype= parameter of np.*array()"""
+    x = da.arange(10, dtype="i2")
+    nx = asarray(x, dtype="i4")
+    assert nx.dtype == np.int32
+
+
+@pytest.mark.parametrize(
+    "asarray",
+    [
+        np.array,
+        pytest.param(
+            np.asarray,
+            marks=pytest.mark.skipif(not NUMPY_GE_200, reason="no copy kwarg"),
+        ),
+        pytest.param(
+            np.asanyarray,
+            marks=pytest.mark.skipif(not NUMPY_GE_210, reason="no copy kwarg"),
+        ),
+    ],
+)
+@pytest.mark.parametrize("chunks", [5, 10])
+def test_numpy_asarray_copy_true(asarray, chunks):
+    """Test np.*array(x, copy=True)"""
+    x = da.asarray(np.arange(10), chunks=chunks)
+
+    nx = asarray(x, copy=True)
+    nx[0] = 42
+    # Did not write back to the buffer held in the Dask graph
+    assert x[0].compute() == 0
+
+
+@pytest.mark.parametrize(
+    "asarray",
+    [
+        pytest.param(
+            np.array,
+            marks=pytest.mark.skipif(
+                not NUMPY_GE_200, reason="copy kwarg not forwarded to dask"
+            ),
+        ),
+        pytest.param(
+            np.asarray,
+            marks=pytest.mark.skipif(not NUMPY_GE_200, reason="no copy kwarg"),
+        ),
+        pytest.param(
+            np.asanyarray,
+            marks=pytest.mark.skipif(not NUMPY_GE_210, reason="no copy kwarg"),
+        ),
+    ],
+)
+@pytest.mark.parametrize("chunks", [5, 10])
+def test_numpy_asarray_copy_false(asarray, chunks):
+    """Test that np.*array(x, copy=False) is forbidden"""
+    x = da.asarray(np.arange(10), chunks=chunks)
+    # Loudly crash if numpy demands for a view of dask's memory
+    with pytest.warns(FutureWarning, match="Can't acquire a memory view"):
+        y = asarray(x, copy=False)
+    assert_eq(y, np.arange(10))
+
+
+@pytest.mark.parametrize(
+    "asarray",
+    [
+        pytest.param(
+            np.array,
+            marks=pytest.mark.skipif(not NUMPY_GE_200, reason="copy=None invalid"),
+        ),
+        pytest.param(
+            np.asarray,
+            marks=pytest.mark.skipif(not NUMPY_GE_200, reason="no copy kwarg"),
+        ),
+        pytest.param(
+            np.asanyarray,
+            marks=pytest.mark.skipif(not NUMPY_GE_210, reason="no copy kwarg"),
+        ),
+    ],
+)
+@pytest.mark.parametrize("chunks", [5, 10])
+def test_numpy_asarray_copy_none(asarray, chunks):
+    """Test np.*array(x, copy=None)"""
+    x = da.asarray(np.arange(10), chunks=chunks)
+
+    nx = asarray(x, copy=None)
+    nx[0] = 42
+    # Did not write back to the buffer held in the Dask graph
+    assert x[0].compute() == 0
+
+
+@pytest.mark.parametrize("asarray", [np.asarray, np.asanyarray, np.array])
+@pytest.mark.parametrize("chunks", [5, 10])
+def test_numpy_asarray_copy_default(asarray, chunks):
+    """Test that np.*array() never returns an object that shares
+    a buffer with the dask graph or a process-local Worker
+    """
+    x = da.asarray(np.arange(10), chunks=chunks)
+
+    # array() defaults to copy=True.
+    # asarray() and asanyarray() default to copy=None.
+    # For Dask, it doesn't make a difference.
+    nx = asarray(x)
+    nx[0] = 42
+    # Did not write back to the buffer held in the Dask graph
+    assert x[0].compute() == 0
+
+
+def test_array_interface_deprecated_kwargs():
+    x = da.ones(10)
+    with pytest.warns(FutureWarning, match="ignored"):
+        x.__array__(something="foo")
+
+
+@pytest.mark.parametrize("chunks", [5, 10])
+def test_compute_copy(chunks):
+    """Test that compute() never returns an object that shares
+    a buffer with the dask graph or a process-local Worker
+    """
+    x = da.asarray(np.arange(10), chunks=chunks)
+
+    nx = x.compute()
+    nx[0] = 42
+    # Did not write back to the buffer held in the Dask graph
+    assert x[0].compute() == 0
+
+    (nx,) = dask.compute(x)
+    nx[0] = 42
+    assert x[0].compute() == 0
+
+
 def test_Array_numpy_gufunc_call__array_ufunc__01():
     x = da.random.default_rng().normal(size=(3, 10, 10), chunks=(2, 10, 10))
     nx = x.compute()