def test_annotation_pack_unpack():
layer = MaterializedLayer({"n": 42}, annotations={"workers": ("alice", )})
packed_anno = layer.__dask_distributed_annotations_pack__()
annotations = {}
Layer.__dask_distributed_annotations_unpack__(annotations, packed_anno,
layer.keys())
assert annotations == {"workers": {"n": ("alice", )}}
def _materialized_layer_pack(
layer: Layer,
all_keys,
known_key_dependencies,
client,
client_keys,
):
from ..client import Future
dsk = dict(layer)
# Find aliases not in `client_keys` and substitute all matching keys
# with its Future
values = {
k: v
for k, v in dsk.items()
if isinstance(v, Future) and k not in client_keys
}
if values:
dsk = subs_multiple(dsk, values)
# Unpack remote data and record its dependencies
dsk = {k: unpack_remotedata(v, byte_keys=True) for k, v in layer.items()}
unpacked_futures = set.union(*[v[1]
for v in dsk.values()]) if dsk else set()
for future in unpacked_futures:
if future.client is not client:
raise ValueError(
"Inputs contain futures that were created by another client.")
if stringify(future.key) not in client.futures:
raise CancelledError(stringify(future.key))
unpacked_futures_deps = {}
for k, v in dsk.items():
if len(v[1]):
unpacked_futures_deps[k] = {f.key for f in v[1]}
dsk = {k: v[0] for k, v in dsk.items()}
# Calculate dependencies without re-calculating already known dependencies
missing_keys = set(dsk.keys()).difference(known_key_dependencies.keys())
dependencies = {
k: keys_in_tasks(all_keys, [dsk[k]], as_list=False)
for k in missing_keys
}
for k, v in unpacked_futures_deps.items():
dependencies[k] = set(dependencies.get(k, ())) | v
# The scheduler expect all keys to be strings
dependencies = {
stringify(k): [stringify(dep) for dep in deps]
for k, deps in dependencies.items()
}
all_keys = all_keys.union(dsk)
dsk = {
stringify(k): stringify(v, exclusive=all_keys)
for k, v in dsk.items()
}
dsk = valmap(dumps_task, dsk)
return {"dsk": dsk, "dependencies": dependencies}
def _materialized_layer_unpack(state, dsk, dependencies, annotations):
dsk.update(state["dsk"])
for k, v in state["dependencies"].items():
dependencies[k] = list(set(dependencies.get(k, ())) | set(v))
if state["annotations"]:
annotations.update(
Layer.expand_annotations(state["annotations"],
state["dsk"].keys()))