gain.task_graph package

Submodules

gain.task_graph.base_executor module

class TaskGraphExecutorBase(task_cache: TaskCache = {}, *, force: bool = False, **kwargs: Any)[source]

Bases: 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 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.

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

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.

gain.task_graph.cache module

class CacheRecord(type: CacheRecordType, result_or_error: Any = None)[source]

Bases: object

The record of one task in a task cache.

Variables:
  • type (CacheRecordType) – The state of the task.

  • result_or_error (Any) – The result of the task, the exception of a failed task, or None when the task cache holds no outcome.

property error: Any

The exception of a task whose record type is ERROR.

property result: Any

The result of a task whose record type is COMPUTED.

result_or_error: Any = None
type: CacheRecordType
class CacheRecordType(*values)[source]

Bases: Enum

The state of a task in a task cache.

Variables:
COMPUTED = 1
ERROR = 2
NEEDS_COMPUTE = 0
class FileTaskCache(cache_dir: str)[source]

Bases: TaskCache

A task cache that keeps one record file for each task.

The record file of a task is <task id>.flag in the task status directory. It holds the record type and the result or the exception of the task. The modification times of the files of a task decide whether a computed record is still valid.

Parameters:

cache_dir (str) – The task status directory. The task cache creates it when it writes the first record.

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

Write the record file of a task.

A failure to write the file is logged, and the run continues.

Parameters:
  • task (Task) – The task that ran.

  • is_error (bool) – True when the task failed. The record type is then ERROR, otherwise COMPUTED.

  • result (Any) – The result of the task, or its exception when it failed. It must be picklable.

get_record(task_desc: TaskDesc) → CacheRecord[source]

Return the record of a task from its record file.

A computed record stays valid while the output files of the task exist and are not older than its input files. A missing input file makes the record invalid. A missing intermediate output file does not, because a later task can consume it. A task without output files compares its input files with its record file. A missing or unreadable record file means that the task must run.

Parameters:

task_desc (TaskDesc) – The description of the task, with its input and output files.

Returns:

The record of the task. Its type is NEEDS_COMPUTE when the record file is missing or unreadable, or when the record is computed but no longer valid. Otherwise it is the record that the file holds, COMPUTED or ERROR.

Return type:

CacheRecord

class NoTaskCache[source]

Bases: dict[Any, Any], TaskCache

A task cache that runs every task and records nothing.

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

Do nothing: the task cache records no outcome.

Parameters:
  • task (Task) – The task that ran.

  • is_error (bool) – True when the task failed.

  • result (Any) – The result of the task, or its exception when it failed.

get_record(task_desc: TaskDesc) → CacheRecord[source]

Return a record whose type is NEEDS_COMPUTE.

Parameters:

task_desc (TaskDesc) – The description of the task. The method does not read it.

Returns:

A record whose type is NEEDS_COMPUTE and that holds no result.

Return type:

CacheRecord

class TaskCache[source]

Bases: object

The interface of a task cache.

A task executor asks the task cache whether a task must run, and records the outcome of each task that runs in it.

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

Record the outcome of a task.

Parameters:
  • task (Task) – The task that ran.

  • is_error (bool) – True when the task failed.

  • result (Any) – The result of the task, or its exception when it failed.

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

Create the task cache for a set of options.

Parameters:
  • force (bool) – When true, the method returns a NoTaskCache.

  • task_progress_mode (bool) – The status mode. When false, the method returns a NoTaskCache.

  • cache_dir (str | None) – The task status directory of a FileTaskCache. None means the current directory.

Returns:

A NoTaskCache when force is true or the status mode is off. Otherwise a FileTaskCache over cache_dir.

Return type:

TaskCache

abstractmethod get_record(task_desc: TaskDesc) → CacheRecord[source]

Return the record of a task.

Parameters:

task_desc (TaskDesc) – The description of the task, from get_task_desc().

Returns:

The record of the task. The task must run unless the record type is COMPUTED. NEEDS_COMPUTE means that the task cache has no valid record of the task.

Return type:

CacheRecord

gain.task_graph.cli_tools module

class TaskGraphCli[source]

Bases: object

Add the task graph options to a command line and run a task graph.

A tool calls add_arguments() on its parser, and then calls process_graph() with its task graph and the parsed arguments as keywords. The page Running in parallel describes the options for the user.

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

Add the task graph options to a command line parser.

