from __future__ import annotations
from abc import abstractmethod
from collections.abc import Generator
from types import TracebackType
from typing import Any
from gain import logging
from gain.task_graph.graph import Task, TaskGraph
logger = logging.getLogger(__name__)
[docs]
class TaskGraphExecutor:
"""Class that executes a task graph."""
[docs]
@abstractmethod
def execute(
self, graph: TaskGraph,
) -> Generator[tuple[Task, Any], None, None]:
"""Start executing the graph.
Return a generator that yields the task in the graph
after they are executed.
This is not necessarily in DFS or BFS order.
This is not even the order in which these tasks are executed.
The only guarantee is that when a task is returned its execution
is already finished.
A generator, not merely an iterator, because abandoning a run is a
supported way to end it -- ``task_graph_run_with_results`` does it on
the first failing task unless ``--keep-going`` -- and closing the
generator is how the run is told to tear itself down. An
implementation that holds resources for the run's lifetime must
release them from that path too, or an abandoned run leaks them for
the life of the process (gain#480).
"""
[docs]
@abstractmethod
def get_completed_tasks(
self, graph: TaskGraph,
) -> Generator[tuple[Task, Any], None, None]:
"""Return an iterator that yields already completed tasks in the graph.
This is not necessarily in DFS or BFS order.
"""
def __enter__(self) -> TaskGraphExecutor:
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
exc_tb: TracebackType | None,
) -> bool:
# Teardown is best-effort: the task graph's results are produced and
# consumed inside the ``with`` body, so a failure while releasing
# executor resources (e.g. a dask worker shutdown timeout) must not
# crash a completed run, nor mask an exception raised in the body.
try:
self.close()
except Exception:
logger.warning("error while closing task graph executor",
exc_info=True)
return exc_type is None
[docs]
@abstractmethod
def close(self) -> None:
"""Clean-up any resources used by the executor."""