gain.task_graph package

Submodules

gain.task_graph.base_executor module

class gain.task_graph.base_executor.TaskGraphExecutorBase(task_cache: TaskCache = {}, *, force: bool = False, **kwargs: Any)

Bases: TaskGraphExecutor

Executor that walks the graph in order that satisfies dependancies.

execute(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).

get_completed_tasks(graph: TaskGraph) → Generator[tuple[Task, Any], None, None]

Return cached tasks and their results.

All tasks that depend on uncomputed tasks are invalidated and will not be returned, even if they have a cached result.

All the tasks that are returned will be preprocessed and removed by the graph internally, so that they are not executed again.

Will not do anything is the executor is in force mode.

gain.task_graph.cache module

class gain.task_graph.cache.CacheRecord(type: CacheRecordType, result_or_error: Any = None)

Bases: object

Encapsulate information about a task in the cache.

property error: Any

invalidate() → CacheRecord

Return a new instance that needs to be recomputed.

property result: Any

result_or_error: Any = None

type: CacheRecordType

class gain.task_graph.cache.CacheRecordType(*values)

Bases: Enum

COMPUTED = 1

ERROR = 2

NEEDS_COMPUTE = 0

class gain.task_graph.cache.FileTaskCache(cache_dir: str)

Bases: TaskCache

Use file modification timestamps to determine if a task needs to run.

cache(task: Task, *, is_error: bool, result: Any) → None

Cache the result or exception of a task.

get_record(task_desc: TaskDesc) → CacheRecord

Get the cache record for a task.

class gain.task_graph.cache.NoTaskCache

Bases: dict[Any, Any], TaskCache

Don’t check any conditions and just run any task.

cache(task: Task, *, is_error: bool, result: Any) → None

Cache the result or exception of a task.

get_record(task_desc: TaskDesc) → CacheRecord

For task in the graph load and yield the cache record.

class gain.task_graph.cache.TaskCache

Bases: object

Store the result of a task in a file and reuse it if possible.

abstractmethod cache(task: Task, *, is_error: bool, result: Any) → None

Cache the result or exception of a task.

static create(*, force: bool = False, task_progress_mode: bool = True, cache_dir: str | None = None) → TaskCache

Create the appropriate task cache.

abstractmethod get_record(task_desc: TaskDesc) → CacheRecord

For task in the graph load and yield the cache record.

gain.task_graph.cli_tools module

class gain.task_graph.cli_tools.TaskGraphCli

Bases: object

Takes care of creating a task graph executor and executing a graph.

static add_arguments(parser: ArgumentParser, *, task_progress_mode: bool = True, default_task_status_dir: str | None = './.task-progress', use_commands: bool = True) → None

Add arguments needed to execute a task graph.

static create_executor(task_cache: TaskCache | None = None, **kwargs: Any) → TaskGraphExecutor

Create a task graph executor according to the args specified.

static process_graph(task_graph: TaskGraph, *, task_progress_mode: bool = True, **kwargs: Any) → bool

Process task_graph in according with the arguments in args.

Return true if the graph get’s successfully processed.

gain.task_graph.cli_tools.task_graph_all_done(task_graph: TaskGraph, task_cache: TaskCache) → bool

Check if the task graph is fully executed.

When all tasks are already computed, the function returns True. If there are tasks, that need to run, the function returns False.

gain.task_graph.cli_tools.task_graph_run(task_graph: TaskGraph, executor: TaskGraphExecutor | None = None, *, keep_going: bool = False) → bool

Execute (runs) the task_graph with the given executor.

gain.task_graph.cli_tools.task_graph_run_with_results(task_graph: TaskGraph, executor: TaskGraphExecutor | None = None, *, keep_going: bool = False) → Generator[Any, None, None]

Run a task graph, yielding the results from each task.

gain.task_graph.cli_tools.task_graph_status(task_graph: TaskGraph, task_cache: TaskCache, verbose: int | None) → bool

Show the status of each task from the task graph.

gain.task_graph.dask_executor module

class gain.task_graph.dask_executor.DaskExecutor(dask_client: Client, task_cache: TaskCache = {}, **kwargs: Any)

Bases: TaskGraphExecutorBase

Dask-based task graph executor.

MAX_RUNNING_TASKS = 700

close() → None

Close the Dask executor.

