Fix Client.persist

Author: fjetter

distributed/client.py

@@ -3763,13 +3763,15 @@ def persist(
         metadata = SpanMetadata(
             collections=[get_collections_metadata(v) for v in collections]
         )
-        expr = collections_to_expr(collections, optimize_graph, **kwargs)
+        expr = collections_to_expr(collections, optimize_graph)
 
-        names = {k for c in collections for k in flatten(c.__dask_keys__())}
+        expr2 = expr.optimize()
+
+        keys = expr2.__dask_keys__()
 
         futures = self._graph_to_futures(
-            expr,
-            names,
+            expr2,
+            list(flatten(keys)),
             workers=workers,
             allow_other_workers=allow_other_workers,
             resources=resources,
@@ -3781,15 +3783,15 @@ def persist(
         )
 
         postpersists = [c.__dask_postpersist__() for c in collections]
-        result = [
-            func({k: futures[k] for k in flatten(c.__dask_keys__())}, *args)
-            for (func, args), c in zip(postpersists, collections)
-        ]
 
         if singleton:
-            return first(result)
+            func, args = postpersists[0]
+            return func({k: futures[k] for k in keys}, *args)
         else:
-            return result
+            return [
+                func({k: futures[k] for k in flatten(ks)}, *args)
+                for (func, args), ks in zip(postpersists, keys)
+            ]
 
     async def _restart(
         self, timeout: str | int | float | NoDefault, wait_for_workers: bool

distributed/tests/test_dask_collections.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import asyncio
+
 import pytest
 
 np = pytest.importorskip("numpy")
@@ -41,6 +43,34 @@ def assert_equal(a, b):
         assert a == b
 
 
+@gen_cluster(client=True)
+async def test_persist(c, s, a, b):
+    df = pd.DataFrame({"x": range(10), "y": range(10, 20)})
+    ddf = dd.from_pandas(df, npartitions=2)
+    df2 = pd.DataFrame({"x": range(20, 30), "y": range(30, 40)})
+    ddf2 = dd.from_pandas(df2, npartitions=2)
+
+    ddfp = await c.persist(ddf)
+    assert s.tasks
+    assert sum(ts.state == "memory" for ts in s.tasks.values()) == 2
+    assert_equal(await c.compute(ddfp), await c.compute(ddf))
+    del ddfp
+
+    while not sum(ts.state == "memory" for ts in s.tasks.values()) == 0:
+        await asyncio.sleep(0.01)
+
+    ddfp1, ddfp2 = c.persist((ddf, ddf2))
+    await wait((ddfp1, ddfp2))
+    assert s.tasks
+    assert sum(ts.state == "memory" for ts in s.tasks.values()) == 4
+
+    assert_equal(await c.compute(ddfp1), await c.compute(ddf))
+    assert_equal(await c.compute(ddfp2), await c.compute(ddf2))
+    del ddfp1, ddfp2
+    while not sum(ts.state == "memory" for ts in s.tasks.values()) == 0:
+        await asyncio.sleep(0.01)
+
+
 @ignore_single_machine_warning
 @gen_cluster(client=True)
 async def test_dataframes(c, s, a, b):