What about having serialize and deserialize methods on these objects that we call instead? That would avoid needing pickling and the alluded to security concerns around it.

What about having serialize and deserialize methods on these objects that we call instead? That would avoid needing pickling and the alluded to security concerns around it.

I think this is essential what we are doing right now just using __reduce__() instead of serialize() ?

Woo, nice green check marks 🎉

I have been working on this PR some time now. Initially, my idea was to implement and use map_tasks() from dask/dask#6689 to handle preprocessing of the high-level graph and use the Layer.__reduce__() methods to handle serialization before sending it to the scheduler. This is a clean generic approach, but it doesn’t work. Distributed requires a very specific protocol when updating the scheduler task graph, which makes it very hard (if not impossible) to implement generically.

I have compiled a list of requirements and thoughts:

• We should support BasicLayer efficiently

• Converting all keys to strings requires domain knowledge in layers like Blockwise where all keys aren't materialized.

• A layer needs ALL keys in the HLG in order to calculate dependencies.

• Dependencies must be calculated before dumps_task() thus dependency calculation must be part of serialization. Some layers can infer key dependencies but consider layers such as Blockwise where kwargs can contain any key. Not to mention BasicLayer, which has no prior dependency knowledge.

• Before serialization, we need to unpack remote data (Futures). This requires map_tasks() to provide key information, which can be problematic: [WIP] Add optional pass_key argument to map_tasks dask#6761 (comment).

• Additionally, we need to search for Futures inside tuples, which map_tasks() cannot support and we need to in cooperate the dependencies of the unpacked Futures into the existing key dependencies without having to recalculate anything.

• Currently, ordering has to be done on the original keys and not after they have been converted to string: Prioritize tasks with their true keys, before stringifying #2006

Because of these requirements, I am exploring a new approach where the serialization and deserialization of layers is specialized. I introduce dumps_highlevelgraph() that do all the preprocessing (unpack Futures, convert keys to strings, calculate dependencies that can cannot be inferred on the scheduler, and dump tasks) and loads_highlevelgraph() that deserialize the required and return a materialized graph and all key dependencies.

The current implementation does all this in Distributed however we should consider move it to the Layer classes in Dask. But notice, these serialization methods will be very Distributed specific and use a lot of functions and classes from Distributed.

NB: the current push is just a rough implementation that always falls back to BasicLayer serialization.

@madsbk if you're around and want to have a high bandwidth conversation I would be interested. Some quick comments/questions

Converting all keys to strings requires domain knowledge in layers like Blockwise where all keys aren't materialized.

Materializing all keys might be ok. Ideally we don't have to materialize the entire graph until we get to the scheduler.

We could also convert to strings after we get to the scheduler and have the full graph.

A layer needs ALL keys in the HLG in order to calculate dependencies.

Again, I think that materializing keys will be ok for now (this is, I think, likely to be faster than the graph). In the future I hope that we can avoid cases where we need to calculate dependencies in the common case.

Dependencies must be calculated before dumps_task() thus dependency calculation must be part of serialization. Some layers can infer key dependencies but consider layers such as Blockwise where kwargs can contain any key. Not to mention BasicLayer, which has no prior dependency knowledge.

Hrm, I'm curious about this. I'll need to think about this more.

Before serialization, we need to unpack remote data (Futures). This requires map_tasks() to provide key information, which can be problematic: dask/dask#6761 (comment).

Also curious.

Additionally, we need to search for Futures inside tuples, which map_tasks() cannot support and we need to in cooperate the dependencies of the unpacked Futures into the existing key dependencies without having to recalculate anything.

We can drop this if we need to

Currently, ordering has to be done on the original keys and not after they have been converted to string: #2006

Ah, right. Well, maybe we can change how we stringify in order to keep ordering. Maybe we do both ordering and stringifying on the scheduler after we've moved everything over.

Materializing all keys might be ok. Ideally we don't have to materialize the entire graph until we get to the scheduler.

Notice, the keys here also refers to the keys inside tasks. Since tasks might not exist yet, it requires domain knowledge to stringify them.

We could also convert to strings after we get to the scheduler and have the full graph.

True, it is properly possible to serialize keys using msgpack in most cases, but dumping tasks before stringify means that the dumped task may contain non-stringified keys.

Moved packing of shuffle layers to Dask: dask/dask#6786 thus CI will fail until it has been merged.

@madsbk is there anything that I can do or should look at to help here?

@madsbk is there anything that I can do or should look at to help here?

The final issue is ordering, which has to be done before converting keys to strings. This is what is failing in CI.

If we do ordering on the scheduler side after stringification I suspect that everything will work the same in almost all cases. I would be willing to take the hit on those cases in order to move forward. Is this an easy option for you?

That seems sensible to me

@mrocklin, I managed to support ordering but it requires a fancy string comparison in order.order(): dask/dask#6807

Let's xfail the test for now, include a link to the order PR in the reason= keyword, and then merge this.

I want to think a bit about prefixing with zeroes before we commit to it.

@mrocklin @jrbourbeau, I think this is ready to be merged. It lacks documentation of the pack/unpack API but let's wait with that until we settle on an API.

I'll take another look in a bit and then merge in.

This looks slick. Thank you for figuring this out @madsbk

@sjperkins this is in. Now would probably be a good time to add in support for annotations. I think that they should go in client.py::Client._graph_to_futures

@sjperkins this is in. Now would probably be a good time to add in support for annotations. I think that they should go in client.py::Client._graph_to_futures

@mrocklin Thanks for the headsup. My schedule is busy till next week Thursday. Would it be OK if I picked it up then?