gain.task_graph.dask_executor.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.

gain.task_graph.dask_run_state module

Single owner of the Dask run loop’s “is anything still outstanding?”.

class gain.task_graph.dask_run_state.GatherBatch(batch_id: int, entries: tuple[tuple[Future, Task], ...])

Bases: object

Finished futures the results worker is collecting from the cluster.

Held by the worker for the whole width of Client.gather(), the mirror image of SubmitBatch.

batch_id: int

entries: tuple[tuple[Future, Task], ...]

property futures: tuple[Future, ...]

Futures in this batch, in order.

property tasks: tuple[Task, ...]

Tasks in this batch, in the same order as futures.

class gain.task_graph.dask_run_state.RunState

Bases: object

All the state one Dask run needs, behind one lock.

A task the run loop takes out of the graph is in exactly one of six states until the run loop yields it:

queued -> in-flight submit -> running
       -> completed -> in-flight gather -> gathered

Every hand-off between two threads – run loop, submit worker, dask callback thread, results worker – is a transition here, and each transition enters the next state and leaves the previous one under a single lock. A task therefore cannot be invisible by being absent from every collection, which is what let a run declare itself finished while a submission was still in flight (gain#365).

The two in-flight states exist for exactly that reason: the workers must not hold a lock across Client.map() or Client.gather(), so a batch that has left one collection and not yet reached the next is represented explicitly, by the worker’s batch handle, rather than by its absence from both.

has_outstanding() is the single query all of this exists to answer, and the run loop’s termination decision is that one call.

The states are not all reachable from every transition, though, and the recovery paths are where that bites: a batch aborted mid-wiring can evict its futures from the collections it holds, but not from the results worker’s hands. So “exactly one result per task” is not left to the collections to imply – _deliver() is the only way into gathered and it enforces the invariant outright (gain#381).

One transition leaves the diagram deliberately: abandon_outstanding() drops tasks without yielding them. It runs only in the run loop’s teardown, once both workers have stopped and the run has been given up on, and it exists so the futures those tasks hold are released rather than left pinning their keys on a client that outlives the run (gain#480). Nothing evaluates has_outstanding() afterwards.

abandon_outstanding() → list[Future]

Take every future the run still owns, emptying what holds them.

The run loop’s teardown calls this once both workers have stopped. Whatever is still in running, completed or the in-flight gather state then belongs to a run that ended without collecting it – a consumer abandoned the generator – and nobody else will ever come for it. The caller releases them, because Future.release() is a dask call and must not run under this lock.

Releasing matters for more than tidiness: an unreleased future keeps its key, and the result that key holds, alive on a client that outlives the run, so an abandoned run’s memory is never reclaimed (gain#480). What it does NOT do is protect the next run – releasing is asynchronous, so it never could. Keys are named per run and per task instead, which is what makes a later run of the same graph independent of whatever this one left behind (gain#531).

Emptying the collections is also what makes the release safe against a dask callback thread that fires afterwards: task_finished() pops from running, finds nothing and returns, so a future cannot be handed to a results worker that has already stopped.

Deliberately not part of shutdown(), which runs before the workers stop and must leave completed work for the results worker to collect – see claim_for_gather().

claim_for_gather() → GatherBatch | None

Take the completed futures into the in-flight gather state.

Blocks until something has completed. Returns None once the run is shutting down and everything completed has been claimed, which is the worker’s cue to stop.

claim_for_submit() → SubmitBatch | None

Take the queued tasks into the in-flight submit state.

Blocks until there is something to submit. Returns None once the run is shutting down, which is the worker’s cue to stop.

enqueue(tasks: Sequence[TaskDesc]) → None

Hand tasks extracted from the graph to the submit worker.

gather_failed(batch: GatherBatch, error: BaseException) → None

Move a batch that could not be gathered out of in-flight gather.

The mirror of gathered() for the failure path: Client.gather() can raise (a lost comm, a dead worker), and errors="skip" suppresses task errors, not transport ones (gain#372). While the batch sits in the in-flight gather state such a failure would be counted as outstanding forever; the failure is delivered as the result of every task in the batch, so the run loop yields it as an error and then terminates, and the batch leaves the gather state in the same lock hold. The caller releases the batch’s futures, as it does after a normal gather – future.release() is a dask call and must not run under this lock.

gathered(batch: GatherBatch, results: Sequence[tuple[Task, Any]]) → None

Move a gathered batch from in-flight gather to results.

A task an aborted submit batch already delivered as an error is not delivered again here – see _deliver() – but the batch leaves the in-flight gather state either way.

has_outstanding() → bool

Answer whether any task is still on its way to being yielded.

The single query, under the single lock: true from the instant a task is enqueued until the instant its result is taken by the run loop, with no gap in between.

shutdown() → None

Tell both workers the run is over, and drop unstarted work.

A run can be shut down with the queue still full – a consumer that stops iterating results abandons the run loop’s generator part way. Those tasks never reached the cluster and nobody will collect them, so they are discarded here, under the same lock that sets the flag. Discarding them is what keeps has_outstanding() truthful: a task left on the queue that no worker will ever claim would be counted as outstanding for as long as this object lived.

Deliberately not what the gather side does – see claim_for_gather(). A completed future holds work the run has already paid for, so it is still handed over; a queued task has cost nothing yet.

submit_aborted(batch: SubmitBatch, futures: Sequence[Future], error: BaseException) → None

Recover a batch whose wiring-up failed after map() returned.

Client.map() handed the futures back, but moving them into running and attaching their completion callbacks raised part way – a client tearing down under Future.add_done_callback, which (unlike release()) does not swallow it (gain#372). Some of the batch’s futures may already sit in running with a callback attached, one of which may even have fired and moved to completed; others never got one. Deliver the whole batch as a per-task error and drop every one of its futures from running (and the batch from submitting, in case submitted raised before it left it), so none lingers there counted as outstanding forever and the run terminates with one result per task. A future a callback already took out of running is simply absent there – the pop shrugs – but it now sits in completed, so it is evicted from there too; otherwise the results worker would gather it and deliver its task a second time, on top of the batch error.

Eviction reaches only as far as the collections this transition can see. The results worker runs in parallel and may have carried a callback-completed future beyond all of them – into the in-flight gather state, into gathered, or out to the run loop, which cannot be taken back at all. Those are _deliver()’s to handle: the error is delivered only for the tasks nothing has delivered yet, so whichever of the two paths arrives first is the one result the task gets (gain#381). If every task in the batch already has a result, the wiring failure cost the run nothing and surfaces only in the log.

submit_failed(batch: SubmitBatch, error: BaseException) → None

Move a batch that could not be submitted out of in-flight submit.

The mirror of submitted() for the failure path: Client.map() can raise (a dead scheduler connection, a serialization error) while the batch sits in the in-flight submit state, where it would be counted as outstanding forever and spin the run loop without end (gain#372). The failure is delivered as the result of every task in the batch – exactly as a task that dies on the worker is delivered – so the run loop yields it as an error and then terminates, and the batch leaves the submit state in the same lock hold.

Delivers through _deliver() like every other path. Nothing can have delivered these tasks already – map() raised, so no future of theirs ever existed to complete – but “every result goes through the ledger” is a rule worth having no exception to: an exception is the kind of thing a later change quietly grows a duplicate behind.

submitted(batch: SubmitBatch, futures: Sequence[Future]) → None

Move a submitted batch from in-flight submit to running.

running knows every future before the batch leaves the in-flight state, so there is no instant at which the batch is in neither – and the caller may only register completion callbacks after this returns, so a future that is already done cannot be reported before its task mapping exists (gain#355).

take_results() → list[tuple[Task, Any]]

Take every gathered result, in completion order.

The results stop being outstanding here, so the caller must feed them back to the graph and yield them before it asks has_outstanding() again.

A task comes out of here at most once for the life of the run, no matter which path delivered it – see _deliver() (gain#381).

task_finished(future: Future) → None

Move a finished future from running to completed.

Called on a dask callback thread, once per future. Futures are handed over only once, but a callback thread is not trusted to guarantee that: a second report of the same future is ignored.

unfinished_count() → int

Count tasks the cluster still owes a result for.

Queued, in-flight submit and running – what the run loop throttles new submissions on. Tasks whose result is already computed but not yet gathered or yielded are not counted: they take no cluster slot.

wait_for_results(timeout: float = 0.05) → None

Block until a result is ready to yield, or timeout elapses.

The timeout bounds how long the run loop goes without re-checking the graph for newly ready tasks.

class gain.task_graph.dask_run_state.SubmitBatch(batch_id: int, tasks: tuple[TaskDesc, ...])

Bases: object

Tasks the submit worker is handing to the cluster.

Held by the worker for the whole width of Client.map(). It is in no collection during that call – being in flight IS its state.

batch_id: int

tasks: tuple[TaskDesc, ...]

gain.task_graph.demo_graphs_cli module

gain.task_graph.demo_graphs_cli.build_demo_graph(graph_type: str, graph_params: list[str] | None) → TaskGraph

Build a demo graph.

gain.task_graph.demo_graphs_cli.main(argv: list[str] | None = None) → None

Entry point for the demo script.

gain.task_graph.demo_graphs_cli.task_part(seconds: str) → str

gain.task_graph.demo_graphs_cli.task_part_b(seconds: str) → str

gain.task_graph.demo_graphs_cli.task_part_c(seconds: str, *_args: str) → str

gain.task_graph.demo_graphs_cli.task_summary(seconds: str) → None

gain.task_graph.demo_graphs_cli.task_summary_b(seconds: str, *args: str) → str

gain.task_graph.demo_graphs_cli.task_summary_c(seconds: str, *args: str) → str

gain.task_graph.demo_graphs_cli.timeout(seconds: str) → float

gain.task_graph.executor module

class gain.task_graph.executor.TaskGraphExecutor

Bases: object

Class that executes a task graph.

abstractmethod close() → None

Clean-up any resources used by the executor.

abstractmethod execute(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).

abstractmethod get_completed_tasks(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.

gain.task_graph.graph module

class gain.task_graph.graph.Task(task_id: str)

Bases: object

Represent one node in a TaskGraph together with its dependencies.

task_id: str

class gain.task_graph.graph.TaskDesc(task: Task, func: Callable[[...], Any], args: list[Any], kwargs: dict[str, Any], deps: list[Task], input_files: list[str], output_files: list[str], intermediate_output_files: list[str])

Bases: object

Represent an immutable full task description with all its properties.

args: list[Any]

deps: list[Task]

func: Callable[[...], Any]

input_files: list[str]

intermediate_output_files: list[str]

kwargs: dict[str, Any]

output_files: list[str]

task: Task

class gain.task_graph.graph.TaskGraph

Bases: object

An object representing a graph of tasks.

add_task(task_desc: TaskDesc) → Task

Add a task to the graph.

add_tasks(task_descs: Sequence[TaskDesc]) → list[Task]

Add multiple tasks to the graph.

as_directed_graph() → DiGraph

Return the task graph as a networkx directed graph.

create_task(task_id: str, func: Callable[[...], Any], *, args: Sequence[Any], kwargs: dict[str, Any] | None = None, deps: Sequence[Task] | None = None, input_files: Sequence[str] | None = None, output_files: Sequence[str] | None = None, intermediate_output_files: Sequence[str] | None = None) → Task

Create a new task and add it to the graph.

Parameters:

• name – Name of the task (used for debugging purposes)

• func – Function to execute

• args – Arguments to that function

• deps – List of TaskNodes on which the current task depends

• input_files – Files that were used to build the graph itself

• output_files – Final output files; if missing the task recomputes

• intermediate_output_files – Pipeline-consumed outputs; if missing falls through to the flag-file check instead of forcing recompute

Return Task:

The newly created task node ID in the graph

empty() → bool

Check if the graph is empty.

extract_tasks(selected_tasks: Sequence[Task]) → Sequence[TaskDesc]

Collects tasks from the task graph and and removes them.

get_task_deps(task: Task) → list[Task]

Get dependancies of a task suitable for dask executor.

get_task_desc(task: Task) → TaskDesc

Get full task description for a given task.

has_task(task: Task) → bool

Check if the graph contains a task.

input_files: list[str]

static make_task(task_id: str, func: Callable[[...], Any], *, args: Sequence[Any], kwargs: dict[str, Any] | None = None, deps: Sequence[Task] | None = None, input_files: Sequence[str] | None = None, output_files: Sequence[str] | None = None, intermediate_output_files: Sequence[str] | None = None) → TaskDesc

Build a task with the given id and function.

process_completed_tasks(task_result: Sequence[tuple[Task, Any]]) → None

Process a completed task.

Parameters:

task – Completed task

prune(tasks_to_keep: Sequence[Task | str]) → None

Prune to keep the specified tasks and their dependencies.

ready_tasks(limit: int = 0) → Sequence[Task]

Return tasks which have no dependencies.

property tasks: Sequence[Task]

Return all tasks in the graph.

topological_order() → Sequence[Task]

Return tasks in topological order.

gain.task_graph.graph.chain_tasks(*tasks: TaskDesc) → TaskDesc

Chain tasks together so that they execute sequentially.

gain.task_graph.graph.sync_tasks() → None

gain.task_graph.logging module

class gain.task_graph.logging.FsspecHandler(logfile: str)

Bases: StreamHandler[Any]

Class to create fsspec based logging handler.

close() → None

Close the stream.

Copied from logging.FileHandler.close().

gain.task_graph.logging.configure_task_logging(log_dir: str | None, task_id: str, verbosity: int) → Handler

Configure and return task logging hadnler.

gain.task_graph.logging.ensure_log_dir(**kwargs: Any) → str

Ensure logging directory exists.

gain.task_graph.logging.safe_task_id(task_id: str) → str

Sanitize a task id into a string usable as a file name.

Must be a pure function of task_id: the forked executor computes the .result file name in the child and recomputes it in the parent, so a non-deterministic result makes the parent read a path that does not exist and silently yield None (gain#573).

gain.task_graph.process_pool_executor module

class gain.task_graph.process_pool_executor.ProcessPoolTaskExecutor(**kwargs: Any)

Bases: TaskGraphExecutorBase

Execute tasks in parallel using Dask to do the heavy lifting.

close() → None

Clean-up any resources used by the executor.

gain.task_graph.sequential_executor module

class gain.task_graph.sequential_executor.SequentialExecutor(task_cache: TaskCache = {}, *, force: bool = False, **kwargs: Any)

Bases: TaskGraphExecutorBase

A Task Graph Executor that executes task in sequential order.

close() → None

Close the executor and release resources.

gain.task_graph.work_dir module

The work-dir convention a task-graph CLI tool follows.

Shared by the annotate tools and binning_tool (gain#1215, gain#1234): the paths the user typed are absolutized before the tool changes into its work directory with chdir, the work directory and the task status/log directories are defaulted from the output, and a work directory the tool created is removed after a clean run.

The parsed-argument keys read here are the task-graph ones gain.task_graph.cli_tools.TaskGraphCli.add_arguments() defines (task_status_dir, task_log_dir, dask_cluster_config_file, command) and the ones the convention itself names, which each tool defines on its own parser: output, work_dir, keep_work_dir and the annotate tools’ keep_parts. Any further path option a tool has (its GRR options) is the caller’s to name; see absolutize_path_args().

gain.task_graph.work_dir.absolutize_path_args(args: dict[str, Any], *, input_key: str, extra_keys: Iterable[str] = ()) → None

Absolutize the tool’s primary input and the common path arguments.

input_key names the tool’s primary input (input for the annotate tools, run_definition for binning_tool); extra_keys names the tool’s own further path options (its GRR options, typically). A key that is absent or empty is left alone.

gain.task_graph.work_dir.apply_work_dir_defaults(args: dict[str, Any]) → None

Default and create the work directory and the task directories.

The convention the annotate tools and binning_tool share: the work directory is the output with its compression suffix and its extension stripped plus _work, the task status and log directories are .task-status and .task-log inside it. work_dir_created records whether the tool created the directory (it did not pre-exist) – maybe_remove_work_dir() reads it – and every non-empty path set here is absolute, so the tool can chdir into the work directory afterwards.

gain.task_graph.work_dir.maybe_remove_work_dir(args: dict[str, Any], *, result: bool) → None

Remove the working directory after a clean run, if the tool made it.

The directory is removed only when every condition holds:

• the tool created it (it did not pre-exist; see work_dir_created),

• the command actually ran annotation (not list/status),

• the run succeeded (result is True – a --keep-going run that finished with task errors returns False and is preserved),

• neither --keep-parts nor --keep-work-dir was requested,

• the output file does not live inside the working directory.

Removal is best-effort: a failure to remove logs a warning and is not fatal, since the annotation has already succeeded.

Module contents