User a layer for p2p shuffle #7180
Conversation
Unit Test ResultsSee test report for an extended history of previous test failures. This is useful for diagnosing flaky tests. 15 files ±0 15 suites ±0 6h 13m 13s ⏱️ - 21m 38s For more details on these failures and errors, see this check. Results for commit dda2ad9. ± Comparison against base commit 8f25111. |
| ) | ||
|
|
||
| def _construct_graph(self, deserializing=None): | ||
| token = tokenize(self.name_input, self.column, self.npartitions, self.parts_out) |
There was a problem hiding this comment.
FWIW putting self.parts_out here is the actual functional change. This causes culled graphs to be unique which matches the way how the shuffle service works very well.
| transferred = df.map_partitions( | ||
| shuffle_transfer, | ||
| id=token, | ||
| npartitions=npartitions, | ||
| column=column, | ||
| meta=df, | ||
| enforce_metadata=False, | ||
| transform_divisions=False, | ||
| ) |
There was a problem hiding this comment.
The one downside of this is that we'll no longer get input task fusion. The previous code here was a BlockWise layer and we'll loose this. However, I consider consistency more important than the input task fusion.
FWIW "fixing"/enabling input (and output) task fusion for graphs of this type would be beneficial for other types of graph as well but I consider this out of scope for now
| dsk = { | ||
| (name, i): (shuffle_unpack, token, i, barrier_key) for i in range(npartitions) | ||
| } | ||
| layer = MaterializedLayer(dsk, annotations={"shuffle": lambda key: key[1]}) |
There was a problem hiding this comment.
cherry on top: no MaterializedLayer
| x = dd.shuffle.shuffle(df, "x", shuffle="p2p") | ||
| full = await x.persist() | ||
| ntasks_full = len(s.tasks) | ||
| del full | ||
| while s.tasks: | ||
| await asyncio.sleep(0) | ||
| partial = await x.tail(compute=False).persist() # Only ask for one key | ||
|
|
||
| assert len(s.tasks) < df.npartitions * 2 | ||
| assert len(s.tasks) < ntasks_full | ||
| del partial |
There was a problem hiding this comment.
Since we're no longer fusing this test was breaking. The new test ensures that the culled graph has fewer tasks instead of asserting on the ratio to input partitions
This moves the graph generation for p2p shuffling to a layer. The benefit is that we can generate a "new graph" when culling.
This specifically helps with problem cases like Case 2 in #6105
as exhibited by test case
test_add_some_results.In the above example,
y, andz, will be treated as unique shuffles which makes releasing keys much more straight forward and decouples the problem significantly.Of course, we may need to submit some data twice but I doubt this is a problem in practice.
FWIW I only inherited from SimpleShuffleLayer because I didn't want to duplicate too much code. It's sufficiently similar to what we need for now to be subclassed