chore: import upstream snapshot with attribution

python/ray/util/dask/__init__.py

@@ -0,0 +1,75 @@
+import dask
+from packaging.version import Version
+
+# Version(dask.__version__) becomes "0" during doc builds.
+if Version(dask.__version__) != Version("0") and Version(dask.__version__) < Version(
+    "2024.11.0"
+):
+    # Dask on Ray doesn't work if Dask version is less than 2024.11.0.
+    raise ImportError(
+        "Dask on Ray requires Dask version 2024.11.0 or later. "
+        "Please upgrade your Dask installation."
+    )
+
+from .callbacks import (
+    ProgressBarCallback,
+    RayDaskCallback,
+    local_ray_callbacks,
+    unpack_ray_callbacks,
+)
+from .optimizations import dataframe_optimize
+from .scheduler import (
+    disable_dask_on_ray,
+    enable_dask_on_ray,
+    ray_dask_get,
+    ray_dask_get_sync,
+)
+
+dask_persist = dask.persist
+
+
+def ray_dask_persist(*args, **kwargs):
+    kwargs["ray_persist"] = True
+    return dask_persist(*args, **kwargs)
+
+
+ray_dask_persist.__doc__ = dask_persist.__doc__
+
+dask_persist_mixin = dask.base.DaskMethodsMixin.persist
+
+
+def ray_dask_persist_mixin(self, **kwargs):
+    kwargs["ray_persist"] = True
+    return dask_persist_mixin(self, **kwargs)
+
+
+ray_dask_persist_mixin.__doc__ = dask_persist_mixin.__doc__
+
+
+# We patch dask in order to inject a kwarg into its `dask.persist()` calls,
+# which the Dask-on-Ray scheduler needs.
+# FIXME(Clark): Monkey patching is bad and we should try to avoid this.
+def patch_dask(ray_dask_persist, ray_dask_persist_mixin):
+    dask.persist = ray_dask_persist
+    dask.base.DaskMethodsMixin.persist = ray_dask_persist_mixin
+
+
+patch_dask(ray_dask_persist, ray_dask_persist_mixin)
+
+__all__ = [
+    # Config
+    "enable_dask_on_ray",
+    "disable_dask_on_ray",
+    # Schedulers
+    "ray_dask_get",
+    "ray_dask_get_sync",
+    # Helpers
+    "ray_dask_persist",
+    # Callbacks
+    "RayDaskCallback",
+    "local_ray_callbacks",
+    "unpack_ray_callbacks",
+    # Optimizations
+    "dataframe_optimize",
+    "ProgressBarCallback",
+]