chore: import upstream snapshot with attribution

python/ray/util/dask/tests/BUILD.bazel

@@ -0,0 +1,90 @@
+# --------------------------------------------------------------------
+# Tests from the python/ray/util/dask/tests directory.
+# Please keep these sorted alphabetically.
+# --------------------------------------------------------------------
+
+load("@rules_python//python:defs.bzl", "py_test")
+
+py_test(
+    name = "test_dask_callback",
+    size = "small",
+    srcs = ["test_dask_callback.py"],
+    tags = [
+        "exclusive",
+        "team:core",
+    ],
+    deps = ["//python/ray/util/dask:dask_lib"],
+)
+
+py_test(
+    name = "test_dask_callback_client_mode",
+    size = "medium",
+    srcs = ["test_dask_callback.py"],
+    main = "test_dask_callback.py",
+    tags = [
+        "client",
+        "exclusive",
+        "team:core",
+    ],
+    deps = ["//python/ray/util/dask:dask_lib"],
+)
+
+py_test(
+    name = "test_dask_optimization",
+    size = "small",
+    srcs = ["test_dask_optimization.py"],
+    tags = [
+        "exclusive",
+        "team:core",
+    ],
+    deps = ["//python/ray/util/dask:dask_lib"],
+)
+
+py_test(
+    name = "test_dask_multi_node",
+    size = "medium",
+    srcs = ["test_dask_multi_node.py"],
+    main = "test_dask_multi_node.py",
+    tags = [
+        "exclusive",
+        "team:core",
+    ],
+    deps = ["//python/ray/util/dask:dask_lib"],
+)
+
+py_test(
+    name = "test_dask_optimization_client_mode",
+    size = "small",
+    srcs = ["test_dask_optimization.py"],
+    main = "test_dask_optimization.py",
+    tags = [
+        "client",
+        "exclusive",
+        "team:core",
+    ],
+    deps = ["//python/ray/util/dask:dask_lib"],
+)
+
+py_test(
+    name = "test_dask_scheduler",
+    size = "small",
+    srcs = ["test_dask_scheduler.py"],
+    tags = [
+        "exclusive",
+        "team:core",
+    ],
+    deps = ["//python/ray/util/dask:dask_lib"],
+)
+
+py_test(
+    name = "test_dask_scheduler_client_mode",
+    size = "small",
+    srcs = ["test_dask_scheduler.py"],
+    main = "test_dask_scheduler.py",
+    tags = [
+        "client",
+        "exclusive",
+        "team:core",
+    ],
+    deps = ["//python/ray/util/dask:dask_lib"],
+)

python/ray/util/dask/tests/test_dask_callback.py

