async def test_gh_8991(c, s, a, b):
# Test illustrating something amiss with HighLevelGraph.key_dependencies.
# The intention is for this to be a cache, so if we clear it, things
# should still work.
# This is a bad test, and we should rethink/remove it as soon as the issue is
# resolved whether, it's fixing the underlying problem or removing
# HighLevelGraph.key_dependencies alltogether.
datasets = pytest.importorskip("dask.datasets")
result = datasets.timeseries().shuffle("x").to_orc("tmp", compute=False)
# Create a dsk and mock sending it to the scheduler.
dsk_opt = result.__dask_optimize__(result.dask, result.key)
unpacked = HighLevelGraph.__dask_distributed_unpack__(
dsk_opt.__dask_distributed_pack__(c, result.key))
deps = unpacked["deps"]
# Create a version of the dsk without the key_dependencies and mock sending it to
# the scheduler as well.
dsk_opt_nokeys = dsk_opt.copy()
dsk_opt_nokeys.key_dependencies.clear()
unpacked_nokeys = HighLevelGraph.__dask_distributed_unpack__(
dsk_opt_nokeys.__dask_distributed_pack__(c, result.key))
deps_nokeys = unpacked_nokeys["deps"]
# The recalculated dependencies should still be the same!
assert deps == deps_nokeys
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"}}
