Annotations for custom graphs in dask

Question

How can I specify resources per task in a custom graph like you can achieve with the dask.annotate context manager for dask collections?

Answer 1

You can convert the custom graph to a High Level Graph object to add annotations.

Taking the example from the custom graph docs that you linked to:

import dask
from dask.highlevelgraph import HighLevelGraph

def load(filename):
    ...

def clean(data):
    ...

def analyze(sequence_of_data):
    ...

def store(result):
    with open(..., 'w') as f:
        f.write(result)

dsk = {'load-1': (load, 'myfile.a.data'),
       'load-2': (load, 'myfile.b.data'),
       'load-3': (load, 'myfile.c.data'),
       'clean-1': (clean, 'load-1'),
       'clean-2': (clean, 'load-2'),
       'clean-3': (clean, 'load-3'),
       'analyze': (analyze, ['clean-%d' % i for i in [1, 2, 3]]),
       'store': (store, 'analyze')}


with dask.annotate(test="value"):
    dsk = HighLevelGraph.from_collections("test", dsk)