Make dask.array.utils functions more generic to other Dask Arrays by mrocklin · Pull Request #10676 · dask/dask

mrocklin · 2023-12-06T01:14:39Z

This will be helpful for dask_expr.array, if we do that soon.

I also ran py.test dask/arrays/tests/test_array_core.py on both this branch and main and notice no slowdown

This will be helpful for dask_expr.array, if we do that soon.

mrocklin · 2023-12-08T22:32:53Z

dask/base.py

+        if "xarray" in type(x).__module__:
+            return x.__dask_graph__() is not None
+        else:
+            return hasattr(x, "__dask_graph__")


@phofl this change is actually kinda important. We were generating the task graph before whenever calling is_dask_collection. Maybe caching was saving us here, but this could be unpleasantly expensive for dask-expr.

The history here is that xarray collections have a __dask_graph__ attribute, but are not always dask collections (they sometimes hold dask arrays but sometimes just have numpy arrays).

Yeah I remember that we discussed this a while back and ended at revisiting if necessary because caching worked fine back then. That change is nice anyway.

This seems pretty expensive for the xarray code path which really should just be something like

any(variable._data for variable in xarray_obj.variables)

Can we add something like __dask_is_present__?

We can add a new protocol. We'll have to do this everywhere, but it seems like an ok thing to do. In the meantime, let's just special-case something for Xarray. I think that it's the only project today where this matters.

I started doing this:

if "xarray" in type(x).__module__: import xarray if isinstance(x, xarray.Dataset): return any(variable._data for variable in x.variables) if isinstance(x, xarray.DataArray): return bool(x._data)

But I'm not sure that _data exists? Maybe we want any(is_dask_collection(variable.data) ...)? Any suggestions?

I suspect this is it:

from dask.base import is_dask_collection def is_dask_xarray(x): import xarray if isinstance(x, xarray.Dataset): return any( is_dask_collection(variable._data) for _, variable in x.variables.items() ) elif isinstance(x, xarray.DataArray): return is_dask_collection(x.variable._data) elif isinstance(x, xarray.Variable): return is_dask_collection(x._data) import xarray as xr ds = xr.tutorial.open_dataset("air_temperature", chunks="auto") computed = ds.compute() assert is_dask_xarray(ds) assert is_dask_xarray(ds.air) assert is_dask_xarray(ds.air.variable) assert not is_dask_xarray(computed) assert not is_dask_xarray(computed.air) assert not is_dask_xarray(computed.air.variable)

I also ran into the issue of materialization on __dask_graph__ when playing with the scheduler integration and I strongly recommend moving dask-expr away from this in favor of an explicit method doing that (that's what I am proposing in dask/dask-expr#294)

@dcherian I've pushed up your changes. Thanks. Using the air_temperature dataset requires an optional dependency pooch, which I was hesitant to add to CI just for this (although maybe I should if it's very lightweight). Alternatively, can you recommend a way to construct a dataset using maybe da.random.random?

@fjetter I'm not sure yet what you mean in the comment above. We need a way to determine if an object is a dask collection. I looked through the PR you mentioned and didn't immediately something there that would solve this problem. Can you help me to understand?

@mrocklin: ds = xr.Dataset({"air": (("time", "lat", "lon"), dask.array.ones((120, 120, 120)))})

mrocklin · 2023-12-15T17:44:56Z

@fjetter can I ask your team to handle this sometime next week?

fjetter · 2023-12-18T15:39:29Z

can I ask your team to handle this sometime next week?

can you state what the expected outcome should be? I'm missing some context about what this change is enabling

mrocklin · 2023-12-18T15:48:54Z

I'm missing some context about what this change is enabling

These changes enable code sharing between dask.array utilities and dask_expr.array (currently in a branch/PR)

They are, I think, innocuous and should be easy to merge.

crusaderky · 2023-12-20T16:40:52Z

dask/base.py

