Fix CI by using client.persist()

Author: hendrikmakait

distributed/comm/tests/test_ws.py

@@ -125,7 +125,7 @@ async def test_collections(c, s, a, b):
     da = pytest.importorskip("dask.array")
     x = da.random.random((1000, 1000), chunks=(100, 100))
     x = x + x.T
-    await x.persist()
+    await c.persist(x)
 
 
 @gen_cluster(client=True, scheduler_kwargs={"protocol": "ws://"})

distributed/dashboard/tests/test_scheduler_bokeh.py

@@ -833,9 +833,9 @@ async def test_TaskGraph(c, s, a, b):
     json.dumps(gp.node_source.data)
 
     da = pytest.importorskip("dask.array")
-    x = da.random.random((20, 20), chunks=(10, 10)).persist()
+    x = c.persist(da.random.random((20, 20), chunks=(10, 10)))
     y = (x + x.T) - x.mean(axis=0)
-    y = y.persist()
+    y = c.persist(y)
     await wait(y)
 
     gp.update()
@@ -912,12 +912,12 @@ async def test_TaskGraph_complex(c, s, a, b):
     da = pytest.importorskip("dask.array")
     gp = TaskGraph(s)
     x = da.random.random((2000, 2000), chunks=(1000, 1000))
-    y = ((x + x.T) - x.mean(axis=0)).persist()
+    y = c.persist((x + x.T) - x.mean(axis=0))
     await wait(y)
     gp.update()
     assert len(gp.layout.index) == len(gp.node_source.data["x"])
     assert len(gp.layout.index) == len(s.tasks)
-    z = (x - y).sum().persist()
+    z = c.persist((x - y).sum())
     await wait(z)
     gp.update()
     assert len(gp.layout.index) == len(gp.node_source.data["x"])
@@ -1177,10 +1177,10 @@ async def test_memory_by_key(c, s, a, b):
     mbk = MemoryByKey(s)
 
     da = pytest.importorskip("dask.array")
-    x = (da.random.random((20, 20), chunks=(10, 10)) + 1).persist(optimize_graph=False)
+    x = c.persist(da.random.random((20, 20), chunks=(10, 10)) + 1, optimize_graph=False)
     await x
 
-    y = await dask.delayed(inc)(1).persist()
+    y = await c.persist(dask.delayed(inc)(1))
 
     mbk.update()
     assert mbk.source.data["name"] == ["add", "inc"]
@@ -1260,10 +1260,10 @@ async def test_aggregate_action(c, s, a, b):
     mbk = AggregateAction(s)
 
     da = pytest.importorskip("dask.array")
-    x = (da.ones((20, 20), chunks=(10, 10)) + 1).persist(optimize_graph=False)
+    x = c.persist(da.ones((20, 20), chunks=(10, 10)) + 1, optimize_graph=False)
 
     await x
-    y = await dask.delayed(inc)(1).persist()
+    y = await c.persist(dask.delayed(inc)(1))
     z = (x + x.T) - x.mean(axis=0)
     await c.compute(z.sum())
 
@@ -1289,9 +1289,9 @@ async def test_compute_per_key(c, s, a, b):
     mbk = ComputePerKey(s)
 
     da = pytest.importorskip("dask.array")
-    x = (da.ones((20, 20), chunks=(10, 10)) + 1).persist(optimize_graph=False)
+    x = c.persist(da.ones((20, 20), chunks=(10, 10)) + 1, optimize_graph=False)
     await x
-    y = await dask.delayed(inc)(1).persist()
+    y = await c.persist(dask.delayed(inc)(1))
     z = (x + x.T) - x.mean(axis=0)
     zsum = z.sum()
     await c.compute(zsum)
@@ -1367,7 +1367,7 @@ async def test_shuffling(c, s, a, b):
         freq="10 s",
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
-        df2 = df.shuffle("x").persist()
+        df2 = c.persist(df.shuffle("x"))
     start = time()
     while not ss.source.data["comm_written"]:
         await asyncio.gather(*[a.heartbeat(), b.heartbeat()])

distributed/shuffle/tests/test_shuffle.py