The options choose the task executor, the number of workers, the task status directory and the command. Pass the parsed arguments to process_graph() with **vars(args). Several keys do not have the name of their option. The options add these keys:

  • -j adds jobs.

  • --process-pool adds use_process_pool.

  • -N adds dask_cluster_name.

  • -c adds dask_cluster_config_file.

  • --task-log-dir adds task_log_dir.

  • The optional command adds command.

  • -t adds task_ids.

  • --keep-going adds keep_going.

  • --fork-tasks adds fork_tasks.

  • --force adds force.

  • -d adds task_status_dir.

The function adds no -v. The commands list and status read verbose, so add it with set_arguments().

Parameters:
  • parser (argparse.ArgumentParser) – The parser that gets the options.

  • task_progress_mode (bool) – When false, the parser gets no --force and no -d, and the parsed arguments have no force and no task_status_dir. Pass the same value to process_graph(): it builds a FileTaskCache by default whatever the parser holds, and without task_status_dir the task cache writes its record files into the current directory.

  • default_task_status_dir (str | None) – The default of -d, the task status directory. None means that a run without -d writes the record files of the task cache into the current directory.

  • use_commands (bool) – When false, the parser gets no optional command run, list or status, and process_graph() always runs the graph.

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

Create the task executor that the parsed arguments choose.

With jobs equal to 1, the function creates a SequentialExecutor. Otherwise, with use_process_pool true, it creates a ProcessPoolTaskExecutor with jobs processes. In every other case, it creates a DaskExecutor on a named cluster.

Parameters:
  • task_cache (TaskCache | None) – The task cache of the task executor. None means a NoTaskCache, which runs every task and records nothing.

  • **kwargs –

    The parsed arguments of add_arguments(). The function reads these keys, and it passes all of them to the task executor too:

    jobs

    The number of workers. A missing key means 1, so a call without keywords creates a SequentialExecutor. A parsed command line always holds jobs, with the value None when -j is absent, so a command line without -j creates a DaskExecutor, or a ProcessPoolTaskExecutor with --process-pool. None means one process for each CPU for the process pool, and the default of the named cluster for dask.

    use_process_pool

    When true, and jobs is not 1, the function creates a ProcessPoolTaskExecutor.

    dask_cluster_name

    The named cluster of the DaskExecutor. None means the default named cluster.

    dask_cluster_config_file

    A cluster config file that describes the cluster of the DaskExecutor. None means that the function uses dask_cluster_name.

Returns:

The new task executor. Use it as a context manager, so that it releases its workers.

Return type:

TaskGraphExecutor

Raises:

AssertionError – When jobs is 1, or use_process_pool is true, and dask_cluster_name or dask_cluster_config_file is set. Also when both dask_cluster_name and dask_cluster_config_file are set.

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

Run a task graph, or report its tasks, by the parsed arguments.

The function prunes the graph to the tasks of task_ids and builds the task cache with TaskCache.create. For the command run, it creates the task executor with create_executor() and runs the graph inside the context of the task executor. For the commands list and status, it prints a table of the task ids and the record type of each task to the standard output, and it runs nothing.

Parameters:
  • task_graph (TaskGraph) – The task graph. A run consumes it.

  • task_progress_mode (bool) – When true, the task cache is a FileTaskCache over task_status_dir. When false, it is a NoTaskCache. This keyword is separate from the keyword of add_arguments(), so pass the same value to both.

  • **kwargs –

    The parsed arguments of add_arguments(). The function reads these keys, and it passes all of them to create_executor():

    command

    run, list or status. None means run.

    task_ids

    The ids of the tasks to keep. The graph keeps the tasks that they depend on too. None means every task.

    task_status_dir

    The task status directory of the FileTaskCache. None, or a missing key, means the current directory.

    force

    When true, the task executor runs every task and writes every record again. The task cache does not change.

    keep_going

    When true, the run continues after a failed task.

    verbose

    When set and not 0, the commands list and status print the exception of each task with an ERROR record. add_arguments() does not add this key. The -v option of VerbosityConfiguration.set_arguments adds it.

Returns:

For run, true when every task succeeds. For list and status, always true.

Return type:

bool

Raises:
  • Exception – The error of the first failed task, for run without keep_going.

  • ValueError – When command is not run, list or status, or when the graph has no task with one of the task_ids.

task_graph_all_done(task_graph: TaskGraph, task_cache: TaskCache) → bool[source]

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.

task_graph_run(task_graph: TaskGraph, executor: TaskGraphExecutor | None = None, *, keep_going: bool = False) → bool[source]

Run a task graph with a task executor.

