chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
Reference in New Issue
Block a user