Remove subgraph callable (#8956)

Author: fjetter

distributed/client.py

@@ -45,7 +45,6 @@
 from dask.core import flatten, validate_key
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import Layer
-from dask.optimization import SubgraphCallable
 from dask.tokenize import tokenize
 from dask.typing import Key, NestedKeys, NoDefault, no_default
 from dask.utils import (
@@ -6120,8 +6119,6 @@ def futures_of(o, client=None):
             stack.extend(x)
         elif type(x) is dict:
             stack.extend(x.values())
-        elif type(x) is SubgraphCallable:
-            stack.extend(x.dsk.values())
         elif isinstance(x, TaskRef):
             if x not in seen:
                 seen.add(x)

distributed/tests/test_client.py

@@ -42,7 +42,6 @@
 import dask
 import dask.bag as db
 from dask import delayed
-from dask.optimization import SubgraphCallable
 from dask.tokenize import tokenize
 from dask.utils import get_default_shuffle_method, parse_timedelta, tmpfile
 
@@ -2626,13 +2625,6 @@ async def test_futures_of_get(c, s, a, b):
     b = db.Bag({("b", i): f for i, f in enumerate([x, y, z])}, "b", 3)
     assert set(futures_of(b)) == {x, y, z}
 
-    sg = SubgraphCallable(
-        {"x": x, "y": y, "z": z, "out": (add, (add, (add, x, y), z), "in")},
-        "out",
-        ("in",),
-    )
-    assert set(futures_of(sg)) == {x, y, z}
-
 
 def test_futures_of_class():
     pytest.importorskip("numpy")
@@ -6191,43 +6183,6 @@ async def test_profile_bokeh(c, s, a, b):
         assert os.path.exists(fn)
 
 
-@gen_cluster(client=True, nthreads=[("", 1)])
-async def test_get_mix_futures_and_SubgraphCallable(c, s, a):
-    future = c.submit(add, 1, 2)
-
-    subgraph = SubgraphCallable(
-        {"_2": (add, "_0", "_1"), "_3": (add, future, "_2")},
-        "_3",
-        ("_0", "_1"),
-    )
-    dsk = {
-        "a": 1,
-        "b": 2,
-        "c": (subgraph, "a", "b"),
-        "d": (subgraph, "c", "b"),
-    }
-
-    future2 = c.get(dsk, "d", sync=False)
-    result = await future2
-    assert result == 11
-
-    # Nested subgraphs
-    subgraph2 = SubgraphCallable(
-        {
-            "_2": (subgraph, "_0", "_1"),
-            "_3": (subgraph, "_2", "_1"),
-            "_4": (add, "_3", future2),
-        },
-        "_4",
-        ("_0", "_1"),
-    )
-
-    dsk2 = {"e": 1, "f": 2, "g": (subgraph2, "e", "f")}
-
-    result = await c.get(dsk2, "g", sync=False)
-    assert result == 22
-
-
 @gen_cluster(client=True)
 async def test_get_mix_futures_and_SubgraphCallable_dask_dataframe(c, s, a, b):
     pd = pytest.importorskip("pandas")

distributed/tests/test_utils_comm.py

@@ -7,7 +7,6 @@
 import pytest
 
 from dask._task_spec import TaskRef
-from dask.optimization import SubgraphCallable
 
 from distributed import wait
 from distributed.compatibility import asyncio_run
@@ -246,18 +245,3 @@ def assert_eq(keys1: set[TaskRef], keys2: set[TaskRef]) -> None:
     res, keys = unpack_remotedata(TaskRef("mykey"))
     assert res == "mykey"
     assert_eq(keys, {TaskRef("mykey")})
-
-    # Check unpack of SC that contains a wrapped key
-    sc = SubgraphCallable({"key": (TaskRef("data"),)}, outkey="key", inkeys=["arg1"])
-    dsk = (sc, "arg1")
-    res, keys = unpack_remotedata(dsk)
-    assert res[0] != sc  # Notice, the first item (the SC) has been changed
-    assert res[1:] == ("arg1", "data")
-    assert_eq(keys, {TaskRef("data")})
-
-    # Check unpack of SC when it takes a wrapped key as argument
-    sc = SubgraphCallable({"key": ("arg1",)}, outkey="key", inkeys=[TaskRef("arg1")])
-    dsk = (sc, "arg1")
-    res, keys = unpack_remotedata(dsk)
-    assert res == (sc, "arg1")  # Notice, the first item (the SC) has NOT been changed
-    assert_eq(keys, set())

distributed/utils_comm.py

@@ -13,7 +13,6 @@
 
 import dask.config
 from dask._task_spec import TaskRef
-from dask.optimization import SubgraphCallable
 from dask.typing import Key
 from dask.utils import is_namedtuple_instance, parse_timedelta
 
@@ -197,30 +196,6 @@ def _unpack_remotedata_inner(
     if typ is tuple:
         if not o:
             return o
-        if type(o[0]) is SubgraphCallable:
-            # Unpack futures within the arguments of the subgraph callable
-            futures: set[TaskRef] = set()
-            args = tuple(_unpack_remotedata_inner(i, byte_keys, futures) for i in o[1:])
-            found_futures.update(futures)
-
-            # Unpack futures within the subgraph callable itself
-            sc: SubgraphCallable = o[0]
-            futures = set()
-            dsk = {
-                k: _unpack_remotedata_inner(v, byte_keys, futures)
-                for k, v in sc.dsk.items()
-            }
-            future_keys: tuple = ()
-            if futures:  # If no futures is in the subgraph, we just use `sc` as-is
-                found_futures.update(futures)
-                future_keys = (
-                    tuple(f.key for f in futures)
-                    if byte_keys
-                    else tuple(f.key for f in futures)
-                )
-                inkeys = tuple(sc.inkeys) + future_keys
-                sc = SubgraphCallable(dsk, sc.outkey, inkeys, sc.name)
-            return (sc,) + args + future_keys
         else:
             return tuple(
                 _unpack_remotedata_inner(item, byte_keys, found_futures) for item in o