Parameters:
  • task_graph (TaskGraph) – The task graph to run. The run consumes it.

  • executor (TaskGraphExecutor | None) – The task executor that runs the graph. None means a new SequentialExecutor.

  • keep_going (bool) – When false, the first failed task raises its error and ends the run. When true, the run prints the error of each failed task to the standard error and continues.

Returns:

True when no task failed.

Return type:

bool

Raises:

Exception – The error of the first failed task, when keep_going is false.

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

Run a task graph and yield the result of each task.

The function yields the results without their tasks, in the order in which the tasks complete. A task that the task cache reports as computed does not run, and the function yields no result for it.

Parameters:
  • task_graph (TaskGraph) – The task graph to run. The run consumes it.

  • executor (TaskGraphExecutor | None) – The task executor that runs the graph. None means a new SequentialExecutor.

  • keep_going (bool) – When false, the first failed task raises its error and ends the run. When true, the function prints the error of each failed task to the standard error and yields the error as the result.

Yields:

Any – The result of each task, or the exception of a failed task.

Raises:

Exception – The error of the first failed task, when keep_going is false.

task_graph_status(task_graph: TaskGraph, task_cache: TaskCache, verbose: int | None) → bool[source]

Show the status of each task from the task graph.

gain.task_graph.dask_executor module

class DaskExecutor(dask_client: Client, task_cache: TaskCache = {}, **kwargs: Any)[source]

Bases: TaskGraphExecutorBase

A task executor that runs the tasks on a dask cluster.

The task executor takes ownership of the dask client: close() shuts down the client and its cluster. A task runs on a worker of the cluster only after the tasks that it depends on.

Parameters:
  • dask_client (Client) – The dask client of the cluster that runs the tasks.

  • task_cache (TaskCache) – The task cache that decides which tasks must run. The default is a NoTaskCache, which runs every task.

  • **kwargs – The other keyword arguments of TaskGraphExecutorBase: force and the options of the task executor.

Variables:

MAX_RUNNING_TASKS (int) – The task executor takes new tasks from the graph only while fewer than this number of its tasks are unfinished.

MAX_RUNNING_TASKS = 700
close() → None[source]

Shut down the dask client and its cluster.

dask_keys(run_id: str, batch: SubmitBatch) → list[str][source]

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.

held_futures(graph: TaskGraph, run_id: str) → list[Future][source]

List the dependency futures run run_id left in the graph.

Only futures of keys the run named (see dask_keys()): a future the caller passed as an ordinary argument is the caller’s to release.

gain.task_graph.dask_run_state module

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

class GatherBatch(batch_id: int, entries: tuple[tuple[Future, Task], ...])[source]

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 RunState[source]

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][source]

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. So does any delivered result’s future the run loop never took. 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[source]

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[source]

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[source]

Hand tasks extracted from the graph to the submit worker.

gather_failed(batch: GatherBatch, error: BaseException) → None[source]

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]]) → list[Future][source]

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

results is in the batch’s order. 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.

The future of each task delivered here with a result is kept for take_result_future(). The rest – failed tasks and dropped duplicates – are returned for the caller to release, outside the lock.

has_outstanding() → bool[source]

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[source]

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[source]

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[source]

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[source]

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_result_future(task: Task) → Future | None[source]

Take the future holding a delivered task’s result.

None for a task delivered as an error, which has no result to pass on.

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

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[source]

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[source]

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[source]

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 SubmitBatch(batch_id: int, tasks: tuple[TaskDesc, ...])[source]

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

build_demo_graph(graph_type: str, graph_params: list[str] | None) → TaskGraph[source]

Build a demo graph.

main(argv: list[str] | None = None) → None[source]

Entry point for the demo script.

task_part(seconds: str) → str[source]
task_part_b(seconds: str) → str[source]
task_part_c(seconds: str, *_args: str) → str[source]
task_summary(seconds: str) → None[source]
timeout(seconds: str) → float[source]

gain.task_graph.executor module

class TaskGraphExecutor[source]

Bases: object

The interface of a task executor.

A task executor runs a task graph. It is a context manager: the with block closes the task executor when it ends, and so releases its workers. A failure in close() is logged and does not hide an error of the with block.

The task executors of GAIn are SequentialExecutor, ProcessPoolTaskExecutor and DaskExecutor.

abstractmethod close() → None[source]

Release the workers and the other resources of the task executor.

The with block of a task executor calls this method when it ends.

abstractmethod execute(graph: TaskGraph) → Generator[tuple[Task, Any], None, None][source]

Run the tasks of a task graph.

The task executor consumes the graph: the graph has no tasks left when the run ends. Build a new graph for each run.