@@ -0,0 +1,236 @@
+import sys
+
+import dask
+import pytest
+
+import ray
+from ray.tests.conftest import *  # noqa: F403, F401
+from ray.util.dask import RayDaskCallback, ray_dask_get
+
+
+@dask.delayed
+def add(x, y):
+    return x + y
+
+
+def test_callback_active():
+    """Test that callbacks are active within context"""
+    assert not RayDaskCallback.ray_active
+
+    with RayDaskCallback():
+        assert RayDaskCallback.ray_active
+
+    assert not RayDaskCallback.ray_active
+
+
+def test_presubmit_shortcircuit(ray_start_regular_shared):
+    """
+    Test that presubmit return short-circuits task submission, and that task's
+    result is set to the presubmit return value.
+    """
+
+    class PresubmitShortcircuitCallback(RayDaskCallback):
+        def _ray_presubmit(self, task, key, deps):
+            return 0
+
+        def _ray_postsubmit(self, task, key, deps, object_ref):
+            pytest.fail(
+                "_ray_postsubmit shouldn't be called when "
+                "_ray_presubmit returns a value"
+            )
+
+    with PresubmitShortcircuitCallback():
+        z = add(2, 3)
+        result = z.compute(scheduler=ray_dask_get)
+
+    assert result == 0
+
+
+def test_pretask_posttask_shared_state(ray_start_regular_shared):
+    """
+    Test that pretask return value is passed to corresponding posttask
+    callback.
+    """
+
+    class PretaskPosttaskCallback(RayDaskCallback):
+        def _ray_pretask(self, key, object_refs):
+            return key
+
+        def _ray_posttask(self, key, result, pre_state):
+            assert pre_state == key
+
+    with PretaskPosttaskCallback():
+        z = add(2, 3)
+        result = z.compute(scheduler=ray_dask_get)
+
+    assert result == 5
+
+
+def test_postsubmit(ray_start_regular_shared):
+    """
+    Test that postsubmit is called after each task.
+    """
+
+    class PostsubmitCallback(RayDaskCallback):
+        def __init__(self, postsubmit_actor):
+            self.postsubmit_actor = postsubmit_actor
+
+        def _ray_postsubmit(self, task, key, deps, object_ref):
+            self.postsubmit_actor.postsubmit.remote(task, key, deps, object_ref)
+
+    @ray.remote
+    class PostsubmitActor:
+        def __init__(self):
+            self.postsubmit_counter = 0
+
+        def postsubmit(self, task, key, deps, object_ref):
+            self.postsubmit_counter += 1
+
+        def get_postsubmit_counter(self):
+            return self.postsubmit_counter
+
+    postsubmit_actor = PostsubmitActor.remote()
+    with PostsubmitCallback(postsubmit_actor):
+        z = add(2, 3)
+        result = z.compute(scheduler=ray_dask_get)
+
+    assert ray.get(postsubmit_actor.get_postsubmit_counter.remote()) == 1
+    assert result == 5
+
+
+def test_postsubmit_all(ray_start_regular_shared):
+    """
+    Test that postsubmit_all is called once.
+    """
+
+    class PostsubmitAllCallback(RayDaskCallback):
+        def __init__(self, postsubmit_all_actor):
+            self.postsubmit_all_actor = postsubmit_all_actor
+
+        def _ray_postsubmit_all(self, object_refs, dsk):
+            self.postsubmit_all_actor.postsubmit_all.remote(object_refs, dsk)
+
+    @ray.remote
+    class PostsubmitAllActor:
+        def __init__(self):
+            self.postsubmit_all_called = False
+
+        def postsubmit_all(self, object_refs, dsk):
+            self.postsubmit_all_called = True
+
+        def get_postsubmit_all_called(self):
+            return self.postsubmit_all_called
+
+    postsubmit_all_actor = PostsubmitAllActor.remote()
+    with PostsubmitAllCallback(postsubmit_all_actor):
+        z = add(2, 3)
+        result = z.compute(scheduler=ray_dask_get)
+
+    assert ray.get(postsubmit_all_actor.get_postsubmit_all_called.remote())
+    assert result == 5
+
+
+def test_finish(ray_start_regular_shared):
+    """
+    Test that finish callback is called once.
+    """
+
+    class FinishCallback(RayDaskCallback):
+        def __init__(self, finish_actor):
+            self.finish_actor = finish_actor
+
+        def _ray_finish(self, result):
+            self.finish_actor.finish.remote(result)
+
+    @ray.remote
+    class FinishActor:
+        def __init__(self):
+            self.finish_called = False
+
+        def finish(self, result):
+            self.finish_called = True
+
+        def get_finish_called(self):
+            return self.finish_called
+
+    finish_actor = FinishActor.remote()
+    with FinishCallback(finish_actor):
+        z = add(2, 3)
+        result = z.compute(scheduler=ray_dask_get)
+
+    assert ray.get(finish_actor.get_finish_called.remote())
+    assert result == 5
+
+
+def test_multiple_callbacks(ray_start_regular_shared):
+    """
+    Test that multiple callbacks are supported.
+    """
+
+    class PostsubmitCallback(RayDaskCallback):
+        def __init__(self, postsubmit_actor):
+            self.postsubmit_actor = postsubmit_actor
+
+        def _ray_postsubmit(self, task, key, deps, object_ref):
+            self.postsubmit_actor.postsubmit.remote(task, key, deps, object_ref)
+
+    @ray.remote
+    class PostsubmitActor:
+        def __init__(self):
+            self.postsubmit_counter = 0
+
+        def postsubmit(self, task, key, deps, object_ref):
+            self.postsubmit_counter += 1
+
+        def get_postsubmit_counter(self):
+            return self.postsubmit_counter
+
+    postsubmit_actor = PostsubmitActor.remote()
+    cb1 = PostsubmitCallback(postsubmit_actor)
+    cb2 = PostsubmitCallback(postsubmit_actor)
+    cb3 = PostsubmitCallback(postsubmit_actor)
+    with cb1, cb2, cb3:
+        z = add(2, 3)
+        result = z.compute(scheduler=ray_dask_get)
+
+    assert ray.get(postsubmit_actor.get_postsubmit_counter.remote()) == 3
+    assert result == 5
+
+
+def test_pretask_posttask_shared_state_multi(ray_start_regular_shared):
+    """
+    Test that pretask return values are passed to the correct corresponding
+    posttask callbacks when multiple callbacks are given.
+    """
+
+    class PretaskPosttaskCallback(RayDaskCallback):
+        def __init__(self, suffix):
+            self.suffix = suffix
+
+        def _ray_pretask(self, key, object_refs):
+            return key + self.suffix
+
+        def _ray_posttask(self, key, result, pre_state):
+            assert pre_state == key + self.suffix
+
+    class PretaskOnlyCallback(RayDaskCallback):
+        def _ray_pretask(self, key, object_refs):
+            return "baz"
+
+    class PosttaskOnlyCallback(RayDaskCallback):
+        def _ray_posttask(self, key, result, pre_state):
+            assert pre_state is None
+
+    cb1 = PretaskPosttaskCallback("foo")
+    cb2 = PretaskOnlyCallback()
+    cb3 = PosttaskOnlyCallback()
+    cb4 = PretaskPosttaskCallback("bar")
+    with cb1, cb2, cb3, cb4:
+        z = add(2, 3)
+        result = z.compute(scheduler=ray_dask_get)
+
+    assert result == 5
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/util/dask/tests/test_dask_multi_node.py