@@ -286,7 +286,7 @@ async def test_bad_disk(c, s, a, b):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = out.persist()
+    out = c.persist(out)
     shuffle_id = await wait_until_new_shuffle_is_initialized(s)
     while not get_active_shuffle_runs(a):
         await asyncio.sleep(0.01)
@@ -1068,7 +1068,7 @@ async def test_heartbeat(c, s, a, b):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = out.persist()
+    out = c.persist(out)
 
     while not s.plugins["shuffle"].heartbeats:
         await asyncio.sleep(0.001)
@@ -1272,7 +1272,7 @@ async def test_head(c, s, a, b):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = await out.head(compute=False).persist()  # Only ask for one key
+    out = await c.persist(out.head(compute=False))  # Only ask for one key
 
     assert list(os.walk(a.local_directory)) == a_files  # cleaned up files?
     assert list(os.walk(b.local_directory)) == b_files
@@ -1380,7 +1380,7 @@ async def test_clean_after_forgotten_early(c, s, a, b):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = out.persist()
+    out = c.persist(out)
     await wait_for_tasks_in_state("shuffle-transfer", "memory", 1, a)
     await wait_for_tasks_in_state("shuffle-transfer", "memory", 1, b)
     del out
@@ -1399,7 +1399,7 @@ async def test_tail(c, s, a, b):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         x = df.shuffle("x")
-    full = await x.persist()
+    full = await c.persist(x)
     ntasks_full = len(s.tasks)
     del full
     while s.tasks:
@@ -1539,7 +1539,7 @@ async def test_crashed_worker_after_shuffle_persisted(c, s, a):
         )
         with dask.config.set({"dataframe.shuffle.method": "p2p"}):
             out = df.shuffle("x")
-        out = out.persist()
+        out = c.persist(out)
 
         await event.wait()
         await out
@@ -1597,7 +1597,7 @@ async def test_new_worker(c, s, a, b):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         shuffled = df.shuffle("x")
-    persisted = shuffled.persist()
+    persisted = c.persist(shuffled)
     while not s.plugins["shuffle"].active_shuffles:
         await asyncio.sleep(0.001)
 
@@ -1646,7 +1646,7 @@ async def test_delete_some_results(c, s, a, b):
         freq="10 s",
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
-        x = df.shuffle("x").persist()
+        x = c.persist(df.shuffle("x"))
     while not s.tasks or not any(ts.state == "memory" for ts in s.tasks.values()):
         await asyncio.sleep(0.01)
 