Parameters:

graph (TaskGraph) – The task graph to run.

Yields:

tuple[Task, Any] – Each task and its result, when the task completes. The result of a failed task is its exception. The order is not the order of the graph, and not the order in which the tasks run. A task is yielded only after it has finished.

gain.task_graph.graph module

class Task(task_id: str)[source]

Bases: object

Identify one task of a task graph.

A Task is the handle that TaskGraph.create_task and TaskGraph.add_task return. Pass it in args, kwargs or deps of another task to make that task depend on it. Two handles with the same id are equal.

Variables:

task_id (str) – The id of the task, unique within its task graph.

task_id: str
class 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])[source]

Bases: object

Describe one task, before or after it is added to a task graph.

TaskGraph.make_task builds a TaskDesc and TaskGraph.add_task adds it to a graph. TaskGraph.get_task_desc returns the description of a task that is in a graph.

Variables:
  • task (Task) – The handle of the task.

  • func (Callable[..., Any]) – The function that the task calls when it runs.

  • args (list[Any]) – The positional arguments of func. When the task runs, each Task in the list is replaced by the result of that task.

  • kwargs (dict[str, Any]) – The keyword arguments of func. When the task runs, each Task value is replaced by the result of that task.

  • deps (list[Task]) – The tasks that must finish before this task runs. The list holds each Task of args and kwargs too.

  • input_files (list[str]) – The files that the task reads. A file task cache runs the task again when one of them is newer than the outputs of the task.

  • output_files (list[str]) – The final output files of the task. A file task cache runs the task again when one of them is missing.

  • intermediate_output_files (list[str]) – The output files that later tasks consume. A missing one does not make a file task cache run the task again.

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 TaskGraph[source]

Bases: object

A directed acyclic graph of tasks.

Build the graph with create_task, or with make_task and add_task. Then give the graph to a task executor, which runs the tasks and consumes the graph.

Variables:

input_files (list[str]) – The files that were used to build the graph. get_task_desc adds them to the input files of each task, so a file task cache runs each task again when one of them changes. The list is empty in a new graph.

add_task(task_desc: TaskDesc) → Task[source]

Add a task description to the graph.

Parameters:

task_desc (TaskDesc) – The description, usually from make_task.

Returns:

The handle of the task.

Return type:

Task

Raises:

ValueError – If the graph already has a task with the same id.

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

Add several task descriptions to the graph.

Parameters:

task_descs (Sequence[TaskDesc]) – The descriptions, for example the result of a create_statistics_build_tasks call.

Returns:

The handles of the tasks, in the order of task_descs.

Return type:

list[Task]

Raises:

ValueError – If the graph already has a task with the id of one of the descriptions, or if two descriptions in task_descs have the same id. In the first case, the graph is unchanged. In the second case, the graph keeps the descriptions that come before the second description with that id.

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[source]

Create a task and add it to the graph.

The call is make_task followed by add_task.

Parameters:
  • task_id (str) – The id of the task, unique within the graph. An id longer than 200 characters is cut to 200.

  • func (Callable[..., Any]) – The function that the task calls. A task executor can run it in another process, so use a module-level function.

  • args (Sequence[Any]) – The positional arguments of func. When the task runs, each Task in args is replaced by the result of that task.

  • kwargs (dict[str, Any] or None) – The keyword arguments of func. When the task runs, each Task value is replaced by the result of that task. None means no keyword arguments.

  • deps (Sequence[Task] or None) – The tasks that must finish before this task runs. The graph adds each Task of args and kwargs to them. A task that is only in deps passes no value. None means no other deps.

  • input_files (Sequence[str] or None) – The files that the task reads. A file task cache runs the task again when one of them is newer than the outputs of the task. None means none.

  • output_files (Sequence[str] or None) – The final output files of the task. A file task cache runs the task again when one of them is missing. None means none.

  • intermediate_output_files (Sequence[str] or None) – The output files that later tasks consume. A missing one does not make a file task cache run the task again. None means none.

Returns:

The handle of the new task.

Return type:

Task

Raises:

ValueError – If the graph already has a task with the same id.

empty() → bool[source]

Check whether the graph has no tasks.

Returns:

True when the graph has no tasks.

Return type:

bool

get_task_desc(task: Task) → TaskDesc[source]

Return the description of a task of the graph.

Parameters:

task (Task) – The handle of the task.

Returns:

The description. Its input_files start with the input_files of the graph.

Return type:

TaskDesc

Raises:

ValueError – If the graph has no such task.

