Update persist in tests for async clients (#9089)

Author: TomAugspurger

distributed/comm/tests/test_ucx.py

@@ -307,16 +307,16 @@ async def test_stress(
         asynchronous=True,
         host=HOST,
     ) as cluster:
-        async with Client(cluster, asynchronous=True):
+        async with Client(cluster, asynchronous=True) as client:
             rs = da.random.RandomState()
             x = rs.random((10000, 10000), chunks=(-1, chunksize))
-            x = x.persist()
+            x = client.persist(x)
             await wait(x)
 
             for _ in range(10):
                 x = x.rechunk((chunksize, -1))
                 x = x.rechunk((-1, chunksize))
-                x = x.persist()
+                x = client.persist(x)
                 await wait(x)
 
 
@@ -377,9 +377,9 @@ async def test_transpose(
     async with LocalCluster(
         protocol="ucx", n_workers=2, threads_per_worker=2, asynchronous=True
     ) as cluster:
-        async with Client(cluster, asynchronous=True):
+        async with Client(cluster, asynchronous=True) as client:
             assert cluster.scheduler_address.startswith("ucx://")
-            x = da.ones((10000, 10000), chunks=(1000, 1000)).persist()
+            x = client.persist(da.ones((10000, 10000), chunks=(1000, 1000)))
             await x
             y = (x + x.T).sum()
             await y

distributed/diagnostics/tests/test_rmm_diagnostics.py

@@ -32,7 +32,7 @@ async def test_rmm_metrics(c, s, *workers):
     assert w.metrics["rmm"]["rmm-used"] == 0
     assert w.metrics["rmm"]["rmm-total"] == parse_bytes("10MiB")
     result = delayed(rmm.DeviceBuffer)(size=10)
-    result = result.persist()
+    result = c.persist(result)
 
     deadline = Deadline.after(5)