various fixes

distributed/dashboard/components/scheduler.py

@@ -64,6 +64,7 @@
     key_split,
     parse_bytes,
     parse_timedelta,
+    stringify,
 )
 
 from distributed.core import Status
@@ -2371,7 +2372,7 @@ def add_new_nodes_edges(self, new, new_edges, update=False):
                     continue
                 xx = x[key]
                 yy = y[key]
-                node_key.append(escape.url_escape(key))
+                node_key.append(escape.url_escape(stringify(key)))
                 node_x.append(xx)
                 node_y.append(yy)
                 node_state.append(task.state)

distributed/shuffle/_shuffle.py

@@ -116,7 +116,7 @@ def rearrange_by_column_p2p(
             f"p2p requires all column names to be str, found: {unsupported}",
         )
 
-    name = f"shuffle-p2p-{token}"
+    name = f"shuffle_p2p-{token}"
     layer = P2PShuffleLayer(
         name,
         column,

distributed/shuffle/tests/test_shuffle.py

@@ -15,6 +15,8 @@
 
 import pytest
 
+from dask.utils import key_split
+
 pd = pytest.importorskip("pandas")
 dd = pytest.importorskip("dask.dataframe")
 
@@ -272,7 +274,13 @@ async def wait_for_tasks_in_state(
         raise TypeError(dask_worker)
 
     while (
-        len([key for key, ts in tasks.items() if prefix in key and ts.state == state])
+        len(
+            [
+                key
+                for key, ts in tasks.items()
+                if prefix in key_split(key) and ts.state == state
+            ]
+        )
         < count
     ):
         await asyncio.sleep(interval)
@@ -752,7 +760,7 @@ async def test_closed_worker_during_unpack(c, s, a, b):
     )
     out = dd.shuffle.shuffle(df, "x", shuffle="p2p")
     x, y = c.compute([df.x.size, out.x.size])
-    await wait_for_tasks_in_state("shuffle-p2p", "memory", 1, b)
+    await wait_for_tasks_in_state("shuffle_p2p", "memory", 1, b)
     await b.close()
 
     x = await x
@@ -779,7 +787,7 @@ async def test_restarting_during_unpack_raises_killed_worker(c, s, a, b):
     )
     out = dd.shuffle.shuffle(df, "x", shuffle="p2p")
     out = c.compute(out.x.size)
-    await wait_for_tasks_in_state("shuffle-p2p", "memory", 1, b)
+    await wait_for_tasks_in_state("shuffle_p2p", "memory", 1, b)
     await b.close()
 
     with pytest.raises(KilledWorker):
@@ -806,7 +814,7 @@ async def test_crashed_worker_during_unpack(c, s, a):
         out = dd.shuffle.shuffle(df, "x", shuffle="p2p")
         y = c.compute(out.x.size)
 
-        await wait_until_worker_has_tasks("shuffle-p2p", killed_worker_address, 1, s)
+        await wait_until_worker_has_tasks("shuffle_p2p", killed_worker_address, 1, s)
         await n.process.process.kill()
 
         y = await y
@@ -1171,7 +1179,7 @@ def block(df, in_event, block_event):
             out = block(out, in_event, block_event)
         out = c.compute(out)
 
-        await wait_until_worker_has_tasks("shuffle-p2p", n.worker_address, 1, s)
+        await wait_until_worker_has_tasks("shuffle_p2p", n.worker_address, 1, s)
         await in_event.wait()
         await n.process.process.kill()
         await block_event.set()
@@ -1199,7 +1207,7 @@ async def test_crashed_worker_after_shuffle_persisted(c, s, a):
         out = dd.shuffle.shuffle(df, "x", shuffle="p2p")
         out = out.persist()
 
-        await wait_until_worker_has_tasks("shuffle-p2p", n.worker_address, 1, s)
+        await wait_until_worker_has_tasks("shuffle_p2p", n.worker_address, 1, s)
         await out
 
         await n.process.process.kill()

distributed/shuffle/tests/utils.py

@@ -74,7 +74,7 @@ def transition(
         assert self.scheduler
         if finish != "waiting":
             return
-        if not key.startswith("shuffle-barrier-"):
+        if not isinstance(key, str) or not key.startswith("shuffle-barrier-"):
             return
         if key in self.seen:
             return

distributed/tests/test_resources.py

@@ -222,9 +222,9 @@ async def test_resources_str(c, s, a, b):
     yy = y.persist(resources={"MyRes": 1})
     await wait(yy)
 
-    ts_first = s.tasks[y.__dask_keys__([0])]
+    ts_first = s.tasks[y.__dask_keys__()[0]]
     assert ts_first.resource_restrictions == {"MyRes": 1}
-    ts_last = s.tasks[y.__dask_keys__([-1])]
+    ts_last = s.tasks[y.__dask_keys__()[-1]]
     assert ts_last.resource_restrictions == {"MyRes": 1}