chore: import upstream snapshot with attribution

doc/source/ray-more-libs/doc_code/dask_on_ray_annotate_example.py

@@ -0,0 +1,52 @@
+import ray
+from ray.util.dask import enable_dask_on_ray
+import dask
+import dask.array as da
+
+# Start Ray.
+# Tip: If connecting to an existing cluster, use ray.init(address="auto").
+ray.init(
+    resources={
+        "custom_resource": 1,
+        "other_custom_resource": 1,
+        "another_custom_resource": 1,
+    }
+)
+
+# Use our Dask config helper to set the scheduler to ray_dask_get globally,
+# without having to specify it on each compute call.
+enable_dask_on_ray()
+
+# All Ray tasks that underly the Dask operations performed in an annotation
+# context will require the indicated resources: 2 CPUs and 0.01 of the custom
+# resource.
+with dask.annotate(
+    ray_remote_args=dict(num_cpus=2, resources={"custom_resource": 0.01})
+):
+    d_arr = da.ones(100)
+
+# Operations on the same collection can have different annotations.
+with dask.annotate(ray_remote_args=dict(resources={"other_custom_resource": 0.01})):
+    d_arr = 2 * d_arr
+
+# This happens outside of the annotation context, so no resource constraints
+# will be attached to the underlying Ray tasks for the sum() operation.
+sum_ = d_arr.sum()
+
+# Compute the result, passing in a default resource request that will be
+# applied to all operations that aren't already annotated with a resource
+# request. In this case, only the sum() operation will get this default
+# resource request.
+# We also give ray_remote_args, which will be given to every Ray task that
+# Dask-on-Ray submits; note that this can also be overridden for individual
+# Dask operations via the dask.annotate API.
+# NOTE: We disable graph optimization since it can break annotations,
+# see this issue: https://github.com/dask/dask/issues/7036.
+result = sum_.compute(
+    ray_remote_args=dict(max_retries=5, resources={"another_custom_resource": 0.01}),
+    optimize_graph=False,
+)
+print(result)
+# 200
+
+ray.shutdown()

doc/source/ray-more-libs/doc_code/dask_on_ray_callbacks.py

@@ -0,0 +1,103 @@
+# flake8: noqa
+import ray
+import dask.array as da
+
+z = da.ones(100)
+
+# fmt: off
+# __timer_callback_begin__
+from ray.util.dask import RayDaskCallback, ray_dask_get
+from timeit import default_timer as timer
+
+
+class MyTimerCallback(RayDaskCallback):
+    def _ray_pretask(self, key, object_refs):
+        # Executed at the start of the Ray task.
+        start_time = timer()
+        return start_time
+
+    def _ray_posttask(self, key, result, pre_state):
+        # Executed at the end of the Ray task.
+        execution_time = timer() - pre_state
+        print(f"Execution time for task {key}: {execution_time}s")
+
+
+with MyTimerCallback():
+    # Any .compute() calls within this context will get MyTimerCallback()
+    # as a Dask-Ray callback.
+    z.compute(scheduler=ray_dask_get)
+# __timer_callback_end__
+# fmt: on
+
+# fmt: off
+# __ray_dask_callback_direct_begin__
+def my_presubmit_cb(task, key, deps):
+    print(f"About to submit task {key}!")
+
+with RayDaskCallback(ray_presubmit=my_presubmit_cb):
+    z.compute(scheduler=ray_dask_get)
+# __ray_dask_callback_direct_end__
+# fmt: on
+
+# fmt: off
+# __ray_dask_callback_subclass_begin__
+class MyPresubmitCallback(RayDaskCallback):
+    def _ray_presubmit(self, task, key, deps):
+        print(f"About to submit task {key}!")
+
+with MyPresubmitCallback():
+    z.compute(scheduler=ray_dask_get)
+# __ray_dask_callback_subclass_end__
+# fmt: on
+
+# fmt: off
+# __multiple_callbacks_begin__
+# The hooks for both MyTimerCallback and MyPresubmitCallback will be
+# called.
+with MyTimerCallback(), MyPresubmitCallback():
+    z.compute(scheduler=ray_dask_get)
+# __multiple_callbacks_end__
+# fmt: on
+
+# fmt: off
+# __caching_actor_begin__
+@ray.remote
+class SimpleCacheActor:
+    def __init__(self):
+        self.cache = {}
+
+    def get(self, key):
+        # Raises KeyError if key isn't in cache.
+        return self.cache[key]
+
+    def put(self, key, value):
+        self.cache[key] = value
+
+
+class SimpleCacheCallback(RayDaskCallback):
+    def __init__(self, cache_actor_handle, put_threshold=10):
+        self.cache_actor = cache_actor_handle
+        self.put_threshold = put_threshold
+
+    def _ray_presubmit(self, task, key, deps):
+        try:
+            return ray.get(self.cache_actor.get.remote(str(key)))
+        except KeyError:
+            return None
+
+    def _ray_pretask(self, key, object_refs):
+        start_time = timer()
+        return start_time
+
+    def _ray_posttask(self, key, result, pre_state):
+        execution_time = timer() - pre_state
+        if execution_time > self.put_threshold:
+            self.cache_actor.put.remote(str(key), result)
+
+
+cache_actor = SimpleCacheActor.remote()
+cache_callback = SimpleCacheCallback(cache_actor, put_threshold=2)
+with cache_callback:
+    z.compute(scheduler=ray_dask_get)
+# __caching_actor_end__
+# fmt: on

