from __future__ import annotations
import contextlib
import multiprocessing as mp
import os
import pickle # ruff: ignore[suspicious-pickle-import]
import time
from abc import abstractmethod
from collections import defaultdict, deque
from collections.abc import Generator, Iterable
from copy import copy
from typing import Any
import fsspec
import networkx
import psutil
from gain import logging
from gain.task_graph.cache import (
CacheRecord,
CacheRecordType,
NoTaskCache,
TaskCache,
)
from gain.task_graph.executor import TaskGraphExecutor
from gain.task_graph.graph import Task, TaskDesc, TaskGraph
from gain.task_graph.logging import (
configure_task_logging,
ensure_log_dir,
safe_task_id,
)
logger = logging.getLogger(__name__)
NO_TASK_CACHE = NoTaskCache()
def _descendants_of_all(
di_graph: networkx.DiGraph, sources: Iterable[Task],
) -> set[Task]:
"""Return every node reachable from any source by one or more edges.
Over an acyclic graph -- which :meth:`TaskGraph.as_directed_graph`
guarantees, refusing a cycle outright -- this is the union of
:func:`networkx.descendants` over ``sources``, computed in a single
breadth-first walk. Each node is expanded at most twice, once as a
source and once when another source reaches it, where a walk per
source expands it once per source it is reachable from -- O(V x E)
over a graph whose tasks are all uncomputed, which is every graph on
a fresh run (gain#1236).
The sources seed the queue but are not themselves in the result: one
is returned only when another source reaches it. That is what makes
this the union of the per-source descendant sets rather than those
sets plus the sources.
"""
reached: set[Task] = set()
queue: deque[Task] = deque(sources)
while queue:
for successor in di_graph.successors(queue.popleft()):
if successor not in reached:
reached.add(successor)
queue.append(successor)
return reached
[docs]
class TaskGraphExecutorBase(TaskGraphExecutor):
"""The base of the task executors of GAIn.
The base runs each task only after the tasks that it depends on, and
it asks a task cache which tasks must run. It records the result or
the error of each task that runs in the task cache.
Parameters
----------
task_cache : TaskCache
The task cache that decides which tasks must run. The default is a
:class:`~gain.task_graph.cache.NoTaskCache`, which runs every task.
force : bool
When true, the task executor runs every task and does not read the
task cache. It still writes the outcome of each task to it.
**kwargs
The options of the task executor. ``task_log_dir`` names the
directory of the task logs. Without it, the task executor writes
them into ``.task-log`` in the current directory. ``verbose`` sets
the verbosity of the task logs. ``fork_tasks`` runs each task in a
new process, and ``task_status_dir`` names the directory of the
result files of those processes.
"""
# NO_TASK_CACHE records nothing, so one shared instance is safe.
def __init__( # pylint: disable=dangerous-default-value
self, task_cache: TaskCache = NO_TASK_CACHE,
*,
force: bool = False,
**kwargs: Any,
):
super().__init__()
self._task_cache = task_cache
self._executing = False
self._force = force
log_dir = ensure_log_dir(**kwargs)
self._params = copy(kwargs)
self._params["task_log_dir"] = log_dir
@staticmethod
def _exec_internal(
task: TaskDesc,
params: dict[str, Any],
) -> Any:
verbose = params.get("verbose")
if verbose is None: # Dont use .get default in case of a Box
verbose = 0
log_dir = params.get("task_log_dir", ".")
task_id = safe_task_id(task.task.task_id)
task_func = task.func
args = task.args
kwargs = task.kwargs
root_logger = logging.getLogger()
handler = configure_task_logging(log_dir, task_id, verbose)
root_logger.addHandler(handler)
task_logger = logging.getLogger("task_executor")
task_logger.info("task <%s> started", task_id)
start = time.time()
process = psutil.Process(os.getpid())
start_memory_mb = process.memory_info().rss / (1024 * 1024)
task_logger.info(
"worker process memory usage: %.2f MB", start_memory_mb)
try:
result = task_func(*args, **kwargs)
except Exception as exp:
# pylint: disable=broad-except
# Through the task's own handler, so the failure is in its
# task log; the exception is still the task's result.
task_logger.exception("task <%s> failed", task_id)
result = exp
elapsed = time.time() - start
task_logger.info("task <%s> finished in %0.2fsec", task_id, elapsed)
finish_memory_mb = process.memory_info().rss / (1024 * 1024)
task_logger.info(
"worker process memory usage: %.2f MB; change: %+0.2f MB",
finish_memory_mb, finish_memory_mb - start_memory_mb)
root_logger.removeHandler(handler)
handler.close()
return result
@staticmethod
def _exec_forked(
task: TaskDesc,
params: dict[str, Any],
) -> None:
task_id = safe_task_id(task.task.task_id)
result_fn = TaskGraphExecutorBase._result_fn(task_id, params)
result = TaskGraphExecutorBase._exec_internal(task, params)
try:
with fsspec.open(result_fn, "wb") as out:
pickle.dump(result, out) # pyright: ignore
except Exception: # pylint: disable=broad-except
logger.exception(
"cannot write result for task %s. Ignoring and continuing.",
result_fn,
)
@staticmethod
def _result_fn(task_id: str, params: dict[str, Any]) -> str:
status_dir = params.get("task_status_dir") or "."
return os.path.join(status_dir, f"{task_id}.result")
@staticmethod
def _exec(
task: TaskDesc,
params: dict[str, Any],
) -> Any:
fork_tasks = params.get("fork_tasks", False)
if not fork_tasks:
return TaskGraphExecutorBase._exec_internal(task, params)
mp.current_process()._config[ # type: ignore # ruff: ignore[private-member-access]
"daemon"] = False
p = mp.Process(
target=TaskGraphExecutorBase._exec_forked,
args=(task, params),
)
p.start()
p.join()
task_id = safe_task_id(task.task.task_id)
result_fn = TaskGraphExecutorBase._result_fn(task_id, params)
try:
with fsspec.open(result_fn, "rb") as infile:
result = pickle.load(infile) # pyright: ignore
except Exception: # pylint: disable=broad-except
logger.exception(
"cannot write result for task %s. Ignoring and continuing.",
result_fn,
)
result = None
return result
def get_completed_tasks(
self, graph: TaskGraph,
) -> Generator[tuple[Task, Any], None, None]:
"""Yield the tasks that the task cache reports as computed.
A task that depends on a task that must run is not yielded, even
when the task cache has a computed record for it. In force mode,
the method yields nothing. :meth:`execute` removes the yielded
tasks from the graph, so they do not run again.
:meta private:
"""
if self._force:
return
cached_tasks: dict[Task, CacheRecord] = {}
uncomputed_tasks: set[Task] = set()
# {file: every task declaring it as an output}, built once so a
# consumer with many missing inputs looks each one up instead of
# walking its ancestors per file (gain#1213). A list, because two
# tasks may declare the same file and both must be invalidated.
producers: dict[str, list[Task]] = defaultdict(list)
with graph as tasks:
di_graph = graph.as_directed_graph()
task_descs = {task: graph.get_task_desc(task) for task in tasks}
for task, task_desc in task_descs.items():
record = self._task_cache.get_record(task_desc)
if record.type != CacheRecordType.COMPUTED:
uncomputed_tasks.add(task)
cached_tasks[task] = record
for output_file in (
task_desc.output_files
+ task_desc.intermediate_output_files
):
producers[output_file].append(task)
intermediates_needing_recompute: set[Task] = set()
for task in uncomputed_tasks:
candidates = [
producer
for input_file in task_descs[task].input_files
if not os.path.exists(input_file)
for producer in producers.get(input_file, ())
if producer not in uncomputed_tasks
and producer not in intermediates_needing_recompute
]
if not candidates:
continue
# Only a producer the consumer depends on is invalidated.
ancestors = networkx.ancestors(di_graph, task)
for producer in candidates:
if producer in ancestors:
cached_tasks[producer] = \
cached_tasks[producer].invalidate()
intermediates_needing_recompute.add(producer)
uncomputed_tasks.update(intermediates_needing_recompute)
for descendant_task in _descendants_of_all(
di_graph, uncomputed_tasks,
):
cached_tasks[descendant_task] = \
cached_tasks[descendant_task].invalidate()
completed_tasks = {
task: record.result_or_error
for task, record in cached_tasks.items()
if record.type == CacheRecordType.COMPUTED
}
for task, result in completed_tasks.items():
yield task, result
[docs]
def execute(
self, graph: TaskGraph,
) -> Generator[tuple[Task, Any], None, None]:
"""Run the tasks of a task graph that must run.
The method first asks the task cache which tasks are computed, and
removes them from the graph. It does not yield them. Then it runs
the other tasks, records the outcome of each one in the task cache
and yields it. When the task cache reports each task as computed,
the method yields nothing.
A task executor runs one graph at a time. Exhaust or close the
generator of a run before you start the next run.
Parameters
----------
graph : TaskGraph
The task graph to run. The method consumes it.
Yields
------
tuple[Task, Any]
Each task that ran and its result, when the task completes.
The result of a failed task is its exception.
"""
assert not self._executing, \
"Cannot execute a new graph while an old one is still running."
self._executing = True
try:
completed_tasks = list(self.get_completed_tasks(graph))
graph.process_completed_tasks(completed_tasks)
if len(graph) == 0:
logger.warning(
"All tasks are already COMPUTED; nothing to compute")
return
# closing(), not a bare `for`: this is a generator, and a consumer
# that abandons it part way -- `task_graph_run_with_results`
# raises out of its `for` loop on the first error unless
# --keep-going -- closes it, throwing GeneratorExit in at the
# yield below. That unwinds this frame, but it does NOT close the
# sub-generator: the iterator a `for` loop holds is released when
# this generator object is deallocated, not when its close()
# returns, and the consumer is still holding a reference to it. So
# `_execute`'s own teardown would run only whenever this object
# happened to be collected -- and for the dask executor that
# teardown is what shuts the run state down and joins both worker
# threads, so until then they spin against the shared client for
# the life of the process (gain#480). Closing it here makes that
# teardown run before close() returns, on every path.
with contextlib.closing(self._execute(graph)) as task_results:
for task_node, result in task_results:
is_error = isinstance(result, BaseException)
self._task_cache.cache(
task_node,
is_error=is_error,
result=result,
)
yield task_node, result
finally:
# In the finally for the same reason: an abandoned run that left
# this flag set made the executor permanently unusable, since the
# assert above then fires on its next graph. Every executor built
# here today is built per CLI invocation, so that is latent -- but
# it turns "the run failed" into "the executor is dead" for any
# caller that reuses one, which is a strange thing to have to
# know. This method sets the flag, so this method clears it.
self._executing = False
@abstractmethod
def _execute(
self, graph: TaskGraph,
) -> Generator[tuple[Task, Any], None, None]:
"""Execute the given task graph.
Must be a generator: :meth:`execute` closes it to tear the run down
when its own consumer abandons it part way (gain#480).
Args:
graph: Task graph to execute.
Yields:
Tuples of (task, result) as tasks complete.
"""