Different behaviour in distributed vs single-machine scheduler: missing task

I’m adding extra tasks on the optimizing step to save some task results to disk. The extra tasks are run on the dask single-machine scheduler, while they are not run on the distributed scheduler. It might not be a bug, but the intended behavior. In that case, how should have I done it?

Thanks!

Minimal Complete Verifiable Example:

Example 1:

import dask

def optimize(dsk, keys):
    print("Optimizing graph")
    dsk = dask.utils.ensure_dict(dsk)
    # Adding an extra task to the graph
    dsk["extra_task"] = (print, "Running extra task")
    return dsk

dask.config.set(delayed_optimize=optimize)

dask.delayed(print)("Running task").compute()

Running example 1 prints:

Optimizing graph
Running extra task
Running task

Example 2:

Start dask-scheduler and a dask-worker:

dask-scheduler & dask-worker localhost:8786 &

Add to example 1 the following:

import dask
import distributed

client = distributed.Client("localhost:8786")

def optimize(dsk, keys):
...

Example 2 prints “Optimizing graph” in the python process running the file, and “Running task” in the dask-worker process, but it doesn’t print “Running extra task”.

Environment:

• Dask and distributed versions: 2021.04.0
• Python version: 3.8.8
• Operating System: Ubuntu 20.04
• Install method (conda, pip, source): conda

New conda environment from environment-3.8.yaml. Cloned and “pip -e” installed dask and distributed, tags 2021.04.0.

[madsbk]

But, are there additional optimizing steps, such as culling, after sending the graph to a scheduler?

Yes if possible, optimizations such as culling and blockwise fusion will be done by the scheduler. Currently, this might not happen if low-level fusion is enabled but we are working on moving most, if not all, optimizations to the scheduler.

I don’t think plugging into the optimizers are an optimal choice here. Maybe introduce a function that parse x.__dask_graph__() and inject load/save instructions instead of relying on the optimize infrastructure?

[maurosilber]

Do you mean something like:

def load_or_save(x):
    """Inject load/save instructions."""
   return ...

x = func(42)

load_or_save(x).compute()  # First time, runs and saves
load_or_save(x).compute()  # Second time, loads

I wanted it to be hidden from the user, as it could be misused as follows:

x = func(42)
x = load_or_save(x)

x,compute()  # Runs and saves
x.compute()  # Runs and saves again