maurosilber/dask-checkpoint

Customizable caching of dask.delayed results.

dask-checkpoint's Issues

Improve optimizer to traverse the dask graph from the requested key

Currently, the optimizer step traverses the full dask graph in no particular order:

https://github.com/maurosilber/pipeline/blob/0cac8b8954b4def43e593040dced79807ac37f3a/pipeline/storage.py#L50-L58

It would be better to traverse it starting from the requested key(s), following through with their dependencies. A task does not need to be checked if all its dependents are already stored and will be loaded.

For instance, consider the following graph, a -> b -> c, where b is already stored.

dsk = {
    "a": (task_a,),
    "b": (task_b, "a"),
    "c": (task_c, "b"),
}

If we request c, which is not stored, then we would need to check b, which is stored and hence loaded. Then, we don't need to check a, which is simply removed from the graph.

optimized_dsk = {
    # "a": (task_a,),
    "b": (load, task_b),
    "c": (task_c, "b"),
}

We could adapt the dask.cull implementation.

Add dataclass_transform

To get type hints for Task subclasses, we currently need to decorate them with @dataclass.

But with dataclass_transform, it wouldn't be necessary anymore: https://github.com/microsoft/pyright/blob/main/specs/dataclass_transforms.md