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"}}