Who's responsible for owning futures in `dask.dataframe.from_delayed`?

[TomAugspurger]

Describe the issue:

Over in rapidsai/dask-upstream-testing#40 (comment), I'm debugging an issue in cuml that cropped up between dask 2025.2.0 and main.

Minimal Complete Verifiable Example:

The following ran with dask / distributed 2025.2.0: uv run --with ipython --with dask[dataframe]==2025.2.0 --with distributed==2025.2.0 --with pyarrow ipython

from distributed import Client
import pandas as pd
import dask.dataframe as dd

client = Client(n_workers=2)

def gen(i):
    return pd.DataFrame({"A": [i]}, index=[i])

futures = [client.submit(gen, i) for i in range(3)]

meta = gen(0)[:0]
df = dd.from_delayed(futures, meta)

df.compute()

import gc

del futures
gc.collect()

df.compute()

With dask / distributed main it fails:

❯ uv run --with ipython  --with "dask[complete] @ git+https://github.com/dask/dask" --with "distributed @ git+https://github.com/dask/distributed" --with pyarrow ipython
 Updated https://github.com/dask/distributed (6691e27d)
 Updated https://github.com/dask/dask (0fa5e18d5)

...

---------------------------------------------------------------------------
FutureCancelledError                      Traceback (most recent call last)
Cell In[1], line 22
     19 del futures
     20 gc.collect()
---> 22 df.compute()

File ~/.cache/uv/archive-v0/n2LAuIxOxeQzXo5nnHTn7/lib/python3.13/site-packages/dask/base.py:373, in DaskMethodsMixin.compute(self, **kwargs)
    349 def compute(self, **kwargs):
    350     """Compute this dask collection
    351
    352     This turns a lazy Dask collection into its in-memory equivalent.
   (...)    371     dask.compute
    372     """
--> 373     (result,) = compute(self, traverse=False, **kwargs)
    374     return result

File ~/.cache/uv/archive-v0/n2LAuIxOxeQzXo5nnHTn7/lib/python3.13/site-packages/dask/base.py:681, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    678     expr = expr.optimize()
    679     keys = list(flatten(expr.__dask_keys__()))
--> 681     results = schedule(expr, keys, **kwargs)
    683 return repack(results)

File ~/.cache/uv/archive-v0/n2LAuIxOxeQzXo5nnHTn7/lib/python3.13/site-packages/distributed/client.py:2395, in Client._gather(self, futures, errors, direct, local_worker)
   2393     exception = st.exception
   2394     traceback = st.traceback
-> 2395     raise exception.with_traceback(traceback)
   2396 if errors == "skip":
   2397     bad_keys.add(key)

FutureCancelledError: ('repartitiontofewer-1204992dafb94ff37742a003f8f0410b', 0) cancelled for reason: lost dependencies.

I'm not quite sure what's going on, but it seems like something changed such that from_delayed no longer captures a reference to the Future object, and so when it's GC'd the task graph is no longer valid. I believe this is related to #8998, but I think they might be different enough to warrant separate issues.

As for what to do, I'm not sure... https://github.com/dask/dask/blob/0fa5e18d511c49f1a9cd5f98c675a9f6cd2fc02f/dask/dataframe/dask_expr/io/_delayed.py#L103-L108 explicitly calls out distributed.Future objects as being supported. So maybe this should continue to work. Alternatively, we can point people to dd.from_graph.

[TomAugspurger]

Looking a bit at the implementation, the Future in that expression looks a bit off.

In [84]: assert fut.key == futures[0].key  # OK!

In [85]:

In [85]: fut.result()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[85], line 1
----> 1 fut.result()

File ~/.cache/uv/archive-v0/n2LAuIxOxeQzXo5nnHTn7/lib/python3.13/site-packages/distributed/client.py:393, in Future.result(self, timeout)
    373 def result(self, timeout=None):
    374     """Wait until computation completes, gather result to local process.
    375
    376     Parameters
   (...)    391         The result of the computation. Or a coroutine if the client is asynchronous.
    392     """
--> 393     self._verify_initialized()
    394     with shorten_traceback():
    395         return self.client.sync(self._result, callback_timeout=timeout)

File ~/.cache/uv/archive-v0/n2LAuIxOxeQzXo5nnHTn7/lib/python3.13/site-packages/distributed/client.py:335, in Future._verify_initialized(self)
    333 def _verify_initialized(self):
    334     if not self.client or not self._state:
--> 335         raise RuntimeError(
    336             f"{type(self)} object not properly initialized. This can happen"
    337             " if the object is being deserialized outside of the context of"
    338             " a Client or Worker."
    339         )

RuntimeError: <class 'distributed.client.Future'> object not properly initialized. This can happen if the object is being deserialized outside of the context of a Client or Worker.

Maybe we just need to ensure that the Future there has the client attached. I'm not sure why the futures in df.expr.delayed_container.operands[0] aren't literally the same future objects though...

[TomAugspurger]

Keeping a ref to whatever the user passed in to from_delayed on the FromDelayed expression seems to fix it...

diff --git a/dask/dataframe/dask_expr/io/_delayed.py b/dask/dataframe/dask_expr/io/_delayed.py
index e948a4a20..7254d6ee0 100644
--- a/dask/dataframe/dask_expr/io/_delayed.py
+++ b/dask/dataframe/dask_expr/io/_delayed.py
@@ -30,6 +30,7 @@ class FromDelayed(PartitionsFiltered, BlockwiseIO):
         "verify_meta",
         "_partitions",
         "prefix",
+        "_ref",
     ]
     _defaults = {
         "meta": None,
@@ -135,6 +136,7 @@ def from_delayed(
         if len(divs) != len(dfs) + 1:  # type: ignore
             raise ValueError("divisions should be a tuple of len(dfs) + 1")
 
+    _ref = dfs
     dfs = [
         delayed(df) if not isinstance(df, Delayed) and hasattr(df, "key") else df
         for df in dfs
@@ -147,7 +149,7 @@ def from_delayed(
     from dask.dataframe.dask_expr._collection import new_collection
 
     result = FromDelayed(
-        DelayedsExpr(*dfs), make_meta(meta), divisions, verify_meta, None, prefix
+        DelayedsExpr(*dfs), make_meta(meta), divisions, verify_meta, None, prefix, _ref
     )
     if pyarrow_strings_enabled() and any(
         pd.api.types.is_object_dtype(dtype)

I'm going to test if this fixes the real problem. I have no idea why that extra reference should be needed, but I think it should be harmless. I don't think we're accidentally keeping something alive that should have had its refcount drop to zero.

cc @fjetter in case you can opine on this / the "fix" on Monday :)

[fjetter]

I can reproduce. I'll have a look.

[fjetter]

A few observations about the reproducer:

• I cannot reproduce it in a Jupyter notebook. I also cannot reproduce it using our default test infrastructure, i.e. I have to explicitly instantiate a LocalCluster instead of relying on any of our test fixtures
• The gc.collect is a red herring. Introducing a brief sleep does the same thing. The reason is that Future.release has to wait for an event loop tick to be able to actually decrease the refs (I have to admit that I don't remember why that is or whether that is still relevant)
• migrate delayed unpack_collections dask#11881 does not resolve this

[fjetter]

But to answer the title of this issue:

Who's responsible for owning futures in dask.dataframe.from_delayed?

The Delayed expression. It should just happen without your ref thing

[fjetter]

I'm almost certain this is a CPython refcounting bug. I seems like I can pass the futures to the expression class (and keep them there) without increasing the python GC refcount. Haven't understood how... but this is probably connected to all the __new__ magic we're doing with expr classes