Repository: dask/dask

Issue:
Performance hit with da.delayed processing in dask 2022080 vs 2022070
**What happened**:
We're using dask.delayed to process Zarr files, see here for details: https://dask.discourse.group/t/using-da-delayed-for-zarr-processing-memory-overhead-how-to-do-it-better/1007
The actual processing may not be optimal and we're looking to make it more memory efficient (separate discussion in the forum post). While doing that, I tested the new `2022080` release. It is about 25x slower for our approach of using a dask.delayed function to process the Zarr array in dask release `2022080` than in dask `2022070` (or than in `2022011`). The full line is `2022080`, the dotted line is `2022070`
![Indexing202207_vs_202208](https://user-images.githubusercontent.com/18033446/185152944-55b2d9ac-34ef-42a0-b2de-b6ebcac892c0.png)

**What you expected to happen**:
I would expect the performance to be the same/similar.

**Minimal Complete Verifiable Example**:

```python
# Generate some synthetic Zarr file of varying size:
import numpy as np
import dask.array as da
import dask


shapes = [
    (2, 2, 16000, 16000),
    (2, 4, 16000, 16000),
    (2, 8, 16000, 16000),
    (2, 16, 16000, 16000),
]

for shape in shapes:
    shape_id = f"{shape}".replace(",", "").replace(" ", "_")[1:-1]
    x = da.random.randint(0, 2 ** 16 - 1,
                          shape,
                          chunks=(1, 1, 2000, 2000),
                          dtype=np.uint16)
    x.to_zarr(f"data_{shape_id}.zarr")

# Our processing functions
SIZE = 2000
# Function which increases value of each index in 
# a SIZExSIZE image by one (=> pseudo-processing)
def shift_img(img):
    return img.copy() + 1

delayed_shift_img = dask.delayed(shift_img)

def process_zarr(input_zarr):
    out_zarr = f"out_{input_zarr}"
    data_old = da.from_zarr(input_zarr)
    data_new = da.empty(data_old.shape,
                        chunks=data_old.chunks,
                        dtype=data_old.dtype)

    n_c, n_z, n_y, n_x = data_old.shape[:]

    print(f"Input file: {input_zarr}")
    print(f"Output file: {out_zarr}")
    print("Array shape:", data_old.shape)
    print("Array chunks:", data_old.chunks)
    print(f"Image size: ({SIZE},{SIZE})")


    for i_c in range(n_c):
        for i_z in range(n_z):
            for i_y in range(0, n_y - 1, SIZE):
                for i_x in range(0, n_x - 1, SIZE):
                    # Need to do this in 2 steps to avoid a
                    # TypeError: Delayed objects of unspecified length have no len()
                    new_img = delayed_shift_img(data_old[i_c, i_z, i_y:i_y+SIZE, i_x:i_x+SIZE])
                    data_new[i_c, i_z, i_y:i_y+SIZE, i_x:i_x+SIZE] = da.from_delayed(new_img,
                                                                                     shape=(SIZE, SIZE),
                                                                                     dtype=data_old.dtype)


    # Write data_new to disk (triggers execution)
    data_new.to_zarr(out_zarr)

    # Clean up output folder
    shutil.rmtree(out_zarr)

# Code to run the analysis
input_zarrs = ['data_2_2_16000_16000.zarr', 'data_2_4_16000_16000.zarr', 'data_2_8_16000_16000.zarr', 'data_2_16_16000_16000.zarr']
for input_zarr in input_zarrs:
    interval = 0.1
    mem = memory_usage((process_zarr, (input_zarr,)), interval=interval)
    time = np.arange(len(mem)) * interval
    mem_file = "log_memory_" + input_zarr.split(".zarr")[0] + "_.dat"
    print(mem_file)
    np.savetxt(mem_file, np.array((time, mem)).T)

```

**Anything else we need to know?**:

**Environment**:

- Dask version: 2022070 vs. 2022080
- Python version: 3.9.12
- Operating System: macOS 12.5
- Install method (conda, pip, source): pip


Discussion:
Thanks for reporting @jluethi. Looking at the `git log` between the `2022.07.0` and `2022.08.0` releases nothing seems like an obvious culprit. Are you able to run `git bisect` to determine which PR lead to a change in behavior?

Also, just to double check, are you using a local scheduler or a distributed cluster? From the example, it looks like a local cluster, but just wanted to confirm 
Thanks @jrbourbeau , will try to narrow it down with `git bisect`. As the only differentiating factor I could find is runtime, it may take a bit. Will report back when I manage to narrow it down.

And yes, just using the local scheduler both for this example as well as in general. 

(We're currently using [parsl](https://parsl-project.org) for orchestration on our cluster to handle some further complexities and use dask within each larger job to parallelise further within each node that we run a job on.)
@jrbourbeau Ok, ran git bisect, thanks a lot for pointing me there, that's a really useful tool!

Using that, I found the issue first at ac74939697c328bd6456aa9fad7d9520328db6e5

Given that this involves some changes to dask graph handling, this could make sense.

Also, again, we may not be using dask optimally here (see discussion about the memory inefficiencies in the forum post linked at the top). But ac74939697c328bd6456aa9fad7d9520328db6e5 makes the code above run significantly slower:
- For the `(2, 4, 16000, 16000)` test, about 27s instead of 7s
- For the `(2, 16, 16000, 16000)` test, about 1800s (~30min) instead of ~80s

Not sure though what other workflows would be affected. At least in that case, performance scaled way worse with larger graphs.

Tagging @rjzamora : would you have an idea why ac74939697c328bd6456aa9fad7d9520328db6e5 would lead to such a performance hit?
>Tagging @rjzamora : would you have an idea why https://github.com/dask/dask/commit/ac74939697c328bd6456aa9fad7d9520328db6e5 would lead to such a performance hit?

Thanks for the ping, and for writing this up!  I will try to figure out if that "fix" caused other graph-related problems.

Resolving patch:
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -959,11 +959,19 @@ def cull(self, keys: Iterable) -> HighLevelGraph:
 
                 # Save the culled layer and its key dependencies
                 ret_layers[layer_name] = culled_layer
-                if isinstance(layer, Blockwise) or (
-                    layer.is_materialized() and (len(layer) == len(culled_deps))
+                if (
+                    isinstance(layer, Blockwise)
+                    or isinstance(layer, MaterializedLayer)
+                    or (layer.is_materialized() and (len(layer) == len(culled_deps)))
                 ):
                     # Don't use culled_deps to update ret_key_deps
-                    # unless they are "direct" key dependencies
+                    # unless they are "direct" key dependencies.
+                    #
+                    # Note that `MaterializedLayer` is "safe", because
+                    # its `cull` method will return a complete dict of
+                    # direct dependencies for all keys in its subgraph.
+                    # See: https://github.com/dask/dask/issues/9389
+                    # for performance motivation
                     ret_key_deps.update(culled_deps)
 
         # Converting dict_keys to a real set lets Python optimise the set


Test patch:
diff --git a/dask/tests/test_layers.py b/dask/tests/test_layers.py
--- a/dask/tests/test_layers.py
+++ b/dask/tests/test_layers.py
@@ -10,6 +10,7 @@
 from distributed import Client, SchedulerPlugin
 from distributed.utils_test import cluster, loop  # noqa F401
 
+from dask.highlevelgraph import HighLevelGraph
 from dask.layers import ArrayChunkShapeDep, ArraySliceDep, fractional_slice
 
 
@@ -250,3 +251,45 @@ def test_dataframe_cull_key_dependencies(op):
     culled_graph = graph.cull(result.__dask_keys__())
 
     assert graph.get_all_dependencies() == culled_graph.get_all_dependencies()
+
+
+def test_dataframe_cull_key_dependencies_materialized():
+    # Test that caching of MaterializedLayer
+    # dependencies during culling doesn't break
+    # the result of ``get_all_dependencies``
+
+    datasets = pytest.importorskip("dask.datasets")
+    dd = pytest.importorskip("dask.dataframe")
+
+    ddf = datasets.timeseries(end="2000-01-15")
+
+    # Build a custom layer to ensure
+    # MaterializedLayer is used
+    name = "custom_graph_test"
+    name_0 = "custom_graph_test_0"
+    dsk = {}
+    for i in range(ddf.npartitions):
+        dsk[(name_0, i)] = (lambda x: x, (ddf._name, i))
+        dsk[(name, i)] = (lambda x: x, (name_0, i))
+    dsk = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
+    result = dd.core.new_dd_object(dsk, name, ddf._meta, ddf.divisions)
+    graph = result.dask
+
+    # HLG cull
+    culled_keys = [k for k in result.__dask_keys__() if k != (name, 0)]
+    culled_graph = graph.cull(culled_keys)
+
+    # Check that culled_deps are cached
+    # See: https://github.com/dask/dask/issues/9389
+    cached_deps = culled_graph.key_dependencies.copy()
+    deps = culled_graph.get_all_dependencies()
+    assert cached_deps == deps
+
+    # Manual cull
+    deps0 = graph.get_all_dependencies()
+    deps0.pop((name, 0))
+    deps0.pop((name_0, 0))
+    deps0.pop((ddf._name, 0))
+
+    # Check that get_all_dependencies results match
+    assert deps0 == deps