+            elif isinstance(x, xarray.Variable):
+                return is_dask_collection(x._data)
+            else:
+                raise TypeError("Unfamiliar with xarray type", type(x))


This is suppressed a few lines below.
Can we reduce the scope of the try...except block to just the type(x).__module__ line?

crusaderky · 2023-12-20T17:02:20Z

dask/base.py

+                and callable(x.__dask_graph__)
+                and not isinstance(x, type)
+            )
+


This PR breaks pint, because just like xarray it violates the dask collection protocol by defining a __dask_graph__() method that may return None:

>>> import dask.array as da, pint >>> q = pint.Quantity(1, "m") >>> q.__dask_graph__() (None) >>> q = pint.Quantity(da.zeros(5), "m") >>> q.__dask_graph__() HighLevelGraph with 1 layers. <dask.highlevelgraph.HighLevelGraph object at 0x7f043c3889a0> 0. zeros_like-902bf3bcabaabb3fedce3f347fb26081

There may be more third-party libraries in the wild that do the same.
A blocklist approach would be more robust:

if ( isinstance(x, type) or not hasattr(x, "__dask_graph__") or not callable(x.__dask_graph__) ): return False pkg_name = getattr(type(x), "__module__", "").split(".")[0] # Add here other third-party packages where you don't want # to call the `__dask_graph__` method - typically because it's expensive if pkg_name == "dask_expr": return True # xarray, pint, and possibly other libraries return None when they wrap a non-dask object return x.__dask_graph__() is not None

The above, however, will mean that we call the __dask_graph__() method of dask_expr.array objects when they're wrapped by xarray or pint. So I would suggest either:

hybrid allowlist+blocklist, with final fallback on calling __dask_graph__()

change the __dask_graph__() method in dask_expr to be trivial (e.g. write a Mapping subclass that is trivial to initialize and is lazy until you actually try to access its contents)

crusaderky · 2023-12-20T17:02:56Z

dask/tests/test_base.py


    x = delayed(1) + 2
    assert is_dask_collection(x)
-    assert not is_dask_collection(2)


Why was this removed?

crusaderky · 2023-12-20T17:13:25Z

dask/base.py

+            if isinstance(x, xarray.Dataset):
+                return any(is_dask_collection(v._data) for _, v in x.variables.items())
+            elif isinstance(x, xarray.DataArray):
+                return is_dask_collection(x.variable._data)


This is returning a false negative when you have an eager variable and lazy non-index coordinates

>>> import dask.array as da, xarray >>> xarray.DataArray([1,2], dims=["x"], coords={"x": [10, 20], "x2": ("x", da.zeros(2))}) <xarray.DataArray (x: 2)> array([1, 2]) Coordinates: * x (x) int64 10 20 x2 (x) float64 dask.array<chunksize=(2,), meta=np.ndarray>

crusaderky · 2023-12-20T17:15:06Z

dask/base.py

+            import xarray
+
+            if isinstance(x, xarray.Dataset):
+                return any(is_dask_collection(v._data) for _, v in x.variables.items())


Suggested change

return any(is_dask_collection(v._data) for _, v in x.variables.items())

return any(is_dask_collection(v._data) for v in x.variables.values())

dcherian · 2023-12-20T17:25:16Z

dask/base.py

+            elif isinstance(x, xarray.DataArray):
+                return is_dask_collection(x.variable._data)


Suggested change

elif isinstance(x, xarray.DataArray):

return is_dask_collection(x.variable._data)

elif isinstance(x, xarray.DataArray):

return is_dask_collection(x.variable._data) or any(is_dask_collection(var._data) for var in x._coords.values())

crusaderky · 2023-12-20T20:11:55Z

Can we go back to the design table for a bit?

What this PR tries to achieve

Chiefly, this PR tries to future-proof is_dask_collection for when we'll migrate dask.array to dask-expr.