doc/source/ray-more-libs/doc_code/dask_on_ray_persist_example.py

@@ -0,0 +1,39 @@
+import ray
+from ray.util.dask import enable_dask_on_ray
+import dask
+import dask.array as da
+
+# Start Ray.
+# Tip: If connecting to an existing cluster, use ray.init(address="auto").
+ray.init()
+
+# Use our Dask config helper to set the scheduler to ray_dask_get globally,
+# without having to specify it on each compute call.
+enable_dask_on_ray()
+
+d_arr = da.ones(100)
+
+# Print the internal Dask graph. Replace this with `print(dask.base.collections_to_dsk([d_arr]))` when dask>=2024.11.0,<2025.4.0.
+print(dask.base.collections_to_expr([d_arr]).dask)
+# {('ones_like-5902a58f37d3b639948dee893f5c4f4a', 0):
+# <Task ('ones_like-5902a58f37d3b639948dee893f5c4f4a', 0)
+# ones_like(...)>}
+
+# This submits all underlying Ray tasks to the cluster and returns
+# a Dask array with the Ray futures inlined.
+d_arr_p = d_arr.persist()
+
+# Notice that the Ray ObjectRef is inlined. The dask.ones() task has
+# been submitted to and is running on the Ray cluster.
+# Replace this in a similar way when dask>=2024.11.0,<2025.4.0.
+print(dask.base.collections_to_expr([d_arr_p]).dask)
+# {('ones_like-5902a58f37d3b639948dee893f5c4f4a', 0):
+# DataNode(ObjectRef(2c329aa28fcae64affffffffffffffffffffffff2c00000001000000))}
+
+# Future computations on this persisted Dask Array will be fast since we
+# already started computing d_arr_p in the background.
+d_arr_p.sum().compute()
+d_arr_p.min().compute()
+d_arr_p.max().compute()
+
+ray.shutdown()

doc/source/ray-more-libs/doc_code/dask_on_ray_scheduler_example.py

@@ -0,0 +1,34 @@
+import ray
+from ray.util.dask import ray_dask_get, enable_dask_on_ray, disable_dask_on_ray
+import dask.array as da
+import dask.dataframe as dd
+import numpy as np
+import pandas as pd
+
+# Start Ray.
+# Tip: If connecting to an existing cluster, use ray.init(address="auto").
+ray.init()
+
+d_arr = da.from_array(np.random.randint(0, 1000, size=(256, 256)))
+
+# The Dask scheduler submits the underlying task graph to Ray.
+d_arr.mean().compute(scheduler=ray_dask_get)
+
+# Use our Dask config helper to set the scheduler to ray_dask_get globally,
+# without having to specify it on each compute call.
+enable_dask_on_ray()
+
+df = dd.from_pandas(
+    pd.DataFrame(np.random.randint(0, 100, size=(1024, 2)), columns=["age", "grade"]),
+    npartitions=2,
+)
+df.groupby(["age"]).mean().compute()
+
+disable_dask_on_ray()
+
+# The Dask config helper can be used as a context manager, limiting the scope
+# of the Dask-on-Ray scheduler to the context.
+with enable_dask_on_ray():
+    d_arr.mean().compute()
+
+ray.shutdown()

doc/source/ray-more-libs/doc_code/dask_on_ray_shuffle_optimization.py

@@ -0,0 +1,30 @@
+import ray
+from ray.util.dask import dataframe_optimize, ray_dask_get
+import dask
+import dask.dataframe as dd
+import numpy as np
+import pandas as pd
+
+# Start Ray.
+# Tip: If connecting to an existing cluster, use ray.init(address="auto").
+ray.init()
+
+# Set the Dask DataFrame optimizer to
+# our custom optimization function, this time using the config setter as a
+# context manager.
+with dask.config.set(scheduler=ray_dask_get, dataframe_optimize=dataframe_optimize):
+    npartitions = 100
+    df = dd.from_pandas(
+        pd.DataFrame(
+            np.random.randint(0, 100, size=(10000, 2)), columns=["age", "grade"]
+        ),
+        npartitions=npartitions,
+    )
+    # We set max_branch to infinity in order to ensure that the task-based
+    # shuffle happens in a single stage, which is required in order for our
+    # optimization to work.
+    df.set_index(["age"], shuffle="tasks", max_branch=float("inf")).head(
+        10, npartitions=-1
+    )
+
+ray.shutdown()