Re-implement shuffle using staging #1030

madsbk · 2022-11-02T14:32:12Z

Introduce staging in explicit-comms. The idea is to "stage" the keys of the input on the workers so that a later explicit-comms task can access and free the data associated with the keys.

Notice, explicit-comms and this new staging approach is still experimental. If or when it gets to a state where it provides a significant performance improvements over a range of workflows, the plan is to tighten up the API.

codecov-commenter · 2022-11-02T14:59:31Z

Codecov Report

❗ No coverage uploaded for pull request base (branch-22.12@f11abe3). Click here to learn what that means.
Patch has no changes to coverable lines.

❗ Current head a721cc6 differs from pull request most recent head 6464eb6. Consider uploading reports for the commit 6464eb6 to get more accurate results

Additional details and impacted files

@@              Coverage Diff               @@
##             branch-22.12   #1030   +/-   ##
==============================================
  Coverage                ?   0.00%           
==============================================
  Files                   ?      18           
  Lines                   ?    2252           
  Branches                ?       0           
==============================================
  Hits                    ?       0           
  Misses                  ?    2252           
  Partials                ?       0

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

wence-

A few minor comments, but no need to act if you feel it doesn't really make sense to.

dask_cuda/explicit_comms/comms.py

dask_cuda/explicit_comms/dataframe/shuffle.py

Co-authored-by: Lawrence Mitchell <[email protected]>

…mm_shuffle_staging

wence-

Mostly just some queries for understanding, with some relatively minor issues.

TBH: This code is quite hard to follow at this point, though that may be my unfamiliarity with it, especially to do with what is executing where.

wence- · 2022-11-22T18:29:53Z

dask_cuda/explicit_comms/comms.py

@@ -147,6 +148,20 @@ async def _stop_ucp_listeners(session_state):
    del session_state["lf"]


+async def _stage_keys(session_state: dict, name: str, keys: set):
+    worker: Worker = session_state["worker"]


Can you add a docstring here?

dask_cuda/explicit_comms/comms.py

wence- · 2022-11-22T18:54:34Z

dask_cuda/explicit_comms/dataframe/shuffle.py

+    for rank, out_part_ids in rank_to_out_part_ids.items():
+        if rank != myrank:
+            msg = {
+                i: to_serialize(out_part_id_to_dataframe.pop(i))
+                for i in (out_part_ids & out_part_id_to_dataframe.keys())
+            }
+            futures.append(eps[rank].write(msg))


OK, loop over people we need to communicate with, rather than all endpoints.

dask_cuda/explicit_comms/dataframe/shuffle.py

wence- · 2022-11-22T19:14:24Z

dask_cuda/explicit_comms/dataframe/shuffle.py

+        recv(eps, myrank, rank_to_out_part_ids, out_part_id_to_dataframe_list, proxify),
+        send(eps, myrank, rank_to_out_part_ids, out_part_id_to_dataframe),


A change here is that you don't collect all the futures for all the tasks and gather them at once, but wrap them up in nested asyncio.gather calls. I don't expect this will really change anything substantive about the performance, but just to note.

wence- · 2022-11-22T19:14:59Z

dask_cuda/explicit_comms/dataframe/shuffle.py

+
+    # Finally, we concatenate the output dataframes into the final output partitions
+    ret: List[DataFrame] = []
+    for out_part_id, dataframe_list in out_part_id_to_dataframe_list.items():


Suggested change

for out_part_id, dataframe_list in out_part_id_to_dataframe_list.items():

for dataframe_list in out_part_id_to_dataframe_list.values():

You never needed the id.

So probably a list comprehension here would be simpler:

ret = [ proxify(dd_concat(dfs, ignore_index=ignore_index)) for dfs in out_part_id_to_dataframe_list.values() ]

Good point, fixed in 6464eb6

wence- · 2022-11-22T19:17:44Z

dask_cuda/explicit_comms/dataframe/shuffle.py

+    rank_to_inkeys = c.stage_keys(name=name, keys=df.__dask_keys__())
+    c.client.cancel(df)  # Notice, since `df` has been staged, nothing is freed here.


