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]]