Moving high level graph pack/unpack Dask by madsbk · Pull Request #7179 · dask/dask

madsbk · 2021-02-05T13:22:25Z

Important: remove CI hack before merge!

This PR moves all pack and unpacking of high level graphs from Distributed to Dask. This work is based on the discussion in dask/distributed#4406.
Beside the move, unpack now doesn't modify the input graph and dependencies, instead the unpacked graph and dependencies a returned, which makes it possible to delegate the pack and unpack of annotations to the base class.

Notice, this PR includes a hack that makes CI install dask/distributed#4489 before testing. This should be removed before merging.

Tests added / passed
Passes black dask / flake8 dask
Remove CI hack

cc. @sjperkins, @ian-r-rose, @jrbourbeau

TODO: remove before PR merge!

ian-r-rose

Thanks for taking this on @madsbk, this looks good to me

dask/highlevelgraph.py

sjperkins · 2021-02-12T10:27:35Z

Thanks for working on this @madsbk. I aim to provide feedback later today.

Co-authored-by: crusaderky <crusaderky@gmail.com>

crusaderky · 2021-02-12T11:52:32Z

dask/highlevelgraph.py

+                expanded[a] = v
+            else:
+                if not keys_stringified:
+                    keys = [stringify(k) for k in keys]
+                    keys_stringified = True
+
+                expanded[a] = dict.fromkeys(keys, v)
+
+        # Merge the expanded annotations with the existing annotations mapping
+        for k, v in expanded.items():
+            v.update(annotations.get(k, {}))
+        annotations.update(expanded)


Suggested change

expanded[a] = v

else:

if not keys_stringified:

keys = [stringify(k) for k in keys]

keys_stringified = True

expanded[a] = dict.fromkeys(keys, v)

# Merge the expanded annotations with the existing annotations mapping

for k, v in expanded.items():

v.update(annotations.get(k, {}))

annotations.update(expanded)

else:

if not keys_stringified:

keys = [stringify(k) for k in keys]

keys_stringified = True

v = dict.fromkeys(keys, v)

try:

v.update(annotations[a])

except KeyError:

pass

annotations[a] = v

In case of collision, the original annotations silently win; is it deliberate?

