Distributed protocol

[mrocklin]

We would like to send expressions up to the Dask scheduler have them optimized and lowered there and then enter the normal scheduler state. To do this I think that we need a new route on client/scheduler communication. Maybe that can be made more general and accept "something picklable that satisfies the .__dask_foo__ protocols" or something similar.

Eventually I'd love it if we held on to the Expr structure in the scheduler to inform future optimizations, but as a first pass I think that we probably just try to end up with the same state that the current update_graph route achieves.

@fjetter I think that this would benefit from someone on your team

[fjetter]

I'm not sure if this is required this early already. While I understand that we might want to have knowledge about the raw input on the scheduler eventually, hooking this up to distributed would be much easier if we did the optimization and materialization on client side first as long as this is still a moving target.

If the expressions can be materialized to an existing HLG or to a low level task / MaterializedGraph we wouldn't need an additional route on the scheduler. At this point I wouldn't know how a protocol like this would even look like and would prefer delaying changes to distributed until we know more about this

[mrocklin]

I'm not sure if this is required this early already.

Yeah, nothing is required yet, we're mostly trying to identify the hard parts of this problem, and make sure that we can solve them well. Then, once we've done that demonstration it'll make sense to do a lot of the tedious work of moving over the bulk of the algorithms.

as long as this is still a moving target

My hope is that by agreeing on some stable protocols we can make everyone happy.

At this point I wouldn't know how a protocol like

Yeah, let's think about this. My recommendation would be that the scheduler receives a pickleable thing that satisfies various __dask_foo__ protocols. For example, it might look like the following:

def update_graph(self, obj, ...):
    graph = obj.__dask_graph__()
    keys = obj.__dask_keys__()
    annotations = obj.__dask_annotations__()

    priority = dask.order(graph)
    ...

[mrocklin]

I'm at a point now where this is the most painful part of my computations (🎉)

@fjetter, I'll probably want to discuss this with you again when you return. My hope is that something very generic, like what I propose above, is feasible.

[mrocklin]

If so, it would be great to have near-instantaneous execution. I suspect that it will viscerally feel nice.

[fjetter]

I poked at this briefly. This is unfortunately a bit of work. HLGs are hard coded in so many places. Introducing this protocol and allowing this to be accepted in Scheduler.update_graph is trivial but on client side there are so many verifications and graph mutations happening that all need to be thrown away / bypassed for a new object. All places that are not immediately assuming a HLG are falling back to materialization via __dask_graph__, __dask_keys__.

FWIW this is not just the distributed client but mostly utility functions from dask. I think we can actually throw away some of this but I haven't gotten so far, yet.

[fjetter]

In this context I would also like to raise the question about key stringification. Right now, stringification happens in a couple of different places but I'm not entirely convinced this is needed.
It happens primarily in distributed but I can see references in Blockwise as well. This matters a little when talking about the protocol. At least existing code is a bit convoluted when it comes to stringified vs non-stringified keys.
I would hope that this protocol uses non-stringified keys entirely (and hope dask-expr has not introduced stringification, yet) such that this becomes an implementation detail of the distributed scheduler and we can possibly remove this down the line.

[mrocklin]

but I can see references in Blockwise as well

Blockwise should generally be much dumber than it is today I think.

I would hope that this protocol uses non-stringified keys entirely (and hope dask-expr has not introduced stringification, yet) such that this becomes an implementation detail of the distributed scheduler and we can possibly remove this down the line.

In general I agree with this. Stringification (or not stringification) should be an implementation detail of the scheduler. The Dask collection code shouldn't think about it.