async def test_annotation_pack_unpack(c, s, a, b):
hlg = HighLevelGraph({"l1": MaterializedLayer({"n": 42})}, {"l1": set()})
packed_hlg = hlg.__dask_distributed_pack__(c, ["n"])
annotations = {"workers": ("alice",)}
unpacked_hlg = HighLevelGraph.__dask_distributed_unpack__(packed_hlg, annotations)
annotations = unpacked_hlg["annotations"]
assert annotations == {"workers": {"n": ("alice",)}}
async def test_pack_MaterializedLayer_handles_futures_in_graph_properly(
c, s, a, b):
fut = c.submit(inc, 1)
hlg = HighLevelGraph(
{
"l1": MaterializedLayer({
"x": fut,
"y": (inc, "x"),
"z": (inc, "y")
})
},
{"l1": set()},
)
# fill hlg.key_dependencies cache. This excludes known futures, so only
# includes a subset of all dependencies. Previously if the cache was present
# the future dependencies would be missing when packed.
hlg.get_all_dependencies()
packed = hlg.__dask_distributed_pack__(c, ["z"], {})
unpacked = HighLevelGraph.__dask_distributed_unpack__(packed)
assert unpacked["deps"] == {"x": {fut.key}, "y": {fut.key}, "z": {"y"}}