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