[REVIEW] Formalization of Computation

[madsbk]

This PR implements graph computations based on the specification in Dask:

A computation may be one of the following:

• Any key present in the Dask graph like 'x'
• Any other value like 1, to be interpreted literally
• A task like (inc, 'x')
• A list of computations, like [1, 'x', (inc, 'x')]

In order to support efficient and flexible task serialization, this PR introduces classes for computations, tasks, data, functions, etc.
This PR continues the existing coarse grained approach where (nested) tasks are pickled. We used to use dumps_function() and warn_dumps() to do this. Now, we use PickledComputation as the outermost wrapper and PickledObjects for already pickled objects. This makes it possible for HLG.unpack(), which runs on the Scheduler, to build new tasks of already pickled objects.

Notable Classes

• PickledObject - An object that are serialized using protocol.pickle. This object isn't a computation by itself instead users can build computations containing them. It is automatically de-serialized by the Worker before execution.

• Computation - A computation that the Worker can execute. The Scheduler sees this as a black box. A computation cannot contain pickled objects but it may contain Serialize and/or Serialized objects, which will be de-serialize when arriving on the Worker automatically.

• PickledComputation - A computation that are serialized using protocol.pickle. The class is derived from Computation but can contain pickled objects. Itself and contained pickled objects will be de-serialized by the Worker before execution.

Notable Functions

• typeset_dask_graph() - Use to typeset a Dask graph, which wrap computations in either the Data or Task class. This should be done before communicating the graph. Note, this replaces the old tlz.valmap(dumps_task, dsk) operation.

---

• Closes [Discussion] Serialize objects within tasks #4673
• Closes [DISCUSSION] Can the scheduler use pickle.dumps? #4890
• Closes [WIP] Distinguish tuples & lists in MsgPack serialization #4575
• Closes Single pass serialization  #4699
• Closes [WIP] Fine grained serialization #4897
• Closes dask.dataframe.read_csv('./filepath/*.csv') returning tuple dask#7777
• Closes dumps_task in SimpleShuffleLayer and BroadcastJoinLayer unpack dask#7650
• Tests added / passed
• Passes black distributed / flake8 distributed / isort distributed

[mrocklin]

@jcrist @TomAugspurger if either of you have time (I suspect not) it would be good to get the perspective of someone who is familiar with Dask broadly, who also isn't familiar with the work that Mads and team have been doing. A sanity check here would be welcome.

[mrocklin]

Should we have PickledObjects? Or should we use the general serialization path for this?

[madsbk]

You mean if we have like a list of objects? That should also just use a single PickledObject to serialize everything in one go.
We use typeset_computation() to look through a computation and wrap individual task functions in PickledCallable.

[mrocklin]

I mean "why would we ever want to pickle an object, rather than give the rest of our serialization machinery a chance to work?" For example, what if the object was a cupy array.

Or maybe this is used very infrequently, and not where user-data is likely to occur?

[madsbk]

This is basically the fine vs coarse grained serialization discussion. This PR continues the existing coarse grained approach where (nested) tasks are pickled. We used to use dumps_function() and warn_dumps() to do this.
Now, we use PickledComputation as the outermost wrapper and PickledObjects for already pickled objects. This makes it possible for HLG.unpack() that runs on the Scheduler to build new tasks of already pickled objects.

[madsbk]

Or maybe this is used very infrequently, and not where user-data is likely to occur?

Yes, including large data in the task directly will raise a warning just like it use to.

[GPUtester]

Can one of the admins verify this patch?

[madsbk]

@jrbourbeau if you can find time to review this PR, it would be great.