wehub-resource-sync
398 lines
12 KiB
Python
398 lines
12 KiB
Python
"""
|
|
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)
|