I think this may be a case of the original annotations (i.e. supplied as kwargs to Client.{submit, persist} overriding layer annotations. @ian-r-rose may know more here.

I think this may be a case of the original annotations (i.e. supplied as kwargs to Client.{submit, persist} overriding layer annotations

Yes, this was also what I was thinking. @ian-r-rose what do you think?

I think this may be a case of the original annotations (i.e. supplied as kwargs to Client.{submit, persist} overriding layer annotations.

Sorry for the slow reply here. Yes, that was my understanding as well. However, I don't have a strong intuition of which should win, and I can think of an argument for both (more specific annotation vs later-applied annotation).

sjperkins · 2021-02-18T13:10:10Z

@madsbk Here's my suggestion, I wasn't able to use your fork as a base for a PR:

sjperkins@7fcd644

Summary of changes:

BasicLayer renamed to MaterialisedLayer.
MaterialisedLayer contains code for packing materialized graphs (a.k.a. python dicts) to the scheduler (previously in Layer)
__dask_distributed_{pack, unpack}__ now calls super() to pack/unpack annotations.
__dask_distributed_{pack, unpack}__ now return dicts.
`_dask_distributed_anno{pack,unpack} removed.
__dask_distributed_{pack, unpack}__ implemented for ParquetSubgraph

More broadly, MaterialisedLayer is now the single point of responsibility for transmitting materialised graphs. Layer returns to a more abstract form (is_materialised return False).

py.test -s -vvv distributed/protocol/tests/test_highlevelgraph.py succeeds for me on the distributed branch.

madsbk · 2021-02-18T14:00:26Z

I have a hard time seeing the advantages of your approach. Personally, I was hoping that we could avoid introducing a new protocol/requirements for the packed state returned by __dask_distributed_pack__(). Now that we have a class hierarchy, why not handle the pack/unpack through inheritance without having to define a shared state dict?

sjperkins · 2021-02-18T15:20:56Z

I have a hard time seeing the advantages of your approach. Personally, I was hoping that we could avoid introducing a new protocol/requirements for the packed state returned by __dask_distributed_pack__(). Now that we have a class hierarchy, why not handle the pack/unpack through inheritance without having to define a shared state dict?

I think I understand that you're asking why the materialised packing behaviour has been moved into MaterialisedLayer? If so, its an attempt to separate interface from implementation. In that sense, Layer is the interface while the derived classes are implementations of that interface. The reason for this is that Implementing materialised behaviour in the interface (base class) could potentially constrain the behaviour of derived classes in future: better to place in the behaviour in a specific implementation class.

I get that it may seem like quibbling, but inheritance hierarchies can get out of hand.

madsbk · 2021-02-18T15:39:09Z

I get that it may seem like quibbling, but inheritance hierarchies can get out of hand.

I agree inheritance hierarchies can get out of hand but now that we have a class hierarchy and we have a perfectly suited default implementation that is compatible with any conceivable derived class, I think it makes sense. Particularly, if the alternative is an introduction of a key scheme for the packed state dict.

sjperkins · 2021-02-18T17:52:36Z

I get that it may seem like quibbling, but inheritance hierarchies can get out of hand.

I agree inheritance hierarchies can get out of hand but now that we have a class hierarchy and we have a perfectly suited default implementation that is compatible with any conceivable derived class, I think it makes sense. Particularly, if the alternative is an introduction of a key scheme for the packed state dict.

OK, I see your point. I'm happy with the PR as it stands then.

@sjperkins

As suggested by @sjperkins, this makes it easier to extend the API in the future.

Layer.__reduce__() isn't used by the client -> scheduler communication any more.

Previously, BasicLayer could be partial materialized, which isn't possible any more. Now BasicLayer is always materialized.

madsbk · 2021-02-19T09:06:07Z

@sjperkins, thanks for the discussion I have added your suggestions to let __dask_distributed_unpack__ return a dict and renamed BasicLayer -> MaterializedLayer.

Also @crusaderky, @ian-r-rose, thanks for the reviews.

@jrbourbeau, I think this is ready to be merged after the removal of the CI hack. Could you take a look and if you agree, we should revert continuous_integration/scripts/install.sh and merge this PR and dask/distributed#4489

See <dask/dask#7179>

madsbk · 2021-02-24T18:22:21Z

@jrbourbeau, I think we should merge this soon. Can you revert continuous_integration/scripts/install.sh and merge this PR and dask/distributed#4489 ?

jrbourbeau

Apologies for the delayed review, thank you for the ping @madsbk. Overall this PR and dask/distributed#4489 look good to me. My primary comment is that there are a few places where I think using the full annotations term instead of anno for shorthand will improve code clarity. Otherwise, I think we're good to revert the temporary CI changes here and merge.

dask/highlevelgraph.py

jrbourbeau · 2021-02-25T04:32:31Z

Hmm actually it looks like there's an issue with how we're handling annotations (see CI over in dask/distributed#4489 -- for example this CI build)

jrbourbeau

Thanks @madsbk! This, and the companion PR in distributed, look good to me : )

Since we may be releasing tomorrow (xref dask/community#129 (comment)), and the changes here are an internal refactor instead of addressing user-facing issues, I'm inclined to wait a day before merging these PRs (until after the release) to give the changes here a little time to simmer in the main branch before pushing them out to users

madsbk · 2021-02-26T08:00:52Z

Since we may be releasing tomorrow (xref dask/community#129 (comment)), and the changes here are an internal refactor instead of addressing user-facing issues, I'm inclined to wait a day before merging these PRs (until after the release) to give the changes here a little time to simmer in the main branch before pushing them out to users

Sounds good!

jrbourbeau · 2021-02-26T17:17:39Z

Note I've opened up a follow-up PR at #7279 to resolve a CI issue

jakirkham · 2021-02-26T17:30:34Z

Thanks James! 😄

madsbk force-pushed the hlg_pack_move_to_dask branch from 8a31de8 to 38413e6 Compare February 5, 2021 14:41

madsbk added 3 commits February 8, 2021 11:02

type cleanup

081af4c

test: added test_combo_of_layer_types()

6ae262e

Moved everything to Dask from Distributed

92fcca4

madsbk force-pushed the hlg_pack_move_to_dask branch 3 times, most recently from 96f4709 to 07af9a2 Compare February 8, 2021 11:41

madsbk added 5 commits February 8, 2021 14:15

CI: pull distributed.git@hlg_pack_move_to_dask

c9bef90

TODO: remove before PR merge!

HLG unpack now doesn't modify input

0b0ae1b

Moved annotation pack and unpack outside of layer

274e2bd

Rename: pack_annotations() => __dask_distributed_anno_unpack__()

4ccd5b8

updated docs

0f4d5d6

madsbk force-pushed the hlg_pack_move_to_dask branch from 3827be1 to 0f4d5d6 Compare February 8, 2021 13:16

madsbk mentioned this pull request Feb 8, 2021

Use the new HLG pack/unpack API in Dask dask/distributed#4489

Merged

2 tasks

madsbk marked this pull request as ready for review February 8, 2021 13:33

madsbk changed the title ~~[WIP] Moving high level graph pack/unpack Dask~~ [REVIEW] Moving high level graph pack/unpack Dask Feb 8, 2021

jakirkham requested review from ian-r-rose, jrbourbeau and sjperkins February 11, 2021 08:07

ian-r-rose reviewed Feb 12, 2021

View reviewed changes

dask/highlevelgraph.py Outdated Show resolved Hide resolved

madsbk added 2 commits February 12, 2021 09:43

Merge branch 'master' of github.com:dask/dask into hlg_pack_move_to_dask

5075f6d

__dask_distributed_unpack__(): returns MutableMapping

7f8d1b3

madsbk requested a review from ian-r-rose February 12, 2021 08:47

crusaderky reviewed Feb 12, 2021

View reviewed changes

madsbk and others added 3 commits February 12, 2021 12:27

Fixing type hints and code clean up

c002390

Co-authored-by: crusaderky <crusaderky@gmail.com>

reformat: black

26d4eea

more clean up

13cc6e8

crusaderky reviewed Feb 12, 2021

View reviewed changes

Merge branch 'master' of github.com:dask/dask into hlg_pack_move_to_dask

b304fbd

madsbk added 4 commits February 19, 2021 09:23

Merge branch 'master' of github.com:dask/dask into hlg_pack_move_to_dask

5d16a4d

__dask_distributed_unpack__: return dict

a3a5e4d

As suggested by @sjperkins, this makes it easier to extend the API in the future.

Layer.__reduce__: updated docstring.

852e305

Layer.__reduce__() isn't used by the client -> scheduler communication any more.

Rename: BasicLayer -> MaterializedLayer

69cf217

Previously, BasicLayer could be partial materialized, which isn't possible any more. Now BasicLayer is always materialized.

madsbk added a commit to madsbk/distributed that referenced this pull request Feb 19, 2021

Use the new high level graph API in Dask

2550bee

See <dask/dask#7179>

madsbk added 2 commits February 19, 2021 10:21

HighLevelGraph.__dask_distributed_unpack__(): also returns a dict now

47823f4

Merge branch 'master' of github.com:dask/dask into hlg_pack_move_to_dask

cbe4d7b

jrbourbeau reviewed Feb 25, 2021

View reviewed changes

dask/highlevelgraph.py Outdated Show resolved Hide resolved

dask/highlevelgraph.py Show resolved Hide resolved

dask/highlevelgraph.py Outdated Show resolved Hide resolved

madsbk added 2 commits February 25, 2021 09:55

Merge branch 'master' of github.com:dask/dask into hlg_pack_move_to_dask

09edb1b

Fixed top-level annotations expand

b6e0d20

madsbk force-pushed the hlg_pack_move_to_dask branch from 4db5b83 to b6e0d20 Compare February 25, 2021 11:19

madsbk and others added 2 commits February 25, 2021 14:53

renamed: anno => annotations

9d91e6a

Revert temporary CI changes

8e360fb

jrbourbeau changed the title ~~[REVIEW] Moving high level graph pack/unpack Dask~~ Moving high level graph pack/unpack Dask Feb 26, 2021

jrbourbeau approved these changes Feb 26, 2021

View reviewed changes

jrbourbeau merged commit 3674876 into dask:master Feb 26, 2021

jrbourbeau mentioned this pull request Feb 26, 2021

Use development version of distributed in CI #7279

Merged

madsbk deleted the hlg_pack_move_to_dask branch March 1, 2021 08:02

jakirkham mentioned this pull request Mar 8, 2021

dask-image imread v0.5.0 not working with dask distributed Client & napari dask/dask-image#194

Closed

Uh oh!

Conversation

madsbk commented Feb 5, 2021 • edited by jrbourbeau Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Important: remove CI hack before merge!

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sjperkins commented Feb 12, 2021

Uh oh!

crusaderky Feb 12, 2021 • 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 Feb 12, 2021

Choose a reason for hiding this comment

Uh oh!

sjperkins Feb 12, 2021

Choose a reason for hiding this comment

Uh oh!

madsbk Feb 16, 2021

Choose a reason for hiding this comment

Uh oh!

ian-r-rose Feb 18, 2021

Choose a reason for hiding this comment

Uh oh!

sjperkins commented Feb 18, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

madsbk commented Feb 18, 2021

Uh oh!

sjperkins commented Feb 18, 2021

Uh oh!

madsbk commented Feb 18, 2021

Uh oh!

sjperkins commented Feb 18, 2021

Uh oh!

madsbk commented Feb 19, 2021

Uh oh!

madsbk commented Feb 24, 2021

Uh oh!

jrbourbeau left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jrbourbeau commented Feb 25, 2021

Uh oh!

jrbourbeau left a comment

Choose a reason for hiding this comment

Uh oh!

madsbk commented Feb 26, 2021

Uh oh!

jrbourbeau commented Feb 26, 2021

Uh oh!

jakirkham commented Feb 26, 2021

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

madsbk commented Feb 5, 2021 •

edited by jrbourbeau

Loading

crusaderky Feb 12, 2021 •

edited

Loading

sjperkins commented Feb 18, 2021 •

edited

Loading