@@ -0,0 +1,91 @@
+import sys
+
+import dask
+import dask.dataframe as dd
+import numpy as np
+import pandas as pd
+import pytest
+
+import ray
+from ray.tests.conftest import *  # noqa: F403, F401
+from ray.util.dask import enable_dask_on_ray
+
+
+@pytest.fixture
+def ray_enable_dask_on_ray():
+    with enable_dask_on_ray():
+        yield
+
+
+def test_ray_dask_resources(ray_start_cluster, ray_enable_dask_on_ray):
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=1)
+    cluster.add_node(num_cpus=1, resources={"other_pin": 1})
+    pinned_node = cluster.add_node(num_cpus=1, num_gpus=1, resources={"pin": 1})
+
+    ray.init(address=cluster.address)
+
+    def get_node_id():
+        return ray._private.worker.global_worker.node.unique_id
+
+    # Test annotations on collection.
+    with dask.annotate(ray_remote_args=dict(num_cpus=1, resources={"pin": 0.01})):
+        c = dask.delayed(get_node_id)()
+    result = c.compute(optimize_graph=False)
+
+    assert result == pinned_node.unique_id
+
+    # Test annotations on compute.
+    c = dask.delayed(get_node_id)()
+    with dask.annotate(ray_remote_args=dict(num_gpus=1, resources={"pin": 0.01})):
+        result = c.compute(optimize_graph=False)
+
+    assert result == pinned_node.unique_id
+
+    # Test compute global Ray remote args.
+    c = dask.delayed(get_node_id)
+    result = c().compute(ray_remote_args={"resources": {"pin": 0.01}})
+
+    assert result == pinned_node.unique_id
+
+    # Test annotations on collection override global resource.
+    with dask.annotate(ray_remote_args=dict(resources={"pin": 0.01})):
+        c = dask.delayed(get_node_id)()
+    result = c.compute(
+        ray_remote_args=dict(resources={"other_pin": 0.01}), optimize_graph=False
+    )
+
+    assert result == pinned_node.unique_id
+
+    # Test top-level resources raises an error.
+    with pytest.raises(ValueError):
+        with dask.annotate(resources={"pin": 0.01}):
+            c = dask.delayed(get_node_id)()
+        result = c.compute(optimize_graph=False)
+    with pytest.raises(ValueError):
+        c = dask.delayed(get_node_id)
+        result = c().compute(resources={"pin": 0.01})
+
+    def get_node_id(row):
+        return pd.Series(ray._private.worker.global_worker.node.unique_id)
+
+    # Test annotations on compute.
+    df = dd.from_pandas(
+        pd.DataFrame(np.random.randint(0, 2, size=(2, 2)), columns=["age", "grade"]),
+        npartitions=2,
+    )
+    c = df.apply(get_node_id, axis=1, meta={0: str})
+    with dask.annotate(ray_remote_args=dict(num_gpus=1, resources={"pin": 0.01})):
+        result = c.compute(optimize_graph=False)
+    assert result[0].iloc[0] == pinned_node.unique_id
+
+    # Test compute global Ray remote args.
+    c = df.apply(get_node_id, axis=1, meta={0: str})
+    result = c.compute(
+        ray_remote_args={"resources": {"pin": 0.01}}, optimize_graph=False
+    )
+    assert result[0].iloc[0] == pinned_node.unique_id
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/util/dask/tests/test_dask_optimization.py

