Use the new high level graph API in Dask

See <dask/dask#7179>

Author: madsbk

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/protocol/highlevelgraph.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
@@ -3635,8 +3635,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():