has_task(task: Task) → bool[source]

Check whether the graph has a task.

Parameters:

task (Task) – The handle of the task.

Returns:

True when the graph has the task.

Return type:

bool

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[source]

Build a task description without adding it to a graph.

The parameters are the parameters of create_task. Add the description to a graph with add_task or add_tasks, or combine several with chain_tasks.

Parameters:
  • task_id (str) – The id of the task. An id longer than 200 characters is cut to 200.

  • func (Callable[..., Any]) – The function that the task calls.

  • args (Sequence[Any]) – The positional arguments of func.

  • kwargs (dict[str, Any] or None) – The keyword arguments of func. None means none.

  • deps (Sequence[Task] or None) – The tasks that must finish first. None means none.

  • input_files (Sequence[str] or None) – The files that the task reads. None means none.

  • output_files (Sequence[str] or None) – The final output files. None means none.

  • intermediate_output_files (Sequence[str] or None) – The output files that later tasks consume. None means none.

Returns:

The description. Its deps hold deps and each Task of args and kwargs.

Return type:

TaskDesc

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

Keep only the named tasks and the tasks that they depend on.

The method removes every other task from the graph.

Parameters:

tasks_to_keep (Sequence[Task | str]) – The tasks to keep, as handles or as task ids.

Raises:

ValueError – If the graph has no task with one of the ids. The graph is then unchanged.

property tasks: Sequence[Task]

The tasks of the graph, in the order they were added.

topological_order() → Sequence[Task][source]

Return the tasks so that each task comes after its deps.

Returns:

The handles of all the tasks. An empty graph gives an empty list.

Return type:

Sequence[Task]

Raises:

ValueError – If the deps of the tasks form a cycle.

chain_tasks(*tasks: TaskDesc) → TaskDesc[source]

Combine task descriptions into one task that runs them in order.

The combined task calls the function of each task in turn. When a task depends on the task before it, it gets the result of that task in place of the Task handle.

Parameters:

*tasks (TaskDesc) – The task descriptions, in the order to run them. Each task after the first depends on no task, or only on the task before it.

Returns:

One task description with the id of the last task, the deps of the first task, and the files of all the tasks. Its result is the result of the last task.

Return type:

TaskDesc

Raises:

ValueError – If tasks is empty, or a task depends on a task other than the one before it.

sync_tasks() → None[source]

gain.task_graph.logging module

class FsspecHandler(logfile: str)[source]

Bases: StreamHandler[Any]

Class to create fsspec based logging handler.

close() → None[source]

Close the stream.

Copied from logging.FileHandler.close().

configure_task_logging(log_dir: str | None, task_id: str, verbosity: int) → Handler[source]

Configure and return task logging hadnler.

ensure_log_dir(**kwargs: Any) → str[source]

Ensure logging directory exists.

safe_task_id(task_id: str) → str[source]

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 ProcessPoolTaskExecutor(**kwargs: Any)[source]

Bases: TaskGraphExecutorBase

A task executor that runs the tasks in a pool of processes.

The constructor creates the pool. The pool starts its worker processes when the task executor submits the first task. A task runs in a worker process only after the tasks that it depends on. The task executor pickles the function and the arguments of each task, so the function must be importable, for example a module-level function.

Parameters:

**kwargs – max_workers is the number of processes in the pool. Without it, the pool has one process for each CPU. The other keyword arguments are those of TaskGraphExecutorBase: task_cache, force and the options of the task executor. The process pool ignores the options task_log_dir, verbose and fork_tasks: it creates the directory of the task logs, but it writes no task log and it does not fork the tasks.

close() → None[source]

Shut the pool down and wait for its processes to exit.

gain.task_graph.sequential_executor module

class SequentialExecutor(task_cache: TaskCache = {}, *, force: bool = False, **kwargs: Any)[source]

Bases: TaskGraphExecutorBase

A task executor that runs the tasks one at a time.

A task runs only after the tasks that it depends on. The tasks run in the current process, unless the fork_tasks option is true: then each task runs in a new process. The constructor takes the keyword arguments of TaskGraphExecutorBase: task_cache, force and the options of the task executor.

close() → None[source]

Do nothing: a sequential task executor holds no workers.

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

absolutize_path_args(args: dict[str, Any], *, input_key: str, extra_keys: Iterable[str] = ()) → None[source]

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.

apply_work_dir_defaults(args: dict[str, Any]) → None[source]

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.

maybe_remove_work_dir(args: dict[str, Any], *, result: bool) → None[source]

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