Culling massive Blockwise graphs is very slow, not constant-time #8570
Description
tl;dr: Could get_all_external_keys and get_output_keys be avoided in HLG culling?
In some workflows, it can be desirable to create a dask Array/DataFrame structure representing some full-size, enormous dataset, then immediately use slicing to sub-select out only a tiny part of it, then work with that. Now that we can use Blockwise for IO (xref #7417), this is an especially appealing pattern, because it should be constant-time to construct the massive graph, since nothing has to be materialized.
I had hoped it would also be linear-time to cull this massive graph, but it appears currently that it's not.
Here's an example where I'm trying to create an xarray representing the Landsat-8 collection at full resolution over the entire continental US. This is a 10PB, 136-million-chunk array that involves 1.3 billion data loading tasks. Here I'm using gjoseph92/stackstac#116 (so the data loading graph is fully blockwise) and #8560 (so we know fuse_roots isn't materializing the graph unnecessarily).
Then I'm sub-selecting a single chunk out of those 136 million. Based on what I know of the graph, this should cull down to 4 tasks.
Here's the HLG for reference
You can see the first layer is materialized with ~100,000 tasks, but the big, 100-million-task one is Blockwise. (#8497 sure would be nice here!)
But when I try to optimize this graph, I see memory usage shoot up until it crashes the kernel on the 32GB machine. Interrupting the kernel after a few seconds makes it pretty clear what's going on: HighLevelGraph.cull is calling get_all_external_keys, which is forcing the generation of all 1.3 billion keys (or 136 million keys? not sure).
Even if it didn't call get_all_external_keys, I see that HighLevelGraph.cull is still calling get_output_keys on every layer. For reference:
Lines 944 to 970 in 358a5e3
Why is it necessary for HLG.cull to ask the layer for all its keys, intersect them itself, then pass that back into the layer? And why is it necessary for cull functions to take all_hlg_keys?
I had thought the interface would be simply: HLG.cull tells each layer the necessary output keys; the layer figures out the rest on its own. If it needs to generate all its keys internally and do that intersection, fine, but for layers that don't need to do this, shouldn't the optimization be available?
@rjzamora @madsbk why does culling work this way? Would it be possible to write Blockwise culling without this, in a way that's truly linear-time to only the number of final keys?