In main, is_dask_collection(x) calls x.__dask_graph__().
For dask.array, dask.dataframe, dask.bag, and dask.delayed this is a trivial attribute fetch.

For xarray, today it's a cheap-ish layer merge:
https://github.com/pydata/xarray/blob/b4444388cb0647c4375d6a364290e4fa5e5f94ba/xarray/core/dataset.py#L872-L881

The problems start if you call it on a dask_expr collection, which triggers a very expensive materialization of the whole graph: https://github.com/dask-contrib/dask-expr/blob/0fab3c19de0e54085119e182de92cbe5ec25f479/dask_expr/_core.py#L429-L445

This PR wants to avoid materialization in dask-expr by changing is_dask_collection not to call __dask_graph__, but just check for its existence.

Why it's not trivial

The problem lies within wrapper libraries - xarray and pint are two, but there may be more I'm not aware of and potentially bespoke, unpublished ones too.

Both xarray and pint define a __dask_graph__ method; however it may return None if they wrap non-dask objects exclusively. Code for xarray is linked above; pint is easier as it wraps around zero or one dask collections: https://github.com/hgrecco/pint/blob/83bffe1df0f18acd451ec5d4622442bdfd2d10f5/pint/facets/dask/__init__.py#L43C1-L48

As of today, the only ways for is_dask_collection to deal with wrapper collections is either

a. call their __dask_graph__ method, which would cause an expensive call to dask_expr.array.Array.__dask_graph__. This is clearly not a viable option.

b. special-case them in dask.base.is_dask_collection, as shown in this PR for xarray only.
I don't like this option because dask should not know about the existence of pint, and it would simply not work with more obscure, possibly closed source, packages that wrap around dask.

I suggest two alternative designs:
c. dask_expr.Expr.__dask_graph__() returns a lazy Mapping subclass, which can be instantiated trivially and only materializes the graph when you call a method such as __getitem__ or __len__. Third party collections that wrap around multiple dask collections, such as xarray, would need to write ad-hoc code to merge multiple such graphs without materializing - much like it already happens with HighLevelGraph: https://github.com/pydata/xarray/blob/b4444388cb0647c4375d6a364290e4fa5e5f94ba/xarray/core/dataset.py#L872-L881

d. define a new API endpoint in the dask protocol, __dask_collections__, which in the trivial case returns [self] but in case of xarray it may return an empty list or multiple underlying dask arrays.

For example, for xarray.Dataset:

def __dask_collections__(self) -> list[DaskCollection]:
    from dask.base import is_dask_collection

    return [v._data for v in self.variables.values() if is_dask_collection(v._data)]

crusaderky · 2023-12-22T12:55:03Z

Offline comment from @fjetter:

I will nuke all of this with

[WIP] Pass non-HLG objects wout materialization distributed#7942

[WIP] Pass non HLG objects wout materialization #10369

[WIP] Allow expressions to be shipped to the scheduler dask-expr#294

xarray, pint will have to migrate to this so we can adjust this accordingly. If xarray truly needs a dynamic switch we probably want to add a dedicated method for this.
Either way, I’m not overly concerned about this PR as long as it doesn’t break anything publicly. The behavior for dask-expr will change anyhow

crusaderky · 2023-12-22T16:05:34Z

In light of the above PRs, I'm switching from allowlist to blocklist, with the expectation that the PRs will be merged before wrapping dask_expr.Array inside xarray or pint becomes important. We can always re-add the allowlist later on.

@mrocklin just wanted to make sure you're happy with this.

dask/tests/test_base.py

phofl

Ideally you shouldn't reach optimise either, that's also expensive if you read from remote storage for example

lgtm otherwise

crusaderky · 2024-01-04T17:00:38Z

@mrocklin I'm merging this. If you have comments happy to open a follow-up PR to address them.

mrocklin · 2024-01-04T17:55:46Z

Grand. Thanks!

…

