Use the new HLG pack/unpack API in Dask

continuous_integration/environment-3.6.yaml

@@ -43,5 +43,5 @@ dependencies:
       # The '.git@master' is a hack needed to work around the failure
       #     ERROR: Package 'dask' requires a different Python: 3.6.12 not in '>=3.6'
       # It only happens in CI; cannot reproduce locally
-      - git+https://github.com/dask/dask.git@master
+      - git+https://github.com/madsbk/dask.git@hlg_pack_move_to_dask
       - keras

continuous_integration/environment-3.7.yaml

@@ -40,6 +40,6 @@ dependencies:
   - zict
   - zstandard
   - pip:
-      - git+https://github.com/dask/dask
+      - git+https://github.com/madsbk/dask.git@hlg_pack_move_to_dask
       - git+https://github.com/jcrist/crick  # Only tested here
       - keras

continuous_integration/environment-3.8.yaml

@@ -44,7 +44,7 @@ dependencies:
   - zict  # overridden by git tip below
   - zstandard
   - pip:
-      - git+https://github.com/dask/dask
+      - git+https://github.com/madsbk/dask.git@hlg_pack_move_to_dask
       - git+https://github.com/dask/s3fs
       - git+https://github.com/dask/zict
       - git+https://github.com/intake/filesystem_spec

distributed/client.py

@@ -61,7 +61,6 @@
 from .metrics import time
 from .protocol import to_serialize
 from .protocol.pickle import dumps, loads
-from .protocol.highlevelgraph import highlevelgraph_pack
 from .publish import Datasets
 from .pubsub import PubSubClientExtension
 from .security import Security
@@ -2540,7 +2539,7 @@ def _graph_to_futures(
             if not isinstance(dsk, HighLevelGraph):
                 dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
 
-            dsk = highlevelgraph_pack(dsk, self, keyset)
+            dsk = dsk.__dask_distributed_pack__(self, keyset)
 
             annotations = {}
             if user_priority:

distributed/scheduler.py

@@ -35,6 +35,7 @@
 from tornado.ioloop import IOLoop, PeriodicCallback
 
 import dask
+from dask.highlevelgraph import HighLevelGraph
 
 from . import profile
 from .batched import BatchedSend
@@ -84,7 +85,6 @@
 from .pubsub import PubSubSchedulerExtension
 from .stealing import WorkStealing
 from .variable import VariableExtension
-from .protocol.highlevelgraph import highlevelgraph_unpack
 
 try:
     from cython import compiled
@@ -3668,8 +3668,10 @@ def update_graph_hlg(
         fifo_timeout=0,
         annotations=None,
     ):
-
-        dsk, dependencies, annotations = highlevelgraph_unpack(hlg, annotations)
+        unpacked_graph = HighLevelGraph.__dask_distributed_unpack__(hlg, annotations)
+        dsk = unpacked_graph["dsk"]
+        dependencies = unpacked_graph["deps"]
+        annotations = unpacked_graph["anno"]
 
         # Remove any self-dependencies (happens on test_publish_bag() and others)
         for k, v in dependencies.items():