chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,10 @@
|
||||
load("@rules_python//python:defs.bzl", "py_library")
|
||||
|
||||
py_library(
|
||||
name = "dask_lib",
|
||||
srcs = glob(
|
||||
["**/*.py"],
|
||||
exclude = ["tests/*.py"],
|
||||
),
|
||||
visibility = ["__subpackages__"],
|
||||
)
|
||||
@@ -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",
|
||||
]
|
||||
@@ -0,0 +1,310 @@
|
||||
import contextlib
|
||||
from collections import defaultdict, namedtuple
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from dask.callbacks import Callback
|
||||
|
||||
import ray
|
||||
|
||||
# The names of the Ray-specific callbacks. These are the kwarg names that
|
||||
# RayDaskCallback will accept on construction, and is considered the
|
||||
# source-of-truth for what Ray-specific callbacks exist.
|
||||
CBS = (
|
||||
"ray_presubmit",
|
||||
"ray_postsubmit",
|
||||
"ray_pretask",
|
||||
"ray_posttask",
|
||||
"ray_postsubmit_all",
|
||||
"ray_finish",
|
||||
)
|
||||
# The Ray-specific callback method names for RayDaskCallback.
|
||||
CB_FIELDS = tuple("_" + field for field in CBS)
|
||||
# The Ray-specific callbacks that we do _not_ wish to drop from RayCallbacks
|
||||
# if not given on a RayDaskCallback instance (will be filled with None
|
||||
# instead).
|
||||
CBS_DONT_DROP = {"ray_pretask", "ray_posttask"}
|
||||
|
||||
# The Ray-specific callbacks for a single RayDaskCallback.
|
||||
RayCallback = namedtuple("RayCallback", " ".join(CBS))
|
||||
|
||||
# The Ray-specific callbacks for one or more RayDaskCallbacks.
|
||||
RayCallbacks = namedtuple("RayCallbacks", " ".join([field + "_cbs" for field in CBS]))
|
||||
|
||||
|
||||
class RayDaskCallback(Callback):
|
||||
"""
|
||||
Extends Dask's `Callback` class with Ray-specific hooks. When instantiating
|
||||
or subclassing this class, both the normal Dask hooks (e.g. pretask,
|
||||
posttask, etc.) and the Ray-specific hooks can be provided.
|
||||
|
||||
See `dask.callbacks.Callback` for usage.
|
||||
|
||||
Caveats: Any Dask-Ray scheduler must bring the Ray-specific callbacks into
|
||||
context using the `local_ray_callbacks` context manager, since the built-in
|
||||
`local_callbacks` context manager provided by Dask isn't aware of this
|
||||
class.
|
||||
"""
|
||||
|
||||
# Set of active Ray-specific callbacks.
|
||||
ray_active = set()
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
for cb in CBS:
|
||||
cb_func = kwargs.pop(cb, None)
|
||||
if cb_func is not None:
|
||||
setattr(self, "_" + cb, cb_func)
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@property
|
||||
def _ray_callback(self):
|
||||
return RayCallback(*[getattr(self, field, None) for field in CB_FIELDS])
|
||||
|
||||
def __enter__(self):
|
||||
self._ray_cm = add_ray_callbacks(self)
|
||||
self._ray_cm.__enter__()
|
||||
super().__enter__()
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
super().__exit__(*args)
|
||||
self._ray_cm.__exit__(*args)
|
||||
|
||||
def register(self):
|
||||
type(self).ray_active.add(self._ray_callback)
|
||||
super().register()
|
||||
|
||||
def unregister(self):
|
||||
type(self).ray_active.remove(self._ray_callback)
|
||||
super().unregister()
|
||||
|
||||
def _ray_presubmit(
|
||||
self, task: Any, key: Any, deps: Dict[Any, Any]
|
||||
) -> Optional[Any]:
|
||||
"""Run before submitting a Ray task.
|
||||
|
||||
If this callback returns a non-`None` value, Ray does _not_ create
|
||||
a task and uses this value as the would-be task's result value.
|
||||
|
||||
Args:
|
||||
task: A Dask task, where the first tuple item is
|
||||
the task function, and the remaining tuple items are
|
||||
the task arguments, which are either the actual argument values,
|
||||
or Dask keys into the deps dictionary whose
|
||||
corresponding values are the argument values.
|
||||
key: The Dask graph key for the given task.
|
||||
deps: The dependencies of this task.
|
||||
|
||||
Returns:
|
||||
Either None, in which case Ray submits a task, or
|
||||
a non-None value, in which case Ray task doesn't submit
|
||||
a task and uses this return value as the
|
||||
would-be task result value.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _ray_postsubmit(
|
||||
self,
|
||||
task: Any,
|
||||
key: Any,
|
||||
deps: Dict[Any, Any],
|
||||
object_ref: ray.ObjectRef,
|
||||
):
|
||||
"""Run after submitting a Ray task.
|
||||
|
||||
Args:
|
||||
task: A Dask task, where the first tuple item is
|
||||
the task function, and the remaining tuple items are
|
||||
the task arguments, which are either the actual argument values,
|
||||
or Dask keys into the deps dictionary whose
|
||||
corresponding values are the argument values.
|
||||
key: The Dask graph key for the given task.
|
||||
deps: The dependencies of this task.
|
||||
object_ref: The object reference for the
|
||||
return value of the Ray task.
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
def _ray_pretask(self, key: Any, object_refs: List[ray.ObjectRef]):
|
||||
"""Run before executing a Dask task within a Ray task.
|
||||
|
||||
This method executes after Ray submits the task within a Ray
|
||||
worker. The return value of this method is passed to the
|
||||
_ray_posttask callback, if provided.
|
||||
|
||||
Args:
|
||||
key: The Dask graph key for the Dask task.
|
||||
object_refs: The object references
|
||||
for the arguments of the Ray task.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _ray_posttask(self, key: Any, result: Any, pre_state: Any):
|
||||
"""Run after executing a Dask task within a Ray task.
|
||||
|
||||
This method executes within a Ray worker. This callback receives the
|
||||
return value of the _ray_pretask callback, if provided.
|
||||
|
||||
Args:
|
||||
key: The Dask graph key for the Dask task.
|
||||
result: The task result value.
|
||||
pre_state: The return value of the corresponding
|
||||
_ray_pretask callback, if said callback is defined.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _ray_postsubmit_all(self, object_refs: List[ray.ObjectRef], dsk: Any):
|
||||
"""Run after Ray submits all tasks.
|
||||
|
||||
Args:
|
||||
object_refs: The object references
|
||||
for the output (leaf) Ray tasks of the task graph.
|
||||
dsk: The Dask graph.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _ray_finish(self, result: Any):
|
||||
"""Run after Ray finishes executing all Ray tasks and returns the final
|
||||
result.
|
||||
|
||||
Args:
|
||||
result: The final result (output) of the Dask
|
||||
computation, before any repackaging is done by
|
||||
Dask collection-specific post-compute callbacks.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class add_ray_callbacks:
|
||||
def __init__(self, *callbacks):
|
||||
self.callbacks = [normalize_ray_callback(c) for c in callbacks]
|
||||
RayDaskCallback.ray_active.update(self.callbacks)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
for c in self.callbacks:
|
||||
RayDaskCallback.ray_active.discard(c)
|
||||
|
||||
|
||||
def normalize_ray_callback(cb):
|
||||
if isinstance(cb, RayDaskCallback):
|
||||
return cb._ray_callback
|
||||
elif isinstance(cb, RayCallback):
|
||||
return cb
|
||||
else:
|
||||
raise TypeError(
|
||||
"Callbacks must be either 'RayDaskCallback' or 'RayCallback' namedtuple"
|
||||
)
|
||||
|
||||
|
||||
def unpack_ray_callbacks(cbs):
|
||||
"""Take an iterable of callbacks, return a list of each callback."""
|
||||
if cbs:
|
||||
# Only drop callback methods that aren't in CBS_DONT_DROP.
|
||||
return RayCallbacks(
|
||||
*(
|
||||
[cb for cb in cbs_ if cb or CBS[idx] in CBS_DONT_DROP] or None
|
||||
for idx, cbs_ in enumerate(zip(*cbs))
|
||||
)
|
||||
)
|
||||
else:
|
||||
return RayCallbacks(*([()] * len(CBS)))
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def local_ray_callbacks(callbacks=None):
|
||||
"""
|
||||
Allows Dask-Ray callbacks to work with nested schedulers.
|
||||
|
||||
Callbacks will only be used by the first started scheduler they encounter.
|
||||
This means that only the outermost scheduler will use global callbacks.
|
||||
"""
|
||||
global_callbacks = callbacks is None
|
||||
if global_callbacks:
|
||||
callbacks, RayDaskCallback.ray_active = (RayDaskCallback.ray_active, set())
|
||||
try:
|
||||
yield callbacks or ()
|
||||
finally:
|
||||
if global_callbacks:
|
||||
RayDaskCallback.ray_active = callbacks
|
||||
|
||||
|
||||
class ProgressBarCallback(RayDaskCallback):
|
||||
def __init__(self):
|
||||
@ray.remote
|
||||
class ProgressBarActor:
|
||||
def __init__(self):
|
||||
self._init()
|
||||
|
||||
def submit(self, key, deps, now):
|
||||
for dep in deps.keys():
|
||||
self.deps[key].add(dep)
|
||||
self.submitted[key] = now
|
||||
self.submission_queue.append((key, now))
|
||||
|
||||
def task_scheduled(self, key, now):
|
||||
self.scheduled[key] = now
|
||||
|
||||
def finish(self, key, now):
|
||||
self.finished[key] = now
|
||||
|
||||
def result(self):
|
||||
return len(self.submitted), len(self.finished)
|
||||
|
||||
def report(self):
|
||||
result = defaultdict(dict)
|
||||
for key, finished in self.finished.items():
|
||||
submitted = self.submitted[key]
|
||||
scheduled = self.scheduled[key]
|
||||
# deps = self.deps[key]
|
||||
result[key]["execution_time"] = (
|
||||
finished - scheduled
|
||||
).total_seconds()
|
||||
# Calculate the scheduling time.
|
||||
# This is inaccurate.
|
||||
# We should subtract scheduled - (last dep completed).
|
||||
# But currently it is not easy because
|
||||
# of how getitem is implemented in dask on ray sort.
|
||||
result[key]["scheduling_time"] = (
|
||||
scheduled - submitted
|
||||
).total_seconds()
|
||||
result["submission_order"] = self.submission_queue
|
||||
return result
|
||||
|
||||
def ready(self):
|
||||
pass
|
||||
|
||||
def reset(self):
|
||||
self._init()
|
||||
|
||||
def _init(self):
|
||||
self.submission_queue = []
|
||||
self.submitted = defaultdict(None)
|
||||
self.scheduled = defaultdict(None)
|
||||
self.finished = defaultdict(None)
|
||||
self.deps = defaultdict(set)
|
||||
|
||||
try:
|
||||
self.pb = ray.get_actor("_dask_on_ray_pb")
|
||||
ray.get(self.pb.reset.remote())
|
||||
except ValueError:
|
||||
self.pb = ProgressBarActor.options(name="_dask_on_ray_pb").remote()
|
||||
ray.get(self.pb.ready.remote())
|
||||
|
||||
def _ray_postsubmit(self, task, key, deps, object_ref):
|
||||
# Indicate the dask task is submitted.
|
||||
self.pb.submit.remote(key, deps, datetime.now())
|
||||
|
||||
def _ray_pretask(self, key, object_refs):
|
||||
self.pb.task_scheduled.remote(key, datetime.now())
|
||||
|
||||
def _ray_posttask(self, key, result, pre_state):
|
||||
# Indicate the dask task is finished.
|
||||
self.pb.finish.remote(key, datetime.now())
|
||||
|
||||
def _ray_finish(self, result):
|
||||
print("All tasks are completed.")
|
||||
@@ -0,0 +1,88 @@
|
||||
import uuid
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Iterator
|
||||
from operator import getitem
|
||||
from typing import Any
|
||||
|
||||
from dask.core import get as get_sync, quote
|
||||
from dask.utils import apply
|
||||
|
||||
import ray
|
||||
|
||||
try:
|
||||
from dataclasses import fields as dataclass_fields, is_dataclass
|
||||
except ImportError:
|
||||
# Python < 3.7
|
||||
def is_dataclass(x):
|
||||
return False
|
||||
|
||||
def dataclass_fields(x):
|
||||
return []
|
||||
|
||||
|
||||
def unpack_object_refs(*args: Any):
|
||||
"""
|
||||
Extract Ray object refs from a set of potentially arbitrarily nested
|
||||
Python objects.
|
||||
|
||||
Intended use is to find all Ray object references in a set of (possibly
|
||||
nested) Python objects, do something to them (get(), wait(), etc.), then
|
||||
repackage them into equivalent Python objects.
|
||||
|
||||
Args:
|
||||
*args: One or more (potentially nested) Python objects that contain
|
||||
Ray object references.
|
||||
|
||||
Returns:
|
||||
A 2-tuple of a flat list of all contained Ray object references, and a
|
||||
function that, when given the corresponding flat list of concrete
|
||||
values, will return a set of Python objects equivalent to that which
|
||||
was given in *args, but with all Ray object references replaced with
|
||||
their corresponding concrete values.
|
||||
"""
|
||||
object_refs = []
|
||||
repack_dsk = {}
|
||||
|
||||
object_refs_token = uuid.uuid4().hex
|
||||
|
||||
def _unpack(expr):
|
||||
if isinstance(expr, ray.ObjectRef):
|
||||
token = expr.hex()
|
||||
repack_dsk[token] = (getitem, object_refs_token, len(object_refs))
|
||||
object_refs.append(expr)
|
||||
return token
|
||||
|
||||
token = uuid.uuid4().hex
|
||||
# Treat iterators like lists
|
||||
typ = list if isinstance(expr, Iterator) else type(expr)
|
||||
if typ in (list, tuple, set):
|
||||
repack_task = (typ, [_unpack(i) for i in expr])
|
||||
elif typ in (dict, OrderedDict):
|
||||
repack_task = (typ, [[_unpack(k), _unpack(v)] for k, v in expr.items()])
|
||||
elif is_dataclass(expr):
|
||||
repack_task = (
|
||||
apply,
|
||||
typ,
|
||||
(),
|
||||
(
|
||||
dict,
|
||||
[
|
||||
[f.name, _unpack(getattr(expr, f.name))]
|
||||
for f in dataclass_fields(expr)
|
||||
],
|
||||
),
|
||||
)
|
||||
else:
|
||||
return expr
|
||||
repack_dsk[token] = repack_task
|
||||
return token
|
||||
|
||||
out = uuid.uuid4().hex
|
||||
repack_dsk[out] = (tuple, [_unpack(i) for i in args])
|
||||
|
||||
def repack(results):
|
||||
dsk = repack_dsk.copy()
|
||||
dsk[object_refs_token] = quote(results)
|
||||
return get_sync(dsk, out)
|
||||
|
||||
return object_refs, repack
|
||||
@@ -0,0 +1,149 @@
|
||||
import warnings
|
||||
|
||||
import dask
|
||||
from dask import core
|
||||
from dask.dataframe.core import _concat
|
||||
from dask.highlevelgraph import HighLevelGraph
|
||||
|
||||
from .scheduler import MultipleReturnFunc, multiple_return_get
|
||||
|
||||
try:
|
||||
from dask.dataframe.optimize import optimize
|
||||
from dask.dataframe.shuffle import SimpleShuffleLayer, shuffle_group
|
||||
except ImportError:
|
||||
# SimpleShuffleLayer doesn't exist in this version of Dask.
|
||||
# This is the case for dask>=2025.1.0.
|
||||
SimpleShuffleLayer = None
|
||||
try:
|
||||
import dask_expr # noqa: F401
|
||||
|
||||
SimpleShuffleLayer = None
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
if SimpleShuffleLayer is not None:
|
||||
|
||||
class MultipleReturnSimpleShuffleLayer(SimpleShuffleLayer):
|
||||
@classmethod
|
||||
def clone(cls, layer: SimpleShuffleLayer):
|
||||
# TODO(Clark): Probably don't need this since SimpleShuffleLayer
|
||||
# implements __copy__() and the shallow clone should be enough?
|
||||
return cls(
|
||||
name=layer.name,
|
||||
column=layer.column,
|
||||
npartitions=layer.npartitions,
|
||||
npartitions_input=layer.npartitions_input,
|
||||
ignore_index=layer.ignore_index,
|
||||
name_input=layer.name_input,
|
||||
meta_input=layer.meta_input,
|
||||
parts_out=layer.parts_out,
|
||||
annotations=layer.annotations,
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f"MultipleReturnSimpleShuffleLayer<name='{self.name}', "
|
||||
f"npartitions={self.npartitions}>"
|
||||
)
|
||||
|
||||
def __reduce__(self):
|
||||
attrs = [
|
||||
"name",
|
||||
"column",
|
||||
"npartitions",
|
||||
"npartitions_input",
|
||||
"ignore_index",
|
||||
"name_input",
|
||||
"meta_input",
|
||||
"parts_out",
|
||||
"annotations",
|
||||
]
|
||||
return (
|
||||
MultipleReturnSimpleShuffleLayer,
|
||||
tuple(getattr(self, attr) for attr in attrs),
|
||||
)
|
||||
|
||||
def _cull(self, parts_out):
|
||||
return MultipleReturnSimpleShuffleLayer(
|
||||
self.name,
|
||||
self.column,
|
||||
self.npartitions,
|
||||
self.npartitions_input,
|
||||
self.ignore_index,
|
||||
self.name_input,
|
||||
self.meta_input,
|
||||
parts_out=parts_out,
|
||||
)
|
||||
|
||||
def _construct_graph(self):
|
||||
"""Construct graph for a simple shuffle operation."""
|
||||
|
||||
shuffle_group_name = "group-" + self.name
|
||||
shuffle_split_name = "split-" + self.name
|
||||
|
||||
dsk = {}
|
||||
n_parts_out = len(self.parts_out)
|
||||
for part_out in self.parts_out:
|
||||
# TODO(Clark): Find better pattern than in-scheduler concat.
|
||||
_concat_list = [
|
||||
(shuffle_split_name, part_out, part_in)
|
||||
for part_in in range(self.npartitions_input)
|
||||
]
|
||||
dsk[(self.name, part_out)] = (_concat, _concat_list, self.ignore_index)
|
||||
for _, _part_out, _part_in in _concat_list:
|
||||
dsk[(shuffle_split_name, _part_out, _part_in)] = (
|
||||
multiple_return_get,
|
||||
(shuffle_group_name, _part_in),
|
||||
_part_out,
|
||||
)
|
||||
if (shuffle_group_name, _part_in) not in dsk:
|
||||
dsk[(shuffle_group_name, _part_in)] = (
|
||||
MultipleReturnFunc(
|
||||
shuffle_group,
|
||||
n_parts_out,
|
||||
),
|
||||
(self.name_input, _part_in),
|
||||
self.column,
|
||||
0,
|
||||
self.npartitions,
|
||||
self.npartitions,
|
||||
self.ignore_index,
|
||||
self.npartitions,
|
||||
)
|
||||
|
||||
return dsk
|
||||
|
||||
def rewrite_simple_shuffle_layer(dsk, keys):
|
||||
if not isinstance(dsk, HighLevelGraph):
|
||||
dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
|
||||
else:
|
||||
dsk = dsk.copy()
|
||||
|
||||
layers = dsk.layers.copy()
|
||||
for key, layer in layers.items():
|
||||
if type(layer) is SimpleShuffleLayer:
|
||||
dsk.layers[key] = MultipleReturnSimpleShuffleLayer.clone(layer)
|
||||
return dsk
|
||||
|
||||
def dataframe_optimize(dsk, keys, **kwargs):
|
||||
if not isinstance(keys, (list, set)):
|
||||
keys = [keys]
|
||||
keys = list(core.flatten(keys))
|
||||
|
||||
if not isinstance(dsk, HighLevelGraph):
|
||||
dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
|
||||
|
||||
dsk = rewrite_simple_shuffle_layer(dsk, keys=keys)
|
||||
return optimize(dsk, keys, **kwargs)
|
||||
|
||||
else:
|
||||
|
||||
def dataframe_optimize(dsk, keys, **kwargs):
|
||||
warnings.warn(
|
||||
"Custom dataframe shuffle optimization only works on "
|
||||
"dask>=2024.11.0,<2025.1.0, you are on version "
|
||||
f"{dask.__version__}."
|
||||
"Doing no additional optimization aside from the default one."
|
||||
)
|
||||
return None
|
||||
@@ -0,0 +1,692 @@
|
||||
import atexit
|
||||
import threading
|
||||
import time
|
||||
import warnings
|
||||
from collections import OrderedDict, defaultdict
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from multiprocessing.pool import ThreadPool
|
||||
from pprint import pprint
|
||||
from typing import Any, Callable, List, Optional
|
||||
|
||||
import dask
|
||||
from dask.core import ishashable, istask
|
||||
|
||||
try:
|
||||
from dask._task_spec import Alias, DataNode, Task, TaskRef, convert_legacy_graph
|
||||
except ImportError:
|
||||
warnings.warn(
|
||||
"Dask on Ray is available only on dask>=2024.11.0, "
|
||||
f"you are on version {dask.__version__}."
|
||||
)
|
||||
from dask.system import CPU_COUNT
|
||||
from dask.threaded import _thread_get_id, pack_exception
|
||||
|
||||
import ray
|
||||
from ray.util.dask.callbacks import local_ray_callbacks, unpack_ray_callbacks
|
||||
from ray.util.dask.common import unpack_object_refs
|
||||
from ray.util.dask.scheduler_utils import apply_sync, get_async
|
||||
|
||||
main_thread = threading.current_thread()
|
||||
default_pool = None
|
||||
pools = defaultdict(dict)
|
||||
pools_lock = threading.Lock()
|
||||
|
||||
TOP_LEVEL_RESOURCES_ERR_MSG = (
|
||||
'Use ray_remote_args={"resources": {...}} instead of resources={...} to specify '
|
||||
"required Ray task resources; see "
|
||||
"https://docs.ray.io/en/master/ray-core/package-ref.html#ray-remote."
|
||||
)
|
||||
|
||||
|
||||
def enable_dask_on_ray(
|
||||
shuffle: Optional[str] = "tasks",
|
||||
use_shuffle_optimization: Optional[bool] = True,
|
||||
) -> dask.config.set:
|
||||
"""
|
||||
Enable Dask-on-Ray scheduler. This helper sets the Dask-on-Ray scheduler
|
||||
as the default Dask scheduler in the Dask config. By default, it will also
|
||||
cause the task-based shuffle to be used for any Dask shuffle operations
|
||||
(required for multi-node Ray clusters, not sharing a filesystem), and will
|
||||
enable a Ray-specific shuffle optimization.
|
||||
|
||||
>>> enable_dask_on_ray()
|
||||
>>> ddf.compute() # <-- will use the Dask-on-Ray scheduler.
|
||||
|
||||
If used as a context manager, the Dask-on-Ray scheduler will only be used
|
||||
within the context's scope.
|
||||
|
||||
>>> with enable_dask_on_ray():
|
||||
... ddf.compute() # <-- will use the Dask-on-Ray scheduler.
|
||||
>>> ddf.compute() # <-- won't use the Dask-on-Ray scheduler.
|
||||
|
||||
Args:
|
||||
shuffle: The shuffle method used by Dask, either "tasks" or
|
||||
"disk". This should be "tasks" if using a multi-node Ray cluster.
|
||||
Defaults to "tasks".
|
||||
use_shuffle_optimization: Enable our custom Ray-specific shuffle
|
||||
optimization. Defaults to True.
|
||||
Returns:
|
||||
The Dask config object, which can be used as a context manager to limit
|
||||
the scope of the Dask-on-Ray scheduler to the corresponding context.
|
||||
"""
|
||||
if use_shuffle_optimization:
|
||||
from ray.util.dask.optimizations import dataframe_optimize
|
||||
else:
|
||||
dataframe_optimize = None
|
||||
# Manually set the global Dask scheduler config.
|
||||
# We also force the task-based shuffle to be used since the disk-based
|
||||
# shuffle doesn't work for a multi-node Ray cluster that doesn't share
|
||||
# the filesystem.
|
||||
return dask.config.set(
|
||||
scheduler=ray_dask_get, shuffle=shuffle, dataframe_optimize=dataframe_optimize
|
||||
)
|
||||
|
||||
|
||||
def disable_dask_on_ray():
|
||||
"""
|
||||
Unsets the scheduler, shuffle method, and DataFrame optimizer.
|
||||
"""
|
||||
return dask.config.set(scheduler=None, shuffle=None, dataframe_optimize=None)
|
||||
|
||||
|
||||
def ray_dask_get(dsk: Any, keys: List[str], **kwargs: Any):
|
||||
"""
|
||||
A Dask-Ray scheduler. This scheduler will send top-level (non-inlined) Dask
|
||||
tasks to a Ray cluster for execution. The scheduler will wait for the
|
||||
tasks to finish executing, fetch the results, and repackage them into the
|
||||
appropriate Dask collections. This particular scheduler uses a threadpool
|
||||
to submit Ray tasks.
|
||||
|
||||
This can be passed directly to `dask.compute()`, as the scheduler:
|
||||
|
||||
>>> dask.compute(obj, scheduler=ray_dask_get)
|
||||
|
||||
You can override the currently active global Dask-Ray callbacks (e.g.
|
||||
supplied via a context manager), the number of threads to use when
|
||||
submitting the Ray tasks, or the threadpool used to submit Ray tasks:
|
||||
|
||||
>>> dask.compute(
|
||||
obj,
|
||||
scheduler=ray_dask_get,
|
||||
ray_callbacks=some_ray_dask_callbacks,
|
||||
num_workers=8,
|
||||
pool=some_cool_pool,
|
||||
)
|
||||
|
||||
Args:
|
||||
dsk: Dask graph, represented as a task DAG dictionary.
|
||||
keys: List of Dask graph keys whose values we wish to
|
||||
compute and return.
|
||||
**kwargs: Optional scheduler overrides. Supported keys include
|
||||
``ray_callbacks`` (Dask-Ray callbacks), ``num_workers`` (number of
|
||||
worker threads to use when traversing the Dask graph), and
|
||||
``pool`` (a multiprocessing threadpool to use to submit Ray
|
||||
tasks).
|
||||
|
||||
Returns:
|
||||
Computed values corresponding to the provided keys.
|
||||
"""
|
||||
num_workers = kwargs.pop("num_workers", None)
|
||||
pool = kwargs.pop("pool", None)
|
||||
# We attempt to reuse any other thread pools that have been created within
|
||||
# this thread and with the given number of workers. We reuse a global
|
||||
# thread pool if num_workers is not given and we're in the main thread.
|
||||
global default_pool
|
||||
thread = threading.current_thread()
|
||||
if pool is None:
|
||||
with pools_lock:
|
||||
if num_workers is None and thread is main_thread:
|
||||
if default_pool is None:
|
||||
default_pool = ThreadPool(CPU_COUNT)
|
||||
atexit.register(default_pool.close)
|
||||
pool = default_pool
|
||||
elif thread in pools and num_workers in pools[thread]:
|
||||
pool = pools[thread][num_workers]
|
||||
else:
|
||||
pool = ThreadPool(num_workers)
|
||||
atexit.register(pool.close)
|
||||
pools[thread][num_workers] = pool
|
||||
|
||||
ray_callbacks = kwargs.pop("ray_callbacks", None)
|
||||
persist = kwargs.pop("ray_persist", False)
|
||||
enable_progress_bar = kwargs.pop("_ray_enable_progress_bar", None)
|
||||
|
||||
# Handle Ray remote args and resource annotations.
|
||||
if "resources" in kwargs:
|
||||
raise ValueError(TOP_LEVEL_RESOURCES_ERR_MSG)
|
||||
ray_remote_args = kwargs.pop("ray_remote_args", {})
|
||||
annotations = dask.get_annotations()
|
||||
if "resources" in annotations:
|
||||
raise ValueError(TOP_LEVEL_RESOURCES_ERR_MSG)
|
||||
|
||||
# Take out the dask graph if it is an Expr for dask>=2025.4.0.
|
||||
if not isinstance(dsk, Mapping):
|
||||
if hasattr(dsk, "_optimized_dsk"):
|
||||
# For Expr with this property
|
||||
dsk = dsk._optimized_dsk
|
||||
else:
|
||||
# For any other Expr
|
||||
dsk = dsk.__dask_graph__()
|
||||
scoped_ray_remote_args = _build_key_scoped_ray_remote_args(
|
||||
dsk, annotations, ray_remote_args
|
||||
)
|
||||
|
||||
with local_ray_callbacks(ray_callbacks) as ray_callbacks:
|
||||
# Unpack the Ray-specific callbacks.
|
||||
(
|
||||
ray_presubmit_cbs,
|
||||
ray_postsubmit_cbs,
|
||||
ray_pretask_cbs,
|
||||
ray_posttask_cbs,
|
||||
ray_postsubmit_all_cbs,
|
||||
ray_finish_cbs,
|
||||
) = unpack_ray_callbacks(ray_callbacks)
|
||||
# Make sure the graph is in the new format
|
||||
dsk = convert_legacy_graph(dsk)
|
||||
# NOTE: We hijack Dask's `get_async` function, injecting a different
|
||||
# task executor.
|
||||
object_refs = get_async(
|
||||
_apply_async_wrapper(
|
||||
pool.apply_async,
|
||||
_rayify_task_wrapper,
|
||||
ray_presubmit_cbs,
|
||||
ray_postsubmit_cbs,
|
||||
ray_pretask_cbs,
|
||||
ray_posttask_cbs,
|
||||
scoped_ray_remote_args,
|
||||
),
|
||||
len(pool._pool),
|
||||
dsk,
|
||||
keys,
|
||||
get_id=_thread_get_id,
|
||||
pack_exception=pack_exception,
|
||||
**kwargs,
|
||||
)
|
||||
if ray_postsubmit_all_cbs is not None:
|
||||
for cb in ray_postsubmit_all_cbs:
|
||||
cb(object_refs, dsk)
|
||||
# NOTE: We explicitly delete the Dask graph here so object references
|
||||
# are garbage-collected before this function returns, i.e. before all
|
||||
# Ray tasks are done. Otherwise, no intermediate objects will be
|
||||
# cleaned up until all Ray tasks are done.
|
||||
del dsk
|
||||
if persist:
|
||||
result = object_refs
|
||||
else:
|
||||
pb_actor = None
|
||||
if enable_progress_bar:
|
||||
pb_actor = ray.get_actor("_dask_on_ray_pb")
|
||||
result = ray_get_unpack(object_refs, progress_bar_actor=pb_actor)
|
||||
if ray_finish_cbs is not None:
|
||||
for cb in ray_finish_cbs:
|
||||
cb(result)
|
||||
|
||||
# cleanup pools associated with dead threads.
|
||||
with pools_lock:
|
||||
active_threads = set(threading.enumerate())
|
||||
if thread is not main_thread:
|
||||
for t in list(pools):
|
||||
if t not in active_threads:
|
||||
for p in pools.pop(t).values():
|
||||
p.close()
|
||||
return result
|
||||
|
||||
|
||||
def _apply_async_wrapper(
|
||||
apply_async: Callable,
|
||||
real_func: Callable,
|
||||
*extra_args: Any,
|
||||
**extra_kwargs: Any,
|
||||
):
|
||||
"""
|
||||
Wraps the given pool `apply_async` function, hotswapping `real_func` in as
|
||||
the function to be applied and adding `extra_args` and `extra_kwargs` to
|
||||
`real_func`'s call.
|
||||
|
||||
Args:
|
||||
apply_async: The pool function to be wrapped.
|
||||
real_func: The real function that we wish the pool apply
|
||||
function to execute.
|
||||
*extra_args: Extra positional arguments to pass to the `real_func`.
|
||||
**extra_kwargs: Extra keyword arguments to pass to the `real_func`.
|
||||
|
||||
Returns:
|
||||
A wrapper function that will ignore it's first `func` argument and
|
||||
pass `real_func` in its place. To be passed to `dask.local.get_async`.
|
||||
"""
|
||||
|
||||
def wrapper(func, args=(), kwds=None, callback=None): # noqa: M511
|
||||
if not kwds:
|
||||
kwds = {}
|
||||
return apply_async(
|
||||
real_func,
|
||||
args=args + extra_args,
|
||||
kwds=dict(kwds, **extra_kwargs),
|
||||
callback=callback,
|
||||
)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def _rayify_task_wrapper(
|
||||
key: Any,
|
||||
task_info: Any,
|
||||
dumps: Callable,
|
||||
loads: Callable,
|
||||
get_id: Callable,
|
||||
pack_exception: Callable,
|
||||
ray_presubmit_cbs: Optional[List[Callable]],
|
||||
ray_postsubmit_cbs: Optional[List[Callable]],
|
||||
ray_pretask_cbs: Optional[List[Callable]],
|
||||
ray_posttask_cbs: Optional[List[Callable]],
|
||||
scoped_ray_remote_args: dict,
|
||||
):
|
||||
"""
|
||||
The core Ray-Dask task execution wrapper, to be given to the thread pool's
|
||||
`apply_async` function. Exactly the same as `execute_task`, except that it
|
||||
calls `_rayify_task` on the task instead of `_execute_task`.
|
||||
|
||||
Args:
|
||||
key: The Dask graph key whose corresponding task we wish to
|
||||
execute.
|
||||
task_info: The task to execute and its dependencies.
|
||||
dumps: A result serializing function.
|
||||
loads: A task_info deserializing function.
|
||||
get_id: An ID generating function.
|
||||
pack_exception: An exception serializing function.
|
||||
ray_presubmit_cbs: Pre-task submission callbacks.
|
||||
ray_postsubmit_cbs: Post-task submission callbacks.
|
||||
ray_pretask_cbs: Pre-task execution callbacks.
|
||||
ray_posttask_cbs: Post-task execution callbacks.
|
||||
scoped_ray_remote_args: Ray task options for each key.
|
||||
|
||||
Returns:
|
||||
A 3-tuple of the task's key, a literal or a Ray object reference for a
|
||||
Ray task's result, and whether the Ray task submission failed.
|
||||
"""
|
||||
try:
|
||||
task, deps = loads(task_info)
|
||||
result = _rayify_task(
|
||||
task,
|
||||
key,
|
||||
deps,
|
||||
ray_presubmit_cbs,
|
||||
ray_postsubmit_cbs,
|
||||
ray_pretask_cbs,
|
||||
ray_posttask_cbs,
|
||||
scoped_ray_remote_args.get(key, {}),
|
||||
)
|
||||
id = get_id()
|
||||
result = dumps((result, id))
|
||||
failed = False
|
||||
except BaseException as e:
|
||||
result = pack_exception(e, dumps)
|
||||
failed = True
|
||||
return key, result, failed
|
||||
|
||||
|
||||
def _rayify_task(
|
||||
task: Any,
|
||||
key: Any,
|
||||
deps: dict,
|
||||
ray_presubmit_cbs: Optional[List[Callable]],
|
||||
ray_postsubmit_cbs: Optional[List[Callable]],
|
||||
ray_pretask_cbs: Optional[List[Callable]],
|
||||
ray_posttask_cbs: Optional[List[Callable]],
|
||||
ray_remote_args: dict,
|
||||
):
|
||||
"""
|
||||
Rayifies the given task, submitting it as a Ray task to the Ray cluster.
|
||||
|
||||
Args:
|
||||
task: A Dask graph value, being either a literal, dependency
|
||||
key, Dask task, or a list thereof.
|
||||
key: The Dask graph key for the given task.
|
||||
deps: The dependencies of this task.
|
||||
ray_presubmit_cbs: Pre-task submission callbacks.
|
||||
ray_postsubmit_cbs: Post-task submission callbacks.
|
||||
ray_pretask_cbs: Pre-task execution callbacks.
|
||||
ray_posttask_cbs: Post-task execution callbacks.
|
||||
ray_remote_args: Ray task options. See :func:`ray.remote` for details.
|
||||
|
||||
Returns:
|
||||
A literal, a Ray object reference representing a submitted task, or a
|
||||
list thereof.
|
||||
"""
|
||||
if isinstance(task, list):
|
||||
# Recursively rayify this list. This will still bottom out at the first
|
||||
# actual task encountered, inlining any tasks in that task's arguments.
|
||||
return [
|
||||
_rayify_task(
|
||||
t,
|
||||
key,
|
||||
deps,
|
||||
ray_presubmit_cbs,
|
||||
ray_postsubmit_cbs,
|
||||
ray_pretask_cbs,
|
||||
ray_posttask_cbs,
|
||||
ray_remote_args,
|
||||
)
|
||||
for t in task
|
||||
]
|
||||
elif istask(task):
|
||||
# Unpacks and repacks Ray object references and submits the task to the
|
||||
# Ray cluster for execution.
|
||||
if ray_presubmit_cbs is not None:
|
||||
alternate_returns = [cb(task, key, deps) for cb in ray_presubmit_cbs]
|
||||
for alternate_return in alternate_returns:
|
||||
# We don't submit a Ray task if a presubmit callback returns
|
||||
# a non-`None` value, instead we return said value.
|
||||
# NOTE: This returns the first non-None presubmit callback
|
||||
# return value.
|
||||
if alternate_return is not None:
|
||||
return alternate_return
|
||||
|
||||
if isinstance(task, Alias):
|
||||
target = task.target
|
||||
if isinstance(target, TaskRef):
|
||||
# for 2024.12.0
|
||||
return deps[target.key]
|
||||
else:
|
||||
# for 2024.12.1+
|
||||
return deps[target]
|
||||
elif isinstance(task, Task):
|
||||
func = task.func
|
||||
else:
|
||||
raise ValueError("Invalid task type: %s" % type(task))
|
||||
|
||||
# If the function's arguments contain nested object references, we must
|
||||
# unpack said object references into a flat set of arguments so that
|
||||
# Ray properly tracks the object dependencies between Ray tasks.
|
||||
arg_object_refs, repack = unpack_object_refs(deps)
|
||||
# Submit the task using a wrapper function.
|
||||
object_refs = dask_task_wrapper.options(
|
||||
name=f"dask:{key!s}",
|
||||
num_returns=(
|
||||
1 if not isinstance(func, MultipleReturnFunc) else func.num_returns
|
||||
),
|
||||
**ray_remote_args,
|
||||
).remote(
|
||||
task,
|
||||
repack,
|
||||
key,
|
||||
ray_pretask_cbs,
|
||||
ray_posttask_cbs,
|
||||
*arg_object_refs,
|
||||
)
|
||||
|
||||
if ray_postsubmit_cbs is not None:
|
||||
for cb in ray_postsubmit_cbs:
|
||||
cb(task, key, deps, object_refs)
|
||||
|
||||
return object_refs
|
||||
elif not ishashable(task):
|
||||
return task
|
||||
elif task in deps:
|
||||
return deps[task]
|
||||
else:
|
||||
return task
|
||||
|
||||
|
||||
@ray.remote
|
||||
def dask_task_wrapper(
|
||||
task: Any,
|
||||
repack: Callable,
|
||||
key: Any,
|
||||
ray_pretask_cbs: Optional[List[Callable]],
|
||||
ray_posttask_cbs: Optional[List[Callable]],
|
||||
*arg_object_refs: ray.ObjectRef,
|
||||
):
|
||||
"""
|
||||
A Ray remote function acting as a Dask task wrapper. This function will
|
||||
repackage the given `arg_object_refs` into its original `deps` using
|
||||
`repack`, and then pass it to the provided Dask Task object , `task`.
|
||||
|
||||
Args:
|
||||
task: The Dask Task class object to execute.
|
||||
repack: A function that repackages the provided args into
|
||||
the original (possibly nested) Python objects.
|
||||
key: The Dask key for this task.
|
||||
ray_pretask_cbs: Pre-task execution callbacks.
|
||||
ray_posttask_cbs: Post-task execution callback.
|
||||
*arg_object_refs: Ray object references representing the dependencies'
|
||||
results.
|
||||
|
||||
Returns:
|
||||
The output of the Dask task. In the context of Ray, a
|
||||
dask_task_wrapper.remote() invocation will return a Ray object
|
||||
reference representing the Ray task's result.
|
||||
"""
|
||||
if ray_pretask_cbs is not None:
|
||||
pre_states = [
|
||||
cb(key, arg_object_refs) if cb is not None else None
|
||||
for cb in ray_pretask_cbs
|
||||
]
|
||||
(repacked_deps,) = repack(arg_object_refs)
|
||||
# De-reference the potentially nested arguments recursively.
|
||||
def _dereference_args(x):
|
||||
if isinstance(x, Task):
|
||||
x.args = _dereference_args(x.args)
|
||||
return x
|
||||
elif isinstance(x, Mapping):
|
||||
return {k: _dereference_args(v) for k, v in x.items()}
|
||||
elif isinstance(x, tuple):
|
||||
return tuple(_dereference_args(x) for x in x)
|
||||
elif isinstance(x, ray.ObjectRef):
|
||||
return ray.get(x)
|
||||
elif isinstance(x, DataNode):
|
||||
if isinstance(x.value, ray.ObjectRef):
|
||||
value = ray.get(x.value)
|
||||
return DataNode(key=x.key, value=value)
|
||||
return x
|
||||
else:
|
||||
return x
|
||||
|
||||
task = _dereference_args(task)
|
||||
result = task(repacked_deps)
|
||||
|
||||
if ray_posttask_cbs is not None:
|
||||
for cb, pre_state in zip(ray_posttask_cbs, pre_states):
|
||||
if cb is not None:
|
||||
cb(key, result, pre_state)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def render_progress_bar(tracker, object_refs):
|
||||
from tqdm import tqdm
|
||||
|
||||
# At this time, every task should be submitted.
|
||||
total, finished = ray.get(tracker.result.remote())
|
||||
reported_finished_so_far = 0
|
||||
pb_bar = tqdm(total=total, position=0)
|
||||
pb_bar.set_description("")
|
||||
|
||||
ready_refs = []
|
||||
|
||||
while finished < total:
|
||||
submitted, finished = ray.get(tracker.result.remote())
|
||||
pb_bar.update(finished - reported_finished_so_far)
|
||||
reported_finished_so_far = finished
|
||||
ready_refs, _ = ray.wait(
|
||||
object_refs, timeout=0, num_returns=len(object_refs), fetch_local=False
|
||||
)
|
||||
if len(ready_refs) == len(object_refs):
|
||||
break
|
||||
time.sleep(0.1)
|
||||
pb_bar.close()
|
||||
submitted, finished = ray.get(tracker.result.remote())
|
||||
if submitted != finished:
|
||||
print("Completed. There was state inconsistency.")
|
||||
|
||||
pprint(ray.get(tracker.report.remote()))
|
||||
|
||||
|
||||
def ray_get_unpack(object_refs: Any, progress_bar_actor: Optional[Any] = None) -> Any:
|
||||
"""
|
||||
Unpacks object references, gets the object references, and repacks.
|
||||
Traverses arbitrary data structures.
|
||||
|
||||
Args:
|
||||
object_refs: A (potentially nested) Python object containing Ray object
|
||||
references.
|
||||
progress_bar_actor: An optional Ray actor used to render a progress bar
|
||||
while waiting on the object references to resolve.
|
||||
|
||||
Returns:
|
||||
The input Python object with all contained Ray object references
|
||||
resolved with their concrete values.
|
||||
"""
|
||||
|
||||
def get_result(object_refs):
|
||||
if progress_bar_actor:
|
||||
render_progress_bar(progress_bar_actor, object_refs)
|
||||
return ray.get(object_refs)
|
||||
|
||||
if isinstance(object_refs, tuple):
|
||||
object_refs = list(object_refs)
|
||||
|
||||
if isinstance(object_refs, list) and any(
|
||||
not isinstance(x, ray.ObjectRef) for x in object_refs
|
||||
):
|
||||
# We flatten the object references before calling ray.get(), since Dask
|
||||
# loves to nest collections in nested tuples and Ray expects a flat
|
||||
# list of object references. We repack the results after ray.get()
|
||||
# completes.
|
||||
object_refs, repack = unpack_object_refs(*object_refs)
|
||||
computed_result = get_result(object_refs)
|
||||
return repack(computed_result)
|
||||
else:
|
||||
return get_result(object_refs)
|
||||
|
||||
|
||||
def ray_dask_get_sync(dsk: Any, keys: List[str], **kwargs: Any):
|
||||
"""
|
||||
A synchronous Dask-Ray scheduler. This scheduler will send top-level
|
||||
(non-inlined) Dask tasks to a Ray cluster for execution. The scheduler will
|
||||
wait for the tasks to finish executing, fetch the results, and repackage
|
||||
them into the appropriate Dask collections. This particular scheduler
|
||||
submits Ray tasks synchronously, which can be useful for debugging.
|
||||
|
||||
This can be passed directly to `dask.compute()`, as the scheduler:
|
||||
|
||||
>>> dask.compute(obj, scheduler=ray_dask_get_sync)
|
||||
|
||||
You can override the currently active global Dask-Ray callbacks (e.g.
|
||||
supplied via a context manager):
|
||||
|
||||
>>> dask.compute(
|
||||
obj,
|
||||
scheduler=ray_dask_get_sync,
|
||||
ray_callbacks=some_ray_dask_callbacks,
|
||||
)
|
||||
|
||||
Args:
|
||||
dsk: Dask graph, represented as a task DAG dictionary.
|
||||
keys: List of Dask graph keys whose values we wish to
|
||||
compute and return.
|
||||
**kwargs: Optional scheduler overrides. Supported keys include
|
||||
``ray_callbacks`` (Dask-Ray callbacks).
|
||||
|
||||
Returns:
|
||||
Computed values corresponding to the provided keys.
|
||||
"""
|
||||
|
||||
ray_callbacks = kwargs.pop("ray_callbacks", None)
|
||||
persist = kwargs.pop("ray_persist", False)
|
||||
|
||||
with local_ray_callbacks(ray_callbacks) as ray_callbacks:
|
||||
# Unpack the Ray-specific callbacks.
|
||||
(
|
||||
ray_presubmit_cbs,
|
||||
ray_postsubmit_cbs,
|
||||
ray_pretask_cbs,
|
||||
ray_posttask_cbs,
|
||||
ray_postsubmit_all_cbs,
|
||||
ray_finish_cbs,
|
||||
) = unpack_ray_callbacks(ray_callbacks)
|
||||
# Make sure the graph is in the new format
|
||||
dsk = convert_legacy_graph(dsk)
|
||||
# NOTE: We hijack Dask's `get_async` function, injecting a different
|
||||
# task executor.
|
||||
object_refs = get_async(
|
||||
_apply_async_wrapper(
|
||||
apply_sync,
|
||||
_rayify_task_wrapper,
|
||||
ray_presubmit_cbs,
|
||||
ray_postsubmit_cbs,
|
||||
ray_pretask_cbs,
|
||||
ray_posttask_cbs,
|
||||
),
|
||||
1,
|
||||
dsk,
|
||||
keys,
|
||||
**kwargs,
|
||||
)
|
||||
if ray_postsubmit_all_cbs is not None:
|
||||
for cb in ray_postsubmit_all_cbs:
|
||||
cb(object_refs, dsk)
|
||||
# NOTE: We explicitly delete the Dask graph here so object references
|
||||
# are garbage-collected before this function returns, i.e. before all
|
||||
# Ray tasks are done. Otherwise, no intermediate objects will be
|
||||
# cleaned up until all Ray tasks are done.
|
||||
del dsk
|
||||
if persist:
|
||||
result = object_refs
|
||||
else:
|
||||
result = ray_get_unpack(object_refs)
|
||||
if ray_finish_cbs is not None:
|
||||
for cb in ray_finish_cbs:
|
||||
cb(result)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class MultipleReturnFunc:
|
||||
func: callable
|
||||
num_returns: int
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
returns = self.func(*args, **kwargs)
|
||||
if isinstance(returns, dict) or isinstance(returns, OrderedDict):
|
||||
returns = [returns[k] for k in range(len(returns))]
|
||||
return returns
|
||||
|
||||
|
||||
def multiple_return_get(multiple_returns, idx):
|
||||
return multiple_returns[idx]
|
||||
|
||||
|
||||
def _build_key_scoped_ray_remote_args(dsk, annotations, ray_remote_args):
|
||||
# Handle per-layer annotations.
|
||||
if not isinstance(dsk, dask.highlevelgraph.HighLevelGraph):
|
||||
dsk = dask.highlevelgraph.HighLevelGraph.from_collections(
|
||||
id(dsk), dsk, dependencies=()
|
||||
)
|
||||
# Build key-scoped annotations.
|
||||
scoped_annotations = {}
|
||||
layers = [(name, dsk.layers[name]) for name in dsk._toposort_layers()]
|
||||
for id_, layer in layers:
|
||||
layer_annotations = layer.annotations
|
||||
if layer_annotations is None:
|
||||
layer_annotations = annotations
|
||||
elif "resources" in layer_annotations:
|
||||
raise ValueError(TOP_LEVEL_RESOURCES_ERR_MSG)
|
||||
for key in layer.get_output_keys():
|
||||
layer_annotations_for_key = annotations.copy()
|
||||
# Layer annotations override global annotations.
|
||||
layer_annotations_for_key.update(layer_annotations)
|
||||
# Let same-key annotations earlier in the topological sort take precedence.
|
||||
layer_annotations_for_key.update(scoped_annotations.get(key, {}))
|
||||
scoped_annotations[key] = layer_annotations_for_key
|
||||
# Build key-scoped Ray remote args.
|
||||
scoped_ray_remote_args = {}
|
||||
for key, annotations in scoped_annotations.items():
|
||||
layer_ray_remote_args = ray_remote_args.copy()
|
||||
# Layer Ray remote args override global Ray remote args given in the compute
|
||||
# call.
|
||||
layer_ray_remote_args.update(annotations.get("ray_remote_args", {}))
|
||||
scoped_ray_remote_args[key] = layer_ray_remote_args
|
||||
return scoped_ray_remote_args
|
||||
@@ -0,0 +1,397 @@
|
||||
"""
|
||||
The following is adapted from Dask release 2021.03.1:
|
||||
https://github.com/dask/dask/blob/2021.03.1/dask/local.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from queue import Empty, Queue
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import dask
|
||||
from dask import config
|
||||
|
||||
try:
|
||||
from dask._task_spec import DataNode, DependenciesMapping
|
||||
except ImportError:
|
||||
warnings.warn(
|
||||
"Dask on Ray is available only on dask>=2024.11.0, "
|
||||
f"you are on version {dask.__version__}."
|
||||
)
|
||||
from dask.callbacks import local_callbacks, unpack_callbacks
|
||||
from dask.core import flatten, get_dependencies, reverse_dict
|
||||
from dask.order import order
|
||||
|
||||
if os.name == "nt":
|
||||
# Python 3 windows Queue.get doesn't handle interrupts properly. To
|
||||
# workaround this we poll at a sufficiently large interval that it
|
||||
# shouldn't affect performance, but small enough that users trying to kill
|
||||
# an application shouldn't care.
|
||||
def queue_get(q):
|
||||
while True:
|
||||
try:
|
||||
return q.get(block=True, timeout=0.1)
|
||||
except Empty:
|
||||
pass
|
||||
|
||||
else:
|
||||
|
||||
def queue_get(q):
|
||||
return q.get()
|
||||
|
||||
|
||||
def start_state_from_dask(
|
||||
dsk: Dict[Any, Any],
|
||||
cache: Optional[Dict[Any, Any]] = None,
|
||||
sortkey: Optional[Callable[[Any], Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Start state from a dask.
|
||||
|
||||
Args:
|
||||
dsk: A dask dictionary specifying a workflow.
|
||||
cache: Temporary storage of results.
|
||||
sortkey: Function to sort keys.
|
||||
|
||||
Returns:
|
||||
Initial scheduler state dict with keys ``dependencies``, ``dependents``,
|
||||
``waiting``, ``waiting_data``, ``cache``, ``ready``, ``running``,
|
||||
``finished``, and ``released``.
|
||||
|
||||
Examples:
|
||||
>>> dsk = {
|
||||
... 'x': 1,
|
||||
... 'y': 2,
|
||||
... 'z': (inc, 'x'),
|
||||
... 'w': (add, 'z', 'y')} # doctest: +SKIP
|
||||
>>> from pprint import pprint # doctest: +SKIP
|
||||
>>> pprint(start_state_from_dask(dsk)) # doctest: +SKIP
|
||||
{'cache': {'x': 1, 'y': 2},
|
||||
'dependencies': {'w': {'z', 'y'}, 'x': set(), 'y': set(), 'z': {'x'}},
|
||||
'dependents': {'w': set(), 'x': {'z'}, 'y': {'w'}, 'z': {'w'}},
|
||||
'finished': set(),
|
||||
'ready': ['z'],
|
||||
'released': set(),
|
||||
'running': set(),
|
||||
'waiting': {'w': {'z'}},
|
||||
'waiting_data': {'x': {'z'}, 'y': {'w'}, 'z': {'w'}}}
|
||||
"""
|
||||
if sortkey is None:
|
||||
sortkey = order(dsk).get
|
||||
if cache is None:
|
||||
cache = config.get("cache", None)
|
||||
if cache is None:
|
||||
cache = dict()
|
||||
|
||||
data_keys = set()
|
||||
for k, v in dsk.items():
|
||||
if isinstance(v, DataNode):
|
||||
cache[k] = v()
|
||||
data_keys.add(k)
|
||||
|
||||
dsk2 = dsk.copy()
|
||||
dsk2.update(cache)
|
||||
|
||||
dependencies = DependenciesMapping(dsk)
|
||||
waiting = {k: set(v) for k, v in dependencies.items() if k not in data_keys}
|
||||
|
||||
dependents = reverse_dict(dependencies)
|
||||
for a in cache:
|
||||
for b in dependents.get(a, ()):
|
||||
waiting[b].remove(a)
|
||||
waiting_data = {k: v.copy() for k, v in dependents.items() if v}
|
||||
|
||||
ready_set = {k for k, v in waiting.items() if not v}
|
||||
ready = sorted(ready_set, key=sortkey, reverse=True)
|
||||
waiting = {k: v for k, v in waiting.items() if v}
|
||||
|
||||
state = {
|
||||
"dependencies": dependencies,
|
||||
"dependents": dependents,
|
||||
"waiting": waiting,
|
||||
"waiting_data": waiting_data,
|
||||
"cache": cache,
|
||||
"ready": ready,
|
||||
"running": set(),
|
||||
"finished": set(),
|
||||
"released": set(),
|
||||
}
|
||||
|
||||
return state
|
||||
|
||||
|
||||
def execute_task(key, task_info, dumps, loads, get_id, pack_exception):
|
||||
"""Compute task and handle all administration.
|
||||
|
||||
See Also:
|
||||
_execute_task : actually execute task
|
||||
"""
|
||||
try:
|
||||
task, data = loads(task_info)
|
||||
result = task(data)
|
||||
id = get_id()
|
||||
result = dumps((result, id))
|
||||
failed = False
|
||||
except BaseException as e:
|
||||
result = pack_exception(e, dumps)
|
||||
failed = True
|
||||
return key, result, failed
|
||||
|
||||
|
||||
def release_data(key, state, delete=True):
|
||||
"""Remove data from temporary storage.
|
||||
|
||||
See Also:
|
||||
finish_task
|
||||
"""
|
||||
if key in state["waiting_data"]:
|
||||
assert not state["waiting_data"][key]
|
||||
del state["waiting_data"][key]
|
||||
|
||||
state["released"].add(key)
|
||||
|
||||
if delete:
|
||||
del state["cache"][key]
|
||||
|
||||
|
||||
DEBUG = False
|
||||
|
||||
|
||||
def finish_task(
|
||||
dsk, key, state, results, sortkey, delete=True, release_data=release_data
|
||||
):
|
||||
"""
|
||||
Update execution state after a task finishes
|
||||
Mutates. This should run atomically (with a lock).
|
||||
"""
|
||||
for dep in sorted(state["dependents"][key], key=sortkey, reverse=True):
|
||||
s = state["waiting"][dep]
|
||||
s.remove(key)
|
||||
if not s:
|
||||
del state["waiting"][dep]
|
||||
state["ready"].append(dep)
|
||||
|
||||
for dep in state["dependencies"][key]:
|
||||
if dep in state["waiting_data"]:
|
||||
s = state["waiting_data"][dep]
|
||||
s.remove(key)
|
||||
if not s and dep not in results:
|
||||
if DEBUG:
|
||||
from chest.core import nbytes
|
||||
|
||||
print(
|
||||
"Key: %s\tDep: %s\t NBytes: %.2f\t Release"
|
||||
% (key, dep, sum(map(nbytes, state["cache"].values()) / 1e6))
|
||||
)
|
||||
release_data(dep, state, delete=delete)
|
||||
elif delete and dep not in results:
|
||||
release_data(dep, state, delete=delete)
|
||||
|
||||
state["finished"].add(key)
|
||||
state["running"].remove(key)
|
||||
|
||||
return state
|
||||
|
||||
|
||||
def nested_get(ind: Union[int, List[Any]], coll: Any) -> Any:
|
||||
"""Get nested index from collection.
|
||||
|
||||
Args:
|
||||
ind: Index or nested list of indices.
|
||||
coll: Collection to index into.
|
||||
|
||||
Returns:
|
||||
Value at the given index, or a nested tuple of values if ``ind`` is a list.
|
||||
|
||||
Examples:
|
||||
>>> nested_get(1, 'abc')
|
||||
'b'
|
||||
>>> nested_get([1, 0], 'abc')
|
||||
('b', 'a')
|
||||
>>> nested_get([[1, 0], [0, 1]], 'abc')
|
||||
(('b', 'a'), ('a', 'b'))
|
||||
"""
|
||||
if isinstance(ind, list):
|
||||
return tuple(nested_get(i, coll) for i in ind)
|
||||
else:
|
||||
return coll[ind]
|
||||
|
||||
|
||||
def default_get_id():
|
||||
"""Default get_id"""
|
||||
return None
|
||||
|
||||
|
||||
def default_pack_exception(e, dumps):
|
||||
raise
|
||||
|
||||
|
||||
def reraise(exc, tb=None):
|
||||
if exc.__traceback__ is not tb:
|
||||
raise exc.with_traceback(tb)
|
||||
raise exc
|
||||
|
||||
|
||||
def identity(x):
|
||||
"""Identity function. Returns x.
|
||||
>>> identity(3)
|
||||
3
|
||||
"""
|
||||
return x
|
||||
|
||||
|
||||
def get_async(
|
||||
apply_async: Callable[..., Any],
|
||||
num_workers: int,
|
||||
dsk: Dict[Any, Any],
|
||||
result: Union[Any, List[Any]],
|
||||
cache: Optional[Dict[Any, Any]] = None,
|
||||
get_id: Callable[[], Any] = default_get_id,
|
||||
rerun_exceptions_locally: Optional[bool] = None,
|
||||
pack_exception: Callable[..., Any] = default_pack_exception,
|
||||
raise_exception: Callable[..., Any] = reraise,
|
||||
callbacks: Optional[Union[Tuple[Any, ...], List[Tuple[Any, ...]]]] = None,
|
||||
dumps: Callable[[Any], Any] = identity,
|
||||
loads: Callable[[Any], Any] = identity,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Asynchronous get function.
|
||||
|
||||
This is a general version of various asynchronous schedulers for dask. It
|
||||
takes a an apply_async function as found on Pool objects to form a more
|
||||
specific ``get`` method that walks through the dask array with parallel
|
||||
workers, avoiding repeat computation and minimizing memory use.
|
||||
|
||||
Args:
|
||||
apply_async: Asynchronous apply function as found on Pool or ThreadPool.
|
||||
num_workers: The number of active tasks we should have at any one time.
|
||||
dsk: A dask dictionary specifying a workflow.
|
||||
result: Keys corresponding to desired data (key or list of keys).
|
||||
cache: Temporary storage of results (dict-like, optional).
|
||||
get_id: Function to return the worker id, takes no arguments. Examples
|
||||
are `threading.current_thread` and `multiprocessing.current_process`.
|
||||
rerun_exceptions_locally: Whether to rerun failing tasks in local
|
||||
process to enable debugging (False by default).
|
||||
pack_exception: Function to take an exception and ``dumps`` method, and
|
||||
return a serialized tuple of ``(exception, traceback)`` to send
|
||||
back to the scheduler. Default is to just raise the exception.
|
||||
raise_exception: Function that takes an exception and a traceback, and
|
||||
raises an error.
|
||||
callbacks: Callbacks are passed in as tuples of length 5. Multiple sets
|
||||
of callbacks may be passed in as a list of tuples. For more
|
||||
information, see the dask.diagnostics documentation.
|
||||
dumps: Function to serialize task data and results to communicate
|
||||
between worker and parent. Defaults to identity.
|
||||
loads: Inverse function of `dumps`. Defaults to identity.
|
||||
**kwargs: Additional keyword arguments (unused).
|
||||
|
||||
Returns:
|
||||
The computed result(s), with the same shape as ``result``.
|
||||
|
||||
See Also:
|
||||
threaded.get
|
||||
"""
|
||||
queue = Queue()
|
||||
|
||||
if isinstance(result, list):
|
||||
result_flat = set(flatten(result))
|
||||
else:
|
||||
result_flat = {result}
|
||||
results = set(result_flat)
|
||||
|
||||
dsk = dict(dsk)
|
||||
with local_callbacks(callbacks) as callbacks:
|
||||
_, _, pretask_cbs, posttask_cbs, _ = unpack_callbacks(callbacks)
|
||||
started_cbs = []
|
||||
succeeded = False
|
||||
# if start_state_from_dask fails, we will have something
|
||||
# to pass to the final block.
|
||||
state = {}
|
||||
try:
|
||||
for cb in callbacks:
|
||||
if cb[0]:
|
||||
cb[0](dsk)
|
||||
started_cbs.append(cb)
|
||||
|
||||
keyorder = order(dsk)
|
||||
|
||||
state = start_state_from_dask(dsk, cache=cache, sortkey=keyorder.get)
|
||||
|
||||
for _, start_state, _, _, _ in callbacks:
|
||||
if start_state:
|
||||
start_state(dsk, state)
|
||||
|
||||
if rerun_exceptions_locally is None:
|
||||
rerun_exceptions_locally = config.get("rerun_exceptions_locally", False)
|
||||
|
||||
if state["waiting"] and not state["ready"]:
|
||||
raise ValueError("Found no accessible jobs in dask")
|
||||
|
||||
def fire_task():
|
||||
"""Fire off a task to the thread pool"""
|
||||
# Choose a good task to compute
|
||||
key = state["ready"].pop()
|
||||
state["running"].add(key)
|
||||
for f in pretask_cbs:
|
||||
f(key, dsk, state)
|
||||
|
||||
# Prep data to send
|
||||
data = {dep: state["cache"][dep] for dep in get_dependencies(dsk, key)}
|
||||
# Submit
|
||||
apply_async(
|
||||
execute_task,
|
||||
args=(
|
||||
key,
|
||||
dumps((dsk[key], data)),
|
||||
dumps,
|
||||
loads,
|
||||
get_id,
|
||||
pack_exception,
|
||||
),
|
||||
callback=queue.put,
|
||||
)
|
||||
|
||||
# Seed initial tasks into the thread pool
|
||||
while state["ready"] and len(state["running"]) < num_workers:
|
||||
fire_task()
|
||||
|
||||
# Main loop, wait on tasks to finish, insert new ones
|
||||
while state["waiting"] or state["ready"] or state["running"]:
|
||||
key, res_info, failed = queue_get(queue)
|
||||
if failed:
|
||||
exc, tb = loads(res_info)
|
||||
if rerun_exceptions_locally:
|
||||
data = {
|
||||
dep: state["cache"][dep]
|
||||
for dep in get_dependencies(dsk, key)
|
||||
}
|
||||
task = dsk[key]
|
||||
task(data) # Re-execute locally
|
||||
else:
|
||||
raise_exception(exc, tb)
|
||||
res, worker_id = loads(res_info)
|
||||
state["cache"][key] = res
|
||||
finish_task(dsk, key, state, results, keyorder.get)
|
||||
for f in posttask_cbs:
|
||||
f(key, res, dsk, state, worker_id)
|
||||
|
||||
while state["ready"] and len(state["running"]) < num_workers:
|
||||
fire_task()
|
||||
|
||||
succeeded = True
|
||||
|
||||
finally:
|
||||
for _, _, _, _, finish in started_cbs:
|
||||
if finish:
|
||||
finish(dsk, state, not succeeded)
|
||||
|
||||
return nested_get(result, state["cache"])
|
||||
|
||||
|
||||
def apply_sync(func, args=(), kwds=None, callback=None):
|
||||
"""A naive synchronous version of apply_async"""
|
||||
if kwds is None:
|
||||
kwds = {}
|
||||
|
||||
res = func(*args, **kwds)
|
||||
if callback is not None:
|
||||
callback(res)
|
||||
@@ -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"],
|
||||
)
|
||||
@@ -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__]))
|
||||
@@ -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__]))
|
||||
@@ -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__]))
|
||||
@@ -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__]))
|
||||
Reference in New Issue
Block a user