wehub-resource-sync
693 lines
25 KiB
Python
693 lines
25 KiB
Python
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
|