Repository navigation
[Discussion] Serialize objects within tasks #4673
Description
When Blockwise materializes its graph on the scheduler, we'd sometimes like to be able to insert bits of data from the client that aren't msgpack-serializable into that graph. Ideally, we could wrap these things in to_serialize on the client, access them as opaque Serialized blobs on the scheduler to insert them into tasks, and when the worker receives the task, the Serialized blobs within would automatically be deserialized.
However, most tasks are pre-serialized by dumps_task, which calls pickle directly. So if a task contains a Serialize object, the normal serialization machinery on the comm will never know, since it's just bytes by the time it reaches dumps. Same thing in the other direction. The reason for this special case is to cache serialized functions—particularly for numba functions, where we want to reuse the same function object instead of re-jitting every time.
I think our requirements are:
- If tasks contain
SerializeorSerializedobjects, those objects are automatically deserialized on the worker, but not on the scheduler. - Serialized functions are cached, like they currently are with
loads_functionanddumps_function. That is,deserialize(serialize(func)) is deserialize(serialize(func))(unlessserialize(func)is huge). Okay if this doesn't apply toserializein general and is specialized to tasks. - Users still see the
warn_dumpswarning when submitting large tasks.
Nice-to-haves:
__dask_distributed_*pack__doesn't usedumps_msgpack/loads_msgpackdirectly; it just produces/consumes the marshalled objects (which may containSerialize/Serialized) and lets comm serialization handle the rest. I think this separation of concerns would be nice for longer-term maintenance.make_blockwise_graphdoesn't need adeserializingargument (doesn't calldumps_function/warn_dumps, or handlestringify_collection_keys/ formatting tasks into a dict)—again, would be nice to separate concerns. I'm wondering ifmake_blockwise_graphcould switch to a generator:dict(make_blockwise_graph(...))would give you a plain dask, butScheduler.update_graph_hlgwould provide the logic to handle the key-stringification, reformatting to a task dict, and possibly even doing something to mark the function as needing to be cached. A generator could help separate concerns without adding the overhead of a full dict traversal.
Ideas:
I think the general idea is to not have dumps_task serialize to bytes anymore, but just return Serialize objects? Maybe even get rid of dumps_task entirely? This would delegate all serialization-to-bytes to the comm level.
- Move function-caching logic into core
serialize/deserialize. Fordeserialize, this will lead to a lot of cache misses. Bad idea. - Register a
dask_serializemethod forFunctionTypethat caches, then just calls pickle. Of course, this won't cache other callables besides functions. Mediocre idea. - Register a dummy serialization family for
"task", which caches, then just calls back intodeserialize. By marking things as"task"when serializing, we can reduce cache misses. Kinda hackey, could work?
Alternatively, the quickest solution would be to leave everything as it is, and have worker._deserialize just call deserialize on the args/kwargs, to unpack any Serialize/Serialized objects they may contain. But I imagine that is undesirable, especially when we already have this nice machinery for handling nested serialization with to_serialize and Serialized.
Pitfalls:
- If tasks are being deserialized at the comm level, what will errors look like to users when unpickling a task fails (such as due to a missing import)? What will happen to the task from the scheduler's perspective?
cc @mrocklin @jrbourbeau @ian-r-rose @jakirkham @madsbk @rjzamora @quasiben: how do these requirements sound? What should be added or removed?
Activity
rjzamora commentedon Apr 5, 2021
Thanks for writing this up @gjoseph92! This seems like a great summary. Just adding some notes in case they are helpful:
Yes - The problem is that
dumps_taskusually callswarn_dumpson the task arguments (which just pickles the arguments together - hiding internalSerializedobjects at deserialization time). When aBlockwiselayer is materialized on the scheduler, we actually don’t usedumps_taskat all, since the function was already serialized by the client. Instead we manually define a dict for each task in (almost) the same waydumps_taskwould have done:{“function": <already-serialized-function>, "args": warn_dumps(args)}.The current
Blockwiseapproach is a problem because: (1)warn_dumpsdoes not properly handle existingSerializedobjects, (2) The manual-dictapproach precludes serialization of in-lined tasks.If we ignore the nested-task case for now, a possible solution is to:
dumps_functionto avoid re-pickling bytes objects (already-serialized functions). Since this utility is specifically designed to serialize functions, it should be fine to avoid serializing abytesobject. (This allows us to go back to usingdumps_taskwhen we construct a graph)warn_serializefunction that does something similar towarn_dumps, but generates aSerializedobject instead of pickled bytes.dumps_task, we add an option to usewarn_serialize(instead ofwarn_dumps) to serialize the “args”/“kwargs” elements of the task dictionary. We can either do this whenSerializedelements are detected in the task args, or we can add an optional argument.serializewill iterate through each task element (and create separate sub-header/frames) whenwarn_serializeis used. For example, this will always be done if something like [WIP] Distinguishtuples &lists in MsgPack serialization #4575 gets in, and can be controlled by aniterate_collectionargument if Add iterate_collection argument to serialize #4641 is integrated._deserializeto avoid calling pickle.loads if the “args”/“kwargs” are notbytesobjects (these elements will already be deserialized ifwarn_serializewas used on the scheduler).I can confirm that the steps above address the current function-caching and
Serialized-object problems. However, to deal with nested (in-lined) tasks, we would need to do something ugly, like wrap the serialized function in a callable object when the graph is constructed on the scheduler. Therefore, in the long run, it may be cleanest use something similar to your last two ideas.add changes discussed in dask#4673
madsbk commentedon Apr 8, 2021
Thanks @gjoseph92, I think you proposal is a reasonable approach but it got me thinking: maybe it is time to standardize tasks.
What if we decouple the arguments of a task from the callable/function? Something like:
Normalize the tasks such that task arguments are either constants (objects not containing any keys) or pure keys, nothing in-between! This normalization will be done automatically and might include wrapping the task in a
SubgraphCallablein order to handle nested keys and/or tasks.Serialize the callable and its constant arguments of the normalized tasks. The pure key arguments doesn't need any serialization.
Wrap the tasks in a class that clearly define what are the serialized callable, what are the serialized constant arguments, and what are the pure input keys.
Now, the scheduler can calculate task dependencies, task annotations, stringify input keys, and modify input keys without having to deserialize anything.
rjzamora commentedon Apr 8, 2021
Thanks for your thoughts here @madsbk ! I do get the impression that this is all further motivation for something like this. I have experimented with multiple ways to address the problem (both within
serializeand with a newserialize_taskutility) and they both feel like workarounds. I am trying to work out aserialize-based approach with minimal changes, but I think the long-term solution may require a new task specification.madsbk commentedon Apr 22, 2021
@gjoseph92, FYI I am working on a PR streamlining serialization: #4699
fjetter commentedon Aug 4, 2023
FYI I have a WIP PR open that is removing
dumps_task#8067 Now that we're much more generously pickling things I don't think this complexity is still warranted