def highlevelgraph_pack(hlg: HighLevelGraph, client, client_keys):
layers = []
# Dump each layer (in topological order)
for layer in (hlg.layers[name] for name in hlg._toposort_layers()):
if not layer.is_materialized():
state = layer.__dask_distributed_pack__(client)
if state is not None:
layers.append({
"__module__": layer.__module__,
"__name__": type(layer).__name__,
"state": state,
})
continue
# Falling back to the default serialization, which will materialize the layer
layers.append({
"__module__":
None,
"__name__":
None,
"state":
_materialized_layer_pack(
layer,
hlg.get_all_external_keys(),
hlg.key_dependencies,
client,
client_keys,
),
})
return msgpack.dumps({"layers": layers}, default=msgpack_encode_default)
def highlevelgraph_pack(hlg: HighLevelGraph, client, client_keys):
"""Pack the high level graph for Scheduler -> Worker communication
The approach is to delegate the packaging to each layer in the high
level graph by calling .__dask_distributed_pack__() on each layer.
If the layer doesn't implement packaging, we materialize the layer
and pack it.
Parameters
----------
hlg: HighLevelGraph
The high level graph to pack
client: distributed.Client
The client calling this function.
client_keys: Iterable
List of keys requested by the client.
Returns
-------
data: bytes
Packed high level graph serialized by msgpack
"""
layers = []
# Dump each layer (in topological order)
for layer in (hlg.layers[name] for name in hlg._toposort_layers()):
if not layer.is_materialized():
state = layer.__dask_distributed_pack__(client)
if state is not None:
layers.append({
"__module__": layer.__module__,
"__name__": type(layer).__name__,
"state": state,
})
continue
# Falling back to the default serialization, which will materialize the layer
layers.append({
"__module__":
None,
"__name__":
None,
"state":
_materialized_layer_pack(
layer,
hlg.get_all_external_keys(),
hlg.key_dependencies,
client,
client_keys,
),
})
return msgpack.dumps({"layers": layers}, default=msgpack_encode_default)
def hlg_layer_topological(hlg: HighLevelGraph, i: int) -> Layer:
"Get the layer from a HighLevelGraph at position ``i``, topologically"
return hlg.layers[hlg._toposort_layers()[i]]