On Thu, Jan 4, 2024 at 9:00 AM crusaderky ***@***.***> wrote: Merged #10676 <#10676> into main. — Reply to this email directly, view it on GitHub <#10676 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTBKC6FOL2KDNJATLL3YM3N4TAVCNFSM6AAAAABAIS2SLSVHI2DSMVQWIX3LMV45UABCJFZXG5LFIV3GK3TUJZXXI2LGNFRWC5DJN5XDWMJRGM4TENJZG42DMMA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

Make dask.array.utils functions more generic to other Dask Arrays

d30b697

This will be helpful for dask_expr.array, if we do that soon.

github-actions bot added the array label Dec 6, 2023

mrocklin mentioned this pull request Dec 6, 2023

Arrays dask/dask-expr#471

Closed

13 tasks

Don't trigger dask_graph call all the time

15c7c0f

mrocklin force-pushed the arrays-generic branch from d0fc005 to 15c7c0f Compare December 8, 2023 22:32

mrocklin commented Dec 8, 2023

View reviewed changes

Merge branch 'main' of github.com:dask/dask into arrays-generic

5412bec

mrocklin added 2 commits December 18, 2023 07:53

Merge branch 'main' of github.com:dask/dask into arrays-generic

aed181b

Fixup is_dask_collection test

e5f2f32

crusaderky requested changes Dec 20, 2023

View reviewed changes

crusaderky reviewed Dec 20, 2023

View reviewed changes

dcherian reviewed Dec 20, 2023

View reviewed changes

crusaderky self-assigned this Dec 20, 2023

fjetter mentioned this pull request Dec 21, 2023

[WIP] Pass non-HLG objects wout materialization dask/distributed#7942

Closed

crusaderky added 2 commits December 22, 2023 12:57

Merge branch 'main' into arrays-generic

8a25075

Blocklist-based implementation

143c6ba

crusaderky self-requested a review December 22, 2023 16:06

crusaderky added 2 commits December 22, 2023 16:23

fix

d5329bd

complete Expr

3d35e88

crusaderky reviewed Dec 22, 2023

View reviewed changes

dask/tests/test_base.py Outdated Show resolved Hide resolved

tweaks

17cb025

crusaderky approved these changes Dec 22, 2023

View reviewed changes

phofl reviewed Dec 22, 2023

View reviewed changes

crusaderky added 2 commits January 2, 2024 14:12

code review

14f7d99

polish

9798fdd

crusaderky merged commit 6d5b994 into dask:main Jan 4, 2024

milesgranger mentioned this pull request Jan 10, 2024

Add extra check in _normalize_lambda for lambda function dask/dask-expr#689

Merged

Illviljan mentioned this pull request Feb 2, 2024

Use strict type hinting for namedarray pydata/xarray#8241

Merged

	return any(is_dask_collection(v._data) for _, v in x.variables.items())
	return any(is_dask_collection(v._data) for v in x.variables.values())

		elif isinstance(x, xarray.DataArray):
		return is_dask_collection(x.variable._data)

Uh oh!

Conversation

mrocklin commented Dec 6, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mrocklin commented Dec 15, 2023

Uh oh!

fjetter commented Dec 18, 2023

Uh oh!

mrocklin commented Dec 18, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

crusaderky Dec 20, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

crusaderky Dec 20, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

crusaderky Dec 20, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

crusaderky commented Dec 20, 2023

What this PR tries to achieve

Why it's not trivial

Uh oh!

crusaderky commented Dec 22, 2023

Uh oh!

crusaderky commented Dec 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

phofl left a comment

Choose a reason for hiding this comment

Uh oh!

crusaderky commented Jan 4, 2024

Uh oh!

mrocklin commented Jan 4, 2024 via email

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

crusaderky Dec 20, 2023 •

edited

Loading

crusaderky Dec 20, 2023 •

edited

Loading

crusaderky Dec 20, 2023 •

edited

Loading

crusaderky commented Dec 22, 2023 •

edited

Loading