@@ -0,0 +1,83 @@
+import sys
+from unittest import mock
+
+import dask
+import dask.dataframe as dd
+import numpy as np
+import pandas as pd
+import pytest
+from packaging.version import Version
+
+from ray.tests.conftest import *  # noqa
+from ray.util.dask import dataframe_optimize
+
+try:
+    import dask_expr  # noqa: F401
+
+    DASK_EXPR_INSTALLED = True
+except ImportError:
+    DASK_EXPR_INSTALLED = False
+    pass
+
+if Version(dask.__version__) < Version("2025.1") and not DASK_EXPR_INSTALLED:
+    from dask.dataframe.shuffle import SimpleShuffleLayer
+
+    from ray.util.dask.optimizations import (
+        MultipleReturnSimpleShuffleLayer,
+        rewrite_simple_shuffle_layer,
+    )
+
+pytestmark = pytest.mark.skipif(
+    Version(dask.__version__) >= Version("2025.1") or DASK_EXPR_INSTALLED,
+    reason="Skip dask tests for Dask 2025.1+",
+)
+
+
+def test_rewrite_simple_shuffle_layer(ray_start_regular_shared):
+    npartitions = 10
+    df = dd.from_pandas(
+        pd.DataFrame(
+            np.random.randint(0, 100, size=(100, 2)), columns=["age", "grade"]
+        ),
+        npartitions=npartitions,
+    )
+    # We set max_branch=npartitions in order to ensure that the task-based
+    # shuffle happens in a single stage, which is required in order for our
+    # optimization to work.
+    a = df.set_index(["age"], shuffle="tasks", max_branch=npartitions)
+
+    dsk = a.__dask_graph__()
+    keys = a.__dask_keys__()
+    assert any(type(v) is SimpleShuffleLayer for k, v in dsk.layers.items())
+    dsk = rewrite_simple_shuffle_layer(dsk, keys)
+    assert all(type(v) is not SimpleShuffleLayer for k, v in dsk.layers.items())
+    assert any(
+        type(v) is MultipleReturnSimpleShuffleLayer for k, v in dsk.layers.items()
+    )
+
+
+@mock.patch("ray.util.dask.optimizations.rewrite_simple_shuffle_layer")
+def test_dataframe_optimize(mock_rewrite, ray_start_regular_shared):
+    def side_effect(dsk, keys):
+        return rewrite_simple_shuffle_layer(dsk, keys)
+
+    mock_rewrite.side_effect = side_effect
+    with dask.config.set(dataframe_optimize=dataframe_optimize):
+        npartitions = 10
+        df = dd.from_pandas(
+            pd.DataFrame(
+                np.random.randint(0, 100, size=(100, 2)), columns=["age", "grade"]
+            ),
+            npartitions=npartitions,
+        )
+        # We set max_branch=npartitions in order to ensure that the task-based
+        # shuffle happens in a single stage, which is required in order for our
+        # optimization to work.
+        a = df.set_index(["age"], shuffle="tasks", max_branch=npartitions).compute()
+
+    assert mock_rewrite.call_count == 2
+    assert a.index.is_monotonic_increasing
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/util/dask/tests/test_dask_scheduler.py