@@ -1669,12 +1669,12 @@ async def test_add_some_results(c, s, a, b):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         x = df.shuffle("x")
-    y = x.partitions[: x.npartitions // 2].persist()
+    y = c.persist(x.partitions[: x.npartitions // 2])
 
     while not s.tasks or not any(ts.state == "memory" for ts in s.tasks.values()):
         await asyncio.sleep(0.01)
 
-    x = x.persist()
+    x = c.persist(x)
 
     await c.compute(x.size)
 
@@ -1697,7 +1697,7 @@ async def test_clean_after_close(c, s, a, b):
 
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = out.persist()
+    out = c.persist(out)
 
     await wait_for_tasks_in_state("shuffle-transfer", "executing", 1, a)
     await wait_for_tasks_in_state("shuffle-transfer", "memory", 1, b)
@@ -2079,7 +2079,7 @@ async def test_deduplicate_stale_transfer(c, s, a, b, wait_until_forgotten):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         shuffled = df.shuffle("x")
-    shuffled = shuffled.persist()
+    shuffled = c.persist(shuffled)
 
     shuffle_extA = a.plugins["shuffle"]
     shuffle_extB = b.plugins["shuffle"]
@@ -2131,7 +2131,7 @@ async def test_handle_stale_barrier(c, s, a, b, wait_until_forgotten):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         shuffled = df.shuffle("x")
-    shuffled = shuffled.persist()
+    shuffled = c.persist(shuffled)
 
     shuffle_extA = a.plugins["shuffle"]
     shuffle_extB = b.plugins["shuffle"]
@@ -2190,7 +2190,7 @@ async def test_shuffle_run_consistency(c, s, a):
     # Initialize first shuffle execution
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = out.persist()
+    out = c.persist(out)
 
     shuffle_id = await wait_until_new_shuffle_is_initialized(s)
     spec = scheduler_ext.get(shuffle_id, a.worker_address)["run_spec"].data
@@ -2214,7 +2214,7 @@ async def test_shuffle_run_consistency(c, s, a):
     # Initialize second shuffle execution
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = out.persist()
+    out = c.persist(out)
 
     new_shuffle_id = await wait_until_new_shuffle_is_initialized(s)
     assert shuffle_id == new_shuffle_id
@@ -2241,7 +2241,7 @@ async def test_shuffle_run_consistency(c, s, a):
     # Create an unrelated shuffle on a different column
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("y")
-    out = out.persist()
+    out = c.persist(out)
     independent_shuffle_id = await wait_until_new_shuffle_is_initialized(s)
     assert shuffle_id != independent_shuffle_id
 
@@ -2284,7 +2284,7 @@ async def test_fail_fetch_race(c, s, a):
     )
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = out.persist()
+    out = c.persist(out)
 
     shuffle_id = await wait_until_new_shuffle_is_initialized(s)
     spec = scheduler_ext.get(shuffle_id, a.worker_address)["run_spec"].data
@@ -2365,7 +2365,7 @@ async def test_replace_stale_shuffle(c, s, a, b):
     # Initialize first shuffle execution
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = out.persist()
+    out = c.persist(out)
 
     shuffle_id = await wait_until_new_shuffle_is_initialized(s)
 
@@ -2391,7 +2391,7 @@ async def test_replace_stale_shuffle(c, s, a, b):
     # Initialize second shuffle execution
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = df.shuffle("x")
-    out = out.persist()
+    out = c.persist(out)
 
     await wait_for_tasks_in_state("shuffle-transfer", "memory", 1, a)
     await run_manager_B.finished_get_shuffle_run.wait()

distributed/tests/test_client.py

@@ -401,7 +401,7 @@ async def test_future_repr(c, s, a, b):
 async def test_future_tuple_repr(c, s, a, b):
     pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
-    y = da.arange(10, chunks=(5,)).persist()
+    y = c.persist(da.arange(10, chunks=(5,)))
     f = futures_of(y)[0]
     for func in [repr, lambda x: x._repr_html_()]:
         for k in f.key:
@@ -4808,7 +4808,7 @@ async def test_restart_workers(c, s, a, b):
     # Persist futures and perform a computation
     size = 100
     x = da.ones(size, chunks=10)
-    x = x.persist()
+    x = c.persist(x)
     assert await c.compute(x.sum()) == size
 
     # Restart a single worker
@@ -5178,7 +5178,7 @@ async def test_serialize_collections(c, s, a, b):
     pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
 
-    x = da.arange(10, chunks=(5,)).persist()
+    x = c.persist(da.arange(10, chunks=(5,)))
 
     def f(x):
         assert isinstance(x, da.Array)
@@ -5360,7 +5360,7 @@ async def test_serialize_collections_of_futures(c, s, a, b):
     from dask.dataframe.utils import assert_eq
 
     df = pd.DataFrame({"x": [1, 2, 3]})
-    ddf = dd.from_pandas(df, npartitions=2).persist()
+    ddf = c.persist(dd.from_pandas(df, npartitions=2))
     future = await c.scatter(ddf)
 
     ddf2 = await future
@@ -5509,7 +5509,7 @@ async def test_call_stack_collections(c, s, a, b):
     pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
 
-    x = da.random.random(100, chunks=(10,)).map_blocks(slowinc, delay=0.5).persist()
+    x = c.persist(da.random.random(100, chunks=(10,)).map_blocks(slowinc, delay=0.5))
     while not a.state.executing_count and not b.state.executing_count:
         await asyncio.sleep(0.001)
     result = await c.call_stack(x)
@@ -5521,7 +5521,7 @@ async def test_call_stack_collections_all(c, s, a, b):
     pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
 
-    x = da.random.random(100, chunks=(10,)).map_blocks(slowinc, delay=0.5).persist()
+    x = c.persist(da.random.random(100, chunks=(10,)).map_blocks(slowinc, delay=0.5))
     while not a.state.executing_count and not b.state.executing_count:
         await asyncio.sleep(0.001)
     result = await c.call_stack()
@@ -6179,7 +6179,7 @@ async def test_get_mix_futures_and_SubgraphCallable_dask_dataframe(c, s, a, b):
     dd = pytest.importorskip("dask.dataframe")
 
     df = pd.DataFrame({"x": range(1, 11)})
-    ddf = dd.from_pandas(df, npartitions=2).persist()
+    ddf = c.persist(dd.from_pandas(df, npartitions=2))
     ddf = ddf.map_partitions(lambda x: x)
     ddf["x"] = ddf["x"].astype("f8")
     ddf = ddf.map_partitions(lambda x: x)
@@ -6239,10 +6239,10 @@ async def test_file_descriptors_dont_leak(Worker):
         async with (
             Worker(s.address),
             Worker(s.address),
-            Client(s.address, asynchronous=True),
+            Client(s.address, asynchronous=True) as c,
         ):
             assert proc.num_fds() > before
-            await df.sum().persist()
+            await c.persist(df.sum())
 
     start = time()
     while proc.num_fds() > before:
@@ -6340,7 +6340,7 @@ async def test_config_inherited_by_subprocess():
 
 @gen_cluster(client=True)
 async def test_futures_of_sorted(c, s, a, b):
-    b = dask.bag.from_sequence(range(10), npartitions=5).persist()
+    b = c.persist(dask.bag.from_sequence(range(10), npartitions=5))
     futures = futures_of(b)
     assert [fut.key for fut in futures] == [k for k in b.__dask_keys__()]
 
@@ -6767,7 +6767,7 @@ async def test_annotations_task_state(c, s, a, b):
         x = da.ones(10, chunks=(5,))
 
     with dask.config.set(optimization__fuse__active=False):
-        x = await x.persist()
+        x = await c.persist(x)
 
     for ts in s.tasks.values():
         assert ts.annotations["qux"] == "bar"
@@ -6822,7 +6822,7 @@ async def test_annotations_priorities(c, s, a, b):
         x = da.ones(10, chunks=(5,))
 
     with dask.config.set(optimization__fuse__active=False):
-        x = await x.persist()
+        x = await c.persist(x)
 
     for ts in s.tasks.values():
         assert ts.priority[0] == -15
@@ -6838,7 +6838,7 @@ async def test_annotations_workers(c, s, a, b):
         x = da.ones(10, chunks=(5,))
 
     with dask.config.set(optimization__fuse__active=False):
-        x = await x.persist()
+        x = await c.persist(x)
 
     for ts in s.tasks.values():
         assert ts.annotations["workers"] == [a.address]
@@ -6857,7 +6857,7 @@ async def test_annotations_retries(c, s, a, b):
         x = da.ones(10, chunks=(5,))
 
     with dask.config.set(optimization__fuse__active=False):
-        x = await x.persist()
+        x = await c.persist(x)
 
     for ts in s.tasks.values():
         assert ts.retries == 2
@@ -6910,7 +6910,7 @@ async def test_annotations_resources(c, s, a, b):
         x = da.ones(10, chunks=(5,))
 
     with dask.config.set(optimization__fuse__active=False):
-        x = await x.persist()
+        x = await c.persist(x)
 
     for ts in s.tasks.values():
         assert ts.resource_restrictions == {"GPU": 1}
@@ -6949,7 +6949,7 @@ async def test_annotations_loose_restrictions(c, s, a, b):
         x = da.ones(10, chunks=(5,))
 
     with dask.config.set(optimization__fuse__active=False):
-        x = await x.persist()
+        x = await c.persist(x)
 
     for ts in s.tasks.values():
         assert not ts.worker_restrictions
@@ -7229,7 +7229,7 @@ async def test_computation_object_code_dask_persist(c, s, a, b):
     da = pytest.importorskip("dask.array")
 
     x = da.ones((10, 10), chunks=(3, 3))
-    future = x.sum().persist()
+    future = c.persist(x.sum())
     await future
 
     test_function_code = inspect.getsource(