import threading
import time
from collections.abc import Generator, Sequence
from copy import copy
from typing import Any
from dask.distributed import Client, Future
from gain import logging
from gain.task_graph.base_executor import TaskGraphExecutorBase
from gain.task_graph.cache import NoTaskCache, TaskCache
from gain.task_graph.dask_run_state import RunState, SubmitBatch
from gain.task_graph.graph import Task, TaskGraph
from gain.task_graph.logging import (
ensure_log_dir,
safe_task_id,
)
NO_TASK_CACHE = NoTaskCache()
logger = logging.getLogger(__name__)
# Named so a leaked run is identifiable in a thread dump -- and so a test
# can assert that a run really did tear its workers down.
SUBMIT_WORKER_THREAD_NAME = "gain-dask-submit-worker"
RESULTS_WORKER_THREAD_NAME = "gain-dask-results-worker"
[docs]
def dask_keys(run_id: str, batch: SubmitBatch) -> list[str]:
"""Name a submit batch's dask keys, one per task.
The key must identify this submission and nothing else. Deriving it from
the task id alone -- as this did until gain#531 -- made every run of the
same graph against the same client submit the same keys, so a run
starting while an earlier one's keys were still known to the scheduler
was deduplicated against them and handed that run's results without
executing anything. Releasing the earlier run's futures cannot close
that window: ``Future.release()`` only queues ``client-releases-keys``,
and the key may still be ``processing`` on a worker.
So the run id is appended: two runs cannot name the same key, whether
the earlier one ended normally, was abandoned, or is still in flight --
including two executors sharing a client concurrently.
The task id alone does not identify a task within a run either.
``safe_task_id`` collapses nine punctuation characters to ``_`` while
``TaskGraph`` enforces uniqueness on the RAW id, so ``annotate chr1``
and ``annotate-chr1`` sanitize to the same string -- and ``Client.map``
builds its layer as a dict, so tasks sharing a key silently collapse
into a single computation and the run reports success having executed
one of them (gain#557). Hence the batch id and the task's position in
the batch: the run state hands out batch ids from a single counter for
the life of the run, so the pair is unique per submitted task by
construction.
The sanitized task id stays at the FRONT so a key remains traceable to
its task in the dashboard and in scheduler logs. This is dask's own
``f"{key}-{token}"`` shape, which is what it would generate itself:
``pure=False`` is already passed at the submit site, so no deduplication
is wanted here in the first place.
"""
keys = [
f"{safe_task_id(task.task.task_id)}-{run_id}-{batch.batch_id}-{index}"
for index, task in enumerate(batch.tasks)
]
# Cheap, and the failure it guards against is silent: duplicate keys are
# not rejected by ``Client.map``, they are merged. Assert rather than
# trust the construction above to stay unique through a later edit.
assert len(set(keys)) == len(keys), "duplicate dask keys in one batch"
return keys
[docs]
class DaskExecutor(TaskGraphExecutorBase):
"""Dask-based task graph executor."""
def __init__(
self, dask_client: Client,
task_cache: TaskCache = NO_TASK_CACHE, **kwargs: Any,
) -> None:
"""Initialize the Dask executor.
Args:
dask_client: Dask client to use for task execution.
"""
super().__init__(task_cache=task_cache, **kwargs)
self._executing = False
self._dask_client = dask_client
log_dir = ensure_log_dir(**kwargs)
self._params = copy(kwargs)
self._params["task_log_dir"] = log_dir
def _submit_worker_func(self, state: RunState) -> None:
"""Hand queued tasks to the cluster until the run shuts down."""
start = time.time()
submit_count = 0
while True:
batch = state.claim_for_submit()
if batch is None:
logger.debug("submit worker stopping")
return
tasks = list(batch.tasks)
# No lock is held across map(). The batch does not need one: it
# sits in the in-flight submit state for the whole width of the
# call, so termination counts it the entire time (gain#365).
#
# map() can raise -- a dead scheduler connection, a
# serialization error -- and this worker is a daemon joined only
# after the run loop, so an escaping exception would kill it with
# the batch stranded in the in-flight submit state:
# has_outstanding() would answer "yes" forever and the run loop
# would spin at its wait timeout without end (gain#372). The
# failure is delivered as the result of every task in the batch
# instead, so the run loop yields it as an error -- exactly as it
# already does for a task that dies on the worker -- and then
# terminates.
try:
futures = self._dask_client.map(
self._exec, tasks,
key=dask_keys(state.run_id, batch),
pure=False,
params=self._params,
)
except BaseException as ex:
# Deliberately broad: in this non-main worker thread a
# transport fault of any type -- including a non-Exception
# BaseException such as a dask CancelledError -- must end the
# run as a delivered task error rather than silently kill the
# worker and strand the batch. KeyboardInterrupt/GeneratorExit
# cannot reach here.
# pylint: disable=broad-except
logger.exception(
"submit worker failed to hand %s task(s) to the "
"cluster; delivering the failure as their result",
len(tasks))
state.submit_failed(batch, ex)
continue
# map() returned, but wiring the batch up -- moving it to running
# and registering completion callbacks -- can raise too, and the
# same stranding applies: submitted() has already emptied the
# in-flight submit state, so a future left in running with no
# callback is outstanding forever and the run loop spins without
# end (gain#372). Future.add_done_callback, unlike release(), does
# not swallow a client tearing down under it. The whole batch is
# delivered as a per-task error and its futures cleared from
# running, so the run yields the failure and terminates.
#
# Callbacks are still registered only after submitted() populates
# running, so a callback firing immediately can never reach the
# run loop before its task mapping exists (gain#355).
try:
state.submitted(batch, futures)
for future in futures:
future.add_done_callback(state.task_finished)
except BaseException as ex:
# Deliberately broad: in this non-main worker thread a fault of
# any type -- including a non-Exception BaseException such as a
# dask CancelledError -- must end the run as a delivered task
# error rather than silently kill the worker and strand the
# batch. KeyboardInterrupt/GeneratorExit cannot reach here.
# pylint: disable=broad-except
logger.exception(
"submit worker failed to wire up %s task(s) after "
"handing them to the cluster; delivering the failure as "
"their result", len(tasks))
state.submit_aborted(batch, futures, ex)
continue
submit_count += len(tasks)
elapsed = time.time() - start
logger.debug(
"submitted %s tasks in %.2f seconds; %.2f tasks/s",
submit_count, elapsed, submit_count / elapsed)
logger.debug(
"total unfinished tasks: %s", state.unfinished_count())
def _results_worker_func(self, state: RunState) -> None:
"""Gather finished futures until the run shuts down."""
processed_results = 0
while True:
batch = state.claim_for_gather()
if batch is None:
break
logger.debug(
"results worker processing %s completed tasks",
len(batch.entries))
# No lock is held across gather() either -- the mirror image of
# the submit worker. The batch is in the in-flight gather state
# for the whole round trip, so a run whose gather outlasts the
# loop's wait cannot be called finished under it (gain#367).
# A *list*, deliberately: ``errors="skip"`` only drops failed
# futures when the container it is handed is a list. Given a
# tuple, ``distributed`` packs the failures back in as ``None``
# and returns the same length -- the check below would always
# match and a crashed task would be delivered as a successful
# ``None`` result, with nothing raised anywhere.
#
# gather() -- and future.result() in the fallback below -- can
# raise a transport error that errors="skip" does not suppress
# (it only skips task errors). This worker is a daemon joined
# only after the run loop, so an escaping exception would kill it
# with the batch stranded in the in-flight gather state and its
# futures never released: the run loop would spin without end
# (gain#372). The failure is delivered as the result of every
# task in the batch instead -- and the batch's futures released
# -- so the run loop yields it as an error and then terminates.
try:
results = self._dask_client.gather(
list(batch.futures), errors="skip")
if len(results) == len(batch.tasks):
gathered = list(zip(batch.tasks, results, strict=True))
else:
logger.error(
"failed to gather results for all %s tasks; "
"looking for exceptions in futures...",
len(batch.tasks))
gathered = []
for future, task in zip(
batch.futures, batch.tasks, strict=True):
try:
result = future.result()
except Exception as ex: # ruff: ignore[blind-except]
# pylint: disable=broad-except
result = ex
gathered.append((task, result))
except BaseException as ex:
# Deliberately broad, as in the submit worker: this non-main
# thread's death would strand the batch, so a transport fault
# of any type -- including a non-Exception BaseException such
# as a dask CancelledError -- must end the run as a delivered
# task error. KeyboardInterrupt/GeneratorExit cannot reach here.
# pylint: disable=broad-except
logger.exception(
"results worker failed to gather %s task(s); delivering "
"the failure as their result", len(batch.tasks))
state.gather_failed(batch, ex)
# Release on the failure path too -- the batch has left the
# gather state, so nothing else will (gain#372). Guarded: this
# runs on the very dead-client condition that caused the
# failure, and a raising release() here would kill this worker
# and re-strand any later batch.
self._release_futures(batch.futures)
continue
state.gathered(batch, gathered)
for future in batch.futures:
future.release()
processed_results += len(gathered)
logger.info("results worker processed %s results", processed_results)
@staticmethod
def _release_futures(futures: Sequence[Future]) -> None:
"""Release each future, surviving a ``release()`` that raises.
Used on the two paths where releasing is itself recovery: the
gather-failure path, whose loop runs on the very dead-client
condition that caused the failure, and the run teardown, which must
finish even for a run abandoned because the cluster is gone.
``distributed``'s ``release()`` is effectively non-throwing, so this
is robustness -- but were one to raise, an unguarded loop would kill
the results worker and re-strand any batch that completes after this
one (the hang gain#372 exists to prevent), or leave the rest of an
abandoned run's keys pinned (gain#480). Each release stands alone.
"""
for future in futures:
try:
future.release()
except BaseException:
# Deliberately broad, and swallowed: this is cleanup on a
# thread whose death would re-strand the run. A fault of any
# type must be logged and stepped over, never let to kill the
# worker. KeyboardInterrupt/GeneratorExit cannot reach here.
# pylint: disable=broad-except
logger.exception(
"failed to release a future after a gather failure; "
"continuing with the rest")
MAX_RUNNING_TASKS = 700
def _execute(
self, graph: TaskGraph,
) -> Generator[tuple[Task, Any], None, None]:
state = RunState()
submit_worker = threading.Thread(
target=self._submit_worker_func, args=(state,),
name=SUBMIT_WORKER_THREAD_NAME, daemon=True)
submit_worker.start()
results_worker = threading.Thread(
target=self._results_worker_func, args=(state,),
name=RESULTS_WORKER_THREAD_NAME, daemon=True)
results_worker.start()
finished_tasks = 0
initial_task_count = len(graph)
# The teardown belongs in a finally: this is a generator, and a
# consumer that stops early -- `task_graph_run_with_results` raises
# out of its `for` loop on the first error unless --keep-going --
# throws GeneratorExit in at the yield below. Without the finally
# that skips the shutdown, both joins and the release, leaking two
# threads and leaving this run's keys pinned on the client.
#
# Someone has to CLOSE this generator for any of that to run, which
# is not automatic: neither closing the generator that iterates this
# one nor raising out of it does so. `TaskGraphExecutorBase.execute`
# and `task_graph_run_with_results` each close what they read for
# exactly that reason (gain#480).
try:
# The run is over when the graph has nothing left to hand out
# and the state has nothing outstanding. One query, one lock,
# one owner: a task is outstanding from the moment it leaves
# the graph until the moment its result is taken below, with no
# gap for it to be invisible in -- so there is nothing for a
# re-read to catch up with. Only this thread takes tasks out of
# the graph or puts results back, so the two reads cannot race
# each other.
while not graph.empty() or state.has_outstanding():
unfinished = state.unfinished_count()
if unfinished < self.MAX_RUNNING_TASKS:
limit = max(self.MAX_RUNNING_TASKS - unfinished, 1)
# Two steps, and the tasks are in neither the graph nor
# the state in between. Safe only because this thread
# is the sole evaluator of the loop condition above:
# nothing can observe the gap. Top the queue up from a
# second thread, or evaluate termination anywhere else,
# and this reopens gain#365 right here -- enqueue and
# extract would then have to happen under one lock.
state.enqueue(
graph.extract_tasks(graph.ready_tasks(limit=limit)))
# Block until a result is ready. The timeout only bounds
# how long we go without re-checking the graph for newly
# ready tasks; unlike the wait() poll it replaced, waking
# up costs nothing per pending future (gain#355).
state.wait_for_results()
for task, result in state.take_results():
graph.process_completed_tasks([(task, result)])
finished_tasks += 1
logger.info(
"finished %s/%s", finished_tasks, initial_task_count)
yield task, result
finally:
state.shutdown()
results_worker.join()
submit_worker.join()
# Both workers have stopped, so anything still held is work this
# run gave up on: the consumer abandoned the generator and the
# results will never be collected. Release it, or every one of
# those keys -- and the result it holds -- stays alive on a
# client that outlives the run, a leak that grows with every
# abandoned run (gain#480). Empty on every normal path, where
# the loop only exits once nothing is outstanding.
#
# No later run depends on this happening, though: keys are named
# per run and per task, so a run cannot be deduplicated against
# another's leftovers whether or not they were released
# (gain#531). Releasing asynchronously, as this does, could never
# have carried that guarantee anyway.
self._release_futures(state.abandon_outstanding())
[docs]
def close(self) -> None:
"""Close the Dask executor."""
logger.info("closing Dask executor")
# shutdown() tears down workers and scheduler gracefully. Retiring /
# closing workers first races the scheduler teardown and floods the
# log with heartbeat failures and "Connection ... closed" lines (#125).
self._dask_client.shutdown()
self._dask_client.close()
logger.info("Dask executor closed")