Collection Protocol #8674

douglasdavis · 2022-02-07T20:25:46Z

PEP 544 introduces the Protocol class to the typing module in Python 3.8 (the soon be the minimum supported version, dask/community#213). Writing new Dask collections for dask-awkward has had me thinking about working on a DaskCollection protocol. I imagine the benefits to be:

usage with static type checkers
- other activity in this area at
- Python supporting IDEs take advantage of typing
self-documenting; some improvements to the custom collections page of the docs. The protocol docs can be autogenerated and added to that page.
purely opt-in feature

The typing.runtime_checkable decorator allows use of isinstance(x, DaskCollection) in any code base
that uses Dask collections; for example:

>>> from dask.typing import DaskCollection
>>> import dask.array as da
>>> x = da.zeros((10, 3))
>>> isinstance(x, DaskCollection)
True

(though this is an order of magnitude slower than dask.base.is_dask_collection which only checks for x.__dask_graph__() is not None; static typing checking & built-in interface documentation are the core benefits IMO)

Something else that came up in the brief discussion on a call last week was having {Scheduler,Worker,Nanny}Plugin protocols in distributed; and perhaps those are better places to start introducing protocols to Dask since on the user side typically more folks would write plugins than new collections.

fjetter

Does the py.typed have any effect? IIUC https://www.python.org/dev/peps/pep-0561/#id13 says we'd need to package this file

douglasdavis · 2022-02-08T15:56:19Z

Yes indeed I was missing the necessary change for it to have an effect (now added as package data in setup.py).

ian-r-rose

Thanks for getting the ball rolling on this @douglasdavis , I'm quite excited to see this get in.

I think it's outside of the scope of this PR, but I'd love to start writing down protocols/aliases for things like "tasks" and "keys", which would help to further specify the type parameters for things like __dask_graph__() and __dask_keys__().

ian-r-rose · 2022-02-17T18:40:29Z

dask/typing.py

+
+@runtime_checkable
+class DaskCollection(Protocol):
+    def __dask_graph__(self) -> Mapping | None:


What is the rationale behind this being None? Is there somewhere where None is a valid graph ("yes" is a perfectly fine answer, I'm just not aware of it)? Since a lot of places check for __dask_graph__() this means a lot of downstream logic around checking for None first

I was taking inspiration from the implementation of is_dask_collection checking for is not None instead of something like hasattr, but since you've pointed it out I feel that anyone implementing a collection would not have a None returning __dask_graph__.. So yea I'd agree it would make sense to just have Mapping here

I noticed that the docs for __dask_graph__() also talk about None: https://docs.dask.org/en/stable/custom-collections.html#dask_graph__ . So the original typing doesn't seem wrong, it's just unfortunate that it would require downstream users to check for it all the time. Perhaps we can remove that?

I'd be in favor of removing the possibility of None. I can't think of a scenario where it would pop up.

Any further thoughts on whether we can remove None? I just took a look at doing that locally, and nothing seemed to blow up.

I also see many places where __dask_graph__() is used without None checks, so certainly usages in this codebase don't really consider it to be likely.

I think if we are removing this, it's worth also deprecating/removing it in is_dask_collection().

Yeah makes sense; I've updated is_dask_collection to use the "old" check by testing hasattr(x, __dask_graph__)

dask/typing.py

douglasdavis · 2022-03-07T16:16:57Z

@ian-r-rose thanks a lot the first round of comments. I've just pushed some updates

added PostComputeCallable and PostPersistCallable protocols
Added a HLGDaskCollection protocol which inherits from DaskCollection but requires definition of a __dask_layers__ method.
also added some docstrings

dask/typing.py

douglasdavis · 2022-04-26T17:27:30Z

dask/base.py

-        return x.__dask_graph__() is not None
-    except (AttributeError, TypeError):
-        return False
+    return hasattr(x, "__dask_graph__") and callable(x.__dask_graph__)


@ian-r-rose there's a test in dask.array that setattrs __dask_graph__ to None:

dask/dask/array/tests/test_dispatch.py

Lines 227 to 236 in 7259f21

def test_direct_deferral_wrapping_override():

"""Directly test Dask defering to an upcast type and the ability to still wrap it."""

a = da.from_array(np.arange(4))

b = WrappedArray(np.arange(4))

assert a.__add__(b) is NotImplemented

# Note: remove dask_graph to be able to wrap b in a dask array

setattr(b, "__dask_graph__", None)

res = a + da.from_array(b)

assert isinstance(res, da.Array)

assert_eq(res, 2 * np.arange(4), check_type=False)

I wonder if it's in the scope of this PR to try to adopt a mechanism for "erasing" an existing __dask_graph__ without setting it to None (adding the callable check allows the test to stay as is).

To me, erasing an existing method on an instance is such specialized (and probably not recommended) behavior that we shouldn't try to capture it in the signature. I don't really even understand what this test is doing, to be honest.

Yea I also can't quite tell what is being accomplished with the test, and see the assignment of the method to None as very specialized. But anyway, I'm satisfied with this definition of is_dask_collection, tests look to be passing and I think it covers what the previous implementation was getting (it while removing the possibility of a None return from the method itself).

ian-r-rose

Thanks @douglasdavis, I think this is close. If some of the self-referential callbacks become too sticky, I'd say we can defer them to follow-ups

dask/typing.py

ian-r-rose · 2022-04-26T23:11:13Z

dask/typing.py

+    def __call__(
+        self,
+        dsk: Mapping,
+        *args: Any,


I think some of the pyright issues that I saw were due to this *args. The implementation for most of the collections don't actually accept arbitrary *args, making the function signatures incompatible.

I'm actually not sure why mypy is okay with it.

Looking at this a bit more: it would be nice if we could type *args with a ParamSpec. However, I think ParamSpec is currently not expressive enough to handle keyword-only args: https://peps.python.org/pep-0612/#concatenating-keyword-parameters.

So I take the above back, I'm not sure we can do better than Any here at the moment.

douglasdavis · 2022-04-27T18:05:22Z

@ian-r-rose I think I'm starting to be pretty content with the status of this PR. Thanks for all of your work on it with me. One thing that I think can improve, but IMO is worth addressing in a future PR, is __dask_optimize__. It's hinted to Any, even hinting it to Callable is causing issues because the core collections are using dask.context.globalmethod to define their __dask_optimize__ implementations. Perhaps a DaskOptimize protocol is something we can work on in the future.

ian-r-rose

Thanks @douglasdavis, I'm happy where this sits right now. I have a couple of very minor nits, but I otherwise think this is ready to be taken for a spin.

What do you think @jrbourbeau and @jsignell?

ian-r-rose · 2022-04-27T18:24:08Z

dask/base.py

+def is_dask_collection(x) -> bool:
    """Returns ``True`` if ``x`` is a dask collection"""
-    try:
-        return x.__dask_graph__() is not None


I think this is a good change, but perhaps it should go through a deprecation cycle. What do you think @jrbourbeau?

I don't think this needs a deprecation cycle.

dask/tests/test_typing.py

ian-r-rose · 2022-04-27T19:48:44Z

dask/typing.py

+    def __call__(
+        self,
+        dsk: Mapping,
+        *args: Any,


Looking at this a bit more: it would be nice if we could type *args with a ParamSpec. However, I think ParamSpec is currently not expressive enough to handle keyword-only args: https://peps.python.org/pep-0612/#concatenating-keyword-parameters.

So I take the above back, I'm not sure we can do better than Any here at the moment.

dask/typing.py

ian-r-rose

Thanks for all your work on this @douglasdavis! I'm happy with this as it stands, my only outstanding question is whether the change to is_dask_collection deserves a deprecation cycle.

jsignell · 2022-04-29T17:44:43Z

Thanks for your work on this @douglasdavis and @ian-r-rose I am having a read through now.

jsignell

This looks really good to me.

ian-r-rose · 2022-05-06T21:14:50Z

Woo! Thanks for your persistence @douglasdavis

jakirkham · 2022-05-06T21:23:10Z

Seeing a test failure in Distributed, which appears may be related to this change. Details in PR ( dask/distributed#6299 ).

[PEP 544](https://www.python.org/dev/peps/pep-0544/) introduces the `Protocol` class to the `typing` module in Python 3.8 (the soon be the minimum supported version, dask/community#213). Writing new Dask collections for [dask-awkward](https://github.com/ContinuumIO/dask-awkward/) has had me thinking about working on a `DaskCollection` protocol. I imagine the benefits to be: - usage with static type checkers - other activity in this area at - dask#8295 - dask#8706 - dask#8854 - Python supporting IDEs take advantage of typing - self-documenting; some improvements to [the custom collections page](https://docs.dask.org/en/latest/custom-collections.html) of the docs. The protocol docs can be autogenerated and added to that page. - purely opt-in feature The `typing.runtime_checkable` decorator allows use of `isinstance(x, DaskCollection)` in any code base that uses Dask collections; for example: ```python >>> from dask.typing import DaskCollection >>> import dask.array as da >>> x = da.zeros((10, 3)) >>> isinstance(x, DaskCollection) True ``` (though this is an order of magnitude slower than `dask.base.is_dask_collection` which only checks for `x.__dask_graph__() is not None`; static typing checking & built-in interface documentation are the core benefits IMO) Something else that came up in the brief discussion on a call last week was having `{Scheduler,Worker,Nanny}Plugin` protocols in `distributed`; and perhaps those are better places to start introducing protocols to Dask since on the user side typically more folks would write plugins than new collections.

fjetter reviewed Feb 8, 2022

View reviewed changes

douglasdavis force-pushed the collection-protocol branch 3 times, most recently from 50d3de8 to ffb767d Compare February 14, 2022 21:03

douglasdavis force-pushed the collection-protocol branch from ffb767d to 9cf3cc9 Compare February 17, 2022 14:36

ian-r-rose reviewed Feb 17, 2022

View reviewed changes

ian-r-rose added the discussion Discussing a topic with no specific actions yet label Feb 17, 2022

douglasdavis force-pushed the collection-protocol branch 2 times, most recently from 52ba882 to fb3677c Compare February 24, 2022 18:34

douglasdavis force-pushed the collection-protocol branch 2 times, most recently from 967415d to fcea3d4 Compare March 1, 2022 16:41

douglasdavis force-pushed the collection-protocol branch from fcea3d4 to 90724c5 Compare March 7, 2022 16:15

douglasdavis force-pushed the collection-protocol branch from 90724c5 to 782be8a Compare March 7, 2022 18:48

ian-r-rose reviewed Mar 7, 2022

View reviewed changes

dask/typing.py Outdated Show resolved Hide resolved

dask/typing.py Outdated Show resolved Hide resolved

dask/typing.py Outdated Show resolved Hide resolved

dask/typing.py Outdated Show resolved Hide resolved

dask/typing.py Outdated Show resolved Hide resolved

douglasdavis force-pushed the collection-protocol branch 5 times, most recently from a71a5d2 to 75c90ff Compare March 15, 2022 15:22

douglasdavis force-pushed the collection-protocol branch 2 times, most recently from 4813edc to c758eb1 Compare March 21, 2022 21:01

douglasdavis force-pushed the collection-protocol branch 2 times, most recently from d51996c to b93503f Compare March 29, 2022 16:05

douglasdavis force-pushed the collection-protocol branch 2 times, most recently from aada9fd to 9786950 Compare April 12, 2022 12:10

douglasdavis force-pushed the collection-protocol branch from 18ff623 to 26b119b Compare April 19, 2022 18:21

douglasdavis added 2 commits April 21, 2022 13:47

skeleton for a Dask collection protocol

71495b6

include py.typed as package data

3f56ef9

old check rely on __dask_graph__ being callable

bff0788

douglasdavis commented Apr 26, 2022

View reviewed changes

Merge remote-tracking branch 'upstream/main' into collection-protocol

ebe056b

ian-r-rose reviewed Apr 26, 2022

View reviewed changes

douglasdavis and others added 5 commits April 27, 2022 10:23

bring back scheduler get protocol; some TypeVar fixes

bc11ed7

fmt

322a23f

fix typo; docstrings; add a test which uses inheritance

6d715af

Merge remote-tracking branch 'upstream/main' into collection-protocol

390d29f

Merge remote-tracking branch 'upstream/main' into collection-protocol

afdebd6

ian-r-rose reviewed Apr 27, 2022

View reviewed changes

ian-r-rose approved these changes Apr 29, 2022

View reviewed changes

jsignell approved these changes Apr 29, 2022

View reviewed changes

Merge remote-tracking branch 'upstream/main' into collection-protocol

022285a

douglasdavis force-pushed the collection-protocol branch from baa5824 to 022285a Compare May 2, 2022 13:30

jsignell approved these changes May 3, 2022

View reviewed changes

Merge remote-tracking branch 'upstream/main' into collection-protocol

ad81e86

douglasdavis force-pushed the collection-protocol branch from b36f571 to ad81e86 Compare May 5, 2022 13:24

jsignell merged commit 1e783d9 into dask:main May 6, 2022

jakirkham mentioned this pull request May 6, 2022

[NO MRG] Test CI dask/distributed#6299

Closed

2 tasks

keewis mentioned this pull request May 7, 2022

⚠️ Nightly upstream-dev CI failed ⚠️ pydata/xarray#6578

Closed

dcherian mentioned this pull request May 9, 2022

New Collections protocol breaks xarray objects #9058

Closed

douglasdavis deleted the collection-protocol branch May 17, 2022 22:18

douglasdavis mentioned this pull request May 17, 2022

Update custom collection documentation (leverage new collection protocol) #9097

Merged

3 tasks

ian-r-rose mentioned this pull request Jun 30, 2022

Only import IPython if type checking #9230

Merged

ian-r-rose mentioned this pull request Oct 13, 2022

Backend library dispatching for IO in Dask-Array and Dask-DataFrame #9475

Merged

	def test_direct_deferral_wrapping_override():
	"""Directly test Dask defering to an upcast type and the ability to still wrap it."""
	a = da.from_array(np.arange(4))
	b = WrappedArray(np.arange(4))
	assert a.__add__(b) is NotImplemented
	# Note: remove dask_graph to be able to wrap b in a dask array
	setattr(b, "__dask_graph__", None)
	res = a + da.from_array(b)
	assert isinstance(res, da.Array)
	assert_eq(res, 2 * np.arange(4), check_type=False)

Uh oh!

Collection Protocol #8674

Collection Protocol #8674

Uh oh!

Conversation

douglasdavis commented Feb 7, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

fjetter left a comment

Choose a reason for hiding this comment

Uh oh!

douglasdavis commented Feb 8, 2022

Uh oh!

ian-r-rose left a comment

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

douglasdavis commented Mar 7, 2022

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

douglasdavis Apr 26, 2022 • 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!

ian-r-rose left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

douglasdavis commented Apr 27, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ian-r-rose left a comment

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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ian-r-rose left a comment

Choose a reason for hiding this comment

douglasdavis commented Feb 7, 2022 •

edited

Loading

douglasdavis Apr 26, 2022 •

edited

Loading

douglasdavis commented Apr 27, 2022 •

edited

Loading