@@ -0,0 +1,98 @@
+import sys
+
+import dask
+import dask.array as da
+import dask.dataframe as dd
+import numpy as np
+import pandas as pd
+import pytest
+
+import ray
+from ray.tests.conftest import *  # noqa: F403, F401
+from ray.util.client.common import ClientObjectRef
+from ray.util.dask import disable_dask_on_ray, enable_dask_on_ray, ray_dask_get
+from ray.util.dask.callbacks import ProgressBarCallback
+
+
+@pytest.fixture
+def ray_enable_dask_on_ray():
+    with enable_dask_on_ray():
+        yield
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
+def test_ray_dask_basic(ray_start_regular_shared):
+    @ray.remote
+    def stringify(x):
+        return "The answer is {}".format(x)
+
+    zero_id = ray.put(0)
+
+    def add(x, y):
+        # Can retrieve ray objects from inside Dask.
+        zero = ray.get(zero_id)
+        # Can call Ray methods from inside Dask.
+        return ray.get(stringify.remote(x + y + zero))
+
+    add = dask.delayed(add)
+
+    expected = "The answer is 6"
+    # Test with explicit scheduler argument.
+    assert add(2, 4).compute(scheduler=ray_dask_get) == expected
+
+    # Test with config setter.
+    enable_dask_on_ray()
+    assert add(2, 4).compute() == expected
+    disable_dask_on_ray()
+
+    # Test with config setter as context manager.
+    with enable_dask_on_ray():
+        assert add(2, 4).compute() == expected
+
+    # Test within Ray task.
+
+    @ray.remote
+    def call_add():
+        z = add(2, 4)
+        with ProgressBarCallback():
+            r = z.compute(scheduler=ray_dask_get)
+        return r
+
+    ans = ray.get(call_add.remote())
+    assert ans == "The answer is 6", ans
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
+def test_ray_dask_persist(ray_start_regular_shared):
+    arr = da.ones(5) + 2
+    result = arr.persist(scheduler=ray_dask_get)
+    assert isinstance(
+        next(iter(result.dask.values())), (ray.ObjectRef, ClientObjectRef)
+    )
+
+
+def test_sort_with_progress_bar(ray_start_regular_shared):
+    npartitions = 10
+    df = dd.from_pandas(
+        pd.DataFrame(
+            np.random.randint(0, 100, size=(100, 2)), columns=["age", "grade"]
+        ),
+        npartitions=npartitions,
+    )
+    # We set max_branch=npartitions in order to ensure that the task-based
+    # shuffle happens in a single stage, which is required in order for our
+    # optimization to work.
+    sorted_with_pb = None
+    sorted_without_pb = None
+    with ProgressBarCallback():
+        sorted_with_pb = df.set_index(
+            ["age"], shuffle_method="tasks", max_branch=npartitions
+        ).compute(scheduler=ray_dask_get, _ray_enable_progress_bar=True)
+    sorted_without_pb = df.set_index(
+        ["age"], shuffle_method="tasks", max_branch=npartitions
+    ).compute(scheduler=ray_dask_get)
+    assert sorted_with_pb.equals(sorted_without_pb)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))