This I do not understand. We just copy some keys around here as far as I can tell (and don't persist anything), so I'm not sure what this even does.

Added more doc in 7ea9f6f

Co-authored-by: Lawrence Mitchell <[email protected]>

madsbk · 2022-11-23T09:43:43Z

Thanks for the review @wence- !
I think I addressed all of your comments?

dask_cuda/explicit_comms/dataframe/shuffle.py

wence-

Thanks for the doc updates!

madsbk · 2022-11-23T11:26:08Z

rerun tests

madsbk · 2022-11-23T13:19:09Z

@gpucibot merge

madsbk added 2 - In Progress Currently a work in progress improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Nov 2, 2022

github-actions bot added the python python code needed label Nov 2, 2022

madsbk force-pushed the xcomm_shuffle_staging branch from 574f766 to daab4b3 Compare November 2, 2022 14:33

wence- reviewed Nov 2, 2022

View reviewed changes

madsbk force-pushed the xcomm_shuffle_staging branch 3 times, most recently from ab2a437 to 169931a Compare November 3, 2022 12:55

madsbk force-pushed the xcomm_shuffle_staging branch from 169931a to de769f5 Compare November 15, 2022 15:44

madsbk and others added 6 commits November 16, 2022 08:38

Re-implement shuffle using staging

21cab61

Simplify code

90255bc

Co-authored-by: Lawrence Mitchell <[email protected]>

Impl. naive multi_shuffle_group()

395cd5f

Impl. single_shuffle_group()

7bf136b

Impl. compute_map_index()

597ac78

Re-impl. shuffle_task

fea22b1

madsbk force-pushed the xcomm_shuffle_staging branch from de769f5 to fea22b1 Compare November 16, 2022 07:39

madsbk added 5 commits November 16, 2022 00:19

recv(): proxify

d3c1c34

doc

746b07d

remove debug

b953fa2

doc

5cb8312

Merge branch 'branch-22.12' of github.com:rapidsai/dask-cuda into xco…

9193832

…mm_shuffle_staging

madsbk changed the title ~~[WIP] Re-implement shuffle using staging~~ Re-implement shuffle using staging Nov 18, 2022

madsbk removed the 2 - In Progress Currently a work in progress label Nov 18, 2022

madsbk marked this pull request as ready for review November 18, 2022 14:23

madsbk requested a review from a team as a code owner November 18, 2022 14:23

madsbk mentioned this pull request Nov 21, 2022

Optimization of explicit-comms-shuffle #1027

Closed

wence- reviewed Nov 22, 2022

View reviewed changes

madsbk and others added 5 commits November 23, 2022 08:47

Typo

d8900ab

Co-authored-by: Lawrence Mitchell <[email protected]>

Impl. pop_staging_area()

7ea9f6f

doc

a721cc6

doc

feac8d5

clean up

6464eb6

wence- reviewed Nov 23, 2022

View reviewed changes

dask_cuda/explicit_comms/dataframe/shuffle.py Show resolved Hide resolved

wence- approved these changes Nov 23, 2022

View reviewed changes

rapids-bot bot merged commit 4d725e3 into rapidsai:branch-22.12 Nov 23, 2022

madsbk deleted the xcomm_shuffle_staging branch November 24, 2022 10:16

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Re-implement shuffle using staging #1030

Re-implement shuffle using staging #1030

madsbk commented Nov 2, 2022 •

edited

Loading

codecov-commenter commented Nov 2, 2022 •

edited

Loading

wence- left a comment

wence- left a comment

wence- Nov 22, 2022

wence- Nov 22, 2022

wence- Nov 22, 2022

wence- Nov 22, 2022

madsbk Nov 23, 2022

wence- Nov 22, 2022

madsbk Nov 23, 2022

madsbk commented Nov 23, 2022

wence- left a comment

madsbk commented Nov 23, 2022

madsbk commented Nov 23, 2022

		recv(eps, myrank, rank_to_out_part_ids, out_part_id_to_dataframe_list, proxify),
		send(eps, myrank, rank_to_out_part_ids, out_part_id_to_dataframe),

	for out_part_id, dataframe_list in out_part_id_to_dataframe_list.items():
	for dataframe_list in out_part_id_to_dataframe_list.values():

		rank_to_inkeys = c.stage_keys(name=name, keys=df.__dask_keys__())
		c.client.cancel(df) # Notice, since `df` has been staged, nothing is freed here.

Re-implement shuffle using staging #1030

Re-implement shuffle using staging #1030

Conversation

madsbk commented Nov 2, 2022 • edited Loading

codecov-commenter commented Nov 2, 2022 • edited Loading

Codecov Report

wence- left a comment

Choose a reason for hiding this comment

wence- left a comment

Choose a reason for hiding this comment

wence- Nov 22, 2022

Choose a reason for hiding this comment

wence- Nov 22, 2022

Choose a reason for hiding this comment

wence- Nov 22, 2022

Choose a reason for hiding this comment

wence- Nov 22, 2022

Choose a reason for hiding this comment

madsbk Nov 23, 2022

Choose a reason for hiding this comment

wence- Nov 22, 2022

Choose a reason for hiding this comment

madsbk Nov 23, 2022

Choose a reason for hiding this comment

madsbk commented Nov 23, 2022

wence- left a comment

Choose a reason for hiding this comment

madsbk commented Nov 23, 2022

madsbk commented Nov 23, 2022

madsbk commented Nov 2, 2022 •

edited

Loading

codecov-commenter commented Nov 2, 2022 •

edited

Loading