diff --git a/doc/manuals/taskvine/vine-graph.md b/doc/manuals/taskvine/vine-graph.md new file mode 100644 index 0000000000..c7cea49dcb --- /dev/null +++ b/doc/manuals/taskvine/vine-graph.md @@ -0,0 +1,296 @@ +# Vine Graph + +Vine Graph describes a directed acyclic graph (DAG) of Python function calls and +runs it with TaskVine. A `Workflow` records the tasks and their dependencies; +`VineGraph` runs the graph and returns the requested results. + +## The programming model + +The main objects are: + +- `Workflow`: owns a graph and all handles created for that graph. +- `TaskHandle`: identifies a task. `workflow.add_task()` returns one. +- `TaskOutputHandle`: represents a task's Python return value, or a selected + part of that value. Obtain one with `task.output()`. +- `FileHandle`: represents an existing frontend file or a file produced in a + task sandbox. +- `VineGraph`: a TaskVine manager that executes a completed workflow. + +`VineGraph.run()` is synchronous: it returns after execution finishes. The +returned dictionary is keyed by the target handles supplied by the caller. + +## First workflow: local execution + +Local execution runs the graph in the manager process and does not need a +worker. It is handy for checking graph construction and task functions. Create +`vine_graph_local.py`: + +```python +from ndcctools.taskvine.vine_graph import VineGraph, Workflow + + +def make_record(value): + return {"values": [value, value + 1], "metadata": {"count": 2}} + + +def scaled_sum(values, count, scale=1): + assert len(values) == count + return sum(values) * scale + + +workflow = Workflow() +record = workflow.add_task(make_record, 20) +answer = workflow.add_task( + scaled_sum, + record.output()["values"], + record.output()["metadata"]["count"], + scale=2, +) + +with VineGraph(port=0) as manager: + results = manager.run( + workflow, + targets=[answer], + params={ + "local-execute": 1, + "output-dir": "./vine-graph-local-output", + }, + ) + +assert results[answer] == 82 +print(results[answer]) +``` + +Run it with: + +```bash +python vine_graph_local.py +``` + +Passing `record` directly would be an error because a `TaskHandle` identifies a +task, not its value. Pass `record.output()` instead. Indexing an output handle, +as in `record.output()["values"]`, selects a value inside the consuming task +without adding another graph node. + +To request every terminal result, pass `targets=workflow.sink_tasks()`. +`sink_tasks()` returns handles for all tasks with no downstream consumers. + +## Passing files between tasks + +Use `workflow.file(path)` for an existing frontend file. Use +`task.file(relative_path)` for a file that the task will create inside its +sandbox. A consumer receives either kind of handle as a local path string. + +For example, `vine_graph_files.py` passes a frontend file to one task and its +output file to another: + +```python +from pathlib import Path + +from ndcctools.taskvine.vine_graph import VineGraph, Workflow + + +def uppercase(source_path): + text = Path(source_path).read_text().upper() + Path("result.txt").write_text(text) + return len(text) + + +def verify(result_path, expected_length): + text = Path(result_path).read_text() + assert len(text) == expected_length + return text + + +Path("input.txt").write_text("vine graph\n") + +workflow = Workflow() +source = workflow.file("input.txt") +producer = workflow.add_task(uppercase, source) +produced_file = producer.file("result.txt") +consumer = workflow.add_task(verify, produced_file, producer.output()) + +with VineGraph(port=0) as manager: + results = manager.run( + workflow, + targets=[consumer, produced_file], + params={ + "local-execute": 1, + "output-dir": "./vine-graph-file-output", + }, + ) + +assert results[consumer] == "VINE GRAPH\n" +assert Path(results[produced_file]).read_text() == "VINE GRAPH\n" +print(results[consumer], end="") +``` + +Task output paths must be non-empty relative paths that stay inside the task +sandbox. Each output path may be declared only once for a given task. Handles +belong to one `Workflow` and cannot be passed into another workflow. + +## Distributed execution with a worker + +Distributed mode is the default. The manager creates a TaskVine task-runner +library, and workers execute the graph nodes. Create +`vine_graph_distributed.py`: + +```python +from ndcctools.taskvine.vine_graph import VineGraph, Workflow + + +def square(value): + return value * value + + +def add_all(*values): + return sum(values) + + +workflow = Workflow() +squares = [workflow.add_task(square, value) for value in range(1, 6)] +total = workflow.add_task(add_all, *(task.output() for task in squares)) + +with VineGraph(port=0, name="vine-graph-example") as manager: + results = manager.run( + workflow, + targets=[total], + params={ + "libcores": 4, + "output-dir": "./vine-graph-distributed-output", + }, + ) + +assert results[total] == 55 +print(results[total]) +``` + +Run the manager in one terminal: + +```bash +python vine_graph_distributed.py +``` + +Start a worker in another terminal: + +```bash +vine_worker -M vine-graph-example --cores 4 +``` + +The manager name is used by remote workers to find the manager through the +TaskVine catalog. Stop the worker with `Ctrl-C`. + +`--cores` on `vine_worker`, `vine_submit_workers`, and `vine_factory` must match +the manager's `libcores`. All examples here use 4. + +## HTCondor workers + +Start the manager: + +```bash +python vine_graph_distributed.py +``` + +In another shell, submit workers to HTCondor: + +```bash +vine_submit_workers -T condor -M vine-graph-example \ + --cores 4 5 +``` + +This submits five workers with four cores each. Use `condor_q` to check their +status. The submit command prints the cluster ID; use `condor_rm CLUSTER_ID` to +stop the workers. + +## Using a factory + +A factory adds and removes workers as the workload changes. Keep the manager +running, then start the factory in another shell: + +```bash +vine_factory -T condor -M vine-graph-example \ + --min-workers 1 --max-workers 10 --cores 4 +``` + +To use the same conda environment on the workers, make a Poncho tarball: + +```bash +poncho_package_create --ignore-editable-packages \ + "$CONDA_PREFIX" vine-graph-env.tar.gz +``` + +Pass it to the factory: + +```bash +vine_factory -T condor -M vine-graph-example \ + --min-workers 1 --max-workers 10 --cores 4 \ + --poncho-env vine-graph-env.tar.gz +``` + +Stop the factory with `Ctrl-C` when the workflow is finished. + +## Dask graphs + +Set `from_dask=True` to convert a Dask-style graph. Low-level task dictionaries +and common Dask collection forms are supported. For example: + +```python +from ndcctools.taskvine.vine_graph import VineGraph + + +def increment(value): + return value + 1 + + +def multiply(left, right): + return left * right + + +dask_graph = { + "incremented": (increment, 1), + "answer": (multiply, "incremented", 10), +} + +with VineGraph(port=0) as manager: + results = manager.run( + dask_graph, + targets=["answer"], + from_dask=True, + params={ + "local-execute": 1, + "output-dir": "./vine-graph-dask-output", + }, + ) + +assert results["answer"] == 20 +print(results["answer"]) +``` + +For Dask collections, pass a dictionary of Dask collection objects. Do not mix +collection and non-collection values in that dictionary. + +## Execution parameters + +Parameters may be supplied through the `params` argument to `run()` or through +`manager.set_params()` before `run()`. + +Useful parameters include: + +| Parameter | Purpose | +| --- | --- | +| `local-execute` | Run in process when set to `1`; use TaskVine workers when `0`. | +| `output-dir` | Store serialized results; local mode also stores task sandboxes here. | +| `checkpoint-dir` | Store executor checkpoints here. | +| `libcores` | Number of cores assigned to the task-runner library. | + +## Run the project regression tests + +From the repository root: + +```bash +cd taskvine/test +./TR_vine_graph_workflow_examples.sh prepare +./TR_vine_graph_workflow_examples.sh run +./TR_vine_graph_dask_adaptor.sh prepare +./TR_vine_graph_dask_adaptor.sh run +``` diff --git a/poncho/src/poncho/library_network_code.py b/poncho/src/poncho/library_network_code.py index a5bedebc30..38a4524309 100755 --- a/poncho/src/poncho/library_network_code.py +++ b/poncho/src/poncho/library_network_code.py @@ -28,6 +28,9 @@ r, w = os.pipe() exec_method = None +# infile load mode for function tasks inside this library +function_infile_load_mode = None + # This class captures how results from FunctionCalls are conveyed from # the library to the manager. @@ -85,6 +88,18 @@ def sigchld_handler(signum, frame): os.write(w, b"a") +# Load the infile for a function task inside this library +def load_function_infile(in_file_path): + if function_infile_load_mode == "cloudpickle": + with open(in_file_path, "rb") as f: + return cloudpickle.load(f) + elif function_infile_load_mode == "json": + with open(in_file_path, "r", encoding="utf-8") as f: + return json.load(f) + else: + raise ValueError(f"invalid infile load mode: {function_infile_load_mode}") + + # Read data from worker, start function, and dump result to `outfile`. def start_function(in_pipe_fd, thread_limit=1): # read length of buffer to read @@ -131,8 +146,7 @@ def start_function(in_pipe_fd, thread_limit=1): os.chdir(function_sandbox) # parameters are represented as infile. - with open("infile", "rb") as f: - event = cloudpickle.load(f) + event = load_function_infile("infile") # output of execution should be dumped to outfile. result = globals()[function_name](event) @@ -158,11 +172,10 @@ def start_function(in_pipe_fd, thread_limit=1): return -1, function_id elif exec_method == "fork": try: - arg_infile = os.path.join(function_sandbox, "infile") - with open(arg_infile, "rb") as f: - event = cloudpickle.load(f) + infile_path = os.path.join(function_sandbox, "infile") + event = load_function_infile(infile_path) except Exception: - stdout_timed_message(f"TASK {function_id} error: can't load the arguments from {arg_infile}") + stdout_timed_message(f"TASK {function_id} error: can't load the arguments from {infile_path}") return -1, function_id p = os.fork() if p == 0: @@ -368,11 +381,16 @@ def main(): global exec_method exec_method = library_info['exec_mode'] + # set infile load mode of functions in this library + global function_infile_load_mode + function_infile_load_mode = library_info['function_infile_load_mode'] + # send configuration of library, just its name for now config = { "name": library_info['library_name'], "taskid": args.task_id, "exec_mode": exec_method, + "function_infile_load_mode": function_infile_load_mode, } send_configuration(config, out_pipe_fd, args.worker_pid) diff --git a/poncho/src/poncho/package_serverize.py b/poncho/src/poncho/package_serverize.py index 4a6e5e7a29..cfc789a11b 100755 --- a/poncho/src/poncho/package_serverize.py +++ b/poncho/src/poncho/package_serverize.py @@ -178,6 +178,7 @@ def pack_library_code(path, envpath): # @param exec_mode The execution mode of functions in this library. # @param hoisting_modules A list of modules imported at the preamble of library, including packages, functions and classes. # @param library_context_info A list containing [library_context_func, library_context_args, library_context_kwargs]. Used to create the library context on remote nodes. +# @param function_infile_load_mode The mode to load infile for function tasks inside this library. # @return A hash value. def generate_library_hash(library_name, function_list, @@ -186,7 +187,8 @@ def generate_library_hash(library_name, add_env, exec_mode, hoisting_modules, - library_context_info): + library_context_info, + function_infile_load_mode): library_info = [library_name] function_list = list(function_list) function_names = set() @@ -234,6 +236,8 @@ def generate_library_hash(library_name, for kwarg in library_context_info[2]: library_info.append(str(kwarg)) library_info.append(str(library_context_info[2][kwarg])) + + library_info.append(str(function_infile_load_mode)) library_info = ''.join(library_info) # linear time complexity msg = hashlib.sha1() @@ -293,6 +297,7 @@ def generate_taskvine_library_code(library_path, hoisting_modules=None): # @param exec_mode execution mode of functions in this library # @param hoisting_modules a list of modules to be imported at the preamble of library # @param library_context_info a list containing a library's context to be created remotely +# @param function_infile_load_mode The mode to load infile for function tasks inside this library. # @return name of the file containing serialized information about the library def generate_library(library_cache_path, library_code_path, @@ -303,7 +308,8 @@ def generate_library(library_cache_path, need_pack=True, exec_mode='fork', hoisting_modules=None, - library_context_info=None + library_context_info=None, + function_infile_load_mode='cloudpickle' ): # create library_info.clpk library_info = {} @@ -313,6 +319,7 @@ def generate_library(library_cache_path, library_info['library_name'] = library_name library_info['exec_mode'] = exec_mode library_info['context_info'] = cloudpickle.dumps(library_context_info) + library_info['function_infile_load_mode'] = function_infile_load_mode with open(library_info_path, 'wb') as f: cloudpickle.dump(library_info, f) diff --git a/taskvine/.gitignore b/taskvine/.gitignore new file mode 100644 index 0000000000..50f0e4da1e --- /dev/null +++ b/taskvine/.gitignore @@ -0,0 +1 @@ +taskvine-insights-blogs \ No newline at end of file diff --git a/taskvine/src/Makefile b/taskvine/src/Makefile index 8f828fd7bf..a4b95f4f50 100644 --- a/taskvine/src/Makefile +++ b/taskvine/src/Makefile @@ -1,17 +1,21 @@ include ../../config.mk include ../../rules.mk -TARGETS=manager worker tools bindings examples +TARGETS=manager worker tools bindings examples vine_graph all: $(TARGETS) worker: manager -bindings: manager +vine_graph: manager +bindings: manager vine_graph tools: manager examples: manager worker tools bindings +# Backward-compatible alias for callers that still use `make graph`. +graph: vine_graph + $(TARGETS): %: - $(MAKE) -C $@ $(MAKECMDGOALS) + $(MAKE) -C $@ $(if $(filter $(TARGETS) graph,$(MAKECMDGOALS)),,$(MAKECMDGOALS)) install: for d in $(TARGETS); do $(MAKE) -C $$d install; done @@ -23,4 +27,4 @@ lint: $(TARGETS) format: $(TARGETS) -.PHONY: all clean install test lint format $(TARGETS) +.PHONY: all clean install test lint format graph $(TARGETS) diff --git a/taskvine/src/bindings/python3/Makefile b/taskvine/src/bindings/python3/Makefile index ca4ca6a52b..cf7797e286 100644 --- a/taskvine/src/bindings/python3/Makefile +++ b/taskvine/src/bindings/python3/Makefile @@ -7,9 +7,10 @@ CCTOOLS_DYNAMIC_SUFFIX = so LOCAL_CCFLAGS = -w -fPIC -DNDEBUG $(CCTOOLS_PYTHON3_CCFLAGS) -I ../../manager LOCAL_LINKAGE = $(CCTOOLS_PYTHON3_LDFLAGS) -lz $(CCTOOLS_OPENSSL_LDFLAGS) $(CCTOOLS_HOME)/taskvine/src/manager/libtaskvine.a $(CCTOOLS_HOME)/dttools/src/libdttools.a -CCTOOLS_FLAKE8_IGNORE_FILES = "cvine.py" +CCTOOLS_FLAKE8_IGNORE_FILES = "cvine.py,vine_graph_capi.py" DSPYTHONSO = ndcctools/taskvine/_cvine.$(CCTOOLS_DYNAMIC_SUFFIX) +VINE_GRAPH_MODULE_DIR = $(CCTOOLS_PYTHON3_PATH)/ndcctools/taskvine/vine_graph LIBRARIES = $(DSPYTHONSO) OBJECTS = vine_wrap.o TARGETS = $(LIBRARIES) @@ -29,9 +30,16 @@ lint: clean: rm -rf $(OBJECTS) $(TARGETS) ndcctools/taskvine/cvine.py vine_wrap.c vine_wrap.o *.pyc __pycache__ + rm -rf ndcctools/taskvine/vine_graph/__pycache__ + find ndcctools/taskvine/vine_graph -name __pycache__ -type d -prune -exec rm -rf {} + install: all mkdir -p $(CCTOOLS_PYTHON3_PATH)/ndcctools/taskvine/compat cp ndcctools/taskvine/*.py $(DSPYTHONSO) $(CCTOOLS_PYTHON3_PATH)/ndcctools/taskvine cp ndcctools/taskvine/compat/*.py $(CCTOOLS_PYTHON3_PATH)/ndcctools/taskvine/compat + mkdir -p $(VINE_GRAPH_MODULE_DIR) + rm -rf $(VINE_GRAPH_MODULE_DIR)/adaptors $(VINE_GRAPH_MODULE_DIR)/task_runner $(VINE_GRAPH_MODULE_DIR)/test $(VINE_GRAPH_MODULE_DIR)/__pycache__ + rm -f $(VINE_GRAPH_MODULE_DIR)/*.py + cd ndcctools/taskvine/vine_graph && find . -path ./test -prune -o -type d -print | while read dir; do mkdir -p "$(VINE_GRAPH_MODULE_DIR)/$$dir"; done + cd ndcctools/taskvine/vine_graph && find . -path ./test -prune -o -name '*.py' -print | while read file; do cp "$$file" "$(VINE_GRAPH_MODULE_DIR)/$$file"; done cp taskvine.py $(CCTOOLS_PYTHON3_PATH)/ diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/__init__.py b/taskvine/src/bindings/python3/ndcctools/taskvine/__init__.py index a2a02ed095..67522cc2de 100644 --- a/taskvine/src/bindings/python3/ndcctools/taskvine/__init__.py +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/__init__.py @@ -77,6 +77,11 @@ class DaskVineWarning(UserWarning): from .compat import DaskVine from .compat import DaskVineDag +try: + from .vine_graph import VineGraph, VineGraphConfig, VineGraphDaskAdaptor # noqa: F401 +except (ImportError, ModuleNotFoundError): + pass + __all__ = [ "Manager", @@ -95,6 +100,13 @@ class DaskVineWarning(UserWarning): "DaskVineDag", ] +if "VineGraph" in globals(): + __all__.append("VineGraph") +if "VineGraphConfig" in globals(): + __all__.append("VineGraphConfig") +if "VineGraphDaskAdaptor" in globals(): + __all__.append("VineGraphDaskAdaptor") + __version__ = cvine.vine_version_string() # vim: set sts=4 sw=4 ts=4 expandtab ft=python: diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/manager.py b/taskvine/src/bindings/python3/ndcctools/taskvine/manager.py index a9c383a085..566e0e88e9 100644 --- a/taskvine/src/bindings/python3/ndcctools/taskvine/manager.py +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/manager.py @@ -1129,8 +1129,9 @@ def check_library_exists(self, library_name): # @param hoisting_modules A list of modules imported at the preamble of library, including packages, functions and classes. # @param exec_mode Execution mode that the library should use to run function calls. Either 'direct' or 'fork' # @param library_context_info A list containing [library_context_func, library_context_args, library_context_kwargs]. Used to create the library context on remote nodes. + # @param function_infile_load_mode The mode to load infile for function tasks inside this library. # @returns A task to be used with @ref ndcctools.taskvine.manager.Manager.install_library. - def create_library_from_functions(self, library_name, *function_list, poncho_env=None, init_command=None, add_env=True, hoisting_modules=None, exec_mode='fork', library_context_info=None): + def create_library_from_functions(self, library_name, *function_list, poncho_env=None, init_command=None, add_env=True, hoisting_modules=None, exec_mode='fork', library_context_info=None, function_infile_load_mode='cloudpickle'): # Delay loading of poncho until here, to avoid bringing in poncho dependencies unless needed. # Ensure poncho python library is available. from ndcctools.poncho import package_serverize @@ -1152,7 +1153,8 @@ def create_library_from_functions(self, library_name, *function_list, poncho_env add_env=add_env, exec_mode=exec_mode, hoisting_modules=hoisting_modules, - library_context_info=library_context_info) + library_context_info=library_context_info, + function_infile_load_mode=function_infile_load_mode) # Create path for caching library code and environment based on function hash. library_cache_dir_name = "vine-library-cache" @@ -1200,7 +1202,8 @@ def create_library_from_functions(self, library_name, *function_list, poncho_env need_pack=need_pack, exec_mode=exec_mode, hoisting_modules=hoisting_modules, - library_context_info=library_context_info) + library_context_info=library_context_info, + function_infile_load_mode=function_infile_load_mode) # enable correct permissions for library code os.chmod(library_code_path, 0o775) diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/__init__.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/__init__.py new file mode 100644 index 0000000000..2e3e8687f3 --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/__init__.py @@ -0,0 +1,18 @@ +# Copyright (C) 2025- The University of Notre Dame +# This software is distributed under the GNU General Public License. +# See the file COPYING for details. + +from .vine_graph import VineGraph, VineGraphConfig +from .adaptors import VineGraphDaskAdaptor, VineGraphGraphedAdaptor +from .workflow import FileHandle, TaskHandle, TaskOutputHandle, Workflow + +__all__ = [ + "VineGraph", + "VineGraphConfig", + "Workflow", + "TaskHandle", + "TaskOutputHandle", + "FileHandle", + "VineGraphDaskAdaptor", + "VineGraphGraphedAdaptor", +] diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/.gitignore b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/.gitignore new file mode 100644 index 0000000000..c18dd8d83c --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/__init__.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/__init__.py new file mode 100644 index 0000000000..c836dd2ac2 --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/__init__.py @@ -0,0 +1,10 @@ +"""Adaptors that lower external workflow formats into VineGraph.""" + +from .adaptor import VineGraphDaskAdaptor, VineGraphGraphedAdaptor, graphed_plan_to_workflow, workflow_to_dask_graph + +__all__ = [ + "VineGraphDaskAdaptor", + "VineGraphGraphedAdaptor", + "graphed_plan_to_workflow", + "workflow_to_dask_graph", +] diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/adaptor.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/adaptor.py new file mode 100644 index 0000000000..9c3cc92b01 --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/adaptor.py @@ -0,0 +1,11 @@ +"""Public adaptor entry points for lowering external graph formats to VineGraph.""" + +from .dask_adaptor import VineGraphDaskAdaptor, workflow_to_dask_graph +from .graphed import VineGraphGraphedAdaptor, graphed_plan_to_workflow + +__all__ = [ + "VineGraphDaskAdaptor", + "VineGraphGraphedAdaptor", + "graphed_plan_to_workflow", + "workflow_to_dask_graph", +] diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/dask_adaptor.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/dask_adaptor.py new file mode 100644 index 0000000000..257200922e --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/dask_adaptor.py @@ -0,0 +1,167 @@ +from collections.abc import Mapping + +try: + import dask +except ImportError: + dask = None + +try: + from dask.base import is_dask_collection +except ImportError: + is_dask_collection = None + +try: + import importlib + + dts = importlib.import_module("dask._task_spec") +except Exception: + dts = None + +from .legacy_dask_adaptor import expand_legacy_subgraph_dsk +from .dask_common import ( + build_task_expr, + identity, + resolve_graph_key_if_task, +) +from .dask_task_spec import DaskTaskSpecConverter +from ..workflow import TaskOutputHandle, Workflow + + +def _apply_with_kwargs_kvlist(func, args_list, kwargs_kvlist): + """Call ``func`` when kwargs are encoded as ``[[key, value], ...]``.""" + return func(*args_list, **{key: value for key, value in kwargs_kvlist}) + + +def workflow_to_dask_graph(workflow): + assert isinstance(workflow, Workflow), "workflow must be a Workflow" + + def ref_to_key(ref): + return ref.task_id + + dsk = {} + for workflow_key, (func_id, args, kwargs) in workflow.task_dict.items(): + func = workflow.callables[func_id] + new_args = workflow._visit_task_output_refs(args, ref_to_key, rewrite=True) + new_kwargs = workflow._visit_task_output_refs(kwargs, ref_to_key, rewrite=True) + + if new_kwargs: + dsk[workflow_key] = (_apply_with_kwargs_kvlist, func, list(new_args), [[k, v] for k, v in new_kwargs.items()]) + else: + dsk[workflow_key] = (func, *new_args) + + return dsk + + +class VineGraphDaskAdaptor: + """Convert Dask graph forms into VineGraph ``Workflow`` task expressions.""" + + def __init__(self, task_dict, expand_subgraphs=False): + if isinstance(task_dict, Workflow): + self.converted = task_dict + return + + self._expand_subgraphs = expand_subgraphs + normalized = self._normalize_task_dict(task_dict) + self.converted = self._convert_to_workflow_tasks(normalized) + + @property + def task_dict(self): + return self.converted + + def _normalize_task_dict(self, task_dict): + if self._is_dask_collection_dict(task_dict): + task_dict = self._dask_collections_to_task_dict(task_dict) + else: + # Plain dicts are VineGraph sexprs; don't let Dask reinterpret kwargs dicts. + task_dict = dict(task_dict) + + if self._expand_subgraphs and not dts and task_dict: + task_dict = expand_legacy_subgraph_dsk(task_dict, dask) + return task_dict + + def _is_dask_collection_dict(self, task_dict): + return bool(is_dask_collection and any(is_dask_collection(value) for value in task_dict.values())) + + def _dask_collections_to_task_dict(self, task_dict): + assert is_dask_collection is not None + from dask.highlevelgraph import HighLevelGraph, ensure_dict + + if not isinstance(task_dict, dict): + raise TypeError("Input must be a dict") + for key, value in task_dict.items(): + if not is_dask_collection(value): + raise TypeError(f"Input must be a dict of DaskCollection, but found {key} with type {type(value)}") + + if dts: + hlg = HighLevelGraph.merge(*(value.dask for value in task_dict.values())).to_dict() + else: + hlg = dask.base.collections_to_dsk(task_dict.values()) + hlg = hlg.to_dict() if hasattr(hlg, "to_dict") else dict(hlg) + return ensure_dict(hlg) + + def _convert_to_workflow_tasks(self, task_dict): + if not task_dict: + return {} + + converted = {} + workflow_keys = set(task_dict.keys()) + task_spec = DaskTaskSpecConverter(dts) if dts else None + + for key, value in task_dict.items(): + if task_spec and task_spec.is_node(value): + converted[key] = task_spec.convert_node(key, value, workflow_keys) + else: + converted[key] = self._convert_legacy_task(value, workflow_keys) + + if task_spec: + while True: + pending = task_spec.pending_nodes(converted) + if not pending: + break + for key, node in pending: + converted[key] = task_spec.convert_node(key, node, workflow_keys) + + return converted + + def _convert_legacy_task(self, sexpr, workflow_keys): + try: + if not isinstance(sexpr, (list, tuple)) and sexpr in workflow_keys: + return build_task_expr(identity, [TaskOutputHandle(sexpr)], {}) + except TypeError: + pass + + if not isinstance(sexpr, (list, tuple)): + return build_task_expr(identity, [sexpr], {}) + if not sexpr: + raise TypeError("Task definition must be a non-empty tuple/list") + + func = sexpr[0] + tail = sexpr[1:] + if tail and isinstance(tail[-1], Mapping): + raw_args, raw_kwargs = tail[:-1], tail[-1] + else: + raw_args, raw_kwargs = tail, {} + + args = tuple(self._wrap_dependency(arg, workflow_keys) for arg in raw_args) + kwargs = {key: self._wrap_dependency(value, workflow_keys) for key, value in raw_kwargs.items()} + return func, args, kwargs + + def _wrap_dependency(self, obj, workflow_keys): + if isinstance(obj, TaskOutputHandle): + return obj + + key = resolve_graph_key_if_task(obj, workflow_keys) + if key is not None: + return TaskOutputHandle(key) + + if isinstance(obj, list): + return [self._wrap_dependency(value, workflow_keys) for value in obj] + if isinstance(obj, tuple): + return tuple(self._wrap_dependency(value, workflow_keys) for value in obj) + if isinstance(obj, Mapping): + return {key: self._wrap_dependency(value, workflow_keys) for key, value in obj.items()} + if isinstance(obj, set): + return {self._wrap_dependency(value, workflow_keys) for value in obj} + if isinstance(obj, frozenset): + return frozenset(self._wrap_dependency(value, workflow_keys) for value in obj) + return obj diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/dask_common.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/dask_common.py new file mode 100644 index 0000000000..7d00e61f5e --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/dask_common.py @@ -0,0 +1,32 @@ +"""Shared helpers for Dask-to-VineGraph adaptor modules.""" + +from ..workflow import TaskOutputHandle + + +def identity(value): + """Return ``value`` unchanged.""" + return value + + +def build_task_expr(func, args, kwargs): + """Build the normalized ``Workflow`` task-expression tuple.""" + return func, tuple(args), dict(kwargs) + + +def resolve_graph_key_if_task(obj, workflow_keys): + """Return the matching graph key when ``obj`` denotes an existing Dask task.""" + if isinstance(obj, TaskOutputHandle): + return None + try: + if obj in workflow_keys: + return obj + except TypeError: + pass + if hasattr(obj, "item") and callable(obj.item): + try: + item = obj.item() + if item in workflow_keys: + return item + except Exception: + pass + return None diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/dask_task_spec.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/dask_task_spec.py new file mode 100644 index 0000000000..aff9ac928e --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/dask_task_spec.py @@ -0,0 +1,328 @@ +import hashlib +from collections.abc import Mapping + +from .dask_common import build_task_expr, identity, resolve_graph_key_if_task +from ..workflow import TaskOutputHandle + + +def _safe_repr(value, limit=800): + """Best-effort repr that stays compact for large graphs.""" + try: + text = repr(value) + except Exception as exc: + text = f"" + if limit and len(text) > limit: + return text[:limit] + "..." + return text + + +def _extract_callable_from_task(node): + for attr in ("function", "op", "callable", "func", "operation", "callable_obj"): + if not hasattr(node, attr): + continue + value = getattr(node, attr) + if value is not None and (callable(value) or hasattr(value, "__call__")): + return value + if hasattr(node, "__call__") and callable(node): + return node + return None + + +def _is_pure_value_op(func): + return func in (dict, list, tuple, set, frozenset) + + +def _is_identity_cast_op(func): + name = getattr(func, "__name__", None) + module = getattr(func, "__module__", None) + return bool(name == "_identity_cast" and module and module.startswith("dask")) + + +def _is_too_large_inline_value(value, *, max_container_len=2000): + try: + return isinstance(value, (list, tuple, set, frozenset, dict)) and len(value) > max_container_len + except Exception: + return False + + +class DaskTaskSpecConverter: + """Convert modern Dask TaskSpec nodes into Workflow task expressions.""" + + def __init__(self, dts_module): + self.dts = dts_module + self.lifted_nodes = {} + self._lift_cache = {} + self._lift_counter = 0 + + def is_node(self, value): + if not self.dts: + return False + try: + return isinstance(value, self.dts.GraphNode) + except AttributeError: + return False + + def pending_nodes(self, converted): + return [(key, node) for key, node in self.lifted_nodes.items() if key not in converted] + + def convert_node(self, key, node, workflow_keys): + if not self.dts: + raise RuntimeError("Dask TaskSpec support unavailable: dask._task_spec is not installed") + + task_cls = getattr(self.dts, "Task", None) + alias_cls = getattr(self.dts, "Alias", None) + literal_cls = getattr(self.dts, "Literal", None) + datanode_cls = getattr(self.dts, "DataNode", None) + nested_cls = getattr(self.dts, "NestedContainer", None) + taskref_cls = getattr(self.dts, "TaskRef", None) + + if task_cls and isinstance(node, task_cls): + return self._convert_task_node(key, node, workflow_keys) + if alias_cls and isinstance(node, alias_cls): + alias_ref = self._extract_alias_target(node, workflow_keys) + if alias_ref is None: + raise ValueError(f"Alias {key} is missing a resolvable upstream task") + return build_task_expr(identity, [alias_ref], {}) + if datanode_cls and isinstance(node, datanode_cls): + return build_task_expr(identity, [node.value], {}) + if literal_cls and isinstance(node, literal_cls): + return build_task_expr(identity, [node.value], {}) + if taskref_cls and isinstance(node, taskref_cls): + return build_task_expr(identity, [TaskOutputHandle(node.key, getattr(node, "path", ()) or ())], {}) + if nested_cls and isinstance(node, nested_cls): + payload = getattr(node, "value", None) + if payload is None: + payload = getattr(node, "data", None) + return build_task_expr(identity, [self.unwrap_operand(payload, workflow_keys, parent_key=key)], {}) + return build_task_expr(identity, [node], {}) + + def _convert_task_node(self, key, node, workflow_keys): + func = _extract_callable_from_task(node) + if func is None: + raise TypeError(f"Task {key} is missing a callable function/op attribute") + + raw_args = getattr(node, "args", ()) or () + raw_kwargs = getattr(node, "kwargs", {}) or {} + + args = [] + try: + for index, arg in enumerate(raw_args): + args.append(self.unwrap_operand(arg, workflow_keys, parent_key=key)) + except Exception as exc: + raise TypeError( + "Failed to adapt TaskSpec node argument while converting to Workflow.\n" + f"- parent_workflow_key: {key!r}\n" + f"- func: {_safe_repr(func)}\n" + f"- arg_index: {index}\n" + f"- arg_value: {_safe_repr(arg)}\n" + f"- raw_args: {_safe_repr(raw_args)}\n" + f"- raw_kwargs: {_safe_repr(raw_kwargs)}" + ) from exc + + kwargs = {} + try: + for kwarg_key, value in raw_kwargs.items(): + kwargs[kwarg_key] = self.unwrap_operand(value, workflow_keys, parent_key=key) + except Exception as exc: + raise TypeError( + "Failed to adapt TaskSpec node kwarg while converting to Workflow.\n" + f"- parent_workflow_key: {key!r}\n" + f"- func: {_safe_repr(func)}\n" + f"- kwarg_key: {kwarg_key!r}\n" + f"- kwarg_value: {_safe_repr(value)}\n" + f"- raw_args: {_safe_repr(raw_args)}\n" + f"- raw_kwargs: {_safe_repr(raw_kwargs)}" + ) from exc + + return build_task_expr(func, args, kwargs) + + def unwrap_operand(self, operand, workflow_keys, *, parent_key=None): + taskref_cls = getattr(self.dts, "TaskRef", None) + if taskref_cls and isinstance(operand, taskref_cls): + return TaskOutputHandle(getattr(operand, "key", None), getattr(operand, "path", ()) or ()) + + alias_cls = getattr(self.dts, "Alias", None) + if alias_cls and isinstance(operand, alias_cls): + alias_ref = self._extract_alias_target(operand, workflow_keys) + if alias_ref is None: + raise ValueError("Alias node is missing a valid upstream source") + return alias_ref + + literal_cls = getattr(self.dts, "Literal", None) + if literal_cls and isinstance(operand, literal_cls): + return getattr(operand, "value", None) + + datanode_cls = getattr(self.dts, "DataNode", None) + if datanode_cls and isinstance(operand, datanode_cls): + return operand.value + + nested_cls = getattr(self.dts, "NestedContainer", None) + if nested_cls and isinstance(operand, nested_cls): + payload = getattr(operand, "value", None) + if payload is None: + payload = getattr(operand, "data", None) + return self.unwrap_operand(payload, workflow_keys, parent_key=parent_key) + + task_cls = getattr(self.dts, "Task", None) + if task_cls and isinstance(operand, task_cls): + return self._unwrap_task_operand(operand, workflow_keys, parent_key=parent_key) + + if isinstance(operand, list): + return [self.unwrap_operand(value, workflow_keys, parent_key=parent_key) for value in operand] + if isinstance(operand, tuple): + return tuple(self.unwrap_operand(value, workflow_keys, parent_key=parent_key) for value in operand) + if isinstance(operand, Mapping): + return {key: self.unwrap_operand(value, workflow_keys, parent_key=parent_key) for key, value in operand.items()} + if isinstance(operand, set): + return {self.unwrap_operand(value, workflow_keys, parent_key=parent_key) for value in operand} + if isinstance(operand, frozenset): + return frozenset(self.unwrap_operand(value, workflow_keys, parent_key=parent_key) for value in operand) + return operand + + def _unwrap_task_operand(self, operand, workflow_keys, *, parent_key=None): + inline_key = getattr(operand, "key", None) + if inline_key is not None and inline_key in workflow_keys: + return TaskOutputHandle(inline_key, ()) + + func = _extract_callable_from_task(operand) + if func is None: + return self._lift_inline_task(operand, workflow_keys, parent_key=parent_key) + if _is_identity_cast_op(func): + return self._unwrap_identity_cast(operand, workflow_keys, parent_key=parent_key) + if _is_pure_value_op(func): + reduced, used_lift = self._reduce_inline_task(operand, workflow_keys, parent_key=parent_key) + if used_lift or _is_too_large_inline_value(reduced): + return self._lift_inline_task(operand, workflow_keys, parent_key=parent_key) + return reduced + return self._lift_inline_task(operand, workflow_keys, parent_key=parent_key) + + def _unwrap_identity_cast(self, operand, workflow_keys, *, parent_key=None): + raw_args = getattr(operand, "args", ()) or () + raw_kwargs = getattr(operand, "kwargs", {}) or {} + typ = raw_kwargs.get("typ", None) + values = [self.unwrap_operand(arg, workflow_keys, parent_key=parent_key) for arg in raw_args] + + if typ in (list, tuple, set, frozenset, dict): + try: + return typ(values) + except Exception: + pass + return self._lift_inline_task(operand, workflow_keys, parent_key=parent_key) + + def _extract_alias_target(self, alias_node, workflow_keys): + fields = getattr(alias_node.__class__, "__dataclass_fields__", {}) if self.dts else {} + path = tuple(getattr(alias_node, "path", ()) or ()) + + for candidate in ("alias_of", "target", "source", "ref"): + if candidate not in fields: + continue + key = resolve_graph_key_if_task(getattr(alias_node, candidate, None), workflow_keys) + if key is not None: + return TaskOutputHandle(key, path) + + deps = getattr(alias_node, "dependencies", None) + if deps: + deps = list(deps) + if len(deps) == 1: + key = resolve_graph_key_if_task(deps[0], workflow_keys) + return TaskOutputHandle(key if key is not None else deps[0], path) + return None + + def _reduce_inline_task(self, task_node, workflow_keys, *, parent_key=None): + func = _extract_callable_from_task(task_node) + raw_args = getattr(task_node, "args", ()) or () + raw_kwargs = getattr(task_node, "kwargs", {}) or {} + used_lift = False + + args = [] + for arg in raw_args: + before = len(self.lifted_nodes) + args.append(self.unwrap_operand(arg, workflow_keys, parent_key=parent_key)) + used_lift = used_lift or (len(self.lifted_nodes) != before) + + kwargs = {} + for key, value in raw_kwargs.items(): + before = len(self.lifted_nodes) + kwargs[key] = self.unwrap_operand(value, workflow_keys, parent_key=parent_key) + used_lift = used_lift or (len(self.lifted_nodes) != before) + + try: + return func(*args, **kwargs), used_lift + except Exception: + return self._lift_inline_task(task_node, workflow_keys, parent_key=parent_key), True + + def _lift_inline_task(self, task_node, workflow_keys, *, parent_key=None): + inline_key = getattr(task_node, "key", None) + if parent_key is not None and inline_key == parent_key: + raise ValueError(f"Refusing to lift Task that would self-reference parent key {parent_key!r}") + + signature = self._structural_signature(task_node, workflow_keys) + cached = self._lift_cache.get(signature) + if cached is not None: + return TaskOutputHandle(cached, ()) + + digest = hashlib.sha1(signature.encode("utf-8")).hexdigest()[:16] + base = f"__lift__{digest}" + new_key = base + while new_key in workflow_keys or new_key in self.lifted_nodes: + self._lift_counter += 1 + new_key = f"{base}_{self._lift_counter}" + + self._lift_cache[signature] = new_key + self.lifted_nodes[new_key] = task_node + workflow_keys.add(new_key) + return TaskOutputHandle(new_key, ()) + + def _structural_signature(self, obj, workflow_keys): + try: + return self._structural_signature_impl(obj, workflow_keys) + except Exception: + return f"fallback({_safe_repr(obj)})" + + def _structural_signature_impl(self, obj, workflow_keys): + taskref_cls = getattr(self.dts, "TaskRef", None) + alias_cls = getattr(self.dts, "Alias", None) + literal_cls = getattr(self.dts, "Literal", None) + datanode_cls = getattr(self.dts, "DataNode", None) + nested_cls = getattr(self.dts, "NestedContainer", None) + task_cls = getattr(self.dts, "Task", None) + + if taskref_cls and isinstance(obj, taskref_cls): + return f"TaskRef({getattr(obj, 'key', None)!r},{tuple(getattr(obj, 'path', ()) or ())!r})" + if alias_cls and isinstance(obj, alias_cls): + ref = self._extract_alias_target(obj, workflow_keys) + return f"Alias({getattr(ref, 'task_id', None)!r},{getattr(ref, 'path', ())!r})" + if literal_cls and isinstance(obj, literal_cls): + return f"Literal({_safe_repr(getattr(obj, 'value', None))})" + if datanode_cls and isinstance(obj, datanode_cls): + return f"DataNode({_safe_repr(getattr(obj, 'value', None))})" + if nested_cls and isinstance(obj, nested_cls): + payload = getattr(obj, "value", None) + if payload is None: + payload = getattr(obj, "data", None) + return f"Nested({self._structural_signature(payload, workflow_keys)})" + if task_cls and isinstance(obj, task_cls): + key = getattr(obj, "key", None) + if key is not None and key in workflow_keys: + return f"TaskKey({key!r})" + func = _extract_callable_from_task(obj) + func_id = (getattr(func, "__module__", None), getattr(func, "__qualname__", None), getattr(func, "__name__", None)) + args = getattr(obj, "args", ()) or () + kwargs = getattr(obj, "kwargs", {}) or {} + arg_sigs = ",".join(self._structural_signature(arg, workflow_keys) for arg in args) + kw_sigs = ",".join(f"{key}={self._structural_signature(value, workflow_keys)}" for key, value in sorted(kwargs.items())) + return f"TaskInline(func={func_id!r},args=[{arg_sigs}],kwargs=[{kw_sigs}])" + if isinstance(obj, list): + return "list(" + ",".join(self._structural_signature(value, workflow_keys) for value in obj) + ")" + if isinstance(obj, tuple): + return "tuple(" + ",".join(self._structural_signature(value, workflow_keys) for value in obj) + ")" + if isinstance(obj, dict): + items = ",".join( + f"{_safe_repr(key)}:{self._structural_signature(value, workflow_keys)}" + for key, value in sorted(obj.items(), key=lambda item: repr(item[0])) + ) + return "dict(" + items + ")" + if isinstance(obj, (set, frozenset)): + items = ",".join(sorted(self._structural_signature(value, workflow_keys) for value in obj)) + return f"{type(obj).__name__}(" + items + ")" + return f"py({_safe_repr(obj)})" diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/graphed.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/graphed.py new file mode 100644 index 0000000000..e542558a46 --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/graphed.py @@ -0,0 +1,103 @@ +import contextlib +from collections import OrderedDict + +from ..workflow import Workflow + + +class _VineGraphGraphedResources: + """Small WorkerResources-compatible cache for one VineGraph task.""" + + def __init__(self, max_open=128): + self._handles = OrderedDict() + self._max_open = max_open + + def open_once(self, uri, opener): + if uri in self._handles: + self._handles.move_to_end(uri) + return self._handles[uri] + handle = opener(uri) + self._handles[uri] = handle + while len(self._handles) > self._max_open: + _, evicted = self._handles.popitem(last=False) + self._close_handle(evicted) + return handle + + def close(self): + for handle in self._handles.values(): + self._close_handle(handle) + self._handles.clear() + + @staticmethod + def _close_handle(handle): + close = getattr(handle, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + +def _graphed_empty(empty_ref): + """Run a graphed Plan empty function.""" + empty = empty_ref[0] + return empty() + + +def _graphed_process(process_ref, partition): + """Run one graphed Plan process task.""" + process = process_ref[0] + resources = _VineGraphGraphedResources() + try: + return process(partition, resources) + finally: + resources.close() + + +def _graphed_combine(combine_ref, left, right): + """Run one graphed Plan combine step.""" + combine = combine_ref[0] + return combine(left, right) + + +def _validate_static_plan(plan): + for attr in ("process", "combine", "empty", "tasks"): + if not hasattr(plan, attr): + raise TypeError(f"graphed plan is missing required attribute {attr!r}") + if getattr(plan, "next_tasks", None) is not None: + raise ValueError("VineGraphGraphedAdaptor only supports static graphed plans") + if getattr(plan, "stop", None) is not None: + raise ValueError("VineGraphGraphedAdaptor does not support graphed StopCondition yet") + + +def graphed_plan_to_workflow(plan, key_prefix="graphed"): + """Convert a static graphed Plan into a VineGraph Workflow.""" + _validate_static_plan(plan) + + workflow = Workflow() + tasks = sorted(tuple(plan.tasks), key=lambda task: task.key) + + empty_task = workflow.add_task(_graphed_empty, [plan.empty]) + + previous_task = empty_task + for task in tasks: + process_task = workflow.add_task(_graphed_process, [plan.process], task.partition) + previous_task = workflow.add_task( + _graphed_combine, + [plan.combine], + previous_task.output(), + process_task.output(), + ) + + workflow.finalize() + return workflow, previous_task + + +class VineGraphGraphedAdaptor: + """Convert a static graphed Plan into a VineGraph Workflow.""" + + def __init__(self, plan, key_prefix="graphed"): + self.plan = plan + self.converted, self.target = graphed_plan_to_workflow(plan, key_prefix=key_prefix) + self.targets = [self.target] + + @property + def task_dict(self): + return self.converted diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/legacy_dask_adaptor.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/legacy_dask_adaptor.py new file mode 100644 index 0000000000..0dee0316d9 --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/adaptors/legacy_dask_adaptor.py @@ -0,0 +1,64 @@ +import hashlib + + +def _legacy_subgraph_key(*parts): + """Stable short key for expanded legacy Dask subgraph tasks.""" + return hashlib.sha256("".join(str(p) for p in parts).encode("utf-8")).hexdigest()[:20] + + +def _rewrite_subgraph_expr(dsk_key, inner_dsk, expr, blockwise_args): + try: + if expr in inner_dsk: + return _legacy_subgraph_key(dsk_key, expr) + except Exception: + pass + + if hasattr(expr, "item") and callable(expr.item): + try: + item = expr.item() + if item in inner_dsk: + return _legacy_subgraph_key(dsk_key, item) + except Exception: + pass + + if isinstance(expr, list): + return [_rewrite_subgraph_expr(dsk_key, inner_dsk, item, blockwise_args) for item in expr] + if isinstance(expr, tuple): + return tuple(_rewrite_subgraph_expr(dsk_key, inner_dsk, item, blockwise_args) for item in expr) + if isinstance(expr, str) and expr.startswith("__dask_blockwise__"): + return blockwise_args[int(expr.split("__")[-1])] + return expr + + +def expand_legacy_subgraph_dsk(task_dict, dask_module): + """Inline legacy ``SubgraphCallable`` layers into a flat dsk.""" + if not task_dict or dask_module is None: + return task_dict + + try: + from dask.optimization import SubgraphCallable + except ImportError: + return task_dict + + expanded = {} + for key, sexpr in task_dict.items(): + if not isinstance(sexpr, (tuple, list)) or not sexpr: + expanded[key] = sexpr + continue + + head = sexpr[0] + tail = sexpr[1:] + if isinstance(head, SubgraphCallable): + expanded[key] = _legacy_subgraph_key(key, head.outkey) + for sub_key, sub_sexpr in head.dsk.items(): + rewritten_key = _legacy_subgraph_key(key, sub_key) + expanded[rewritten_key] = _rewrite_subgraph_expr(key, head.dsk, sub_sexpr, tail) + elif callable(head): + expanded[key] = sexpr + else: + raise TypeError( + f"Legacy dsk task {key!r} has non-callable head {type(head).__name__}; " + "expected SubgraphCallable or a callable." + ) + + return expanded diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/capi_bridge.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/capi_bridge.py new file mode 100644 index 0000000000..74a107b6fd --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/capi_bridge.py @@ -0,0 +1,146 @@ +# Copyright (C) 2025 The University of Notre Dame +# This software is distributed under the GNU General Public License. +# See the file COPYING for details. + +"""Bridge from Python VineGraph objects to the C vine_graph API.""" + +import ctypes +import importlib +import os +import sys + +# SWIG-generated vine_graph_capi.py imports "cvine"; wire that top-level name +# to the real TaskVine Python module before importing the generated bindings. +# Also expose _cvine's C symbols globally so _vine_graph_capi reuses the same +# TaskVine/dttools runtime instead of loading a second copy of those globals. +cvine = importlib.import_module("ndcctools.taskvine.cvine") +_rtld_now = getattr(os, "RTLD_NOW", 0) +_rtld_global = getattr(os, "RTLD_GLOBAL", ctypes.RTLD_GLOBAL) +ctypes.CDLL(cvine._cvine.__file__, mode=_rtld_now | _rtld_global) +sys.modules.setdefault("cvine", cvine) + +from . import vine_graph_capi # noqa: E402 + + +class VineGraphCapiBridge: + """Thin bridge around the SWIG bindings.""" + + def __init__(self, c_taskvine): + """Create the backing C vine_graph objects.""" + self._c_graph = vine_graph_capi.vine_graph_executor_create_graph(c_taskvine) + self._c_executor = vine_graph_capi.vine_graph_executor_create(c_taskvine, self._c_graph) + self._workflow_key_to_scheduler_key = {} + self._scheduler_key_to_workflow_key = {} + + def tune(self, name, value): + """Forward a tuning parameter to the C vine_graph executor.""" + if vine_graph_capi.vine_graph_executor_tune(self._c_executor, name, value) != 0: + raise RuntimeError(f"Failed to tune executor parameter {name!r}={value!r}") + + def add_node(self, workflow_key, is_target=None): + """Create a C node and record its workflow key.""" + node_id = vine_graph_capi.vine_graph_executor_add_node(self._c_executor) + self._workflow_key_to_scheduler_key[workflow_key] = node_id + self._scheduler_key_to_workflow_key[node_id] = workflow_key + if is_target is not None and bool(is_target): + vine_graph_capi.vine_graph_set_target(self._c_graph, node_id) + return node_id + + def set_target(self, workflow_key): + """Mark a node as a target.""" + node_id = self._workflow_key_to_scheduler_key.get(workflow_key) + if node_id is None: + raise KeyError(f"Workflow key not found: {workflow_key}") + vine_graph_capi.vine_graph_set_target(self._c_graph, node_id) + + def add_dependency(self, parent_workflow_key, child_workflow_key): + """Add an edge between two existing nodes.""" + wk2sk = self._workflow_key_to_scheduler_key + if parent_workflow_key not in wk2sk or child_workflow_key not in wk2sk: + raise KeyError("parent or child workflow_key missing in mapping; call add_node() first") + vine_graph_capi.vine_graph_add_dependency( + self._c_graph, wk2sk[parent_workflow_key], wk2sk[child_workflow_key] + ) + + def group_chain_like_tasks(self): + """Merge maximal singleton linear chains into supernodes (C vine_graph_group_chain_like_tasks).""" + return vine_graph_capi.vine_graph_group_chain_like_tasks(self._c_graph) + + def compute_topology_metrics(self): + """Finalize the C graph and compute topology metrics.""" + vine_graph_capi.vine_graph_executor_finalize(self._c_executor) + + def get_node_outfile_remote_name(self, workflow_key): + """Return the output path assigned by the C graph.""" + if workflow_key not in self._workflow_key_to_scheduler_key: + raise KeyError(f"Workflow key not found: {workflow_key}") + return vine_graph_capi.vine_graph_get_node_outfile_remote_name( + self._c_graph, self._workflow_key_to_scheduler_key[workflow_key] + ) + + def get_task_runner_library_name(self): + """Return the generated task runner library name.""" + return vine_graph_capi.vine_graph_get_task_runner_library_name(self._c_graph) + + def set_task_runner_function(self, task_runner_function): + """Set the worker-side task runner entry point.""" + vine_graph_capi.vine_graph_set_task_runner_function_name( + self._c_graph, task_runner_function.__name__ + ) + + def declare_input_file(self, file_id, source_path): + """Declare one frontend file.""" + if vine_graph_capi.vine_graph_executor_declare_input_file( + self._c_executor, file_id, source_path + ) != 0: + raise RuntimeError(f"failed to declare input file {file_id}: {source_path}") + + def add_task_input_file(self, workflow_key, file_id, task_path): + """Mount a declared FileHandle into a task.""" + task_id = self._workflow_key_to_scheduler_key.get(workflow_key) + if task_id is None: + raise KeyError(f"Workflow key not found: {workflow_key}") + if vine_graph_capi.vine_graph_executor_add_task_input_file( + self._c_executor, task_id, file_id, task_path + ) != 0: + raise RuntimeError(f"failed to mount input file {file_id} on task {workflow_key}") + + def add_task_output_file(self, workflow_key, file_id, task_path, is_target=False): + """Declare and mount a task-produced FileHandle.""" + task_id = self._workflow_key_to_scheduler_key.get(workflow_key) + if task_id is None: + raise KeyError(f"Workflow key not found: {workflow_key}") + if vine_graph_capi.vine_graph_executor_add_task_output_file( + self._c_executor, task_id, file_id, task_path, int(bool(is_target)) + ) != 0: + raise RuntimeError(f"failed to declare output file {file_id} on task {workflow_key}") + + def get_file_target_path(self, file_id): + """Return the manager-side path of a retrieved output file.""" + path = vine_graph_capi.vine_graph_executor_get_file_target_path(self._c_executor, file_id) + if not path: + raise RuntimeError(f"file {file_id} has no manager-side target path") + return path + + def execute(self): + """Execute the graph.""" + vine_graph_capi.vine_graph_executor_execute(self._c_executor) + + def get_makespan_us(self): + """Return the graph makespan in microseconds.""" + return vine_graph_capi.vine_graph_executor_get_makespan_us(self._c_executor) + + def get_total_recovery_tasks(self): + """Return the total number of submitted recovery tasks.""" + return vine_graph_capi.vine_graph_executor_get_total_recovery_tasks(self._c_executor) + + def get_completed_recovery_tasks(self): + """Return the number of completed recovery tasks.""" + return vine_graph_capi.vine_graph_executor_get_completed_recovery_tasks(self._c_executor) + + def delete(self): + """Delete the backing C graph.""" + vine_graph_capi.vine_graph_executor_delete(self._c_executor) + self._c_executor = None + vine_graph_capi.vine_graph_delete(self._c_graph) + self._c_graph = None diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/.gitignore b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/.gitignore new file mode 100644 index 0000000000..c18dd8d83c --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/__init__.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/__init__.py new file mode 100644 index 0000000000..4105c37eb7 --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/__init__.py @@ -0,0 +1,11 @@ +from .registration import TaskRunnerRegistration +from .execution import ( + compute_task, + run_scheduler_keys, +) + +__all__ = [ + "TaskRunnerRegistration", + "run_scheduler_keys", + "compute_task", +] diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/execution.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/execution.py new file mode 100644 index 0000000000..0f3a9f2085 --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/execution.py @@ -0,0 +1,250 @@ +# Copyright (C) 2025- The University of Notre Dame +# This software is distributed under the GNU General Public License. +# See the file COPYING for details. + + +import sys +import time +import copy +import dataclasses +from collections import defaultdict, deque + +from ndcctools.taskvine.utils import load_variable_from_library +from ..workflow import _TaskOutputAttribute + + +def _resolve_nested_legacy_tasks(obj, memo=None): + """Evaluate nested legacy Dask tasks encoded as plain ``(func, *args)`` tuples.""" + if memo is None: + memo = {} + + if obj is None or isinstance(obj, (str, bytes, bytearray, memoryview, int, float, bool)): + return obj + + oid = id(obj) + if oid in memo: + return memo[oid] + + if type(obj) is tuple and obj and callable(obj[0]): + func = obj[0] + args = tuple(_resolve_nested_legacy_tasks(v, memo) for v in obj[1:]) + result = func(*args) + memo[oid] = result + return result + + if isinstance(obj, list): + out = [] + memo[oid] = out + out.extend(_resolve_nested_legacy_tasks(v, memo) for v in obj) + return out + + if isinstance(obj, deque): + out = deque(maxlen=obj.maxlen) + memo[oid] = out + out.extend(_resolve_nested_legacy_tasks(v, memo) for v in obj) + return out + + if type(obj) is tuple: + out = tuple(_resolve_nested_legacy_tasks(v, memo) for v in obj) + memo[oid] = out + return out + + if isinstance(obj, tuple) and hasattr(obj, "_fields"): + out = obj.__class__(*(_resolve_nested_legacy_tasks(v, memo) for v in obj)) + memo[oid] = out + return out + + if isinstance(obj, dict): + try: + out = copy.copy(obj) + out.clear() + except Exception: + out = {} + memo[oid] = out + for key, value in obj.items(): + out[key] = _resolve_nested_legacy_tasks(value, memo) + return out + + if isinstance(obj, set): + out = set() + memo[oid] = out + out.update(_resolve_nested_legacy_tasks(v, memo) for v in obj) + return out + + if isinstance(obj, frozenset): + out = frozenset(_resolve_nested_legacy_tasks(v, memo) for v in obj) + memo[oid] = out + return out + + if dataclasses.is_dataclass(obj) and not isinstance(obj, type): + out = copy.copy(obj) + memo[oid] = out + for field in dataclasses.fields(obj): + object.__setattr__(out, field.name, _resolve_nested_legacy_tasks(getattr(obj, field.name), memo)) + return out + + return obj + + +def compute_task(workflow, task_expr): + func_id, args, kwargs = task_expr + func = workflow.callables[func_id] + cache = {} + + def _follow_path(value, path): + current = value + for token in path: + if isinstance(token, _TaskOutputAttribute): + current = getattr(current, token.name) + elif isinstance(current, (list, tuple)): + current = current[token] + elif isinstance(current, dict): + current = current[token] + else: + current = getattr(current, token) + return current + + def on_ref(r): + if r.task_id not in cache: + x = workflow.load_task_output(r.task_id) + cache[r.task_id] = x + else: + x = cache[r.task_id] + if r.path: + return _follow_path(x, r.path) + return x + + def on_file(f): + if f.workflow_id != workflow._workflow_id: + raise ValueError("file belongs to a different Workflow") + return workflow.file_input_path(f.file_id) + + r_args, r_kwargs = workflow._visit_task_output_refs( + (args, kwargs), on_ref, rewrite=True, on_file=on_file + ) + + r_args, r_kwargs = _resolve_nested_legacy_tasks((r_args, r_kwargs)) + + return func(*r_args, **r_kwargs) + + +def topo_sort_group_scheduler_keys(workflow, member_scheduler_keys): + """Return a topological order of scheduler keys for one batched task-runner call. + + Edges count only when both endpoints lie in the batch. Ties are broken by the order keys + appear in the argument list which usually matches the CSV written by C with the leader first. + """ + # Kahn tie-break uses the order keys were listed because plain sets forget that order. + ordered = [] + seen = set() + for k in member_scheduler_keys: + kk = int(k) + if kk not in seen: + seen.add(kk) + ordered.append(kk) + + nodes = set(ordered) + order_index = {k: i for i, k in enumerate(ordered)} + adj = defaultdict(set) + indeg = {k: 0 for k in nodes} + + # Restrict Workflow to this batch: edge parent_sk -> child_sk only when both endpoints are members. + # Aligns with the C executor DAG via the same parents_of as Python Workflow. + for child_sk in nodes: + wk_child = workflow.scheduler_key_to_task_id[child_sk] + for wk_parent in workflow.parents_of.get(wk_child, ()): + parent_sk = workflow.task_id_to_scheduler_key[wk_parent] + if parent_sk in nodes and child_sk not in adj[parent_sk]: + adj[parent_sk].add(child_sk) + indeg[child_sk] += 1 + + # Kahn: serial execution order within this task-runner call. + q = deque(k for k in ordered if indeg[k] == 0) + out = [] + while q: + u = q.popleft() + out.append(u) + ready_next = [] + for v in adj[u]: + indeg[v] -= 1 + if indeg[v] == 0: + ready_next.append(v) + ready_next.sort(key=lambda x: order_index[x]) + for v in ready_next: + q.append(v) + if len(out) != len(nodes): + raise ValueError("task_group: cycle among members per Workflow graph") + return out + + +def run_single_workflow_node(workflow, scheduler_key): + """Run one node: scheduler_key -> workflow_key, execute, write outfile for downstream refs.""" + workflow_key = workflow.scheduler_key_to_task_id[scheduler_key] + task_expr = workflow.task_dict[workflow_key] + + output = compute_task(workflow, task_expr) + + time.sleep(workflow.extra_task_sleep_time[workflow_key]) + + workflow.save_task_output(workflow_key, output) + + +def _scheduler_keys_spec_to_list(scheduler_keys_spec): + """ + Normalize task-runner input payload to a list of integer scheduler keys. + Primary wire form is one comma-separated string (e.g. ``\"1,2,3\"``). + Also accepts a bare int (legacy JSON) or a list of ints/strings from json. + """ + if isinstance(scheduler_keys_spec, str): + parts = [p.strip() for p in scheduler_keys_spec.split(",") if p.strip()] + if not parts: + raise ValueError("run_scheduler_keys: empty scheduler key list after parsing") + return [int(p, 10) for p in parts] + if isinstance(scheduler_keys_spec, int): + return [scheduler_keys_spec] + if isinstance(scheduler_keys_spec, list): + if not scheduler_keys_spec: + raise ValueError("run_scheduler_keys: empty list") + return [int(x) for x in scheduler_keys_spec] + raise TypeError( + f"run_scheduler_keys: expected str, int, or list of keys, got {type(scheduler_keys_spec).__name__}" + ) + + +def _workflow_from_task_runner_context(): + """ + Resolve the Workflow from the TaskVine function context. + + The normal path is ``ndcctools.taskvine.utils.load_variable_from_library``. + Keep a ``__main__`` lookup first so generated function scripts that inject + context directly into their own module namespace also work. + """ + main = sys.modules.get("__main__") + g = getattr(main, "graph", None) if main is not None else None + if g is not None: + return g + return load_variable_from_library("graph") + + +def run_scheduler_keys(scheduler_keys_spec): + """Task runner entry that parses keys, orders them, runs each node, and writes outfiles.""" + workflow = _workflow_from_task_runner_context() + keys = _scheduler_keys_spec_to_list(scheduler_keys_spec) + ordered = topo_sort_group_scheduler_keys(workflow, keys) + leader_sk = keys[0] + # The infile from C lists the leader first. Kahn might pick another indeg-zero key first which + # would reorder writes and confuse executor validation, so move the leader to the front + # when it has no parent that is also inside this batch. + if ordered[0] != leader_sk: + batch = set(ordered) + wk_leader = workflow.scheduler_key_to_task_id[leader_sk] + for wkp in workflow.parents_of.get(wk_leader, ()): + psk = workflow.task_id_to_scheduler_key[wkp] + if psk in batch: + raise ValueError( + f"run_scheduler_keys: leader {leader_sk} must run first but has intra-batch parent {psk}" + ) + ordered = [leader_sk] + [sk for sk in ordered if sk != leader_sk] + + for sk in ordered: + run_single_workflow_node(workflow, sk) diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/registration.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/registration.py new file mode 100644 index 0000000000..9dab23c5af --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/task_runner/registration.py @@ -0,0 +1,89 @@ +# Copyright (C) 2025- The University of Notre Dame +# This software is distributed under the GNU General Public License. +# See the file COPYING for details. + +import os +import uuid +import cloudpickle +import copy +import dataclasses +import types +import time +import random +import hashlib +import collections +import collections.abc + +from ..workflow import FileHandle, TaskHandle, TaskOutputHandle, TaskOutputWrapper, Workflow, _TaskOutputAttribute +from .execution import run_scheduler_keys +from ndcctools.taskvine.utils import load_variable_from_library + + +class TaskRunnerRegistration: + def __init__(self, vine_graph): + self.vine_graph = vine_graph + + self.name = None + self.cores = None + + self.task = None + + # These modules are included in the generated function context so task calls can execute directly. + self.hoisting_modules = [ + os, cloudpickle, copy, dataclasses, uuid, hashlib, random, types, collections, collections.abc, time, + Workflow, FileHandle, TaskHandle, TaskOutputHandle, TaskOutputWrapper, _TaskOutputAttribute, + load_variable_from_library, run_scheduler_keys + ] + + # Environment files are sent with the task runner context and exposed under their remote paths. + self.env_files = {} + + # The context loader rebuilds the Workflow object in the remote function context. + self.context_loader_func = None + self.context_loader_args = [] + self.context_loader_kwargs = {} + + self.local_path = None + self.remote_path = None + + def set_cores(self, cores): + self.cores = cores + + def set_name(self, name): + self.name = name + + def add_hoisting_modules(self, new_modules): + assert isinstance(new_modules, list), "new_modules must be a list of modules" + self.hoisting_modules.extend(new_modules) + + def add_env_files(self, new_env_files): + assert isinstance(new_env_files, dict), "new_env_files must be a dictionary" + self.env_files.update(new_env_files) + + def set_context_loader(self, context_loader_func, context_loader_args=[], context_loader_kwargs={}): + self.context_loader_func = context_loader_func + self.context_loader_args = context_loader_args + self.context_loader_kwargs = context_loader_kwargs + + def install(self): + assert self.name is not None, "Task runner name must be set before installing (use set_name method)" + assert self.cores is not None, "Task runner cores must be set before installing (use set_cores method)" + + self.task = self.vine_graph.create_library_from_functions( + self.name, + run_scheduler_keys, + library_context_info=[self.context_loader_func, self.context_loader_args, self.context_loader_kwargs], + add_env=False, + function_infile_load_mode="json", + hoisting_modules=self.hoisting_modules, + ) + for local, remote in self.env_files.items(): + if not os.path.exists(local): + raise FileNotFoundError(f"Local file {local} not found") + self.task.add_input(self.vine_graph.declare_file(local, cache=True, peer_transfer=True), remote) + self.task.set_cores(self.cores) + self.task.set_function_slots(self.cores) + self.vine_graph.install_library(self.task) + + def uninstall(self): + self.vine_graph.remove_library(self.name) diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/utils.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/utils.py new file mode 100644 index 0000000000..31e156434d --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/utils.py @@ -0,0 +1,25 @@ +import cloudpickle +import os + + +def context_loader_func(graph_pkl): + graph = cloudpickle.loads(graph_pkl) + return {"graph": graph} + + +def remove_tree_contents(root_dir): + """Remove files under the run-info template directory.""" + if not os.path.exists(root_dir): + return + for dirpath, dirnames, filenames in os.walk(root_dir): + for filename in filenames: + file_path = os.path.join(dirpath, filename) + try: + os.remove(file_path) + except FileNotFoundError: + print(f"Failed to delete file {file_path}") + + +def color_text(text, color_code): + """Return text wrapped in an ANSI color code.""" + return f"\033[{color_code}m{text}\033[0m" diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/vine_graph.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/vine_graph.py new file mode 100644 index 0000000000..0edbff9b3b --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/vine_graph.py @@ -0,0 +1,442 @@ +# Copyright (C) 2025 The University of Notre Dame +# This software is distributed under the GNU General Public License. +# See the file COPYING for details. + +from ndcctools.taskvine.manager import Manager + +from .adaptors import VineGraphDaskAdaptor +from .task_runner import TaskRunnerRegistration, compute_task, run_scheduler_keys +from .workflow import FileHandle, Workflow, TaskHandle, TaskOutputHandle, TaskOutputWrapper +from .capi_bridge import VineGraphCapiBridge +from .utils import color_text, context_loader_func, remove_tree_contents + +import cloudpickle +import os +import random +import signal +import sys +import time + + +class VineGraphConfig: + def __init__(self): + """Store VineGraph configuration by the layer that consumes it.""" + self.manager_tuning = { + "worker-source-max-transfers": 100, + "max-retrievals": -1, + "prefer-dispatch": 1, + "transient-error-interval": 1, + "attempt-schedule-depth": 10000, + "temp-replica-count": 1, + "enforce-worker-eviction-interval": -1, + "shift-disk-load": 0, + "clean-redundant-replicas": 0, + } + self.executor_tuning = { + "failure-injection-step-percent": -1, + "task-priority-mode": "largest-input-first", + "prune-depth": 1, + "output-dir": "./outputs", + "checkpoint-dir": "./checkpoints", + "checkpoint-fraction": 0, + "progress-bar-update-interval-sec": 0.1, + "print-graph-details": 0, + } + self.task_runner = { + "libcores": 16, + } + self.execution = { + "schedule": "worst", + "extra-task-output-size-mb": [0, 0], + "extra-task-sleep-time": [0, 0], + # 1 = run Workflow in-process (topological order), no workers / no task runner library; stdout stays on frontend. + "local-execute": 0, + # 1 = merge maximal linear chains into supernodes before finalize (vine_graph_group_chain_like_tasks). + "task-group": 0, + } + + def _sections(self): + return ( + self.manager_tuning, + self.executor_tuning, + self.task_runner, + self.execution, + ) + + def update_param(self, param_name, new_value): + """Update one parameter.""" + for section in self._sections(): + if param_name in section: + section[param_name] = new_value + return + # Unknown parameters are assumed to be TaskVine manager tuning knobs. + self.manager_tuning[param_name] = new_value + + def get_value_of(self, param_name): + """Return the current value for a parameter.""" + for section in self._sections(): + if param_name in section: + return section[param_name] + raise ValueError(f"Invalid param name: {param_name}") + + +class VineGraph(Manager): + def __init__(self, + *args, + **kwargs): + """Create a VineGraph manager.""" + + signal.signal(signal.SIGINT, self._on_sigint) + + self.params = VineGraphConfig() + + run_info_path = kwargs.get("run_info_path", None) + run_info_template = kwargs.get("run_info_template", None) + + self.run_info_template_path = None + if run_info_path and run_info_template: + self.run_info_template_path = os.path.join(run_info_path, run_info_template) + if self.run_info_template_path: + remove_tree_contents(self.run_info_template_path) + + # Manager lifetime is tied to this object. + super().__init__(*args, **kwargs) + + print(f"=== Manager name: {color_text(self.name, 92)}") + print(f"=== Manager port: {color_text(self.port, 92)}") + print(f"=== Runtime directory: {color_text(self.runtime_directory, 92)}") + self._sigint_received = False + + def get_param(self, param_name): + """Return a parameter value.""" + return self.params.get_value_of(param_name) + + def set_params(self, new_params): + """Apply a batch of parameter overrides.""" + assert isinstance(new_params, dict), "new_params must be a dict" + for k, new_v in new_params.items(): + self.params.update_param(k, new_v) + + def tune_manager(self): + """Apply manager-side tuning.""" + for k, v in self.params.manager_tuning.items(): + try: + self.tune(k, v) + except Exception: + raise ValueError(f"Unrecognized parameter: {k}") + + def tune_capi_bridge(self, bridge): + """Apply C API bridge tuning.""" + for k, v in self.params.executor_tuning.items(): + bridge.tune(k, str(v)) + + def _rep_key(self, k, r): + return k if r == 0 else ("__rep", r, k) + + def _replicate_graph(self, task_dict, target_keys, repeats): + if repeats <= 1: + return task_dict, target_keys + if isinstance(task_dict, Workflow): + old_workflow = task_dict + if old_workflow.input_files or old_workflow.output_files: + raise ValueError("repeats cannot be combined with FileHandle dependencies") + new_workflow = Workflow() + new_workflow.callables = list(old_workflow.callables) + new_workflow._callable_index = dict(old_workflow._callable_index) + for r in range(repeats): + def rewriter(ref): + return TaskOutputHandle( + self._rep_key(ref.task_id, r), + ref.path, + workflow_id=new_workflow._workflow_id, + ) + + def _rewrite(obj): + return old_workflow._visit_task_output_refs(obj, rewriter, rewrite=True) + + for k, (func_id, args, kwargs) in old_workflow.task_dict.items(): + new_args, new_kwargs = _rewrite((args, kwargs)) + new_workflow._add_task_with_key( + self._rep_key(k, r), + old_workflow.callables[func_id], + *new_args, + **new_kwargs, + ) + new_workflow.finalize() + executor_targets = list(target_keys) + for r in range(1, repeats): + executor_targets.extend(self._rep_key(k, r) for k in target_keys if k in old_workflow.task_dict) + return new_workflow, executor_targets + else: + temp_workflow = Workflow() + expanded = {} + for r in range(repeats): + def rewriter(ref): + return TaskOutputHandle(self._rep_key(ref.task_id, r), ref.path) + + def _rewrite(obj): + return temp_workflow._visit_task_output_refs(obj, rewriter, rewrite=True) + + for k, v in task_dict.items(): + func, args, kwargs = v + new_args, new_kwargs = _rewrite((args, kwargs)) + expanded[self._rep_key(k, r)] = (func, new_args, new_kwargs) + executor_targets = list(target_keys) + for r in range(1, repeats): + executor_targets.extend(self._rep_key(k, r) for k in target_keys if k in task_dict) + return expanded, executor_targets + + def build_workflow(self, task_dict): + if isinstance(task_dict, Workflow): + workflow = task_dict + else: + workflow = Workflow() + + for k, v in task_dict.items(): + func, args, kwargs = v + assert callable(func), f"Task {k} does not have a callable" + workflow._add_task_with_key(k, func, *args, **kwargs) + + workflow.finalize() + + return workflow + + def build_capi_bridge(self, py_graph, target_keys): + """Build the C vine_graph mirror from the Python graph.""" + assert py_graph is not None, "Python graph must be built before building the VineGraphCapiBridge" + + bridge = VineGraphCapiBridge(self._taskvine) + + bridge.set_task_runner_function(run_scheduler_keys) + + self.tune_manager() + self.tune_capi_bridge(bridge) + + topo_order = py_graph.get_topological_order() + + for k in topo_order: + node_id = bridge.add_node(k) + py_graph.task_id_to_scheduler_key[k] = node_id + py_graph.scheduler_key_to_task_id[node_id] = k + for pk in py_graph.parents_of[k]: + bridge.add_dependency(pk, k) + + for k in target_keys: + bridge.set_target(k) + + return bridge + + def build_workflow_and_capi_bridge(self, task_dict, target_keys, file_target_ids=()): + """Build the Python graph and its C mirror.""" + py_graph = self.build_workflow(task_dict) + + # Ignore requested targets that are not in the graph. + missing_keys = [k for k in target_keys if k not in py_graph.task_dict] + if missing_keys: + print(f"=== Warning: the following target keys are not in the graph: {','.join(map(str, missing_keys))}") + target_keys = list(set(target_keys) - set(missing_keys)) + + bridge = self.build_capi_bridge(py_graph, target_keys) + + # Declare each FileHandle once, then mount it on its producer/consumers. + for file_id, source_path in py_graph.input_files.items(): + bridge.declare_input_file(file_id, source_path) + file_target_ids = set(file_target_ids) + for file_id, (workflow_key, task_path) in py_graph.output_files.items(): + bridge.add_task_output_file( + workflow_key, file_id, task_path, is_target=file_id in file_target_ids + ) + for file_id, consumers in py_graph.file_consumers.items(): + task_path = py_graph.file_input_path(file_id) + for workflow_key in consumers: + bridge.add_task_input_file(workflow_key, file_id, task_path) + + # Matches --task-group on the Python side so the C layer knows whether merging is allowed. + bridge.tune( + "chain-grouping-enabled", + "1" if int(self.get_param("task-group")) else "0", + ) + + if int(self.get_param("task-group")): + bridge.group_chain_like_tasks() + + bridge.compute_topology_metrics() + + # Save output locations back into the Python graph after finalize may adjust checkpoint paths. + for k in py_graph.task_id_to_scheduler_key: + outfile_remote_name = bridge.get_node_outfile_remote_name(k) + py_graph.outfile_remote_name[k] = outfile_remote_name + + return py_graph, bridge + + def build_task_runner_registration(self, py_graph, bridge, hoisting_modules, env_files): + """Build the TaskVine task runner registration.""" + task_runner_registration = TaskRunnerRegistration(self) + task_runner_registration.add_hoisting_modules(hoisting_modules) + task_runner_registration.add_env_files(env_files) + task_runner_registration.set_context_loader(context_loader_func, context_loader_args=[cloudpickle.dumps(py_graph)]) + task_runner_registration.set_cores(self.get_param("libcores")) + task_runner_registration.set_name(bridge.get_task_runner_library_name()) + + return task_runner_registration + + def _print_local_progress(self, done, total, started_at): + """Print a simple local-execute progress line without external dependencies.""" + bar_width = 24 + filled = int(bar_width * done / total) if total else bar_width + bar = "#" * filled + "-" * (bar_width - filled) + percent = 100.0 * done / total if total else 100.0 + elapsed = time.time() - started_at + sys.stdout.write(f"\rExecuting Tasks [{bar}] {done}/{total} {percent:5.1f}% elapsed {elapsed:.1f}s") + if done == total: + sys.stdout.write("\n") + sys.stdout.flush() + + def _execute_workflow_local(self, py_graph): + """Run the workflow locally in topological order.""" + out_dir = os.path.abspath(self.get_param("output-dir")) + os.makedirs(out_dir, exist_ok=True) + prev_cwd = os.getcwd() + py_graph._local_execute = True + for k, remote_name in list(py_graph.outfile_remote_name.items()): + if not os.path.isabs(remote_name): + py_graph.outfile_remote_name[k] = os.path.join(out_dir, remote_name) + t0 = time.time() + try: + order = py_graph.get_topological_order() + interval = float(self.get_param("progress-bar-update-interval-sec")) + if interval <= 0: + interval = 0.1 + + n = len(order) + if n == 0: + return time.time() - t0 + + self._print_local_progress(0, n, t0) + last_update = time.time() + for i, k in enumerate(order, 1): + task_dir = os.path.join(out_dir, ".vine_graph_tasks", f"task-{py_graph.task_id_to_scheduler_key[k]}") + os.makedirs(task_dir, exist_ok=True) + remove_tree_contents(task_dir) + for task_path in py_graph.output_files_by_task.get(k, {}): + parent = os.path.dirname(os.path.join(task_dir, task_path)) + os.makedirs(parent, exist_ok=True) + os.chdir(task_dir) + try: + out = compute_task(py_graph, py_graph.task_dict[k]) + finally: + os.chdir(prev_cwd) + for task_path, file_id in py_graph.output_files_by_task.get(k, {}).items(): + local_path = os.path.abspath(os.path.join(task_dir, task_path)) + if not os.path.isfile(local_path): + raise FileNotFoundError( + f"task {k} did not produce declared output file {task_path!r}" + ) + py_graph._local_file_paths[file_id] = local_path + py_graph.save_task_output(k, out) + now = time.time() + if now - last_update >= interval or i == n: + self._print_local_progress(i, n, t0) + last_update = now + finally: + os.chdir(prev_cwd) + return time.time() - t0 + + def run( + self, + task_dict, + targets=None, + params=None, + hoisting_modules=None, + env_files=None, + from_dask=False, + expand_subgraphs=False, + repeats=1, + ): + """Build the graph, run it, and return the requested results.""" + requested_targets = list(targets or []) + params = {} if params is None else params + hoisting_modules = [] if hoisting_modules is None else hoisting_modules + env_files = {} if env_files is None else env_files + self.set_params(params) + + if from_dask: + task_dict = VineGraphDaskAdaptor(task_dict, expand_subgraphs=expand_subgraphs).converted + + file_target_ids = set() + if isinstance(task_dict, Workflow): + result_items = [] + for target in requested_targets: + if isinstance(target, TaskHandle): + result_items.append(("task", target, task_dict._task_key(target))) + elif isinstance(target, FileHandle): + if target.workflow_id != task_dict._workflow_id: + raise ValueError("file target belongs to a different Workflow") + if target.file_id not in task_dict.input_files and target.file_id not in task_dict.output_files: + raise ValueError("file target does not belong to this Workflow") + result_items.append(("file", target, target.file_id)) + if target.file_id in task_dict.output_files: + file_target_ids.add(target.file_id) + else: + raise TypeError("Workflow targets must be TaskHandle or FileHandle objects") + else: + if any(isinstance(target, (TaskHandle, FileHandle)) for target in requested_targets): + raise TypeError("TaskHandle and FileHandle targets require a Workflow") + result_items = [("task", target, target) for target in requested_targets] + + scheduler_targets = [key for kind, _, key in result_items if kind == "task"] + task_dict, scheduler_targets = self._replicate_graph(task_dict, scheduler_targets, repeats) + + py_graph, bridge = self.build_workflow_and_capi_bridge( + task_dict, scheduler_targets, file_target_ids + ) + # Optional synthetic output size / sleep for testing. + for k in py_graph.task_dict: + py_graph.extra_task_output_size_mb[k] = random.uniform(*self.get_param("extra-task-output-size-mb")) + py_graph.extra_task_sleep_time[k] = random.uniform(*self.get_param("extra-task-sleep-time")) + + local_execute = bool(self.get_param("local-execute")) + task_runner_registration = None + + try: + if local_execute: + print("=== local-execute: running Workflow in process (no workers)", flush=True) + makespan_s = self._execute_workflow_local(py_graph) + completed_recovery_tasks = 0 + else: + task_runner_registration = self.build_task_runner_registration(py_graph, bridge, hoisting_modules, env_files) + task_runner_registration.install() + bridge.execute() + makespan_s = round(bridge.get_makespan_us() / 1e6, 6) + completed_recovery_tasks = bridge.get_completed_recovery_tasks() + + total_tasks_completed = len(py_graph.task_dict) + completed_recovery_tasks + throughput_tps = round(total_tasks_completed / makespan_s, 6) if makespan_s > 0 else 0.0 + print(f"=== Makespan: {makespan_s:.6f} seconds") + print(f"=== Total tasks completed: {total_tasks_completed}") + print(f"=== Throughput: {throughput_tps:.6f} tasks/s") + + results = {} + for kind, public_target, key in result_items: + if kind == "task": + if key not in py_graph.task_dict: + continue + outfile_path = os.path.join(self.get_param("output-dir"), py_graph.outfile_remote_name[key]) + results[public_target] = TaskOutputWrapper.load_from_path(outfile_path) + elif key in py_graph.input_files: + results[public_target] = py_graph.input_files[key] + elif local_execute: + results[public_target] = py_graph._local_file_paths[key] + else: + results[public_target] = os.path.abspath(bridge.get_file_target_path(key)) + return results + finally: + try: + if task_runner_registration is not None: + task_runner_registration.uninstall() + finally: + bridge.delete() + + def _on_sigint(self, signum, frame): + self._sigint_received = True + raise KeyboardInterrupt diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/vine_graph_capi.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/vine_graph_capi.py new file mode 100644 index 0000000000..9fbb59ae1b --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/vine_graph_capi.py @@ -0,0 +1,204 @@ +# This file was automatically generated by SWIG (https://www.swig.org). +# Version 4.4.1 +# +# Do not make changes to this file unless you know what you are doing - modify +# the SWIG interface file instead. + +from sys import version_info as _swig_python_version_info +# Import the low-level C/C++ module +if getattr(globals().get("__spec__"), "parent", None) or __package__ or "." in __name__: + from . import _vine_graph_capi +else: + import _vine_graph_capi + +try: + import builtins as __builtin__ +except ImportError: + import __builtin__ + +def _swig_repr(self): + try: + strthis = "proxy of " + self.this.__repr__() + except __builtin__.Exception: + strthis = "" + return "<%s.%s; %s >" % (self.__class__.__module__, self.__class__.__name__, strthis,) + + +def _swig_setattr_nondynamic_instance_variable(set): + def set_instance_attr(self, name, value): + if name == "this": + set(self, name, value) + elif name == "thisown": + self.this.own(value) + elif hasattr(self, name) and isinstance(getattr(type(self), name), property): + set(self, name, value) + else: + raise AttributeError("You cannot add instance attributes to %s" % self) + return set_instance_attr + + +def _swig_setattr_nondynamic_class_variable(set): + def set_class_attr(cls, name, value): + if hasattr(cls, name) and not isinstance(getattr(cls, name), property): + set(cls, name, value) + else: + raise AttributeError("You cannot add class attributes to %s" % cls) + return set_class_attr + + +def _swig_add_metaclass(metaclass): + """Class decorator for adding a metaclass to a SWIG wrapped class - a slimmed down version of six.add_metaclass""" + def wrapper(cls): + return metaclass(cls.__name__, cls.__bases__, cls.__dict__.copy()) + return wrapper + + +class _SwigNonDynamicMeta(type): + """Meta class to enforce nondynamic attributes (no new attributes) for a class""" + __setattr__ = _swig_setattr_nondynamic_class_variable(type.__setattr__) + + +import cvine +class vine_graph(object): + thisown = property(lambda x: x.this.own(), lambda x, v: x.this.own(v), doc="The membership flag") + __repr__ = _swig_repr + nodes = property(_vine_graph_capi.vine_graph_nodes_get, _vine_graph_capi.vine_graph_nodes_set) + super_leader_to_members = property(_vine_graph_capi.vine_graph_super_leader_to_members_get, _vine_graph_capi.vine_graph_super_leader_to_members_set) + outfile_cachename_to_node = property(_vine_graph_capi.vine_graph_outfile_cachename_to_node_get, _vine_graph_capi.vine_graph_outfile_cachename_to_node_set) + file_id_to_file = property(_vine_graph_capi.vine_graph_file_id_to_file_get, _vine_graph_capi.vine_graph_file_id_to_file_set) + supernode_leader_child_to_input_source = property(_vine_graph_capi.vine_graph_supernode_leader_child_to_input_source_get, _vine_graph_capi.vine_graph_supernode_leader_child_to_input_source_set) + checkpoint_dir = property(_vine_graph_capi.vine_graph_checkpoint_dir_get, _vine_graph_capi.vine_graph_checkpoint_dir_set) + output_dir = property(_vine_graph_capi.vine_graph_output_dir_get, _vine_graph_capi.vine_graph_output_dir_set) + task_runner_library_name = property(_vine_graph_capi.vine_graph_task_runner_library_name_get, _vine_graph_capi.vine_graph_task_runner_library_name_set) + task_runner_function_name = property(_vine_graph_capi.vine_graph_task_runner_function_name_get, _vine_graph_capi.vine_graph_task_runner_function_name_set) + checkpoint_fraction = property(_vine_graph_capi.vine_graph_checkpoint_fraction_get, _vine_graph_capi.vine_graph_checkpoint_fraction_set) + prune_depth = property(_vine_graph_capi.vine_graph_prune_depth_get, _vine_graph_capi.vine_graph_prune_depth_set) + print_graph_details = property(_vine_graph_capi.vine_graph_print_graph_details_get, _vine_graph_capi.vine_graph_print_graph_details_set) + chain_grouping_enabled = property(_vine_graph_capi.vine_graph_chain_grouping_enabled_get, _vine_graph_capi.vine_graph_chain_grouping_enabled_set) + + def __init__(self): + _vine_graph_capi.vine_graph_swiginit(self, _vine_graph_capi.new_vine_graph()) + __swig_destroy__ = _vine_graph_capi.delete_vine_graph + +# Register vine_graph in _vine_graph_capi: +_vine_graph_capi.vine_graph_swigregister(vine_graph) + +def vine_graph_input_producer_node(g, parent, child): + return _vine_graph_capi.vine_graph_input_producer_node(g, parent, child) + +def vine_graph_create(runtime_dir): + return _vine_graph_capi.vine_graph_create(runtime_dir) + +def vine_graph_add_node(g): + return _vine_graph_capi.vine_graph_add_node(g) + +def vine_graph_set_target(g, node_id): + return _vine_graph_capi.vine_graph_set_target(g, node_id) + +def vine_graph_add_dependency(g, parent_id, child_id): + return _vine_graph_capi.vine_graph_add_dependency(g, parent_id, child_id) + +def vine_graph_finalize(g): + return _vine_graph_capi.vine_graph_finalize(g) + +def vine_graph_get_node_heavy_score(g, node_id): + return _vine_graph_capi.vine_graph_get_node_heavy_score(g, node_id) + +def vine_graph_get_node_outfile_remote_name(g, node_id): + return _vine_graph_capi.vine_graph_get_node_outfile_remote_name(g, node_id) + +def vine_graph_delete(g): + return _vine_graph_capi.vine_graph_delete(g) + +def vine_graph_get_task_runner_library_name(g): + return _vine_graph_capi.vine_graph_get_task_runner_library_name(g) + +def vine_graph_set_task_runner_function_name(g, task_runner_function_name): + return _vine_graph_capi.vine_graph_set_task_runner_function_name(g, task_runner_function_name) + +def vine_graph_tune(g, name, value): + return _vine_graph_capi.vine_graph_tune(g, name, value) + +def vine_graph_node_is_supernode_leader(n): + return _vine_graph_capi.vine_graph_node_is_supernode_leader(n) + +def vine_graph_supernode_leader_node(g, n): + return _vine_graph_capi.vine_graph_supernode_leader_node(g, n) + +def vine_graph_group_chain_like_tasks(g): + return _vine_graph_capi.vine_graph_group_chain_like_tasks(g) +TASK_PRIORITY_MODE_RANDOM = _vine_graph_capi.TASK_PRIORITY_MODE_RANDOM +TASK_PRIORITY_MODE_DEPTH_FIRST = _vine_graph_capi.TASK_PRIORITY_MODE_DEPTH_FIRST +TASK_PRIORITY_MODE_BREADTH_FIRST = _vine_graph_capi.TASK_PRIORITY_MODE_BREADTH_FIRST +TASK_PRIORITY_MODE_FIFO = _vine_graph_capi.TASK_PRIORITY_MODE_FIFO +TASK_PRIORITY_MODE_LIFO = _vine_graph_capi.TASK_PRIORITY_MODE_LIFO +TASK_PRIORITY_MODE_LARGEST_INPUT_FIRST = _vine_graph_capi.TASK_PRIORITY_MODE_LARGEST_INPUT_FIRST +TASK_PRIORITY_MODE_LARGEST_STORAGE_FOOTPRINT_FIRST = _vine_graph_capi.TASK_PRIORITY_MODE_LARGEST_STORAGE_FOOTPRINT_FIRST +class vine_graph_executor(object): + thisown = property(lambda x: x.this.own(), lambda x, v: x.this.own(v), doc="The membership flag") + __repr__ = _swig_repr + graph = property(_vine_graph_capi.vine_graph_executor_graph_get, _vine_graph_capi.vine_graph_executor_graph_set) + manager = property(_vine_graph_capi.vine_graph_executor_manager_get, _vine_graph_capi.vine_graph_executor_manager_set) + task_id_to_node = property(_vine_graph_capi.vine_graph_executor_task_id_to_node_get, _vine_graph_capi.vine_graph_executor_task_id_to_node_set) + resubmit_queue = property(_vine_graph_capi.vine_graph_executor_resubmit_queue_get, _vine_graph_capi.vine_graph_executor_resubmit_queue_set) + time_first_task_dispatched = property(_vine_graph_capi.vine_graph_executor_time_first_task_dispatched_get, _vine_graph_capi.vine_graph_executor_time_first_task_dispatched_set) + time_last_task_retrieved = property(_vine_graph_capi.vine_graph_executor_time_last_task_retrieved_get, _vine_graph_capi.vine_graph_executor_time_last_task_retrieved_set) + makespan_us = property(_vine_graph_capi.vine_graph_executor_makespan_us_get, _vine_graph_capi.vine_graph_executor_makespan_us_set) + time_spent_on_cut_propagation = property(_vine_graph_capi.vine_graph_executor_time_spent_on_cut_propagation_get, _vine_graph_capi.vine_graph_executor_time_spent_on_cut_propagation_set) + completed_recovery_tasks = property(_vine_graph_capi.vine_graph_executor_completed_recovery_tasks_get, _vine_graph_capi.vine_graph_executor_completed_recovery_tasks_set) + pfs_usage_bytes = property(_vine_graph_capi.vine_graph_executor_pfs_usage_bytes_get, _vine_graph_capi.vine_graph_executor_pfs_usage_bytes_set) + total_preprocessing_time_us = property(_vine_graph_capi.vine_graph_executor_total_preprocessing_time_us_get, _vine_graph_capi.vine_graph_executor_total_preprocessing_time_us_set) + total_postprocessing_time_us = property(_vine_graph_capi.vine_graph_executor_total_postprocessing_time_us_get, _vine_graph_capi.vine_graph_executor_total_postprocessing_time_us_set) + task_priority_mode = property(_vine_graph_capi.vine_graph_executor_task_priority_mode_get, _vine_graph_capi.vine_graph_executor_task_priority_mode_set) + failure_injection_step_percent = property(_vine_graph_capi.vine_graph_executor_failure_injection_step_percent_get, _vine_graph_capi.vine_graph_executor_failure_injection_step_percent_set) + progress_bar_update_interval_sec = property(_vine_graph_capi.vine_graph_executor_progress_bar_update_interval_sec_get, _vine_graph_capi.vine_graph_executor_progress_bar_update_interval_sec_set) + + def __init__(self): + _vine_graph_capi.vine_graph_executor_swiginit(self, _vine_graph_capi.new_vine_graph_executor()) + __swig_destroy__ = _vine_graph_capi.delete_vine_graph_executor + +# Register vine_graph_executor in _vine_graph_capi: +_vine_graph_capi.vine_graph_executor_swigregister(vine_graph_executor) + +def vine_graph_executor_create(manager, graph): + return _vine_graph_capi.vine_graph_executor_create(manager, graph) + +def vine_graph_executor_create_graph(manager): + return _vine_graph_capi.vine_graph_executor_create_graph(manager) + +def vine_graph_executor_delete(e): + return _vine_graph_capi.vine_graph_executor_delete(e) + +def vine_graph_executor_add_node(e): + return _vine_graph_capi.vine_graph_executor_add_node(e) + +def vine_graph_executor_finalize(e): + return _vine_graph_capi.vine_graph_executor_finalize(e) + +def vine_graph_executor_declare_input_file(e, file_id, source_path): + return _vine_graph_capi.vine_graph_executor_declare_input_file(e, file_id, source_path) + +def vine_graph_executor_add_task_output_file(e, task_id, file_id, task_path, is_target): + return _vine_graph_capi.vine_graph_executor_add_task_output_file(e, task_id, file_id, task_path, is_target) + +def vine_graph_executor_add_task_input_file(e, task_id, file_id, task_path): + return _vine_graph_capi.vine_graph_executor_add_task_input_file(e, task_id, file_id, task_path) + +def vine_graph_executor_get_file_target_path(e, file_id): + return _vine_graph_capi.vine_graph_executor_get_file_target_path(e, file_id) + +def vine_graph_executor_tune(e, name, value): + return _vine_graph_capi.vine_graph_executor_tune(e, name, value) + +def vine_graph_executor_execute(e): + return _vine_graph_capi.vine_graph_executor_execute(e) + +def vine_graph_executor_get_makespan_us(e): + return _vine_graph_capi.vine_graph_executor_get_makespan_us(e) + +def vine_graph_executor_get_total_recovery_tasks(e): + return _vine_graph_capi.vine_graph_executor_get_total_recovery_tasks(e) + +def vine_graph_executor_get_completed_recovery_tasks(e): + return _vine_graph_capi.vine_graph_executor_get_completed_recovery_tasks(e) + diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/workflow.py b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/workflow.py new file mode 100644 index 0000000000..d9d99c7642 --- /dev/null +++ b/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph/workflow.py @@ -0,0 +1,478 @@ +# Copyright (C) 2025- The University of Notre Dame +# This software is distributed under the GNU General Public License. +# See the file COPYING for details. + +import collections +import collections.abc +import copy +import dataclasses +import cloudpickle +import os +import uuid + + +# Lightweight wrapper around task results that optionally pads the payload. The +# padding lets tests model large outputs without altering the logical result. +class TaskOutputWrapper: + def __init__(self, result, extra_size_mb=None): + self.result = result + self.extra_obj = bytearray(int(extra_size_mb * 1024 * 1024)) if extra_size_mb and extra_size_mb > 0 else None + + @staticmethod + def load_from_path(path): + try: + with open(path, "rb") as f: + result_obj = cloudpickle.load(f) + assert isinstance(result_obj, TaskOutputWrapper), "Loaded object is not of type TaskOutputWrapper" + return result_obj.result + except FileNotFoundError: + raise FileNotFoundError(f"Task result file not found at {path}") + + +# A reference to a task output. This is used to represent the output of a task as a dependency of another task. +class TaskOutputHandle: + """Symbolic handle to a task result or a value projected from it.""" + + __slots__ = ("_workflow_id", "_task_id", "_path") + + def __init__(self, task_id, path=(), *, workflow_id=None): + self._workflow_id = workflow_id + self._task_id = task_id + self._path = tuple(path) + + @property + def workflow_id(self): + return self._workflow_id + + @property + def task_id(self): + return self._task_id + + @property + def path(self): + return self._path + + def __getitem__(self, key): + # One [] operation always appends one path token. Consequently, + # output["a"]["b"] is a two-level lookup while output[("a", "b")] + # looks up one tuple dictionary key, matching normal Python semantics. + return TaskOutputHandle( + self._task_id, + self._path + (key,), + workflow_id=self._workflow_id, + ) + + def attr(self, name): + """Select an attribute from this output without overloading dictionary lookup.""" + if not isinstance(name, str): + raise TypeError("attribute name must be a string") + return TaskOutputHandle( + self._task_id, + self._path + (_TaskOutputAttribute(name),), + workflow_id=self._workflow_id, + ) + + +@dataclasses.dataclass(frozen=True) +class _TaskOutputAttribute: + name: str + + +class FileHandle: + """Symbolic handle to an external input or a task-produced file.""" + + __slots__ = ("_workflow_id", "_file_id") + + def __init__(self, workflow_id, file_id): + self._workflow_id = workflow_id + self._file_id = file_id + + @property + def workflow_id(self): + return self._workflow_id + + @property + def file_id(self): + return self._file_id + + def __repr__(self): + return f"FileHandle(id={self._file_id})" + + +class TaskHandle: + """User-facing symbolic handle returned by Workflow.add_task().""" + + __slots__ = ("_workflow", "_task_id") + + def __init__(self, workflow, task_id): + self._workflow = workflow + self._task_id = task_id + + def output(self): + """Return a symbolic reference to this task's Python result.""" + return TaskOutputHandle(self._task_id, workflow_id=self._workflow._workflow_id) + + def file(self, path): + """Return a handle for a file written in this task's sandbox.""" + return self._workflow._declare_task_file(self, path) + + def __repr__(self): + return f"TaskHandle(id={self._task_id})" + + +# The Workflow is a directed acyclic graph (DAG) that represents the logical dependencies between tasks. +# It is used to build the C executor graph. +class Workflow: + + _LEAF_TYPES = (str, bytes, bytearray, memoryview, int, float, bool, type(None)) + + def __init__(self): + self._workflow_id = uuid.uuid4().hex + self._next_task_id = 1 + self._next_file_id = 1 + self.callables = [] + self._callable_index = {} + + self.task_dict = {} + + self.parents_of = collections.defaultdict(set) # workflow_key -> set of workflow_keys + self.children_of = collections.defaultdict(set) # workflow_key -> set of workflow_keys + + self.input_files = {} # file_id -> absolute frontend path + self.output_files = {} # file_id -> (producer task id, task-relative path) + self.output_files_by_task = collections.defaultdict(dict) # task id -> relative path -> file_id + self.file_consumers = collections.defaultdict(set) # file_id -> consumer task ids + self._local_execute = False + self._local_file_paths = {} + + self.outfile_remote_name = collections.defaultdict(lambda: None) # workflow_key -> remote outfile name, will be set by the executor graph + + self.task_id_to_scheduler_key = {} # workflow_key -> scheduler key (C node id) + self.scheduler_key_to_task_id = {} # scheduler key -> workflow_key + + self.extra_task_output_size_mb = {} # workflow_key -> extra size in MB + self.extra_task_sleep_time = {} # workflow_key -> extra sleep time in seconds + + def _intern_callable(self, func): + idx = self._callable_index.get(func) + if idx is None: + idx = len(self.callables) + self.callables.append(func) + self._callable_index[func] = idx + return idx + + def _visit_task_output_refs(self, obj, on_ref, *, rewrite: bool, on_file=None): + memo = {} + active_immutable = set() + + def dict_key_contains_ref(key, seen): + if isinstance(key, (TaskOutputHandle, TaskHandle, FileHandle)): + return True + if key is None or isinstance(key, self._LEAF_TYPES): + return False + oid = id(key) + if oid in seen: + return False + seen.add(oid) + if dataclasses.is_dataclass(key) and not isinstance(key, type): + return any(dict_key_contains_ref(getattr(key, f.name), seen) for f in dataclasses.fields(key)) + if isinstance(key, collections.abc.Mapping): + return any( + dict_key_contains_ref(k, seen) or dict_key_contains_ref(v, seen) + for k, v in key.items() + ) + if isinstance(key, (list, tuple, set, frozenset, collections.deque)): + return any(dict_key_contains_ref(v, seen) for v in key) + try: + state = vars(key) + except TypeError: + state = None + if state and any(dict_key_contains_ref(v, seen) for v in state.values()): + return True + for slot in getattr(type(key), "__slots__", ()): + if isinstance(slot, str) and hasattr(key, slot): + if dict_key_contains_ref(getattr(key, slot), seen): + return True + return False + + def rec(x): + if isinstance(x, TaskOutputHandle): + return on_ref(x) + + if isinstance(x, FileHandle): + if on_file is None: + raise TypeError("FileHandle is not supported in this graph conversion") + return on_file(x) + + if isinstance(x, TaskHandle): + raise TypeError("TaskHandle cannot be used as an argument directly; use task.output()") + + if x is None or isinstance(x, self._LEAF_TYPES): + return x if rewrite else None + + oid = id(x) + if oid in memo: + return memo[oid] if rewrite else None + + if not rewrite: + memo[oid] = None + + if isinstance(x, collections.abc.Mapping): + for k in x.keys(): + if dict_key_contains_ref(k, set()): + raise ValueError("dependency handles cannot be used as dict keys") + if not rewrite: + for v in x.values(): + rec(v) + return None + + # copy+clear preserves dict subclasses and defaultdict factories, + # while publishing the empty object first preserves aliases/cycles. + try: + out = copy.copy(x) + out.clear() + except Exception: + out = {} + memo[oid] = out + for k, v in x.items(): + out[k] = rec(v) + return out + + if isinstance(x, list): + if not rewrite: + for v in x: + rec(v) + return None + out = [] + memo[oid] = out + out.extend(rec(v) for v in x) + return out + + if isinstance(x, collections.deque): + if not rewrite: + for v in x: + rec(v) + return None + out = collections.deque(maxlen=x.maxlen) + memo[oid] = out + out.extend(rec(v) for v in x) + return out + + if isinstance(x, set): + if not rewrite: + for v in x: + rec(v) + return None + out = set() + memo[oid] = out + out.update(rec(v) for v in x) + return out + + if isinstance(x, tuple) and hasattr(x, "_fields"): # namedtuple + if not rewrite: + for v in x: + rec(v) + return None + if oid in active_immutable: + raise ValueError("cyclic dependency containers involving namedtuple are not supported") + active_immutable.add(oid) + try: + out = x.__class__(*(rec(v) for v in x)) + finally: + active_immutable.remove(oid) + memo[oid] = out + return out + + if isinstance(x, (tuple, frozenset)): + if not rewrite: + for v in x: + rec(v) + return None + if oid in active_immutable: + raise ValueError("cyclic dependency containers involving immutable containers are not supported") + active_immutable.add(oid) + try: + values = [rec(v) for v in x] + out = tuple(values) if isinstance(x, tuple) else frozenset(values) + finally: + active_immutable.remove(oid) + memo[oid] = out + return out + + if dataclasses.is_dataclass(x) and not isinstance(x, type): + if not rewrite: + for field in dataclasses.fields(x): + rec(getattr(x, field.name)) + return None + out = copy.copy(x) + memo[oid] = out + for field in dataclasses.fields(x): + object.__setattr__(out, field.name, rec(getattr(x, field.name))) + return out + + if dict_key_contains_ref(x, set()): + raise TypeError( + "dependency handle inside an arbitrary custom object is not supported; " + "use a dataclass or a supported container" + ) + + return x if rewrite else None + + return rec(obj) + + def _find_dependencies(self, obj): + parents = set() + files = set() + + def on_ref(r): + if r.workflow_id is not None and r.workflow_id != self._workflow_id: + raise ValueError("task output belongs to a different Workflow") + parents.add(r.task_id) + return None + + def on_file(f): + if f.workflow_id != self._workflow_id: + raise ValueError("file belongs to a different Workflow") + if f.file_id not in self.input_files and f.file_id not in self.output_files: + raise ValueError("file does not belong to this Workflow") + files.add(f.file_id) + if f.file_id in self.output_files: + producer_task_id = self.output_files[f.file_id][0] + parents.add(producer_task_id) + return None + + self._visit_task_output_refs(obj, on_ref, rewrite=False, on_file=on_file) + return parents, files + + def _allocate_task_id(self): + while self._next_task_id in self.task_dict: + self._next_task_id += 1 + task_id = self._next_task_id + self._next_task_id += 1 + return task_id + + def _add_task_with_key(self, workflow_key, func, *args, **kwargs): + if workflow_key in self.task_dict: + raise ValueError(f"Task {workflow_key} already exists") + if not callable(func): + raise TypeError("task function must be callable") + + func_id = self._intern_callable(func) + self.task_dict[workflow_key] = (func_id, args, kwargs) + + parents, files = self._find_dependencies((args, kwargs)) + + for parent in parents: + self.parents_of[workflow_key].add(parent) + self.children_of[parent].add(workflow_key) + + for file_id in files: + self.file_consumers[file_id].add(workflow_key) + + return TaskHandle(self, workflow_key) + + def add_task(self, func, *args, **kwargs): + """Add a task with a private generated id and return its TaskHandle.""" + if not callable(func): + raise TypeError("add_task expects a callable followed by its arguments") + return self._add_task_with_key(self._allocate_task_id(), func, *args, **kwargs) + + def _task_key(self, task): + if not isinstance(task, TaskHandle): + raise TypeError("expected a TaskHandle from this Workflow") + if task._workflow is not self: + raise ValueError("task belongs to a different Workflow") + return task._task_id + + def _task_handle(self, task_id): + """Return a handle for an internal/adaptor task id.""" + if task_id not in self.task_dict: + raise KeyError(f"Task {task_id} does not exist") + return TaskHandle(self, task_id) + + def _allocate_file_id(self): + file_id = self._next_file_id + self._next_file_id += 1 + return file_id + + def file(self, path): + """Return a handle for an existing file in the frontend filesystem.""" + source_path = os.fspath(path) + if not isinstance(source_path, str): + raise TypeError("input file path must be a string or path-like object") + if "\0" in source_path: + raise ValueError("input file path contains a null byte") + source_path = os.path.abspath(source_path) + if not os.path.isfile(source_path): + raise FileNotFoundError(source_path) + file_id = self._allocate_file_id() + self.input_files[file_id] = source_path + return FileHandle(self._workflow_id, file_id) + + def _declare_task_file(self, task, path): + task_id = self._task_key(task) + path = os.fspath(path) + if not isinstance(path, str): + raise TypeError("output file path must be a string or path-like object") + if "\0" in path: + raise ValueError("output file path contains a null byte") + normalized = os.path.normpath(path) + if not path or normalized in ("", ".") or os.path.isabs(path) or normalized == ".." or normalized.startswith("../"): + raise ValueError("output file path must be a non-empty relative path inside the task sandbox") + if normalized in self.output_files_by_task[task_id]: + raise ValueError(f"task already declares output file {normalized!r}") + file_id = self._allocate_file_id() + self.output_files[file_id] = (task_id, normalized) + self.output_files_by_task[task_id][normalized] = file_id + return FileHandle(self._workflow_id, file_id) + + def file_input_path(self, file_id): + """Resolve a FileHandle to the path visible to the current task.""" + if self._local_execute: + if file_id in self._local_file_paths: + return self._local_file_paths[file_id] + if file_id in self.input_files: + return self.input_files[file_id] + raise RuntimeError(f"file {file_id} is not available yet") + if file_id in self.input_files: + base = os.path.basename(self.input_files[file_id]) + else: + base = os.path.basename(self.output_files[file_id][1]) + return f"vine-graph-file-{file_id}-{base}" + + def save_task_output(self, workflow_key, output): + with open(self.outfile_remote_name[workflow_key], "wb") as f: + wrapped_output = TaskOutputWrapper(output, extra_size_mb=self.extra_task_output_size_mb[workflow_key]) + cloudpickle.dump(wrapped_output, f) + + def load_task_output(self, workflow_key): + return TaskOutputWrapper.load_from_path(self.outfile_remote_name[workflow_key]) + + def get_topological_order(self): + indegree = {} + for workflow_key in self.task_dict: + indegree[workflow_key] = len(self.parents_of.get(workflow_key, ())) + + q = collections.deque(t for t, d in indegree.items() if d == 0) + order = [] + + while q: + u = q.popleft() + order.append(u) + + for v in self.children_of.get(u, ()): + indegree[v] -= 1 + if indegree[v] == 0: + q.append(v) + + if len(order) != len(self.task_dict): + raise ValueError("Graph has a cycle or missing dependencies") + + return order + + def sink_tasks(self): + """Return handles for tasks with no downstream children.""" + return [self._task_handle(key) for key in self.task_dict if not self.children_of.get(key)] + + def finalize(self): + """Finalize the workflow. Dependencies are recorded when tasks are added.""" diff --git a/taskvine/src/manager/taskvine.h b/taskvine/src/manager/taskvine.h index 205024a085..1d156490b8 100644 --- a/taskvine/src/manager/taskvine.h +++ b/taskvine/src/manager/taskvine.h @@ -282,6 +282,12 @@ This may be called on tasks after they are returned from @ref vine_wait. */ void vine_task_delete(struct vine_task *t); +/** Add a reference to an existing task object, return the same object. +@param t A task object. +@return The same task object, or null. +*/ +struct vine_task *vine_task_addref(struct vine_task *t); + /** Indicate the command to be executed. @param t A task object. @param cmd The command to be executed. This string will be duplicated by this call, so the argument may be freed or @@ -491,6 +497,12 @@ int vine_task_set_monitor_output(struct vine_task *t, const char *monitor_output const char *vine_task_get_state(struct vine_task *t); +/** Return the original task id that a recovery task is restoring. +@param t A task object. +@return The source task id for a recovery task, or zero for non-recovery tasks. +*/ +int vine_task_get_recovery_source_task_id(struct vine_task *t); + /** Get the command line of the task. @param t A task object. @return The command line set by @ref vine_task_create. @@ -717,6 +729,14 @@ has previously been called on this object. */ const char *vine_file_contents(struct vine_file *f); +/** Release a reference to a file object. +Most declared files should be released with @ref vine_undeclare_file; this function is for file references owned outside +the manager declaration table. +@param f A file object. +@return The remaining reference count, or zero. +*/ +int vine_file_delete(struct vine_file *f); + /** Get the length of a vine file. @param f A file object. @return The length of the file, or zero if unknown. @@ -736,6 +756,18 @@ const char *vine_file_source(struct vine_file *f); */ vine_file_type_t vine_file_type(struct vine_file *f); +/** Get the manager-side cached name of a file. +@param f A file object. +@return The cached name, or null. +*/ +const char *vine_file_cached_name(struct vine_file *f); + +/** Return non-zero if a file is currently being recovered. +@param f A file object. +@return Non-zero if the file is currently being recovered. +*/ +int vine_file_is_recovering(struct vine_file *f); + /** Get the number of replicas of a file. @param m A manager object @param f A file object. @@ -938,6 +970,13 @@ but is still available on the manager's site, and can be recovered by submitting */ int vine_prune_file(struct vine_manager *m, struct vine_file *f); +/** Return a declared file by its cached name, or NULL if it is unknown to the manager. +@param m A manager object. +@param cached_name The file cache name. +@return The declared file object, or null. +*/ +struct vine_file *vine_manager_lookup_file(struct vine_manager *m, const char *cached_name); + //@} /** @name Functions - Managers */ @@ -1130,6 +1169,14 @@ int vine_enable_peer_transfers(struct vine_manager *m); /** Disable taskvine peer transfers to be scheduled by the manager **/ int vine_disable_peer_transfers(struct vine_manager *m); +/** Enable external recovery handling by returning recovery tasks from vine_wait. +By default, recovery tasks are handled internally by the manager. **/ +int vine_enable_external_recovery_handling(struct vine_manager *m); + +/** Disable external recovery handling. +Recovery tasks will be handled internally by the manager. **/ +int vine_disable_external_recovery_handling(struct vine_manager *m); + /** When enabled, resources to tasks in are assigned in proportion to the size of the worker. If a resource is specified (e.g. with @ref vine_task_set_cores), proportional resources never go below explicit specifications. This mode is most diff --git a/taskvine/src/manager/vine_file.c b/taskvine/src/manager/vine_file.c index bb21bf303a..863e1aedda 100644 --- a/taskvine/src/manager/vine_file.c +++ b/taskvine/src/manager/vine_file.c @@ -373,6 +373,20 @@ const char *vine_file_source(struct vine_file *f) return f->source; } +const char *vine_file_cached_name(struct vine_file *f) +{ + return f ? f->cached_name : NULL; +} + +int vine_file_is_recovering(struct vine_file *f) +{ + if (!f || f->type != VINE_TEMP || !f->recovery_task) { + return 0; + } + + return f->recovery_task->state != VINE_TASK_INITIAL && f->recovery_task->state != VINE_TASK_DONE; +} + void vine_file_set_mode(struct vine_file *f, int mode) { /* The mode must contain, at a minimum, owner-rw (0600) (so that we can delete it) */ diff --git a/taskvine/src/manager/vine_manager.c b/taskvine/src/manager/vine_manager.c index 5d14cba017..81fe4b8d12 100644 --- a/taskvine/src/manager/vine_manager.c +++ b/taskvine/src/manager/vine_manager.c @@ -4345,6 +4345,20 @@ int vine_disable_peer_transfers(struct vine_manager *q) return 1; } +int vine_enable_external_recovery_handling(struct vine_manager *q) +{ + debug(D_VINE, "External recovery handling enabled"); + q->external_recovery_handling = 1; + return 1; +} + +int vine_disable_external_recovery_handling(struct vine_manager *q) +{ + debug(D_VINE, "External recovery handling disabled"); + q->external_recovery_handling = 0; + return 1; +} + int vine_enable_proportional_resources(struct vine_manager *q) { debug(D_VINE, "Proportional resources enabled"); @@ -5310,6 +5324,10 @@ struct vine_task *find_task_to_return(struct vine_manager *q, const char *tag, i break; case VINE_TASK_TYPE_RECOVERY: + /* If configured for external recovery handling, return it to the user. */ + if (q->external_recovery_handling) { + return t; + } /* If this is a recovery task, it is owned by the manager, @@ -6087,10 +6105,12 @@ int vine_tune(struct vine_manager *q, const char *name, double value) } else if (!strcmp(name, "max-library-retries")) { q->max_library_retries = MIN(1, value); + } else if (!strcmp(name, "disk-proportion-available-to-task")) { if (value < 1 && value > 0) { q->disk_proportion_available_to_task = value; } + } else if (!strcmp(name, "enforce-worker-eviction-interval")) { q->enforce_worker_eviction_interval = (timestamp_t)(MAX(0, (int)value) * ONE_SECOND); diff --git a/taskvine/src/manager/vine_manager.h b/taskvine/src/manager/vine_manager.h index 10bc220cdf..d5a8237690 100644 --- a/taskvine/src/manager/vine_manager.h +++ b/taskvine/src/manager/vine_manager.h @@ -222,6 +222,7 @@ struct vine_manager { int temp_replica_count; /* Number of replicas per temp file */ int clean_redundant_replicas; /* If true, remove redundant replicas of temp files to save disk space. */ int shift_disk_load; /* If true, shift storage burden to more available workers to minimize disk usage peaks. */ + int external_recovery_handling; /* If true, recovery tasks are returned by vine_wait to the user. By default they are handled internally. */ double resource_submit_multiplier; /* Factor to permit overcommitment of resources at each worker. */ double bandwidth_limit; /* Artificial limit on bandwidth of manager<->worker transfers. */ diff --git a/taskvine/src/manager/vine_task.c b/taskvine/src/manager/vine_task.c index c401324b9f..dd8783f6b4 100644 --- a/taskvine/src/manager/vine_task.c +++ b/taskvine/src/manager/vine_task.c @@ -829,6 +829,27 @@ const char *vine_task_get_state(struct vine_task *t) return vine_task_state_to_string(t->state); } +int vine_task_get_recovery_source_task_id(struct vine_task *t) +{ + if (!t || t->type != VINE_TASK_TYPE_RECOVERY) { + return 0; + } + + /* + Recovery tasks are copies of original tasks and keep the same output file objects. + The file records the original producer so the task does not need duplicate state. + */ + struct vine_mount *m; + LIST_ITERATE(t->output_mounts, m) + { + if (m && m->file && m->file->original_producer_task_id > 0) { + return m->file->original_producer_task_id; + } + } + + return 0; +} + #define METRIC(x) \ if (!strcmp(name, #x)) \ return t->x; diff --git a/taskvine/src/vine_graph/.gitignore b/taskvine/src/vine_graph/.gitignore new file mode 100644 index 0000000000..1576058b30 --- /dev/null +++ b/taskvine/src/vine_graph/.gitignore @@ -0,0 +1,6 @@ +*.o +*.py[cod] +*.so +__pycache__/ +vine_graph_capi.py +vine_graph_wrap.c diff --git a/taskvine/src/vine_graph/Makefile b/taskvine/src/vine_graph/Makefile new file mode 100644 index 0000000000..481c715dc0 --- /dev/null +++ b/taskvine/src/vine_graph/Makefile @@ -0,0 +1,74 @@ +include ../../../config.mk +include ../../../rules.mk + +CCTOOLS_DYNAMIC_SUFFIX = so + +PROJECT_NAME := vine_graph +SOURCE_DIR := $(CCTOOLS_HOME)/taskvine/src/vine_graph +PY_MODULE_DIR := $(CCTOOLS_HOME)/taskvine/src/bindings/python3/ndcctools/taskvine/vine_graph +MODULE_DIR := $(CCTOOLS_PYTHON3_PATH)/ndcctools/taskvine/vine_graph +TEST_MODULE_DIR := $(CCTOOLS_HOME)/test_support/python_modules/$(CCTOOLS_PYTHON_TEST_DIR)/ndcctools/taskvine/vine_graph +MODULE_NAME := vine_graph_capi + +SOURCES := vine_graph_node.c vine_graph.c vine_graph_executor.c +OBJECTS := $(SOURCES:%.c=%.o) +SWIG_WRAP := vine_graph_wrap.c +WRAP_OBJ := vine_graph_wrap.o +PYMODULE := _$(MODULE_NAME).$(CCTOOLS_DYNAMIC_SUFFIX) +PYMODULE_BINDING := $(PY_MODULE_DIR)/$(PYMODULE) + +LOCAL_LINKAGE += $(CCTOOLS_HOME)/dttools/src/progress_bar.o +LOCAL_CCFLAGS += -I $(CCTOOLS_HOME)/taskvine/src/manager -I $(SOURCE_DIR) + +.PHONY: all install clean lint format test + +all: $(PYMODULE_BINDING) + +$(PYMODULE_BINDING): $(PYMODULE) + cp $(PYMODULE) $(PYMODULE_BINDING) + +$(SWIG_WRAP): vine_graph.i vine_graph.h vine_graph_executor.h + $(CCTOOLS_SWIG) -python -threads -relativeimport \ + -I$(CCTOOLS_HOME)/taskvine/src/manager \ + -I$(CCTOOLS_HOME)/dttools/src \ + -I$(SOURCE_DIR) \ + -outdir $(SOURCE_DIR) -o $@ $< + +$(WRAP_OBJ): $(SWIG_WRAP) + $(CCTOOLS_CC) -o $@ -c $(CCTOOLS_INTERNAL_CCFLAGS) $(LOCAL_CCFLAGS) $(CCTOOLS_PYTHON3_CCFLAGS) -w -fPIC -DNDEBUG $< + +%.o: %.c vine_graph.h vine_graph_executor.h vine_graph_node.h + $(CCTOOLS_CC) -o $@ -c $(CCTOOLS_INTERNAL_CCFLAGS) $(LOCAL_CCFLAGS) -fPIC -DNDEBUG $< + +$(PYMODULE): $(WRAP_OBJ) $(OBJECTS) + $(CCTOOLS_LD) -o $@ $(CCTOOLS_DYNAMIC_FLAG) $(CCTOOLS_INTERNAL_LDFLAGS) $(LOCAL_LDFLAGS) $^ $(LOCAL_LINKAGE) $(CCTOOLS_PYTHON3_LDFLAGS) $(CCTOOLS_EXTERNAL_LINKAGE) + +install: all + mkdir -p $(CCTOOLS_INSTALL_DIR)/$(PROJECT_NAME)/include + cp $(CCTOOLS_HOME)/taskvine/src/manager/taskvine.h $(CCTOOLS_INSTALL_DIR)/$(PROJECT_NAME)/include/ + mkdir -p $(MODULE_DIR) + rm -f $(MODULE_DIR)/graph.py $(MODULE_DIR)/graph_capi.py $(MODULE_DIR)/_graph_capi.$(CCTOOLS_DYNAMIC_SUFFIX) + rm -f $(MODULE_DIR)/vine_graph_capi.py $(MODULE_DIR)/_vine_graph_capi.$(CCTOOLS_DYNAMIC_SUFFIX) + rm -f $(MODULE_DIR)/capi_bridge.py + cp $(PYMODULE) $(MODULE_DIR) + cp $(MODULE_NAME).py $(MODULE_DIR) + cp $(PY_MODULE_DIR)/capi_bridge.py $(MODULE_DIR) + mkdir -p $(TEST_MODULE_DIR) + rm -f $(TEST_MODULE_DIR)/vine_graph_capi.py $(TEST_MODULE_DIR)/_vine_graph_capi.$(CCTOOLS_DYNAMIC_SUFFIX) + rm -f $(TEST_MODULE_DIR)/capi_bridge.py + cp $(PYMODULE) $(TEST_MODULE_DIR) + cp $(MODULE_NAME).py $(TEST_MODULE_DIR) + cp $(MODULE_DIR)/capi_bridge.py $(TEST_MODULE_DIR) + +clean: + rm -f *.o + rm -rf __pycache__ + rm -f $(SWIG_WRAP) $(WRAP_OBJ) $(PYMODULE) $(PYMODULE_BINDING) $(MODULE_NAME).py + +lint: + clang-format -Werror --dry-run --style='file:$(CCTOOLS_HOME)/.clang-format' $(SOURCE_DIR)/*.[ch] + +format: + clang-format -i $(SOURCE_DIR)/*.[ch] + +test: diff --git a/taskvine/src/vine_graph/vine_graph.c b/taskvine/src/vine_graph/vine_graph.c new file mode 100644 index 0000000000..e31e8a516d --- /dev/null +++ b/taskvine/src/vine_graph/vine_graph.c @@ -0,0 +1,1035 @@ +#include +#include +#include +#include +#include +#include + +#include "debug.h" +#include "vine_graph.h" +#include "priority_queue.h" +#include "set.h" +#include "stringtools.h" +#include "uuid.h" +#include "xxmalloc.h" + +/*************************************************************/ +/* Private Functions */ +/*************************************************************/ + +/** + * Compute a topological ordering of the executor graph. + * Call only after all nodes, edges, and metrics have been populated. + * @param g Reference to the executor graph. + * @return Nodes in topological order. + */ +static struct list *vine_graph_compute_topological_order(struct vine_graph *g) +{ + if (!g) { + return NULL; + } + + int total_nodes = itable_size(g->nodes); + struct list *topo_order = list_create(); + struct itable *in_degree_map = itable_create(0); + struct priority_queue *pq = priority_queue_create(total_nodes); + + uint64_t nid; + struct vine_graph_node *node; + int iteration; + ITABLE_ITERATE(g->nodes, iteration, nid, node) + { + int deg = list_size(node->parents); + itable_insert(in_degree_map, nid, (void *)(intptr_t)deg); + if (deg == 0) { + priority_queue_push(pq, node, -(double)node->node_id); + } + } + + while (priority_queue_size(pq) > 0) { + struct vine_graph_node *current = priority_queue_pop(pq); + list_push_tail(topo_order, current); + + struct vine_graph_node *child; + LIST_ITERATE(current->children, child) + { + intptr_t raw_deg = (intptr_t)itable_lookup(in_degree_map, child->node_id); + int deg = (int)raw_deg - 1; + itable_insert(in_degree_map, child->node_id, (void *)(intptr_t)deg); + + if (deg == 0) { + priority_queue_push(pq, child, -(double)child->node_id); + } + } + } + + if (list_size(topo_order) != total_nodes) { + debug(D_ERROR, "Error: executor graph contains cycles or is malformed."); + debug(D_ERROR, "Expected %d nodes, but only sorted %d.", total_nodes, list_size(topo_order)); + + uint64_t id; + ITABLE_ITERATE(g->nodes, iteration, id, node) + { + intptr_t raw_deg = (intptr_t)itable_lookup(in_degree_map, id); + int deg = (int)raw_deg; + if (deg > 0) { + debug(D_ERROR, " Node %" PRIu64 " has in-degree %d. Parents:", id, deg); + struct vine_graph_node *p; + LIST_ITERATE(node->parents, p) + { + debug(D_ERROR, " -> %" PRIu64, p->node_id); + } + } + } + + list_delete(topo_order); + itable_delete(in_degree_map); + priority_queue_delete(pq); + exit(1); + } + + itable_delete(in_degree_map); + priority_queue_delete(pq); + return topo_order; +} + +/** + * Extract weakly connected components of the executor graph. + * Currently used for debugging and instrumentation only. + * @param g Reference to the executor graph. + * @return List of weakly connected components. + */ +static struct list *vine_graph_extract_weak_components(struct vine_graph *g) +{ + if (!g) { + return NULL; + } + + struct set *visited = set_create(0); + struct list *components = list_create(); + + uint64_t nid; + struct vine_graph_node *node; + int iteration; + ITABLE_ITERATE(g->nodes, iteration, nid, node) + { + if (set_lookup(visited, node)) { + continue; + } + + struct list *component = list_create(); + struct list *queue = list_create(); + + list_push_tail(queue, node); + set_insert(visited, node); + list_push_tail(component, node); + + while (list_size(queue) > 0) { + struct vine_graph_node *curr = list_pop_head(queue); + + struct vine_graph_node *p; + LIST_ITERATE(curr->parents, p) + { + if (!set_lookup(visited, p)) { + list_push_tail(queue, p); + set_insert(visited, p); + list_push_tail(component, p); + } + } + + struct vine_graph_node *c; + LIST_ITERATE(curr->children, c) + { + if (!set_lookup(visited, c)) { + list_push_tail(queue, c); + set_insert(visited, c); + list_push_tail(component, c); + } + } + } + + list_push_tail(components, component); + list_delete(queue); + } + + set_delete(visited); + return components; +} + +/** + * Compute the heavy score of a node in the executor graph. + * @param node Reference to the node. + * @return Heavy score. + */ +static double vine_graph_node_compute_heavy_score(struct vine_graph_node *node) +{ + if (!node) { + return 0; + } + + double up_score = node->depth * node->upstream_subgraph_size * node->fan_in; + double down_score = node->height * node->downstream_subgraph_size * node->fan_out; + + return up_score / (down_score + 1); +} + +static void vine_graph_node_set_local_output_file(struct vine_graph_node *node) +{ + node->outfile_type = VINE_GRAPH_NODE_OUTFILE_TYPE_LOCAL; +} + +static void vine_graph_node_set_temp_output_file(struct vine_graph_node *node) +{ + node->outfile_type = VINE_GRAPH_NODE_OUTFILE_TYPE_TEMP; +} + +/** + * Compute upstream/downstream subgraph sizes and heavy scores for each node. + * This is expensive (can approach transitive-closure cost) and should only be + * invoked when heavy-score-based checkpoint selection is enabled. + */ +static void vine_graph_compute_topology_scores(struct vine_graph *g, struct list *topo_order) +{ + if (!g || !topo_order) { + return; + } + + struct vine_graph_node *node; + struct vine_graph_node *parent_node; + struct vine_graph_node *child_node; + + struct itable *upstream_map = itable_create(0); // reachable ancestors per node + struct itable *downstream_map = itable_create(0); // reachable descendants per node + uint64_t nid_tmp; + int iteration; + ITABLE_ITERATE(g->nodes, iteration, nid_tmp, node) + { + struct set *upstream = set_create(0); + struct set *downstream = set_create(0); + itable_insert(upstream_map, node->node_id, upstream); + itable_insert(downstream_map, node->node_id, downstream); + } + + LIST_ITERATE(topo_order, node) + { + struct set *upstream = itable_lookup(upstream_map, node->node_id); + LIST_ITERATE(node->parents, parent_node) + { + struct set *parent_upstream = itable_lookup(upstream_map, parent_node->node_id); + set_insert_set(upstream, parent_upstream); // in-place union, not set_union + set_insert(upstream, parent_node); + } + } + + LIST_ITERATE_REVERSE(topo_order, node) + { + struct set *downstream = itable_lookup(downstream_map, node->node_id); + LIST_ITERATE(node->children, child_node) + { + struct set *child_downstream = itable_lookup(downstream_map, child_node->node_id); + set_insert_set(downstream, child_downstream); // in-place union, not set_union + set_insert(downstream, child_node); + } + } + + LIST_ITERATE(topo_order, node) + { + node->upstream_subgraph_size = set_size(itable_lookup(upstream_map, node->node_id)); + node->downstream_subgraph_size = set_size(itable_lookup(downstream_map, node->node_id)); + node->fan_in = list_size(node->parents); + node->fan_out = list_size(node->children); + set_delete(itable_lookup(upstream_map, node->node_id)); + set_delete(itable_lookup(downstream_map, node->node_id)); + } + + itable_delete(upstream_map); + itable_delete(downstream_map); + + LIST_ITERATE(topo_order, node) + { + node->heavy_score = vine_graph_node_compute_heavy_score(node); // ranks checkpoint candidates + } +} + +/*************************************************************/ +/* Public APIs */ +/*************************************************************/ + +/** Tune the executor graph. + * @param g Reference to the executor graph. + * @param name Reference to the name of the parameter to tune. + * @param value Reference to the value of the parameter to tune. + * @return 0 on success, -1 on failure. + */ +int vine_graph_tune(struct vine_graph *g, const char *name, const char *value) +{ + if (!g || !name || !value) { + return -1; + } + + if (strcmp(name, "output-dir") == 0) { + if (mkdir(value, 0777) != 0 && errno != EEXIST) { + debug(D_ERROR, "failed to mkdir %s (errno=%d)", value, errno); + return -1; + } + free(g->output_dir); + g->output_dir = xxstrdup(value); + + } else if (strcmp(name, "prune-depth") == 0) { + int k = atoi(value); + if (k < 0) { + debug(D_ERROR, "invalid prune-depth: %s (must be >= 0; 0 disables prune-depth release)", value); + return -1; + } + g->prune_depth = k; + + } else if (strcmp(name, "checkpoint-fraction") == 0) { + double fraction = atof(value); + if (fraction < 0.0 || fraction > 1.0) { + debug(D_ERROR, "invalid checkpoint fraction: %s (must be between 0.0 and 1.0)", value); + return -1; + } + g->checkpoint_fraction = fraction; + + } else if (strcmp(name, "checkpoint-dir") == 0) { + if (mkdir(value, 0777) != 0 && errno != EEXIST) { + debug(D_ERROR, "failed to mkdir %s (errno=%d)", value, errno); + return -1; + } + free(g->checkpoint_dir); + g->checkpoint_dir = xxstrdup(value); + + } else if (strcmp(name, "print-graph-details") == 0) { + g->print_graph_details = (atoi(value) == 1) ? 1 : 0; + } else if (strcmp(name, "chain-grouping-enabled") == 0) { + /* Stays aligned with Python's --task-group. When zero the executor never merges chain members. */ + g->chain_grouping_enabled = (atoi(value) != 0) ? 1 : 0; + } else { + debug(D_ERROR, "invalid parameter name: %s", name); + return -1; + } + + return 0; +} + +/** + * Get the outfile remote name of a node in the executor graph. + * @param g Reference to the executor graph. + * @param node_id Reference to the node id. + * @return The outfile remote name. + */ +const char *vine_graph_get_node_outfile_remote_name(const struct vine_graph *g, uint64_t node_id) +{ + if (!g) { + return NULL; + } + + struct vine_graph_node *node = itable_lookup(g->nodes, node_id); + if (!node) { + return NULL; + } + + return node->outfile_remote_name; +} + +/** + * Get the task runner library name of the executor graph. + * @param g Reference to the executor graph. + * @return The task runner library name. + */ +const char *vine_graph_get_task_runner_library_name(const struct vine_graph *g) +{ + if (!g) { + return NULL; + } + + return g->task_runner_library_name; +} + +/** + * Set the task runner function name of the executor graph. + * @param g Reference to the executor graph. + * @param task_runner_function_name Reference to the task runner function name. + */ +void vine_graph_set_task_runner_function_name(struct vine_graph *g, const char *task_runner_function_name) +{ + if (!g || !task_runner_function_name) { + return; + } + + if (g->task_runner_function_name) { + free(g->task_runner_function_name); + } + + g->task_runner_function_name = xxstrdup(task_runner_function_name); +} + +/** + * Get the heavy score of a node in the executor graph. + * @param g Reference to the executor graph. + * @param node_id Reference to the node id. + * @return The heavy score. + */ +double vine_graph_get_node_heavy_score(const struct vine_graph *g, uint64_t node_id) +{ + if (!g) { + return -1; + } + + struct vine_graph_node *node = itable_lookup(g->nodes, node_id); + if (!node) { + return -1; + } + + return node->heavy_score; +} + +/** + * Compute the topology metrics of the executor graph, including depth, height, upstream and downstream counts, + * heavy scores, and weakly connected components. Must be called after all nodes and dependencies are added. + * @param g Reference to the executor graph. + */ +void vine_graph_finalize(struct vine_graph *g) +{ + if (!g) { + return; + } + + struct list *topo_order = vine_graph_compute_topological_order(g); // required for all metric passes + if (!topo_order) { + return; + } + + struct vine_graph_node *node; + struct vine_graph_node *parent_node; + struct vine_graph_node *child_node; + + /* Longest path from any source in topo order. */ + LIST_ITERATE(topo_order, node) + { + node->depth = 0; + LIST_ITERATE(node->parents, parent_node) + { + if (node->depth < parent_node->depth + 1) { + node->depth = parent_node->depth + 1; + } + } + } + + /* Longest path to any sink in reverse topo order. */ + LIST_ITERATE_REVERSE(topo_order, node) + { + node->height = 0; + LIST_ITERATE(node->children, child_node) + { + if (node->height < child_node->height + 1) { + node->height = child_node->height + 1; + } + } + } + + int total_nodes = list_size(topo_order); + int total_target_nodes = 0; + LIST_ITERATE(topo_order, node) + { + if (node->is_target) { + total_target_nodes++; + } + } + + /* + * Pick how many non-target nodes become shared-filesystem checkpoints. + * If zero, skip heavy-score passes entirely. + */ + int checkpoint_count = (int)((total_nodes - total_target_nodes) * g->checkpoint_fraction); + if (checkpoint_count < 0) { + checkpoint_count = 0; + } + + if (checkpoint_count > 0) { + vine_graph_compute_topology_scores(g, topo_order); // expensive, only if ranking needed + + struct priority_queue *sorted_nodes = priority_queue_create(total_nodes); + LIST_ITERATE(topo_order, node) + { + priority_queue_push(sorted_nodes, node, node->heavy_score); + } + + int assigned_checkpoint_count = 0; + while ((node = priority_queue_pop(sorted_nodes))) { + if (node->is_target) { + vine_graph_node_set_local_output_file(node); // targets keep managed local returns + continue; + } + if (assigned_checkpoint_count < checkpoint_count) { + /* + * Top heavy_score nodes checkpoint to shared storage under checkpoint_dir. + * No vine_file handle for that mode. + */ + node->outfile_type = VINE_GRAPH_NODE_OUTFILE_TYPE_SHARED_FILE_SYSTEM; + char *shared_file_system_outfile_path = string_format("%s/%s", g->checkpoint_dir, node->outfile_remote_name); + free(node->outfile_remote_name); + node->outfile_remote_name = shared_file_system_outfile_path; + assigned_checkpoint_count++; + } else { + vine_graph_node_set_temp_output_file(node); // remaining nodes use temp storage + } + } + priority_queue_delete(sorted_nodes); + } else { + LIST_ITERATE(topo_order, node) + { + if (node->is_target) { + vine_graph_node_set_local_output_file(node); + } else { + vine_graph_node_set_temp_output_file(node); // no checkpoint budget, all non-targets are temp + } + } + } + + if (g->print_graph_details) { + // weakly connected components and vine_graph_node_debug_print, debug only + struct list *weakly_connected_components = vine_graph_extract_weak_components(g); + struct list *component; + int component_index = 0; + debug(D_VINE, "graph has %d weakly connected components\n", list_size(weakly_connected_components)); + LIST_ITERATE(weakly_connected_components, component) + { + debug(D_VINE, "component %d size: %d\n", component_index, list_size(component)); + list_delete(component); + component_index++; + } + list_delete(weakly_connected_components); + + LIST_ITERATE(topo_order, node) + { + vine_graph_node_debug_print(node); + } + } + + list_delete(topo_order); + + return; +} + +/** + * Create a new node and track it in the executor graph. + * @param g Reference to the executor graph. + * @return The auto-assigned node id. + */ +uint64_t vine_graph_add_node(struct vine_graph *g) +{ + if (!g) { + return 0; + } + + uint64_t candidate_id = itable_size(g->nodes); + candidate_id += 1; // skip zero, search upward until unused + while (itable_lookup(g->nodes, candidate_id)) { + candidate_id++; + } + uint64_t node_id = candidate_id; + + struct vine_graph_node *node = vine_graph_node_create(node_id); // defaults to non-target + + if (!node) { + debug(D_ERROR, "failed to create node %" PRIu64, node_id); + vine_graph_delete(g); + exit(1); + } + + itable_insert(g->nodes, node_id, node); + + return node_id; +} + +/** + * Mark a node as a retrieval target. + */ +void vine_graph_set_target(struct vine_graph *g, uint64_t node_id) +{ + if (!g) { + return; + } + struct vine_graph_node *node = itable_lookup(g->nodes, node_id); + if (!node) { + debug(D_ERROR, "node %" PRIu64 " not found", node_id); + exit(1); + } + + node->is_target = 1; +} + +/** + * Create a new executor graph using graph-owned path configuration. + * @param runtime_dir Runtime directory used as the default path root. + * @return A new executor graph instance. + */ +struct vine_graph *vine_graph_create(const char *runtime_dir) +{ + if (!runtime_dir) { + return NULL; + } + + struct vine_graph *g = xxmalloc(sizeof(struct vine_graph)); + + g->checkpoint_dir = xxstrdup(runtime_dir); // default to current working directory + g->output_dir = xxstrdup(runtime_dir); // default to current working directory + + g->nodes = itable_create(0); + g->super_leader_to_members = itable_create(0); + g->outfile_cachename_to_node = hash_table_create(0, 0); + g->file_id_to_file = itable_create(0); + g->supernode_leader_child_to_input_source = hash_table_create(0, 0); + + cctools_uuid_t task_runner_library_name_id; + cctools_uuid_create(&task_runner_library_name_id); + g->task_runner_library_name = xxstrdup(task_runner_library_name_id.str); + + g->task_runner_function_name = NULL; + g->checkpoint_fraction = 0.0; + + /* Default prune-depth: release a TEMP node as soon as all of its + * direct children have completed. Set to 0 via tune("prune-depth") to + * disable and rely exclusively on cut-propagation. */ + g->prune_depth = 1; + + g->print_graph_details = 0; + g->chain_grouping_enabled = 0; + + return g; +} + +/** + * Add a dependency between two nodes in the executor graph. Note that the input-output file relationship + * is not handled here, because their file names might not have been determined yet. + * @param g Reference to the executor graph. + * @param parent_id Reference to the parent node id. + * @param child_id Reference to the child node id. + */ +void vine_graph_add_dependency(struct vine_graph *g, uint64_t parent_id, uint64_t child_id) +{ + if (!g) { + return; + } + + struct vine_graph_node *parent_node = itable_lookup(g->nodes, parent_id); + struct vine_graph_node *child_node = itable_lookup(g->nodes, child_id); + if (!parent_node) { + debug(D_ERROR, "parent node %" PRIu64 " not found", parent_id); + exit(1); + } + if (!child_node) { + debug(D_ERROR, "child node %" PRIu64 " not found", child_id); + exit(1); + } + + vine_graph_node_ensure_dependency(parent_node, child_node); + + return; +} + +/** + * Copy a list of struct vine_graph_node * into a new list in the same order (pointer values only, not deep copy). + * Snapshots parents/children before mutating edges during supernode rewire. Caller frees the new list. + */ +static struct list *vine_graph_node_list_copy(struct list *src) +{ + struct list *dst = list_create(); + if (!src) { + return dst; + } + struct vine_graph_node *x; + LIST_ITERATE(src, x) + { + list_push_tail(dst, x); + } + return dst; +} + +/** + * After supernode registration: for each node in mset, remove internal edges (both ends in mset). + * Rewire edges that cross the group boundary through leader. mset contains leader plus all non-leader members. + */ +static void vine_graph_supernode_rewire(struct vine_graph *g, struct vine_graph_node *leader, struct set *mset) +{ + struct vine_graph_node *m; + uint64_t nid; + int iteration; + ITABLE_ITERATE(g->nodes, iteration, nid, m) + { + if (!set_lookup(mset, m)) { + continue; + } + + struct list *psnap = vine_graph_node_list_copy(m->parents); + struct vine_graph_node *p; + while ((p = list_pop_head(psnap))) { + if (set_lookup(mset, p)) { + vine_graph_node_remove_dependency(p, m); + } else { + vine_graph_node_remove_dependency(p, m); + vine_graph_node_ensure_dependency(p, leader); + } + } + list_delete(psnap); + + struct list *csnap = vine_graph_node_list_copy(m->children); + struct vine_graph_node *c; + while ((c = list_pop_head(csnap))) { + if (set_lookup(mset, c)) { + vine_graph_node_remove_dependency(m, c); + } else { + vine_graph_node_remove_dependency(m, c); + vine_graph_node_ensure_dependency(leader, c); + /* + * Scheduling sees leader->c; data still flows from member m (e.g. chain tail). + * Record so materialize mounts m->outfile, not the leader's primary outfile. + */ + if (m != leader) { + char *lckey = string_format("%" PRIu64 ",%" PRIu64, leader->node_id, c->node_id); + /* + * The table keeps the first insert if the same key is stored twice. Drop the + * old mapping explicitly so the node that truly produces the file for c wins. + */ + hash_table_remove(g->supernode_leader_child_to_input_source, lckey); + hash_table_insert(g->supernode_leader_child_to_input_source, lckey, m); + free(lckey); + } + } + } + list_delete(csnap); + } +} + +/** + * Resolve the leader struct vine_graph_node * for any node in a supernode (singleton maps to itself). + * Returns NULL if g or n is NULL or the leader id is missing from g. + */ +struct vine_graph_node *vine_graph_supernode_leader_node(struct vine_graph *g, struct vine_graph_node *n) +{ + if (!g || !n) { + return NULL; + } + return itable_lookup(g->nodes, n->super_leader_id); +} + +/** + * List of non-leader member nodes for leader_id. Owned by the graph; do not free. + * Items are struct vine_graph_node *. NULL if g is NULL or there is no entry. + */ +struct list *vine_graph_supernode_nonleader_members(struct vine_graph *g, uint64_t leader_id) +{ + if (!g) { + return NULL; + } + return itable_lookup(g->super_leader_to_members, leader_id); +} + +/** + * Register a supernode: validate members, rewire externals to leader_id, set super_leader_id on all members, + * clear fired_parents, and store non-leaders in super_leader_to_members. Nodes and plain deps must exist first. + * Returns 0 on success, -1 on error (see debug log). + */ +int vine_graph_supernode_register(struct vine_graph *g, uint64_t leader_id, const uint64_t *member_ids, int n_member_ids) +{ + if (!g || member_ids == NULL || n_member_ids < 1) { + debug(D_ERROR, "vine_graph_supernode_register: invalid arguments"); + return -1; + } + + struct vine_graph_node *leader = itable_lookup(g->nodes, leader_id); + if (!leader) { + debug(D_ERROR, "vine_graph_supernode_register: leader %" PRIu64 " not found", leader_id); + return -1; + } + + struct set *seen = set_create(0); + struct set *mset = set_create(0); + set_insert(mset, leader); + + for (int i = 0; i < n_member_ids; i++) { + uint64_t mid = member_ids[i]; + if (mid == leader_id) { + debug(D_ERROR, "vine_graph_supernode_register: member list must not contain leader %" PRIu64, leader_id); + set_delete(seen); + set_delete(mset); + return -1; + } + if (set_lookup(seen, (void *)(uintptr_t)mid)) { + debug(D_ERROR, "vine_graph_supernode_register: duplicate member %" PRIu64, mid); + set_delete(seen); + set_delete(mset); + return -1; + } + + struct vine_graph_node *mn = itable_lookup(g->nodes, mid); + if (!mn) { + debug(D_ERROR, "vine_graph_supernode_register: member node %" PRIu64 " not found", mid); + set_delete(seen); + set_delete(mset); + return -1; + } + if (mn->super_leader_id != mn->node_id) { + debug(D_ERROR, + "vine_graph_supernode_register: node %" PRIu64 " already belongs to leader %" PRIu64, + mid, + mn->super_leader_id); + set_delete(seen); + set_delete(mset); + return -1; + } + set_insert(seen, (void *)(uintptr_t)mid); + set_insert(mset, mn); + } + set_delete(seen); + + if (leader->super_leader_id != leader->node_id) { + debug(D_ERROR, + "vine_graph_supernode_register: leader %" PRIu64 " already belongs to group %" PRIu64, + leader_id, + leader->super_leader_id); + set_delete(mset); + return -1; + } + + if (itable_lookup(g->super_leader_to_members, leader_id)) { + debug(D_ERROR, "vine_graph_supernode_register: leader %" PRIu64 " already registered", leader_id); + set_delete(mset); + return -1; + } + + vine_graph_supernode_rewire(g, leader, mset); + + struct vine_graph_node *x; + uint64_t xid; + int iteration; + ITABLE_ITERATE(g->nodes, iteration, xid, x) + { + if (set_lookup(mset, x)) { + x->super_leader_id = leader_id; + vine_graph_node_clear_fired_parents(x); + } + } + + struct list *nonleaders = list_create(); + for (int i = 0; i < n_member_ids; i++) { + struct vine_graph_node *mn = itable_lookup(g->nodes, member_ids[i]); + list_push_tail(nonleaders, mn); + } + itable_insert(g->super_leader_to_members, leader_id, nonleaders); + + set_delete(mset); + return 0; +} + +/** Exactly one child pointer, or NULL if not exactly one. */ +static struct vine_graph_node *vine_graph_node_only_child(struct vine_graph_node *n) +{ + if (!n || list_size(n->children) != 1) { + return NULL; + } + struct vine_graph_node *c; + LIST_ITERATE(n->children, c) + { + return c; + } + return NULL; +} + +/** Exactly one parent pointer, or NULL if not exactly one. */ +static struct vine_graph_node *vine_graph_node_only_parent(struct vine_graph_node *n) +{ + if (!n || list_size(n->parents) != 1) { + return NULL; + } + struct vine_graph_node *p; + LIST_ITERATE(n->parents, p) + { + return p; + } + return NULL; +} + +/* File mounts need separate TaskVine sandboxes, so they form chain-group boundaries. */ +static int vine_graph_node_allows_chain_grouping(struct vine_graph_node *n) +{ + return n && list_size(n->extra_inputs) == 0 && list_size(n->extra_outputs) == 0; +} + +/** + * Head of a maximal linear chain: singleton, and not strictly inside such a chain from the left + * (no parent, fan-in > 1, or unique parent has fan-out > 1). + */ +static int vine_graph_node_is_chain_head(struct vine_graph_node *n) +{ + if (!vine_graph_node_is_supernode_leader(n) || !vine_graph_node_allows_chain_grouping(n)) { + return 0; + } + if (list_size(n->parents) == 0) { + return 1; + } + if (list_size(n->parents) > 1) { + return 1; + } + struct vine_graph_node *p = vine_graph_node_only_parent(n); + return p && (!vine_graph_node_allows_chain_grouping(p) || list_size(p->children) > 1); +} + +struct pending_chain_group { + uint64_t leader_id; + uint64_t *member_ids; + int n_members; +}; + +int vine_graph_group_chain_like_tasks(struct vine_graph *g) +{ + if (!g) { + return -1; + } + if (!g->chain_grouping_enabled) { + return 0; + } + + /* + * Registering a supernode rewires the graph. Find every chain on the original adjacency + * first, then register them, so a half-updated graph does not confuse later chain walks. + */ + struct list *pending = list_create(); + + uint64_t nid; + struct vine_graph_node *n; + int iteration; + ITABLE_ITERATE(g->nodes, iteration, nid, n) + { + if (!vine_graph_node_is_chain_head(n)) { + continue; + } + if (list_size(n->children) != 1) { + continue; + } + + struct list *chain = list_create(); + struct vine_graph_node *cur = n; + for (;;) { + list_push_tail(chain, cur); + if (list_size(cur->children) != 1) { + break; + } + struct vine_graph_node *c = vine_graph_node_only_child(cur); + if (!c || !vine_graph_node_is_supernode_leader(c) || !vine_graph_node_allows_chain_grouping(c)) { + break; + } + if (c->is_target) { + /* Leave the retrieval or output consumer as its own task outside the merged chain. */ + break; + } + if (list_size(c->parents) != 1) { + break; + } + if (vine_graph_node_only_parent(c) != cur) { + break; + } + cur = c; + } + + int L = list_size(chain); + if (L < 2) { + list_delete(chain); + continue; + } + + struct pending_chain_group *pc = xxmalloc(sizeof(*pc)); + pc->leader_id = n->node_id; + pc->n_members = L - 1; + pc->member_ids = xxmalloc((size_t)(L - 1) * sizeof(uint64_t)); + int mi = 0; + int first = 1; + struct vine_graph_node *x; + LIST_ITERATE(chain, x) + { + if (first) { + first = 0; + continue; + } + pc->member_ids[mi++] = x->node_id; + } + list_push_tail(pending, pc); + list_delete(chain); + } + + int groups = 0; + struct pending_chain_group *pc; + while ((pc = (struct pending_chain_group *)list_pop_head(pending))) { + if (vine_graph_supernode_register(g, pc->leader_id, pc->member_ids, pc->n_members) == 0) { + groups++; + } + free(pc->member_ids); + free(pc); + } + list_delete(pending); + + return groups; +} + +struct vine_graph_node *vine_graph_input_producer_node(struct vine_graph *g, struct vine_graph_node *parent, struct vine_graph_node *child) +{ + if (!g || !parent || !child) { + return parent; + } + if (!g->chain_grouping_enabled || !g->supernode_leader_child_to_input_source) { + return parent; + } + + char *k = string_format("%" PRIu64 ",%" PRIu64, parent->node_id, child->node_id); + struct vine_graph_node *src = (struct vine_graph_node *)hash_table_lookup(g->supernode_leader_child_to_input_source, k); + free(k); + return src ? src : parent; +} + +/** + * Delete an executor graph instance. + * @param g Reference to the executor graph. + */ +void vine_graph_delete(struct vine_graph *g) +{ + if (!g) { + return; + } + + uint64_t lid; + struct list *mems; + int iteration; + if (g->super_leader_to_members) { + ITABLE_ITERATE(g->super_leader_to_members, iteration, lid, mems) + { + (void)lid; + if (mems) { + list_delete(mems); + } + } + itable_delete(g->super_leader_to_members); + g->super_leader_to_members = NULL; + } + + uint64_t nid; + struct vine_graph_node *node; + ITABLE_ITERATE(g->nodes, iteration, nid, node) + { + vine_graph_node_delete(node); + } + + free(g->task_runner_library_name); + free(g->task_runner_function_name); + free(g->checkpoint_dir); + free(g->output_dir); + + itable_delete(g->nodes); + hash_table_delete(g->outfile_cachename_to_node); + + hash_table_delete(g->supernode_leader_child_to_input_source); + + itable_delete(g->file_id_to_file); + + free(g); +} diff --git a/taskvine/src/vine_graph/vine_graph.h b/taskvine/src/vine_graph/vine_graph.h new file mode 100644 index 0000000000..fe58833eda --- /dev/null +++ b/taskvine/src/vine_graph/vine_graph.h @@ -0,0 +1,158 @@ +#ifndef VINE_GRAPH_H +#define VINE_GRAPH_H + +#include "hash_table.h" +#include "itable.h" + +#include "vine_graph_node.h" + +struct vine_graph { + struct itable *nodes; + /** Maps each supernode leader id to the list of non-leader member nodes. Empty when no groups exist. */ + struct itable *super_leader_to_members; + struct hash_table *outfile_cachename_to_node; + /** Maps public FileHandle ids to their single declared vine_file. */ + struct itable *file_id_to_file; + /** + * Maps leader id and downstream child id to the member that actually wrote the outfile the + * child should read. After a rewire the DAG edge may show the leader while the tail member + * still owns the bytes. Keys are decimal node ids joined with a comma. + */ + struct hash_table *supernode_leader_child_to_input_source; + + char *checkpoint_dir; + char *output_dir; + char *task_runner_library_name; + char *task_runner_function_name; + + double checkpoint_fraction; + int prune_depth; + + int print_graph_details; + /* + * Zero unless the user turned on chain grouping. When zero the executor does not merge + * members onto one task, does not list several scheduler keys in one runner infile, and does + * not remap which node's outfile a consumer should mount. + */ + int chain_grouping_enabled; +}; + +// Public graph API (declarations below) + +/** + * Resolve which producer node's outfile a child should consume for a scheduling parent edge. + * After supernode rewire, parent may be the leader while the file is produced by a member (e.g. tail). + * @return The node that owns the vine_file inputs should mount, or @p parent if there is no override. + */ +struct vine_graph_node *vine_graph_input_producer_node(struct vine_graph *g, struct vine_graph_node *parent, struct vine_graph_node *child); + +/** Create an executor graph and return it. +@param runtime_dir Runtime directory used for default graph output paths. +@return A new executor graph. +*/ +struct vine_graph *vine_graph_create(const char *runtime_dir); + +/** Create a new node in the executor graph. +@param g Reference to the executor graph. +@return The auto-assigned node id. +*/ +uint64_t vine_graph_add_node(struct vine_graph *g); + +/** Mark a node as a retrieval target. +@param g Reference to the executor graph. +@param node_id Identifier of the node to mark as target. +*/ +void vine_graph_set_target(struct vine_graph *g, uint64_t node_id); + +/** Add a dependency between two nodes in the executor graph. +@param g Reference to the executor graph. +@param parent_id Identifier of the parent node. +@param child_id Identifier of the child node. +*/ +void vine_graph_add_dependency(struct vine_graph *g, uint64_t parent_id, uint64_t child_id); + +/** Finalize the metrics of the executor graph. +@param g Reference to the executor graph. +*/ +void vine_graph_finalize(struct vine_graph *g); + +/** Get the heavy score of a node in the executor graph. +@param g Reference to the executor graph. +@param node_id Identifier of the node. +@return The heavy score. +*/ +double vine_graph_get_node_heavy_score(const struct vine_graph *g, uint64_t node_id); + +/** Get the outfile remote name of a node in the executor graph. +@param g Reference to the executor graph. +@param node_id Identifier of the node. +@return The outfile remote name. +*/ +const char *vine_graph_get_node_outfile_remote_name(const struct vine_graph *g, uint64_t node_id); + +/** Delete an executor graph. +@param g Reference to the executor graph. +*/ +void vine_graph_delete(struct vine_graph *g); + +/** Get the task runner library name of the executor graph. +@param g Reference to the executor graph. +@return The task runner library name. +*/ +const char *vine_graph_get_task_runner_library_name(const struct vine_graph *g); + +/** Set the task runner function name of the executor graph. +@param g Reference to the executor graph. +@param task_runner_function_name Reference to the task runner function name. +*/ +void vine_graph_set_task_runner_function_name(struct vine_graph *g, const char *task_runner_function_name); + +/** Tune the executor graph. +@param g Reference to the executor graph. +@param name Reference to the name of the parameter to tune. +@param value Reference to the value of the parameter to tune. +@return 0 on success, -1 on failure. +*/ +int vine_graph_tune(struct vine_graph *g, const char *name, const char *value); + +/** + * True if this node may submit a TaskVine task: singletons and supernode leaders only. + * Non-leader members run inside the leader's runner and must not be scheduled separately. + * Returns non-zero when n is non-NULL and node_id equals super_leader_id. + */ +static inline int vine_graph_node_is_supernode_leader(const struct vine_graph_node *n) +{ + return n && n->node_id == n->super_leader_id; +} + +/** + * Look up the leader node for n using n->super_leader_id (identity when n is already the leader). + * Returns the leader struct vine_graph_node *, or NULL on invalid input or if the leader id is missing from g. + */ +struct vine_graph_node *vine_graph_supernode_leader_node(struct vine_graph *g, struct vine_graph_node *n); + +/** + * Merge nodes into one supernode: rewire external edges to leader_id and set super_leader_id on all members. + * Call after all plain vine_graph_add_dependency edges exist and before vine_graph_executor_finalize. + * Every involved node must currently be its own group (super_leader_id == node_id); member_ids must list + * only non-leader members (not leader_id) without duplicates. Returns 0 on success, -1 on error. + */ +int vine_graph_supernode_register(struct vine_graph *g, uint64_t leader_id, const uint64_t *member_ids, int n_member_ids); + +/** + * Non-leader members for a registered supernode as a list of struct vine_graph_node *. + * Do not free the list; it is owned by g. Returns NULL if g is NULL or no group was registered for that leader. + */ +struct list *vine_graph_supernode_nonleader_members(struct vine_graph *g, uint64_t leader_id); + +/** + * Collapse each maximal singleton linear chain into one supernode (leader = chain head). + * A chain is a path n0->n1->... where each n_i (i > 0) has exactly one parent (n_{i-1}) and + * each n_i (i < last) has exactly one child (n_{i+1}); n0 is not preceded by such an edge from a + * single-child parent (head: no parent, multiple parents, or parent with multiple children). + * Call after all vine_graph_add_dependency edges and before vine_graph_executor_finalize. + * Returns the number of supernodes registered, or -1 if g is NULL. + */ +int vine_graph_group_chain_like_tasks(struct vine_graph *g); + +#endif // VINE_GRAPH_H diff --git a/taskvine/src/vine_graph/vine_graph.i b/taskvine/src/vine_graph/vine_graph.i new file mode 100644 index 0000000000..e546ec57c4 --- /dev/null +++ b/taskvine/src/vine_graph/vine_graph.i @@ -0,0 +1,21 @@ +/* SWIG interface for local executor graph API bindings */ +%module vine_graph_capi + +%{ +#include "int_sizes.h" +#include "vine_graph.h" +#include "vine_graph_executor.h" +%} + +%include "stdint.i" +%include "int_sizes.h" + +/* uint64_t[] + length: not mapped yet; expose via VineGraphExecutor when needed. */ +%ignore vine_graph_supernode_register; +%ignore vine_graph_supernode_nonleader_members; + +/* Import existing SWIG interface for type information (do not wrap again) */ +%import "../bindings/python3/taskvine.i" + +%include "vine_graph.h" +%include "vine_graph_executor.h" diff --git a/taskvine/src/vine_graph/vine_graph_executor.c b/taskvine/src/vine_graph/vine_graph_executor.c new file mode 100644 index 0000000000..6244edcc1c --- /dev/null +++ b/taskvine/src/vine_graph/vine_graph_executor.c @@ -0,0 +1,1441 @@ +#include +#include +#include +#include +#include +#include + +#include "buffer.h" +#include "debug.h" +#include "vine_graph_executor.h" +#include "macros.h" +#include "progress_bar.h" +#include "random.h" +#include "set.h" +#include "stringtools.h" +#include "xxmalloc.h" + +#include "taskvine.h" + +static volatile sig_atomic_t interrupted = 0; + +static void vine_graph_executor_submit_node(struct vine_graph_executor *e, struct vine_graph_node *node); +static struct vine_task *vine_graph_executor_make_vine_task(struct vine_graph_executor *e); +static void vine_graph_executor_materialize_node(struct vine_graph_executor *e, struct vine_graph_node *node); +static void vine_graph_executor_run_completion_postprocess(struct vine_graph_executor *e, struct vine_graph_node *node); + +static uint64_t vine_graph_executor_count_completed_user_nodes(const struct vine_graph *g) +{ + uint64_t n = 0; + uint64_t nid; + struct vine_graph_node *nd; + int iteration; + + if (!g) { + return 0; + } + ITABLE_ITERATE(g->nodes, iteration, nid, nd) + { + if (nd->completed) { + n++; + } + } + return n; +} + +/* + * The leader's task has passed validation. Mark that node completed and, when chain grouping + * applies, mark every non-leader in the same group. Those members did not receive their own + * vine tasks because they ran inside the leader's single submission. + */ +static void vine_graph_executor_mark_user_node_completed_after_success(struct vine_graph *g, struct vine_graph_node *leader) +{ + if (!g || !leader) { + return; + } + + leader->completed = 1; + + if (!g->chain_grouping_enabled) { + return; + } + + struct list *smems = vine_graph_supernode_nonleader_members(g, leader->node_id); + if (!smems) { + return; + } + + struct vine_graph_node *m; + LIST_ITERATE(smems, m) + { + if (!m->completed) { + m->completed = 1; + } + } +} + +/* + * Build the JSON payload for the runner infile argument. Field fn_args[0] names the scheduler + * keys that run_scheduler_keys should execute. A merged chain passes one comma-separated string + * such as "1,2,3". A singleton passes a single id. The payload must be one JSON string in that + * slot instead of an array that mixes strings and bare numbers. + */ +static char *vine_graph_executor_format_runner_infile_json(struct vine_graph *g, struct vine_graph_node *node) +{ + if (!g || !node) { + return NULL; + } + + if (!g->chain_grouping_enabled) { + return string_format("{\"fn_args\":[\"%" PRIu64 "\"],\"fn_kwargs\":{}}", node->node_id); + } + + struct list *mems = vine_graph_supernode_nonleader_members(g, node->node_id); + if (!mems || list_size(mems) == 0) { + return string_format("{\"fn_args\":[\"%" PRIu64 "\"],\"fn_kwargs\":{}}", node->node_id); + } + + buffer_t buf; + buffer_init(&buf); + buffer_printf(&buf, "{\"fn_args\":[\"%" PRIu64 "", node->node_id); + struct vine_graph_node *m; + LIST_ITERATE(mems, m) + { + buffer_printf(&buf, ",%" PRIu64 "", m->node_id); + } + buffer_printf(&buf, "\"],\"fn_kwargs\":{}}"); + + char *s = xxstrdup(buffer_tostring(&buf)); + buffer_free(&buf); + return s; +} + +static void vine_graph_io_mount_add(struct list *lst, struct vine_file *f, const char *remote_name) +{ + struct vine_graph_io_mount *m = xxmalloc(sizeof(*m)); + m->file = f; + m->remote_name = xxstrdup(remote_name); + list_push_tail(lst, m); +} + +/* Undeclare runner infile buffer (before discarding the vine_task). */ +static void vine_graph_executor_clear_node_runner_arg(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + if (!e || !node || !node->task_runner_arg_file) { + return; + } + vine_undeclare_file(e->manager, node->task_runner_arg_file); + node->task_runner_arg_file = NULL; +} + +/* Initialize runtime fields and default tuning values for a new executor. */ +static void vine_graph_executor_init_runtime(struct vine_graph_executor *e) +{ + if (!e) { + return; + } + + e->task_id_to_node = itable_create(0); + e->resubmit_queue = list_create(); + e->time_first_task_dispatched = UINT64_MAX; // sentinel until first task commit time + e->time_last_task_retrieved = 0; + e->makespan_us = 0; + e->completed_recovery_tasks = 0; + e->time_spent_on_cut_propagation = 0; + e->pfs_usage_bytes = 0; + e->total_preprocessing_time_us = 0; + e->total_postprocessing_time_us = 0; + e->task_priority_mode = TASK_PRIORITY_MODE_LARGEST_INPUT_FIRST; + e->failure_injection_step_percent = -1.0; + e->progress_bar_update_interval_sec = 0.1; +} + +static int vine_graph_task_not_submitted(struct vine_task *task) +{ + return !task || vine_task_get_id(task) <= 0; +} + +/* Release the task-id lookup table and the resubmit queue. */ +static void vine_graph_executor_clear_runtime(struct vine_graph_executor *e) +{ + if (!e) { + return; + } + if (e->task_id_to_node) { + itable_delete(e->task_id_to_node); + e->task_id_to_node = NULL; + } + if (e->resubmit_queue) { + list_delete(e->resubmit_queue); + e->resubmit_queue = NULL; + } +} + +/* Allocate an executor bound to the given manager and graph. */ +struct vine_graph_executor *vine_graph_executor_create(struct vine_manager *manager, struct vine_graph *graph) +{ + if (!manager || !graph) { + return NULL; + } + + struct vine_graph_executor *e = malloc(sizeof(*e)); + if (!e) { + return NULL; + } + + e->graph = graph; + e->manager = manager; + vine_graph_executor_init_runtime(e); + return e; +} + +/* Create a new graph for the manager's runtime directory. */ +struct vine_graph *vine_graph_executor_create_graph(struct vine_manager *manager) +{ + if (!manager) { + return NULL; + } + + const char *runtime_dir = vine_get_runtime_directory(manager); + return vine_graph_create(runtime_dir); +} + +/* Undeclare managed files, remove local outputs, and free the executor. */ +void vine_graph_executor_delete(struct vine_graph_executor *e) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (g && e->manager) { + uint64_t nid; + struct vine_graph_node *node; + int iteration; + ITABLE_ITERATE(g->nodes, iteration, nid, node) + { + if (node->task_runner_arg_file) { + vine_undeclare_file(e->manager, node->task_runner_arg_file); // before graph free to avoid double free + node->task_runner_arg_file = NULL; + } + switch (node->outfile_type) { + case VINE_GRAPH_NODE_OUTFILE_TYPE_TEMP: + break; + case VINE_GRAPH_NODE_OUTFILE_TYPE_SHARED_FILE_SYSTEM: + if (node->outfile_remote_name) { + unlink(node->outfile_remote_name); + } + break; + case VINE_GRAPH_NODE_OUTFILE_TYPE_LOCAL: + if (node->outfile && vine_file_source(node->outfile)) { + unlink(vine_file_source(node->outfile)); + } + break; + } + if (node->outfile) { + hash_table_remove(g->outfile_cachename_to_node, vine_file_cached_name(node->outfile)); + vine_undeclare_file(e->manager, node->outfile); + node->outfile = NULL; + } + } + + uint64_t file_id; + struct vine_file *file; + ITABLE_ITERATE(g->file_id_to_file, iteration, file_id, file) + { + vine_undeclare_file(e->manager, file); + } + itable_clear(g->file_id_to_file, NULL); + } + vine_graph_executor_clear_runtime(e); + free(e); +} + +/* + * Create a new library task (not yet published on the node). The caller attaches IO, then sets + * node->task only when the task is fully configured (atomic materialize). + */ +static struct vine_task *vine_graph_executor_make_vine_task(struct vine_graph_executor *e) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g) { + return NULL; + } + + if (!g->task_runner_function_name) { + debug(D_ERROR, "task runner function name is not set"); + vine_graph_delete(g); + exit(1); + } + if (!g->task_runner_library_name) { + debug(D_ERROR, "task runner library name is not set"); + vine_graph_delete(g); + exit(1); + } + + struct vine_task *t = vine_task_create(g->task_runner_function_name); + vine_task_set_library_required(t, g->task_runner_library_name); + vine_task_addref(t); // keep alive across vine_submit and vine_wait + return t; +} + +/* + * Attach inputs, outputs, and the infile buffer to a new vine_task at submit time. The node's + * task pointer stays NULL until that bundle is complete and ready for vine_submit. + */ +static void vine_graph_executor_materialize_node(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !node) { + return; + } + + if (node->task && !vine_graph_task_not_submitted(node->task)) { + return; + } + /* INITIAL task implies mounts + runner infile are already complete. */ + if (node->task) { + return; + } + + vine_graph_executor_clear_node_runner_arg(e, node); + + struct vine_task *t = vine_graph_executor_make_vine_task(e); + if (!t) { + return; + } + + if (node->outfile) { + vine_task_add_output(t, node->outfile, node->outfile_remote_name, VINE_TRANSFER_ALWAYS); + } + + void *item; + LIST_ITERATE(node->extra_outputs, item) + { + struct vine_graph_io_mount *m = (struct vine_graph_io_mount *)item; + vine_task_add_output(t, m->file, m->remote_name, VINE_TRANSFER_ALWAYS); + } + + /* + * When chains are merged, only the leader submits a task, but that task must list every + * member's declared outputs so the manager can track each outfile by node id. + */ + if (g->chain_grouping_enabled && vine_graph_node_is_supernode_leader(node)) { + struct list *mems = vine_graph_supernode_nonleader_members(g, node->node_id); + if (mems) { + struct vine_graph_node *m; + LIST_ITERATE(mems, m) + { + if (m->outfile) { + vine_task_add_output(t, m->outfile, m->outfile_remote_name, VINE_TRANSFER_ALWAYS); + } + LIST_ITERATE(m->extra_outputs, item) + { + struct vine_graph_io_mount *em = (struct vine_graph_io_mount *)item; + vine_task_add_output(t, em->file, em->remote_name, VINE_TRANSFER_ALWAYS); + } + } + } + } + + struct vine_graph_node *parent_node; + LIST_ITERATE(node->parents, parent_node) + { + struct vine_graph_node *src = vine_graph_input_producer_node(g, parent_node, node); + if (src && src->outfile) { + vine_task_add_input(t, src->outfile, src->outfile_remote_name, VINE_TRANSFER_ALWAYS); + } + } + + LIST_ITERATE(node->extra_inputs, item) + { + struct vine_graph_io_mount *m = (struct vine_graph_io_mount *)item; + vine_task_add_input(t, m->file, m->remote_name, VINE_TRANSFER_ALWAYS); + } + + char *task_arguments = vine_graph_executor_format_runner_infile_json(g, node); + if (!task_arguments) { + goto fail_task; + } + struct vine_file *arg_file = + vine_declare_buffer(e->manager, task_arguments, strlen(task_arguments), VINE_CACHE_LEVEL_TASK, VINE_UNLINK_WHEN_DONE); + free(task_arguments); + if (!arg_file) { + goto fail_task; + } + vine_task_add_input(t, arg_file, "infile", VINE_TRANSFER_ALWAYS); + + node->task = t; + node->task_runner_arg_file = arg_file; + return; + +fail_task: + vine_task_delete(t); +} + +/* + * Declare the output vine_file for this node from outfile_type. + * Shared filesystem outputs may leave outfile unset. + */ +static void vine_graph_executor_declare_node_outfile(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !node || node->outfile) { + return; + } + + switch (node->outfile_type) { + case VINE_GRAPH_NODE_OUTFILE_TYPE_LOCAL: { + char *local_outfile_path = string_format("%s/%s", g->output_dir, node->outfile_remote_name); + node->outfile = vine_declare_file(e->manager, local_outfile_path, VINE_CACHE_LEVEL_WORKFLOW, 0); + free(local_outfile_path); + break; + } + case VINE_GRAPH_NODE_OUTFILE_TYPE_TEMP: + node->outfile = vine_declare_temp(e->manager); + break; + case VINE_GRAPH_NODE_OUTFILE_TYPE_SHARED_FILE_SYSTEM: + break; + } +} + +/* + * Allocate the next graph node. Per-node vine_task and I/O mounts appear later during + * vine_graph_executor_materialize_node at submit time. + */ +uint64_t vine_graph_executor_add_node(struct vine_graph_executor *e) +{ + if (!e || !e->graph) { + return 0; + } + + uint64_t node_id = vine_graph_add_node(e->graph); + return node_id; +} + +/* + * Finalize the graph: declare outputs with cached_name registration, + * attach parent inputs, and set remaining parent counts for scheduling. + */ +void vine_graph_executor_finalize(struct vine_graph_executor *e) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g) { + return; + } + + vine_graph_finalize(g); + + /* + * Two passes. Declare outputs and cached_name map first so parent + * vine_file objects exist. Task-level input/output mounts are applied + * in vine_graph_executor_materialize_node at submit time. + */ + uint64_t nid; + struct vine_graph_node *node; + int iteration; + ITABLE_ITERATE(g->nodes, iteration, nid, node) + { + vine_graph_executor_declare_node_outfile(e, node); + if (node->outfile) { + hash_table_insert(g->outfile_cachename_to_node, vine_file_cached_name(node->outfile), node); + } + } + + ITABLE_ITERATE(g->nodes, iteration, nid, node) + { + node->remaining_parents_count = list_size(node->parents); + } +} + +int vine_graph_executor_declare_input_file(struct vine_graph_executor *e, uint64_t file_id, const char *source_path) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !file_id || !source_path || itable_lookup(g->file_id_to_file, file_id)) { + return -1; + } + + struct vine_file *file = vine_declare_file(e->manager, source_path, VINE_CACHE_LEVEL_WORKFLOW, 0); + if (!file) { + return -1; + } + itable_insert(g->file_id_to_file, file_id, file); + return 0; +} + +int vine_graph_executor_add_task_output_file(struct vine_graph_executor *e, uint64_t task_id, uint64_t file_id, const char *task_path, int is_target) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !task_id || !file_id || !task_path || itable_lookup(g->file_id_to_file, file_id)) { + return -1; + } + + struct vine_graph_node *node = itable_lookup(g->nodes, task_id); + if (!node) { + return -1; + } + + struct vine_file *file = NULL; + if (is_target) { + const char *base = strrchr(task_path, '/'); + base = base ? base + 1 : task_path; + char *target_path = string_format("%s/file-%" PRIu64 "-%s", g->output_dir, file_id, base); + file = vine_declare_file(e->manager, target_path, VINE_CACHE_LEVEL_WORKFLOW, 0); + free(target_path); + } else { + file = vine_declare_temp(e->manager); + } + if (!file) { + return -1; + } + + itable_insert(g->file_id_to_file, file_id, file); + vine_graph_io_mount_add(node->extra_outputs, file, task_path); + return 0; +} + +int vine_graph_executor_add_task_input_file(struct vine_graph_executor *e, uint64_t task_id, uint64_t file_id, const char *task_path) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !task_id || !file_id || !task_path) { + return -1; + } + + struct vine_graph_node *node = itable_lookup(g->nodes, task_id); + struct vine_file *file = itable_lookup(g->file_id_to_file, file_id); + if (!node || !file) { + return -1; + } + + vine_graph_io_mount_add(node->extra_inputs, file, task_path); + return 0; +} + +const char *vine_graph_executor_get_file_target_path(struct vine_graph_executor *e, uint64_t file_id) +{ + struct vine_graph *g = e ? e->graph : NULL; + struct vine_file *file = g ? itable_lookup(g->file_id_to_file, file_id) : NULL; + return file ? vine_file_source(file) : NULL; +} + +/* Apply executor-level tuning. Unknown keys are forwarded to vine_graph_tune. */ +int vine_graph_executor_tune(struct vine_graph_executor *e, const char *name, const char *value) +{ + if (!e || !name || !value) { + return -1; + } + + if (strcmp(name, "failure-injection-step-percent") == 0) { + e->failure_injection_step_percent = atof(value); + + } else if (strcmp(name, "task-priority-mode") == 0) { + if (strcmp(value, "random") == 0) { + e->task_priority_mode = TASK_PRIORITY_MODE_RANDOM; + } else if (strcmp(value, "depth-first") == 0) { + e->task_priority_mode = TASK_PRIORITY_MODE_DEPTH_FIRST; + } else if (strcmp(value, "breadth-first") == 0) { + e->task_priority_mode = TASK_PRIORITY_MODE_BREADTH_FIRST; + } else if (strcmp(value, "fifo") == 0) { + e->task_priority_mode = TASK_PRIORITY_MODE_FIFO; + } else if (strcmp(value, "lifo") == 0) { + e->task_priority_mode = TASK_PRIORITY_MODE_LIFO; + } else if (strcmp(value, "largest-input-first") == 0) { + e->task_priority_mode = TASK_PRIORITY_MODE_LARGEST_INPUT_FIRST; + } else if (strcmp(value, "largest-storage-footprint-first") == 0) { + e->task_priority_mode = TASK_PRIORITY_MODE_LARGEST_STORAGE_FOOTPRINT_FIRST; + } else { + debug(D_ERROR, "invalid priority mode: %s", value); + return -1; + } + + } else if (strcmp(name, "progress-bar-update-interval-sec") == 0) { + double val = atof(value); + e->progress_bar_update_interval_sec = (val > 0.0) ? val : 0.1; + + } else { + return vine_graph_tune(e->graph, name, value); + } + + return 0; +} + +/* Set the interrupted flag when SIGINT is received. */ +static void vine_graph_executor_handle_sigint(int signal) +{ + interrupted = 1; +} + +/* Compute submission priority for a node using the configured scheduling policy. */ +static double vine_graph_executor_calculate_task_priority(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!node || !g) { + return 0; + } + + double priority = 0; + timestamp_t current_time = timestamp_get(); + struct vine_graph_node *parent_node; + + switch (e->task_priority_mode) { + case TASK_PRIORITY_MODE_RANDOM: + priority = random_double(); + break; + case TASK_PRIORITY_MODE_DEPTH_FIRST: + priority = (double)node->depth; + break; + case TASK_PRIORITY_MODE_BREADTH_FIRST: + priority = -(double)node->depth; + break; + case TASK_PRIORITY_MODE_FIFO: + priority = -(double)current_time; // earlier time yields higher priority + break; + case TASK_PRIORITY_MODE_LIFO: + priority = (double)current_time; + break; + case TASK_PRIORITY_MODE_LARGEST_INPUT_FIRST: + LIST_ITERATE(node->parents, parent_node) + { + struct vine_graph_node *src = vine_graph_input_producer_node(g, parent_node, node); + if (!src || !src->outfile) { + continue; + } + priority += (double)vine_file_size(src->outfile); + } + break; + case TASK_PRIORITY_MODE_LARGEST_STORAGE_FOOTPRINT_FIRST: + LIST_ITERATE(node->parents, parent_node) + { + struct vine_graph_node *src = vine_graph_input_producer_node(g, parent_node, node); + if (!src || !src->outfile) { + continue; + } + if (!parent_node->task) { + continue; + } + timestamp_t parent_task_completion_time = vine_task_get_metric(parent_node->task, "time_workers_execute_last"); + priority += (double)vine_file_size(src->outfile) * (double)parent_task_completion_time; + } + break; + } + + return priority; +} + +/* Submit the node task if it is still initial, and record the manager task id for later lookup. */ +static void vine_graph_executor_submit_node(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !node) { + return; + } + + timestamp_t t_pre = timestamp_get(); + + vine_graph_executor_materialize_node(e, node); + + if (!node->task) { + debug(D_ERROR, "vine_graph_executor_submit_node: node %" PRIu64 " has no task after materialize", node->node_id); + goto record_preprocessing; + } + + if (!vine_graph_task_not_submitted(node->task)) { + debug(D_VINE, + "vine_graph_executor_submit_node: skipping node %" PRIu64 " (task already submitted, state=%s, task_id=%d)", + node->node_id, + vine_task_get_state(node->task), + vine_task_get_id(node->task)); + goto record_preprocessing; + } + + double priority = vine_graph_executor_calculate_task_priority(e, node); + vine_task_set_priority(node->task, priority); + + int task_id = vine_submit(e->manager, node->task); + + if (task_id <= 0) { + debug(D_ERROR, "vine_graph_executor_submit_node: failed to submit node %" PRIu64 " (returned task_id=%d)", node->node_id, task_id); + goto record_preprocessing; + } + + itable_insert(e->task_id_to_node, (uint64_t)task_id, node); // reverse lookup from vine_wait + debug(D_VINE, "submitted node %" PRIu64 " with task id %d", node->node_id, task_id); + +record_preprocessing: { + uint64_t dt = (uint64_t)(timestamp_get() - t_pre); + node->preprocessing_time_us = dt; + e->total_preprocessing_time_us += dt; + debug(D_VINE, + "node %" PRIu64 " preprocessing %" PRIu64 " us, graph cumulative %" PRIu64 " us", + node->node_id, + dt, + e->total_preprocessing_time_us); +} +} + +/* + * Return true when this node is allowed to submit. If chain grouping is active, only the + * supernode leader may submit because other members are executed inside the leader's task. + */ +static int vine_graph_node_ready_for_submission(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!e || !node || node->remaining_parents_count != 0 || node->completed) { + return 0; + } + if (g && g->chain_grouping_enabled && !vine_graph_node_is_supernode_leader(node)) { + return 0; + } + if (node->in_resubmit_queue) { + return 0; + } + if (node->task && !vine_graph_task_not_submitted(node->task)) { + return 0; + } + return 1; +} + +/* Submit ready source nodes and enable delivery of recovery tasks to the application. */ +static void vine_graph_executor_submit_initial_ready_nodes(struct vine_graph_executor *e) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !e->manager) { + return; + } + + uint64_t nid; + struct vine_graph_node *node; + int iteration; + ITABLE_ITERATE(g->nodes, iteration, nid, node) + { + if (vine_graph_node_ready_for_submission(e, node)) { + vine_graph_executor_submit_node(e, node); + } + } + + vine_enable_external_recovery_handling(e->manager); // driver must observe recovery completions for cut and prune +} + +/* After one parent completes, decrement remaining parents and submit children that become ready. */ +static void vine_graph_executor_submit_unblocked_children(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !node) { + return; + } + + struct vine_graph_node *child_node; + LIST_ITERATE(node->children, child_node) + { + if (!child_node) { + continue; + } + + if (!child_node->fired_parents) { + child_node->fired_parents = set_create(0); + } + if (set_lookup(child_node->fired_parents, node)) { + continue; + } + set_insert(child_node->fired_parents, node); + + if (child_node->remaining_parents_count > 0) { + child_node->remaining_parents_count--; + } + + if (vine_graph_node_ready_for_submission(e, child_node)) { + vine_graph_executor_submit_node(e, child_node); + } + } +} + +/* Map a completed vine_task to the corresponding graph node, including recovery tasks. */ +static struct vine_graph_node *vine_graph_executor_node_from_task(struct vine_graph_executor *e, struct vine_task *task) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !task) { + return NULL; + } + + /* Recovery completions map to the original producer; regular completions map to themselves. */ + int lookup_task_id = vine_task_get_recovery_source_task_id(task); + if (lookup_task_id <= 0) { + lookup_task_id = vine_task_get_id(task); + } + if (lookup_task_id > 0) { + return itable_lookup(e->task_id_to_node, (uint64_t)lookup_task_id); + } + + debug(D_ERROR, "task %d has no graph node mapping", vine_task_get_id(task)); + return NULL; +} + +/* Update shared-filesystem byte counters when a node's credited output size changes. */ +static void vine_graph_executor_account_pfs_write(struct vine_graph_executor *e, struct vine_graph_node *n, size_t new_size) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !n) { + return; + } + + size_t prev = n->pfs_credited_bytes; + if (new_size == prev) { + return; + } + + if (new_size > prev) { + e->pfs_usage_bytes += (new_size - prev); + } else { + e->pfs_usage_bytes -= (prev - new_size); // retry may produce a smaller file + } + n->pfs_credited_bytes = new_size; + + debug(D_VINE, + "pfs write: node %" PRIu64 " size=%zu (prev=%zu) usage=%" PRIu64, + n->node_id, + new_size, + prev, + e->pfs_usage_bytes); +} + +/* Remove a node's credited bytes from the shared-filesystem usage total. */ +static void vine_graph_executor_account_pfs_delete(struct vine_graph_executor *e, struct vine_graph_node *n) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !n) { + return; + } + + size_t credited = n->pfs_credited_bytes; + if (credited == 0) { + return; + } + + e->pfs_usage_bytes -= credited; + n->pfs_credited_bytes = 0; + + debug(D_VINE, + "pfs delete: node %" PRIu64 " size=%zu usage=%" PRIu64, + n->node_id, + credited, + e->pfs_usage_bytes); +} + +/* Return non-zero if the completed node retains output on local disk or a shared filesystem path. */ +static int vine_graph_node_is_anchored(const struct vine_graph_node *n) +{ + if (!n || !n->completed) { + return 0; + } + return n->outfile_type == VINE_GRAPH_NODE_OUTFILE_TYPE_SHARED_FILE_SYSTEM || n->outfile_type == VINE_GRAPH_NODE_OUTFILE_TYPE_LOCAL; +} + +/* Return non-zero when a temporary output file has a recovery task that is neither initial nor finished. */ +static int vine_graph_node_is_mid_recovery(const struct vine_graph_node *n) +{ + if (!n || !n->outfile || vine_file_type(n->outfile) != VINE_TEMP) { + return 0; + } + return vine_file_is_recovering(n->outfile); +} + +/* Remove or prune the node's result file according to its output storage mode. */ +static void vine_graph_executor_delete_node_output(struct vine_graph_executor *e, struct vine_graph_node *n) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !n) { + return; + } + + switch (n->outfile_type) { + case VINE_GRAPH_NODE_OUTFILE_TYPE_TEMP: + if (n->outfile) { + vine_prune_file(e->manager, n->outfile); + } + break; + case VINE_GRAPH_NODE_OUTFILE_TYPE_SHARED_FILE_SYSTEM: + vine_graph_executor_account_pfs_delete(e, n); + if (n->outfile_remote_name) { + unlink(n->outfile_remote_name); + } + break; + case VINE_GRAPH_NODE_OUTFILE_TYPE_LOCAL: + if (n->outfile && vine_file_source(n->outfile)) { + unlink(vine_file_source(n->outfile)); + } + break; + } +} + +/* Attempt to mark a completed node as cut and delete its return file when all children permit release. */ +static int vine_graph_executor_try_cut_node(struct vine_graph_executor *e, struct vine_graph_node *n) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !n || n->cut || !n->completed) { + return 0; + } + + struct vine_graph_node *c; + LIST_ITERATE(n->children, c) + { + if ((!vine_graph_node_is_anchored(c) && !c->cut) || vine_graph_node_is_mid_recovery(c)) { + return 0; // wait for anchored, cut, or non-recovery children + } + } + + n->cut = 1; + debug(D_VINE, "cut: node %" PRIu64 " outfile_type=%d is_target=%d", n->node_id, n->outfile_type, n->is_target); + + if (!n->is_target) { + vine_graph_executor_delete_node_output(e, n); // targets keep data for retrieval + } + + return 1; +} + +/* Walk upstream from a completed node and apply cut propagation along the worklist. */ +static void vine_graph_executor_propagate_cut_from(struct vine_graph_executor *e, struct vine_graph_node *start) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !start || !start->completed) { + return; + } + + timestamp_t t0 = timestamp_get(); + vine_graph_executor_try_cut_node(e, start); + + /* Upstream BFS: when a node is cut, enqueue its parents for the same check. */ + struct list *worklist = list_create(); + struct vine_graph_node *p; + LIST_ITERATE(start->parents, p) + { + list_push_tail(worklist, p); + } + + while (list_size(worklist) > 0) { + struct vine_graph_node *m = list_pop_head(worklist); + if (vine_graph_executor_try_cut_node(e, m)) { + LIST_ITERATE(m->parents, p) + { + list_push_tail(worklist, p); + } + } + } + + list_delete(worklist); + e->time_spent_on_cut_propagation += timestamp_get() - t0; +} + +/* Return non-zero if every descendant within the given depth bound is complete and not mid-recovery. */ +static int vine_graph_node_descendants_completed_within_depth(struct vine_graph_node *a, int depth) +{ + if (!a || depth <= 0) { + return 1; + } + + struct set *visited = set_create(0); + struct list *current = list_create(); + list_push_tail(current, a); + set_insert(visited, a); + + int ok = 1; + /* Expand one child frontier per iteration up to depth hops from a. */ + for (int d = 0; d < depth && ok; d++) { + struct list *next = list_create(); + struct vine_graph_node *n; + LIST_ITERATE(current, n) + { + struct vine_graph_node *c; + LIST_ITERATE(n->children, c) + { + if (set_lookup(visited, c)) { + continue; + } + set_insert(visited, c); + if (!c->completed || vine_graph_node_is_mid_recovery(c)) { + ok = 0; + break; + } + list_push_tail(next, c); + } + if (!ok) { + break; + } + } + list_delete(current); + current = next; + } + + list_delete(current); + set_delete(visited); + return ok; +} + +/* Release a temporary output when prune-depth constraints and descendant completion are satisfied. */ +static void vine_graph_executor_try_prune_depth_release(struct vine_graph_executor *e, struct vine_graph_node *a) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !a) { + return; + } + if (a->released_by_prune_depth) { + return; + } + if (a->outfile_type != VINE_GRAPH_NODE_OUTFILE_TYPE_TEMP) { + return; + } + if (a->is_target) { + return; + } + if (!a->completed || !a->outfile) { + return; + } + if (!vine_graph_node_descendants_completed_within_depth(a, g->prune_depth)) { + return; // wait until descendants within prune_depth layers are settled + } + + vine_graph_executor_delete_node_output(e, a); + a->released_by_prune_depth = 1; + + debug(D_VINE, "prune-depth release: node %" PRIu64 " depth=%d", a->node_id, g->prune_depth); +} + +/* Apply prune-depth release starting at a node and extending up to k ancestor levels. */ +static void vine_graph_executor_apply_prune_depth_from(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !node) { + return; + } + int k = g->prune_depth; + if (k <= 0) { + return; + } + + vine_graph_executor_try_prune_depth_release(e, node); + + struct set *visited = set_create(0); + struct list *current = list_create(); + list_push_tail(current, node); + set_insert(visited, node); + + /* Visit new parents up to k levels, trying prune release on each. */ + for (int d = 1; d <= k; d++) { + struct list *next = list_create(); + struct vine_graph_node *n; + LIST_ITERATE(current, n) + { + struct vine_graph_node *p; + LIST_ITERATE(n->parents, p) + { + if (set_lookup(visited, p)) { + continue; + } + set_insert(visited, p); + list_push_tail(next, p); + vine_graph_executor_try_prune_depth_release(e, p); + } + } + list_delete(current); + current = next; + } + + list_delete(current); + set_delete(visited); +} + +/* + * Completion hook after a node finishes: cut propagation, prune-depth handling, and timing. + * Postprocessing wall time is charged to the node that triggered this call. + */ +static void vine_graph_executor_run_completion_postprocess(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + if (!e || !node) { + return; + } + + timestamp_t t0 = timestamp_get(); + vine_graph_executor_propagate_cut_from(e, node); + vine_graph_executor_apply_prune_depth_from(e, node); + uint64_t dt = (uint64_t)(timestamp_get() - t0); + node->postprocessing_time_us = dt; + e->total_postprocessing_time_us += dt; + debug(D_VINE, + "node %" PRIu64 " postprocessing %" PRIu64 " us, graph cumulative %" PRIu64 " us", + node->node_id, + dt, + e->total_postprocessing_time_us); +} + +#define RESUBMIT_SCAN_LIMIT 100 +#define RESUBMIT_COOLDOWN_USECS ((timestamp_t)1000000) + +/* + * Queue a retry after failure. With chain grouping the queue stores the leader so one retry + * resubmits the whole merged task. Without grouping the failing node itself is the leader. + */ +static void vine_graph_executor_queue_node_retry(struct vine_graph_executor *e, struct vine_graph_node *node) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g || !node) { + return; + } + + struct vine_graph_node *leader = node; + if (g->chain_grouping_enabled) { + struct vine_graph_node *mapped = vine_graph_supernode_leader_node(g, node); + if (mapped) { + leader = mapped; + } + } + if (!leader) { + return; + } + + if (leader->in_resubmit_queue) { + return; + } + + leader->last_failure_time = timestamp_get(); + list_push_tail(e->resubmit_queue, leader); + leader->in_resubmit_queue = 1; +} + +/* + * Process the resubmit queue after cooldown. A ready head is popped, its failed task is torn down, + * and vine_graph_executor_submit_node runs again. If the head is still cooling off, rotate it to the tail so + * other leaders behind it are not stuck forever while the driver waits. + */ +static void vine_graph_executor_drain_resubmit_queue(struct vine_graph_executor *e) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g) { + return; + } + + timestamp_t now = timestamp_get(); + int queued = list_size(e->resubmit_queue); + if (queued == 0) { + return; + } + + int budget = queued < RESUBMIT_SCAN_LIMIT ? queued : RESUBMIT_SCAN_LIMIT; + int resubmits = 0; + int rotations_without_resubmit = 0; + + while (resubmits < budget && list_size(e->resubmit_queue) > 0) { + struct vine_graph_node *node = list_peek_head(e->resubmit_queue); + if (!node) { + break; + } + if (now - node->last_failure_time >= RESUBMIT_COOLDOWN_USECS) { + list_pop_head(e->resubmit_queue); + node->in_resubmit_queue = 0; + + debug(D_VINE, "Resubmitting node %" PRIu64, node->node_id); + vine_graph_executor_clear_node_runner_arg(e, node); + if (node->task) { + vine_task_delete(node->task); + node->task = NULL; + } + vine_graph_executor_submit_node(e, node); + resubmits++; + rotations_without_resubmit = 0; + } else { + list_pop_head(e->resubmit_queue); + list_push_tail(e->resubmit_queue, node); + rotations_without_resubmit++; + if (rotations_without_resubmit >= list_size(e->resubmit_queue)) { + break; + } + } + } +} + +/* + * Verify status and on-disk outputs (primary outfile only). + * On failure enqueue a retry for the node. + */ +static int vine_graph_executor_validate_node_outputs_or_retry(struct vine_graph_executor *e, struct vine_graph_node *retry_node, struct vine_graph_node *output_node, struct vine_task *task) +{ + switch (output_node->outfile_type) { + case VINE_GRAPH_NODE_OUTFILE_TYPE_SHARED_FILE_SYSTEM: { + struct stat info; + // shared path may lag behind a successful task result code + if (stat(output_node->outfile_remote_name, &info) < 0) { + debug(D_VINE, "Task %d succeeded but missing sharedfs output %s", vine_task_get_id(task), output_node->outfile_remote_name); + vine_graph_executor_queue_node_retry(e, retry_node); + return 0; + } + output_node->outfile_size_bytes = info.st_size; + vine_graph_executor_account_pfs_write(e, output_node, (size_t)info.st_size); + break; + } + case VINE_GRAPH_NODE_OUTFILE_TYPE_LOCAL: + case VINE_GRAPH_NODE_OUTFILE_TYPE_TEMP: + if (output_node->outfile) { + output_node->outfile_size_bytes = vine_file_size(output_node->outfile); + } + break; + } + + return 1; +} + +/* + * Verify one extra output mount after a successful task (VINE_FILE paths must exist). + */ +static int vine_graph_executor_validate_io_mount_or_retry( + struct vine_graph_executor *e, struct vine_graph_node *retry_node, struct vine_graph_io_mount *mount, struct vine_task *task) +{ + if (!mount || !mount->file) { + return 1; + } + struct vine_file *f = mount->file; + + switch (vine_file_type(f)) { + case VINE_TEMP: + case VINE_BUFFER: + break; + case VINE_FILE: + if (vine_file_source(f)) { + struct stat info; + if (stat(vine_file_source(f), &info) < 0) { + debug(D_VINE, + "Task %d succeeded but missing extra output file %s (%s)", + vine_task_get_id(task), + mount->remote_name ? mount->remote_name : "?", + vine_file_source(f)); + vine_graph_executor_queue_node_retry(e, retry_node); + return 0; + } + } + break; + default: + break; + } + return 1; +} + +/* + * Primary outfile (per node outfile_type) plus every extra_outputs mount declared for that graph node. + */ +static int vine_graph_executor_validate_all_declared_outputs_or_retry( + struct vine_graph_executor *e, struct vine_graph_node *retry_node, struct vine_graph_node *output_node, struct vine_task *task) +{ + if (!vine_graph_executor_validate_node_outputs_or_retry(e, retry_node, output_node, task)) { + return 0; + } + void *item; + LIST_ITERATE(output_node->extra_outputs, item) + { + if (!vine_graph_executor_validate_io_mount_or_retry(e, retry_node, (struct vine_graph_io_mount *)item, task)) { + return 0; + } + } + return 1; +} + +static int vine_graph_executor_validate_task_or_retry(struct vine_graph_executor *e, struct vine_graph_node *node, struct vine_task *task) +{ + struct vine_graph *g = e ? e->graph : NULL; + /* + * Returning zero means the completion is rejected, a retry is enqueued, and the progress + * bar must not advance because the batch did not truly succeed yet. + */ + if (vine_task_get_result(task) != VINE_RESULT_SUCCESS || vine_task_get_exit_code(task) != 0) { + debug(D_VINE, + "Task %d failed (result=%d, exit=%d)", + vine_task_get_id(task), + vine_task_get_result(task), + vine_task_get_exit_code(task)); + vine_graph_executor_queue_node_retry(e, node); + return 0; + } + + if (!vine_graph_executor_validate_all_declared_outputs_or_retry(e, node, node, task)) { + return 0; + } + + /* + * One vine task may carry outputs for the whole supernode. Validate every grouped member's + * declared files against that same task, not only the leader's primary outfile. + */ + if (g && g->chain_grouping_enabled) { + struct list *mems = vine_graph_supernode_nonleader_members(g, node->node_id); + if (mems) { + struct vine_graph_node *m; + LIST_ITERATE(mems, m) + { + if (!vine_graph_executor_validate_all_declared_outputs_or_retry(e, node, m, task)) { + return 0; + } + } + } + } + + return 1; +} + +/* Return the recorded user-task makespan in microseconds. */ +uint64_t vine_graph_executor_get_makespan_us(const struct vine_graph_executor *e) +{ + if (!e) { + return 0; + } + + return (uint64_t)e->makespan_us; +} + +/* Return the manager's cumulative count of recovery tasks submitted. */ +uint64_t vine_graph_executor_get_total_recovery_tasks(const struct vine_graph_executor *e) +{ + if (!e || !e->manager) { + return 0; + } + + struct vine_stats stats; + vine_get_stats(e->manager, &stats); + return (uint64_t)stats.tasks_recovery; +} + +/* Return how many recovery tasks have completed in the current executor run. */ +uint64_t vine_graph_executor_get_completed_recovery_tasks(const struct vine_graph_executor *e) +{ + if (!e) { + return 0; + } + + return e->completed_recovery_tasks; +} + +/* Main loop: submit work, wait, handle recovery, update graph state. */ +void vine_graph_executor_execute(struct vine_graph_executor *e) +{ + struct vine_graph *g = e ? e->graph : NULL; + if (!g) { + return; + } + + interrupted = 0; + void (*previous_sigint_handler)(int) = signal(SIGINT, vine_graph_executor_handle_sigint); + + debug(D_VINE, "start executing executor graph"); + + vine_graph_executor_submit_initial_ready_nodes(e); + + struct ProgressBar *pbar = progress_bar_init("Executing Tasks"); + progress_bar_set_update_interval(pbar, e->progress_bar_update_interval_sec); + e->completed_recovery_tasks = 0; + + struct ProgressBarPart *user_tasks_part = progress_bar_create_part("User", itable_size(g->nodes)); + struct ProgressBarPart *recovery_tasks_part = progress_bar_create_part("Recovery", 0); + progress_bar_bind_part(pbar, user_tasks_part); + progress_bar_bind_part(pbar, recovery_tasks_part); + + const uint64_t user_node_total = itable_size(g->nodes); + + double next_failure_threshold = -1.0; + if (e->failure_injection_step_percent > 0) { + next_failure_threshold = e->failure_injection_step_percent / 100.0; + } + + int wait_timeout = 1; // short timeout after a result, longer when idle + + /* + * Stop the main loop once every graph node reports completed. Count completed nodes directly + * instead of relying on progress bar ticks because grouped members can flip to completed in + * a single completion event and the bar would miss intermediate steps. + */ + while (vine_graph_executor_count_completed_user_nodes(g) < user_node_total) { + if (interrupted) { + break; + } + + vine_graph_executor_drain_resubmit_queue(e); + progress_bar_set_part_total(pbar, recovery_tasks_part, vine_graph_executor_get_total_recovery_tasks(e)); + + struct vine_task *task = vine_wait(e->manager, wait_timeout); + if (task) { + wait_timeout = 0; + + struct vine_graph_node *node = vine_graph_executor_node_from_task(e, task); + if (!node) { + debug(D_ERROR, "fatal: task %d could not be mapped to a task node, this indicates a serious bug.", vine_task_get_id(task)); + exit(1); + } + + timestamp_t commit_end = vine_task_get_metric(task, "time_when_commit_end"); + if (commit_end > 0) { + e->time_first_task_dispatched = MIN(e->time_first_task_dispatched, commit_end); // makespan start + } + + /* + * User and recovery progress advances only after outputs validate. Failed tasks enter + * the retry path and leave the bar unchanged until a later successful completion. + */ + if (!vine_graph_executor_validate_task_or_retry(e, node, task)) { + continue; + } + + int first_completion = 0; + + if (vine_task_get_recovery_source_task_id(task) > 0) { + e->completed_recovery_tasks++; + progress_bar_update_part( + pbar, + recovery_tasks_part, + e->completed_recovery_tasks - recovery_tasks_part->current); + + /* Reset cut and prune-depth flags for recovery tasks. */ + node->cut = 0; + node->released_by_prune_depth = 0; + + /* Only postprocess recovery tasks. */ + vine_graph_executor_run_completion_postprocess(e, node); + } else { + timestamp_t retrieval_time = (timestamp_t)vine_task_get_metric(task, "time_when_retrieval"); + e->time_last_task_retrieved = MAX(e->time_last_task_retrieved, retrieval_time); + e->makespan_us = e->time_last_task_retrieved - e->time_first_task_dispatched; + + first_completion = !node->completed; + vine_graph_executor_mark_user_node_completed_after_success(g, node); + + if (first_completion) { + if (user_tasks_part->current == 0) { + progress_bar_set_start_time(pbar, vine_task_get_metric(task, "time_when_commit_start")); + } + + vine_graph_node_update_critical_path_time(node, vine_task_get_metric(task, "time_workers_execute_last")); + } + + uint64_t completed_user_nodes = vine_graph_executor_count_completed_user_nodes(g); + if (completed_user_nodes > user_tasks_part->current) { + /* Count graph nodes, then advance by the delta to avoid double-counting retries. */ + progress_bar_update_part(pbar, user_tasks_part, completed_user_nodes - user_tasks_part->current); + } + + if (e->failure_injection_step_percent > 0) { + // test hook, drop workers at stepped progress thresholds + double progress = (double)user_tasks_part->current / (double)user_tasks_part->total; + if (progress >= next_failure_threshold && vine_manager_release_random_worker(e->manager)) { + debug(D_VINE, "released a random worker at %.2f%% (threshold %.2f%%)", progress * 100, next_failure_threshold * 100); + next_failure_threshold += e->failure_injection_step_percent / 100.0; + } + } + + /* Postprocess the node and submit its children. */ + vine_graph_executor_run_completion_postprocess(e, node); + vine_graph_executor_submit_unblocked_children(e, node); + } + } else { + wait_timeout = 1; // no task ready, wait with default blocking timeout + } + } + + progress_bar_finish(pbar); + progress_bar_delete(pbar); + + debug(D_VINE, "total time spent on cut propagation: %.6f seconds\n", e->time_spent_on_cut_propagation / 1e6); + + signal(SIGINT, previous_sigint_handler); + if (interrupted) { + raise(SIGINT); // restore handler first, then honor prior interrupt + } +} diff --git a/taskvine/src/vine_graph/vine_graph_executor.h b/taskvine/src/vine_graph/vine_graph_executor.h new file mode 100644 index 0000000000..ae96d851eb --- /dev/null +++ b/taskvine/src/vine_graph/vine_graph_executor.h @@ -0,0 +1,57 @@ +#ifndef VINE_GRAPH_EXECUTOR_H +#define VINE_GRAPH_EXECUTOR_H + +#include "vine_graph.h" + +#include "taskvine.h" + +typedef enum { + TASK_PRIORITY_MODE_RANDOM = 0, + TASK_PRIORITY_MODE_DEPTH_FIRST, + TASK_PRIORITY_MODE_BREADTH_FIRST, + TASK_PRIORITY_MODE_FIFO, + TASK_PRIORITY_MODE_LIFO, + TASK_PRIORITY_MODE_LARGEST_INPUT_FIRST, + TASK_PRIORITY_MODE_LARGEST_STORAGE_FOOTPRINT_FIRST +} task_priority_mode_t; + +struct vine_graph_executor { + struct vine_graph *graph; // DAG executed by this executor + struct vine_manager *manager; // TaskVine runtime + + struct itable *task_id_to_node; // maps vine task id to graph node after submit + struct list *resubmit_queue; // nodes waiting for retry + + timestamp_t time_first_task_dispatched; // earliest dispatch time among user tasks + timestamp_t time_last_task_retrieved; // latest user task retrieval time + timestamp_t makespan_us; // workflow span in microseconds + timestamp_t time_spent_on_cut_propagation; // time spent in cut propagation + uint64_t completed_recovery_tasks; // recovery completions seen this run + uint64_t pfs_usage_bytes; // bytes credited for shared filesystem outputs + /** Sum of @c vine_graph_executor_submit_node preprocessing intervals across all nodes (microseconds). */ + uint64_t total_preprocessing_time_us; + /** Sum of @c vine_graph_executor_run_completion_postprocess intervals across all completions (microseconds). */ + uint64_t total_postprocessing_time_us; + + task_priority_mode_t task_priority_mode; // schedule order before submit + double failure_injection_step_percent; // optional worker release steps for tests + double progress_bar_update_interval_sec; +}; + +struct vine_graph_executor *vine_graph_executor_create(struct vine_manager *manager, struct vine_graph *graph); +struct vine_graph *vine_graph_executor_create_graph(struct vine_manager *manager); +void vine_graph_executor_delete(struct vine_graph_executor *e); +uint64_t vine_graph_executor_add_node(struct vine_graph_executor *e); +void vine_graph_executor_finalize(struct vine_graph_executor *e); +int vine_graph_executor_declare_input_file(struct vine_graph_executor *e, uint64_t file_id, const char *source_path); +int vine_graph_executor_add_task_output_file(struct vine_graph_executor *e, uint64_t task_id, uint64_t file_id, const char *task_path, int is_target); +int vine_graph_executor_add_task_input_file(struct vine_graph_executor *e, uint64_t task_id, uint64_t file_id, const char *task_path); +const char *vine_graph_executor_get_file_target_path(struct vine_graph_executor *e, uint64_t file_id); + +int vine_graph_executor_tune(struct vine_graph_executor *e, const char *name, const char *value); +void vine_graph_executor_execute(struct vine_graph_executor *e); +uint64_t vine_graph_executor_get_makespan_us(const struct vine_graph_executor *e); +uint64_t vine_graph_executor_get_total_recovery_tasks(const struct vine_graph_executor *e); +uint64_t vine_graph_executor_get_completed_recovery_tasks(const struct vine_graph_executor *e); + +#endif // VINE_GRAPH_EXECUTOR_H diff --git a/taskvine/src/vine_graph/vine_graph_node.c b/taskvine/src/vine_graph/vine_graph_node.c new file mode 100644 index 0000000000..85bbb0a7fc --- /dev/null +++ b/taskvine/src/vine_graph/vine_graph_node.c @@ -0,0 +1,276 @@ +#include +#include + +#include "debug.h" +#include "list.h" +#include "stringtools.h" +#include "xxmalloc.h" + +#include "vine_graph_node.h" + +/*************************************************************/ +/* Public APIs */ +/*************************************************************/ + +/** + * Update the critical path time of a node. + * @param node Reference to the node object. + * @param execution_time Reference to the execution time of the node. + */ +void vine_graph_node_update_critical_path_time(struct vine_graph_node *node, timestamp_t execution_time) +{ + timestamp_t max_parent_critical_path_time = 0; + struct vine_graph_node *parent_node; + LIST_ITERATE(node->parents, parent_node) + { + if (parent_node->critical_path_time > max_parent_critical_path_time) { + max_parent_critical_path_time = parent_node->critical_path_time; + } + } + node->critical_path_time = max_parent_critical_path_time + execution_time; +} + +/** + * Create a new node owned by the C-side graph. + * @param node_id Graph-assigned identifier that keeps C and Python in sync. + * @return Newly allocated node. + */ +struct vine_graph_node *vine_graph_node_create(uint64_t node_id) +{ + struct vine_graph_node *node = xxmalloc(sizeof(struct vine_graph_node)); + + node->is_target = 0; + node->node_id = node_id; + node->super_leader_id = node_id; + + node->task = NULL; + node->task_runner_arg_file = NULL; + node->outfile = NULL; + node->outfile_remote_name = string_format("outfile_node_%" PRIu64, node->node_id); + node->outfile_type = VINE_GRAPH_NODE_OUTFILE_TYPE_TEMP; + + node->parents = list_create(); + node->children = list_create(); + node->extra_outputs = list_create(); + node->extra_inputs = list_create(); + node->remaining_parents_count = 0; + node->fired_parents = NULL; + node->completed = 0; + node->cut = 0; + node->released_by_prune_depth = 0; + node->outfile_size_bytes = 0; + node->pfs_credited_bytes = 0; + node->in_resubmit_queue = 0; + node->last_failure_time = 0; + + node->depth = -1; + node->height = -1; + node->upstream_subgraph_size = -1; + node->downstream_subgraph_size = -1; + node->fan_in = -1; + node->fan_out = -1; + node->heavy_score = -1; + + node->critical_path_time = -1; + node->preprocessing_time_us = 0; + node->postprocessing_time_us = 0; + + return node; +} + +/** Non-zero if parent->child is already in the adjacency lists (checked via parent's children). */ +static int vine_graph_node_dependency_exists(struct vine_graph_node *parent, struct vine_graph_node *child) +{ + struct vine_graph_node *x; + if (!parent || !child) { + return 0; + } + LIST_ITERATE(parent->children, x) + { + if (x == child) { + return 1; + } + } + return 0; +} + +void vine_graph_node_remove_dependency(struct vine_graph_node *parent, struct vine_graph_node *child) +{ + if (!parent || !child) { + return; + } + list_remove(child->parents, parent); + list_remove(parent->children, child); +} + +void vine_graph_node_ensure_dependency(struct vine_graph_node *parent, struct vine_graph_node *child) +{ + if (!parent || !child || vine_graph_node_dependency_exists(parent, child)) { + return; + } + list_push_tail(child->parents, parent); + list_push_tail(parent->children, child); +} + +/** + * Drop fired_parents so executor scheduling can recount parents (e.g. after supernode merge rewires edges). + */ +void vine_graph_node_clear_fired_parents(struct vine_graph_node *n) +{ + if (!n || !n->fired_parents) { + return; + } + set_delete(n->fired_parents); + n->fired_parents = NULL; +} + +/** + * Construct the task arguments for the node. + * @param node Reference to the node object. + * @return The task arguments in JSON format: {"fn_args": ["node_id"], "fn_kwargs": {}} (string for run_scheduler_keys). + */ +char *vine_graph_node_construct_task_arguments(struct vine_graph_node *node) +{ + if (!node) { + return NULL; + } + return string_format("{\"fn_args\":[\"%" PRIu64 "\"],\"fn_kwargs\":{}}", node->node_id); +} + +/** + * Print the info of the node. + * @param node Reference to the node object. + */ +void vine_graph_node_debug_print(struct vine_graph_node *node) +{ + if (!node) { + return; + } + + debug(D_VINE, "---------------- Node Info ----------------"); + debug(D_VINE, "node_id: %" PRIu64, node->node_id); + debug(D_VINE, "preprocessing_time_us (last): %" PRIu64, node->preprocessing_time_us); + debug(D_VINE, "postprocessing_time_us (last): %" PRIu64, node->postprocessing_time_us); + + if (!node->task) { + debug(D_VINE, "task: (none yet)"); + debug(D_VINE, "-------------------------------------------"); + return; + } + + debug(D_VINE, "task_id: %d", vine_task_get_id(node->task)); + debug(D_VINE, "depth: %d", node->depth); + debug(D_VINE, "height: %d", node->height); + + if (node->outfile_remote_name) { + debug(D_VINE, "outfile_remote_name: %s", node->outfile_remote_name); + } + + if (node->outfile) { + const char *type_str = "UNKNOWN"; + switch (vine_file_type(node->outfile)) { + case VINE_FILE: + type_str = "VINE_FILE"; + break; + case VINE_TEMP: + type_str = "VINE_TEMP"; + break; + case VINE_URL: + type_str = "VINE_URL"; + break; + case VINE_BUFFER: + type_str = "VINE_BUFFER"; + break; + case VINE_MINI_TASK: + type_str = "VINE_MINI_TASK"; + break; + } + const char *cached_name = vine_file_cached_name(node->outfile); + debug(D_VINE, "outfile_type: %s", type_str); + debug(D_VINE, "outfile_cached_name: %s", cached_name ? cached_name : "(null)"); + } else { + debug(D_VINE, "outfile_type: SHARED_FILE_SYSTEM or none"); + } + + char *parent_ids = NULL; // comma separated parent ids for logging + struct vine_graph_node *p; + LIST_ITERATE(node->parents, p) + { + if (!parent_ids) { + parent_ids = string_format("%" PRIu64, p->node_id); + } else { + char *tmp = string_format("%s, %" PRIu64, parent_ids, p->node_id); + free(parent_ids); + parent_ids = tmp; + } + } + + char *child_ids = NULL; // comma separated child ids for logging + struct vine_graph_node *c; + LIST_ITERATE(node->children, c) + { + if (!child_ids) { + child_ids = string_format("%" PRIu64, c->node_id); + } else { + char *tmp = string_format("%s, %" PRIu64, child_ids, c->node_id); + free(child_ids); + child_ids = tmp; + } + } + + debug(D_VINE, "parents: %s", parent_ids ? parent_ids : "(none)"); + debug(D_VINE, "children: %s", child_ids ? child_ids : "(none)"); + + free(parent_ids); + free(child_ids); + + debug(D_VINE, "-------------------------------------------"); +} + +/** + * Delete the node and all of its associated resources. + * @param node Reference to the node object. + */ +void vine_graph_node_delete(struct vine_graph_node *node) +{ + if (!node) { + return; + } + + if (node->outfile_remote_name) { + free(node->outfile_remote_name); + } + + vine_task_delete(node->task); + node->task = NULL; + + if (node->task_runner_arg_file) { + vine_file_delete(node->task_runner_arg_file); + node->task_runner_arg_file = NULL; + } + if (node->outfile) { + vine_file_delete(node->outfile); + node->outfile = NULL; + } + + list_delete(node->parents); + list_delete(node->children); + + while (list_size(node->extra_inputs) > 0) { + struct vine_graph_io_mount *m = list_pop_head(node->extra_inputs); + free(m->remote_name); + free(m); + } + list_delete(node->extra_inputs); + while (list_size(node->extra_outputs) > 0) { + struct vine_graph_io_mount *m = list_pop_head(node->extra_outputs); + free(m->remote_name); + free(m); + } + list_delete(node->extra_outputs); + + if (node->fired_parents) { + set_delete(node->fired_parents); + } + free(node); +} diff --git a/taskvine/src/vine_graph/vine_graph_node.h b/taskvine/src/vine_graph/vine_graph_node.h new file mode 100644 index 0000000000..bcbd8304a1 --- /dev/null +++ b/taskvine/src/vine_graph/vine_graph_node.h @@ -0,0 +1,128 @@ +#ifndef VINE_GRAPH_NODE_H +#define VINE_GRAPH_NODE_H + +#include + +#include "set.h" +#include "timestamp.h" + +#include "list.h" +#include "taskvine.h" + +/** + * One element of @c extra_outputs or @c extra_inputs: a logical filename plus its @c vine_file + * (declared during graph build; attached to @c vine_task in @c vine_graph_executor_materialize_node). + */ +struct vine_graph_io_mount { + struct vine_file *file; + char *remote_name; +}; + +/** The storage type of the node's output file. */ +typedef enum { + VINE_GRAPH_NODE_OUTFILE_TYPE_LOCAL = 0, // staged file under graph output_dir + VINE_GRAPH_NODE_OUTFILE_TYPE_TEMP, // TaskVine temp blob + VINE_GRAPH_NODE_OUTFILE_TYPE_SHARED_FILE_SYSTEM, // path on shared storage, no vine_file +} vine_graph_node_outfile_type_t; + +/** The node object. */ +struct vine_graph_node { + uint64_t node_id; // graph assigned id + /** + * Supernode leader id for scheduling: equals @c node_id for a single-node group. + * After @c vine_graph_supernode_register, every member shares the same leader id. + */ + uint64_t super_leader_id; + int is_target; // if set, output is retrieved when the task completes + + struct vine_task *task; + struct vine_file *task_runner_arg_file; // JSON args buffer for the runner + struct vine_file *outfile; // NULL when output is PFS only + char *outfile_remote_name; + size_t outfile_size_bytes; + vine_graph_node_outfile_type_t outfile_type; + size_t pfs_credited_bytes; // contribution to executor pfs_usage_bytes + + struct list *parents; + struct list *children; + /** + * Files tracked by TaskHandle.file(), beyond this node's primary + * Python-result outfile. Filled before + * @c node->task exists; consumed when building the task at submit / materialize time. + */ + struct list *extra_outputs; + /** + * FileHandle inputs beyond those implied by Python-result dependencies. Same lifecycle + * as @c extra_outputs: queued at graph build, wired on @c vine_task at materialize. + */ + struct list *extra_inputs; + + int remaining_parents_count; // parents not yet satisfied for scheduling + struct set *fired_parents; // parents already counted toward that count + int completed; + int cut; // return released by cut, cleared if recovery restores file + /** Non-zero after this node's temp output was released under @c graph->prune_depth; cleared on recovery. */ + int released_by_prune_depth; + int in_resubmit_queue; + timestamp_t last_failure_time; // last enqueue to resubmit queue + + int depth; + int height; + int upstream_subgraph_size; + int downstream_subgraph_size; + int fan_in; + int fan_out; + double heavy_score; + + timestamp_t critical_path_time; + /** Latest @c vine_graph_executor_submit_node interval for this node (microseconds); graph total is on @c struct vine_graph_executor. */ + uint64_t preprocessing_time_us; + /** Latest @c vine_graph_executor_run_completion_postprocess interval for this node (microseconds); graph total on executor. */ + uint64_t postprocessing_time_us; +}; + +/** Create a new node. +@param node_id Unique node identifier supplied by the owning graph. +@return Newly allocated node instance. +*/ +struct vine_graph_node *vine_graph_node_create(uint64_t node_id); + +/** + * Remove parent->child from both endpoints' parents/children lists (no-op if NULL). + * Used when rewiring supernodes so stale edges do not corrupt fan-in/out. + */ +void vine_graph_node_remove_dependency(struct vine_graph_node *parent, struct vine_graph_node *child); + +/** + * Add parent->child if that edge is not already present (idempotent). + */ +void vine_graph_node_ensure_dependency(struct vine_graph_node *parent, struct vine_graph_node *child); + +/** + * Drop fired_parents so executor scheduling can recount parents (e.g. after supernode merge rewires edges). + */ +void vine_graph_node_clear_fired_parents(struct vine_graph_node *n); + +/** Create the task arguments for a node. +@param node Reference to the node. +@return The task arguments in JSON format: {"fn_args": ["node_id"], "fn_kwargs": {}} (string id for run_scheduler_keys). +*/ +char *vine_graph_node_construct_task_arguments(struct vine_graph_node *node); + +/** Delete a node and release owned resources. +@param node Reference to the node. +*/ +void vine_graph_node_delete(struct vine_graph_node *node); + +/** Print information about a node. +@param node Reference to the node. +*/ +void vine_graph_node_debug_print(struct vine_graph_node *node); + +/** Update the critical path time of a node. +@param node Reference to the node. +@param execution_time Reference to the execution time of the node. +*/ +void vine_graph_node_update_critical_path_time(struct vine_graph_node *node, timestamp_t execution_time); + +#endif // VINE_GRAPH_NODE_H diff --git a/taskvine/src/worker/vine_worker.c b/taskvine/src/worker/vine_worker.c index 83bb87fead..d615cd105e 100644 --- a/taskvine/src/worker/vine_worker.c +++ b/taskvine/src/worker/vine_worker.c @@ -1658,6 +1658,7 @@ static void check_libraries_ready(struct link *manager) uint64_t library_task_id; struct vine_process *library_process; int iteration; + struct list *failed_libraries = list_create(); struct link_info library_link_info; library_link_info.events = LINK_READ; @@ -1678,12 +1679,12 @@ static void check_libraries_ready(struct link *manager) debug(D_VINE, "Library %s reports ready to execute functions.", library_process->task->provides_library); library_process->library_ready = 1; } else { - /* Kill library if it fails the startup check. */ + /* Collect failed libraries so procs_running is not re-iterated here. */ debug(D_VINE, "Library %s task id %" PRIu64 " verification failed (unexpected response). Killing it.", library_process->task->provides_library, library_task_id); - handle_failed_library_process(library_process, manager); + list_push_tail(failed_libraries, library_process); } } else { /* The library is running and the link has no readable data, do nothing until the @@ -1692,6 +1693,12 @@ static void check_libraries_ready(struct link *manager) library_link_info.revents = 0; } + + while ((library_process = list_pop_head(failed_libraries))) { + handle_failed_library_process(library_process, manager); + } + + list_delete(failed_libraries); } /* Start working for the (newly connected) manager on this given link. */ diff --git a/taskvine/test/TR_vine_graph_dask_adaptor.sh b/taskvine/test/TR_vine_graph_dask_adaptor.sh new file mode 100755 index 0000000000..2e22a63a49 --- /dev/null +++ b/taskvine/test/TR_vine_graph_dask_adaptor.sh @@ -0,0 +1,48 @@ +#!/bin/sh +set -e + +. ../../dttools/test/test_runner_common.sh + +import_config_val CCTOOLS_PYTHON_TEST_EXEC +import_config_val CCTOOLS_PYTHON_TEST_DIR + +export PYTHONPATH=$(pwd)/../../test_support/python_modules/${CCTOOLS_PYTHON_TEST_DIR}:$PYTHONPATH + +STATUS_FILE=vine_graph_dask_adaptor.status + +check_needed() +{ + [ -n "${CCTOOLS_PYTHON_TEST_EXEC}" ] || return 1 + "${CCTOOLS_PYTHON_TEST_EXEC}" -c "import dask" || return 1 + return 0 +} + +prepare() +{ + rm -f $STATUS_FILE + return 0 +} + +run() +{ + ${CCTOOLS_PYTHON_TEST_EXEC} vine_graph_dask_adaptor.py + echo $? > $STATUS_FILE + + status=$(cat $STATUS_FILE) + if [ $status -ne 0 ] + then + exit 1 + fi + + exit 0 +} + +clean() +{ + rm -f $STATUS_FILE + exit 0 +} + +dispatch "$@" + +# vim: set noexpandtab tabstop=4: diff --git a/taskvine/test/TR_vine_graph_task_group.sh b/taskvine/test/TR_vine_graph_task_group.sh new file mode 100755 index 0000000000..7553b3e0ab --- /dev/null +++ b/taskvine/test/TR_vine_graph_task_group.sh @@ -0,0 +1,86 @@ +#!/bin/sh +set -e + +. ../../dttools/test/test_runner_common.sh + +import_config_val CCTOOLS_PYTHON_TEST_EXEC +import_config_val CCTOOLS_PYTHON_TEST_DIR + +export PYTHONPATH=$(pwd)/../../test_support/python_modules/${CCTOOLS_PYTHON_TEST_DIR}:$PYTHONPATH +export PATH=$(dirname "${CCTOOLS_PYTHON_TEST_EXEC}"):$PATH + +CASE=chain-branches:6 + +check_needed() +{ + [ -n "${CCTOOLS_PYTHON_TEST_EXEC}" ] || return 1 + "${CCTOOLS_PYTHON_TEST_EXEC}" -c "import cloudpickle" || return 1 + return 0 +} + +prepare() +{ + rm -f vine_graph_task_group_*.status vine_graph_task_group_*.port vine_graph_task_group_*.result vine_graph_task_group_*.log worker.*.log + return 0 +} + +run_one() +{ + task_group=$1 + port_file=vine_graph_task_group_${task_group}.port + status_file=vine_graph_task_group_${task_group}.status + result_file=vine_graph_task_group_${task_group}.result + run_log=vine_graph_task_group_${task_group}.log + worker_log=worker.${task_group}.log + + rm -f $port_file $status_file $result_file $run_log $worker_log + + ( ${CCTOOLS_PYTHON_TEST_EXEC} vine_graph_workflow_examples.py $port_file --case $CASE --task-group $task_group --result-file $result_file --no-print-results --timeout 90 > $run_log 2>&1; echo $? > $status_file ) & + + wait_for_file_creation $port_file 15 + + cores=16 + memory=2000 + disk=2000 + run_taskvine_worker $port_file $worker_log + + wait_for_file_creation $status_file 60 + + status=$(cat $status_file) + if [ $status -ne 0 ] + then + exit 1 + fi + + test -s $result_file +} + +print_throughput_comparison() +{ + without_group=$(awk '/^=== Throughput:/ {print $3 " " $4; exit}' vine_graph_task_group_0.log) + with_group=$(awk '/^=== Throughput:/ {print $3 " " $4; exit}' vine_graph_task_group_1.log) + + echo "Task-group throughput comparison:" + echo " task-group=0: ${without_group:-unknown}" + echo " task-group=1: ${with_group:-unknown}" +} + +run() +{ + run_one 0 + run_one 1 + print_throughput_comparison + require_identical_files vine_graph_task_group_0.result vine_graph_task_group_1.result + exit 0 +} + +clean() +{ + rm -f vine_graph_task_group_*.status vine_graph_task_group_*.port vine_graph_task_group_*.result vine_graph_task_group_*.log worker.*.log + rm -rf vine-run-info + exit 0 +} + +dispatch "$@" + +# vim: set noexpandtab tabstop=4: diff --git a/taskvine/test/TR_vine_graph_workflow_examples.sh b/taskvine/test/TR_vine_graph_workflow_examples.sh new file mode 100755 index 0000000000..61cab64170 --- /dev/null +++ b/taskvine/test/TR_vine_graph_workflow_examples.sh @@ -0,0 +1,61 @@ +#!/bin/sh +set -e + +. ../../dttools/test/test_runner_common.sh + +import_config_val CCTOOLS_PYTHON_TEST_EXEC +import_config_val CCTOOLS_PYTHON_TEST_DIR + +export PYTHONPATH=$(pwd)/../../test_support/python_modules/${CCTOOLS_PYTHON_TEST_DIR}:$PYTHONPATH +export PATH=$(dirname "${CCTOOLS_PYTHON_TEST_EXEC}"):$PATH + +STATUS_FILE=vine_graph_workflow_examples.status +PORT_FILE=vine_graph_workflow_examples.port +RESULT_FILE=vine_graph_workflow_examples.result + +check_needed() +{ + [ -n "${CCTOOLS_PYTHON_TEST_EXEC}" ] || return 1 + "${CCTOOLS_PYTHON_TEST_EXEC}" -c "import cloudpickle" || return 1 + return 0 +} + +prepare() +{ + rm -f $STATUS_FILE $PORT_FILE $RESULT_FILE worker.log + return 0 +} + +run() +{ + ( ${CCTOOLS_PYTHON_TEST_EXEC} vine_graph_workflow_examples.py $PORT_FILE --case corner-cases --result-file $RESULT_FILE --no-print-results --timeout 90; echo $? > $STATUS_FILE ) & + + wait_for_file_creation $PORT_FILE 15 + + cores=16 + memory=2000 + disk=2000 + run_taskvine_worker $PORT_FILE worker.log + + wait_for_file_creation $STATUS_FILE 60 + + status=$(cat $STATUS_FILE) + if [ $status -ne 0 ] + then + exit 1 + fi + + test -s $RESULT_FILE + exit 0 +} + +clean() +{ + rm -f $STATUS_FILE $PORT_FILE $RESULT_FILE worker.log + rm -rf vine-run-info + exit 0 +} + +dispatch "$@" + +# vim: set noexpandtab tabstop=4: diff --git a/taskvine/test/vine_graph_dask_adaptor.py b/taskvine/test/vine_graph_dask_adaptor.py new file mode 100644 index 0000000000..a541b8c3a4 --- /dev/null +++ b/taskvine/test/vine_graph_dask_adaptor.py @@ -0,0 +1,145 @@ +import warnings + +from ndcctools.taskvine.vine_graph import TaskOutputHandle +from ndcctools.taskvine.vine_graph.adaptors import VineGraphDaskAdaptor + + +MIN_TASKS = 10 + + +def require_module(name): + """Import an optional test dependency, or skip that case if it is unavailable.""" + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + module = __import__(name, fromlist=["*"]) + except Exception: + return None + return module + + +def get_delayed(): + """Return Dask's delayed constructor across common Dask versions.""" + try: + from dask import delayed + + return delayed + except ImportError: + import dask + + return dask.delayed + + +def inc(value): + return value + 1 + + +def add_many(*values): + return sum(values) + + +def assert_has_task_output_handle(value): + """Return true if a converted argument tree contains a VineGraph dependency.""" + if isinstance(value, TaskOutputHandle): + return True + if isinstance(value, (list, tuple, set, frozenset)): + return any(assert_has_task_output_handle(item) for item in value) + if isinstance(value, dict): + return any(assert_has_task_output_handle(item) for item in value.values()) + return False + + +def assert_converts(name, graph, min_tasks=1, expect_dependency=True): + """Convert a Dask graph and check the basic VineGraph task-expression shape.""" + converted = VineGraphDaskAdaptor(graph).converted + assert len(converted) >= min_tasks, f"{name}: expected at least {min_tasks} tasks, got {len(converted)}" + + has_dependency = False + for _func, args, kwargs in converted.values(): + has_dependency = has_dependency or any(assert_has_task_output_handle(arg) for arg in args) + has_dependency = has_dependency or any(assert_has_task_output_handle(value) for value in kwargs.values()) + + if expect_dependency: + assert has_dependency, f"{name}: expected at least one TaskOutputHandle dependency" + + return converted + + +def check_low_level_dict(): + """Check Dask's explicit low-level task dictionary form.""" + graph = {f"x-{i}": (inc, i) for i in range(MIN_TASKS)} + graph["total"] = (add_many, *graph.keys()) + + assert_converts( + "low-level-dict", + graph, + min_tasks=MIN_TASKS, + ) + + +def check_delayed(): + """Check the common dask.delayed graph construction API.""" + delayed = get_delayed() + + leaves = [delayed(inc, pure=True)(i) for i in range(MIN_TASKS)] + total = delayed(add_many, pure=True)(*leaves) + + converted = assert_converts("delayed", {"result": total}, min_tasks=MIN_TASKS) + + add_many_tasks = [ + task_expr + for task_expr in converted.values() + if getattr(task_expr[0], "__name__", None) == "add_many" + ] + assert len(add_many_tasks) == 1 + + _func, args, kwargs = add_many_tasks[0] + assert kwargs == {} + assert len(args) == MIN_TASKS + assert all(isinstance(arg, TaskOutputHandle) for arg in args) + + +def check_array_collection(): + """Check a Dask Array collection when the optional array package is available.""" + da = require_module("dask.array") + if da is None: + return + + values = da.arange(2 * MIN_TASKS, chunks=2) + total = (values + 1).sum() + assert_converts("array", {"result": total}, min_tasks=MIN_TASKS) + + +def check_bag_collection(): + """Check a Dask Bag collection when the optional bag package is available.""" + db = require_module("dask.bag") + if db is None: + return + + total = db.from_sequence(list(range(MIN_TASKS)), npartitions=MIN_TASKS).map(inc).sum() + assert_converts("bag", {"result": total}, min_tasks=MIN_TASKS) + + +def check_dataframe_collection(): + """Check a Dask DataFrame collection when dataframe dependencies are available.""" + dd = require_module("dask.dataframe") + pd = require_module("pandas") + if dd is None or pd is None: + return + + frame = pd.DataFrame({"value": list(range(MIN_TASKS))}) + ddf = dd.from_pandas(frame, npartitions=MIN_TASKS) + total = (ddf.value + 1).sum() + assert_converts("dataframe", {"result": total}, min_tasks=MIN_TASKS) + + +def main(): + check_low_level_dict() + check_delayed() + check_array_collection() + check_bag_collection() + check_dataframe_collection() + + +if __name__ == "__main__": + main() diff --git a/taskvine/test/vine_graph_workflow_examples.py b/taskvine/test/vine_graph_workflow_examples.py new file mode 100644 index 0000000000..e53b7d2823 --- /dev/null +++ b/taskvine/test/vine_graph_workflow_examples.py @@ -0,0 +1,622 @@ +import argparse +from collections import deque, namedtuple +from dataclasses import dataclass +import json +import os +import shutil +import signal +import sys +import tempfile +from pathlib import Path + +import cloudpickle +import ndcctools.taskvine.vine_graph.vine_graph as vine_graph_mod +from ndcctools.taskvine.vine_graph import FileHandle, TaskHandle, TaskOutputHandle, VineGraph, Workflow + +TEST_DIR = Path(__file__).resolve().parent +MAX_BRANCHES = 32 + + +def add(*args): + return sum(args) + + +def make_simple_graph(): + bg = Workflow() + bg.add_task(add, 1, 5) + return bg + + +def make_chain_graph(chain_len=1, branches=1): + chain_len = max(1, int(chain_len)) + branches = max(1, int(branches)) + bg = Workflow() + for b in range(branches): + prev = bg.add_task(add, 1) + for i in range(1, chain_len): + prev = bg.add_task(add, prev.output()) + return bg + + +def make_chain_rich(n=1): + n = max(1, int(n)) + bg = Workflow() + if n == 1: + bg.add_task(add, 1) + return bg + branch_count = min(MAX_BRANCHES, max(1, n // 8)) + base, extra = divmod(n, branch_count) + for b in range(branch_count): + size = base + (1 if b < extra else 0) + prev = bg.add_task(add, 1) + for i in range(1, size): + prev = bg.add_task(add, prev.output()) + return bg + + +def make_individuals(n=1): + n = max(1, int(n)) + bg = Workflow() + for _ in range(n): + bg.add_task(add, 1) + return bg + + +def make_trivial(n=1): + bg = Workflow() + for _ in range(max(1, int(n))): + bg.add_task(add, 1) + return bg + + +def _add_binary_tree(bg, n): + tasks = [None] * n + last = (n - 2) // 2 + for i in range(last + 1, n): + tasks[i] = bg.add_task(add, 1) + for i in range(last, -1, -1): + deps = [tasks[2 * i + 1].output()] + if 2 * i + 2 < n: + deps.append(tasks[2 * i + 2].output()) + tasks[i] = bg.add_task(add, *deps) + + +def make_binary_tree(n=1): + n = max(1, int(n)) + bg = Workflow() + _add_binary_tree(bg, n) + return bg + + +def make_binary_forest(n=None, *, branches=5, level=8): + bg = Workflow() + if n is not None: + n = max(1, int(n)) + branches = max(1, min(n, MAX_BRANCHES)) + base, extra = divmod(n, branches) + for b in range(branches): + size = base + (1 if b < extra else 0) + _add_binary_tree(bg, size) + else: + branches, level = max(1, branches), max(1, level) + tree_size = 2**level - 1 + for _ in range(branches): + _add_binary_tree(bg, tree_size) + return bg + + +CornerRecord = namedtuple("CornerRecord", ["left", "right"]) + + +@dataclass +class CornerArguments: + value: object + nested: object + + +@dataclass(frozen=True) +class FrozenCornerArguments: + value: object + + +@dataclass +class FileCornerArguments: + left: object + nested: object + + +class CornerObject: + def __init__(self): + self.metadata = {"count": 3} + + +class UnsupportedCornerArguments: + def __init__(self, value): + self.value = value + + +def make_corner_payload(): + return { + "metadata": {"count": 3}, + "values": [10, 20, 30], + ("tuple", "key"): 7, + "record": CornerRecord(4, 5), + } + + +def make_corner_object(): + return CornerObject() + + +def check_shared_alias(left, right, *, keyword): + assert left is right is keyword + assert left == [3] + return "shared-ok" + + +def check_cycle(value): + assert value[0] is value + assert value[1] == 3 + return "cycle-ok" + + +def check_dataclasses(mutable, frozen): + assert isinstance(mutable, CornerArguments) + assert mutable.value == 3 + assert mutable.nested == {"value": [20]} + assert isinstance(frozen, FrozenCornerArguments) + assert frozen.value == 7 + return "dataclass-ok" + + +def check_containers(value): + assert value["list"] == [3, 20] + assert value["tuple"] == (3, 20) + assert value["set"] == {3, 20} + assert value["frozenset"] == frozenset({3, 20}) + assert value["deque"] == deque([3, 20]) + assert value["namedtuple"] == CornerRecord(3, 20) + return "containers-ok" + + +def write_corner_file(value): + with open("shared-name.txt", "w") as stream: + stream.write(value) + return f"producer-{value}" + + +def read_corner_files(files, external, repeated=None): + with open(files.left) as stream: + left = stream.read() + with open(files.nested["right"]) as stream: + right = stream.read() + with open(repeated) as stream: + repeated_value = stream.read() + with open(external) as stream: + external_prefix = stream.read(15) + assert left == repeated_value == "left" + assert right == "right" + assert "import argparse" in external_prefix + assert files.left != files.nested["right"] + return "files-ok" + + +def read_one_corner_file(path): + with open(path) as stream: + return stream.read() + + +def write_multiple_corner_files(): + os.makedirs("nested", exist_ok=True) + Path("nested/first.txt").write_text("first") + Path("second.txt").write_text("second") + return "multiple-producer" + + +def read_multiple_corner_files(first, second): + assert Path(first).read_text() == "first" + assert Path(second).read_text() == "second" + return "multiple-files-ok" + + +def verify_corner_results(*values, keyword_value=None): + assert values == ( + 3, + 20, + 7, + 4, + 3, + "shared-ok", + "cycle-ok", + "dataclass-ok", + "containers-ok", + "files-ok", + "right", + "producer-left", + "producer-right", + "multiple-files-ok", + ) + assert keyword_value == 30 + return "corner-cases-ok" + + +def make_corner_cases_graph(): + workflow = Workflow() + payload = workflow.add_task(make_corner_payload) + obj = workflow.add_task(make_corner_object) + + count = payload.output()["metadata"]["count"] + second = payload.output()["values"][1] + tuple_key = payload.output()[("tuple", "key")] + namedtuple_index = payload.output()["record"][0] + object_attribute = obj.output().attr("metadata")["count"] + + shared = [count] + shared_check = workflow.add_task(check_shared_alias, shared, shared, keyword=shared) + + cyclic = [] + cyclic.append(cyclic) + cyclic.append(count) + cycle_check = workflow.add_task(check_cycle, cyclic) + + mutable_box = CornerArguments(count, {"value": [second]}) + frozen_box = FrozenCornerArguments(tuple_key) + dataclass_check = workflow.add_task(check_dataclasses, mutable_box, frozen_box) + + container_check = workflow.add_task( + check_containers, + { + "list": [count, second], + "tuple": (count, second), + "set": {count, second}, + "frozenset": frozenset({count, second}), + "deque": deque([count, second]), + "namedtuple": CornerRecord(count, second), + }, + ) + + external = workflow.file(__file__) + left_producer = workflow.add_task(write_corner_file, "left") + right_producer = workflow.add_task(write_corner_file, "right") + left_file = left_producer.file("shared-name.txt") + right_file = right_producer.file("shared-name.txt") + file_consumer = workflow.add_task( + read_corner_files, + FileCornerArguments(left_file, {"right": right_file}), + external, + repeated=left_file, + ) + second_consumer = workflow.add_task(read_one_corner_file, right_file) + multiple_producer = workflow.add_task(write_multiple_corner_files) + first_file = multiple_producer.file("nested/first.txt") + second_file = multiple_producer.file("second.txt") + multiple_consumer = workflow.add_task(read_multiple_corner_files, first_file, second_file) + + final = workflow.add_task( + verify_corner_results, + count, + second, + tuple_key, + namedtuple_index, + object_attribute, + shared_check.output(), + cycle_check.output(), + dataclass_check.output(), + container_check.output(), + file_consumer.output(), + second_consumer.output(), + left_producer.output(), + right_producer.output(), + multiple_consumer.output(), + keyword_value=payload.output()["values"][2], + ) + assert isinstance(left_file, FileHandle) + workflow._corner_target_id = final._task_id + workflow._corner_file_target = left_file + return workflow + + +def check_rejected_corner_cases(): + left = Workflow() + right = Workflow() + left_task = left.add_task(add, 1) + assert isinstance(left_task.output(), TaskOutputHandle) + + try: + left.add_task("legacy-key", add, 1) + except TypeError as exc: + assert "callable" in str(exc) + else: + raise AssertionError("legacy add_task(key, func, ...) was accepted") + + try: + right.add_task(add, left_task.output()) + except ValueError as exc: + assert "different Workflow" in str(exc) + else: + raise AssertionError("cross-workflow dependency was accepted") + + try: + left.add_task(add, left_task) + except TypeError as exc: + assert "task.output()" in str(exc) + else: + raise AssertionError("bare TaskHandle argument was accepted") + + try: + left.add_task(add, {left_task.output(): 1}) + except ValueError as exc: + assert "dict key" in str(exc) + else: + raise AssertionError("TaskOutputHandle dictionary key was accepted") + + try: + left.add_task(add, {("nested", left_task.output()): 1}) + except ValueError as exc: + assert "dict key" in str(exc) + else: + raise AssertionError("nested TaskOutputHandle dictionary key was accepted") + + try: + left.add_task(add, UnsupportedCornerArguments(left_task.output())) + except TypeError as exc: + assert "custom object" in str(exc) + else: + raise AssertionError("hidden dependency in arbitrary object was accepted") + + left_file = left.file(__file__) + assert isinstance(left_file, FileHandle) + + try: + right.add_task(add, left_file) + except ValueError as exc: + assert "different Workflow" in str(exc) + else: + raise AssertionError("cross-workflow FileHandle was accepted") + + try: + left.add_task(add, {left_file: 1}) + except ValueError as exc: + assert "dict key" in str(exc) + else: + raise AssertionError("FileHandle dictionary key was accepted") + + try: + left.add_task(add, UnsupportedCornerArguments(left_file)) + except TypeError as exc: + assert "custom object" in str(exc) + else: + raise AssertionError("hidden FileHandle was accepted") + + try: + left.add_task(add, FileHandle(left._workflow_id, 999999)) + except ValueError as exc: + assert "does not belong" in str(exc) + else: + raise AssertionError("unknown FileHandle was accepted") + + producer = left.add_task(add, 1) + producer.file("same.txt") + try: + producer.file("same.txt") + except ValueError as exc: + assert "already declares" in str(exc) + else: + raise AssertionError("duplicate output path on one task was accepted") + + for invalid in ("", ".", "../escape", "/absolute"): + try: + producer.file(invalid) + except ValueError: + pass + else: + raise AssertionError(f"invalid output path was accepted: {invalid!r}") + + assert not hasattr(producer, "produces") + assert not hasattr(producer, "consumes") + assert not hasattr(left, "declare_input_file") + assert not hasattr(producer, "declare_output_file") + assert not hasattr(left, "declare_file") + assert not hasattr(producer, "declare_file") + + +def build_graph(name, n=None): + if name == "simple": + return make_simple_graph() + if name == "chain": + return make_chain_graph(max(1, n or 8)) + if name == "chain-branches": + return make_chain_graph(max(1, n or 8), branches=4) + if name == "chain-rich": + return make_chain_rich(max(1, n or 1000)) + if name == "binary-forest": + return make_binary_forest(n) + if name == "individuals": + return make_individuals(max(1, n or 1000)) + if name == "trivial": + return make_trivial(max(1, n or 1000)) + if name == "binary-tree": + return make_binary_tree(max(1, n or 1000)) + if name == "corner-cases": + return make_corner_cases_graph() + raise ValueError(name) + + +def parse_cases(specs): + out = [] + for s in specs: + name, _, n = s.strip().partition(":") + out.append((name, None if not n else int(n))) + return out + + +def _sink_tasks(workflow): + return workflow.sink_tasks() + + +def _run_vine_graph( + graph, n, task_group, port, port_file, logs, tag, out_dir, ckpt_dir, + priority, manager_name, libcores, +): + run_info = logs / tag + if run_info.exists(): + shutil.rmtree(run_info) + + wf = build_graph(graph, n) + corner_target = TaskHandle(wf, wf._corner_target_id) if graph == "corner-cases" else None + targets = [corner_target, wf._corner_file_target] if corner_target is not None else _sink_tasks(wf) + + def context_loader(graph_pkl): + cwd = os.getcwd() + if cwd not in sys.path: + sys.path.insert(0, cwd) + return {"graph": cloudpickle.loads(graph_pkl)} + + vine_graph_mod.context_loader_func = context_loader + try: + cloudpickle.register_pickle_by_value(sys.modules[__name__]) + except Exception: + pass + + with VineGraph(port=port, name=manager_name, run_info_path=str(logs), run_info_template=tag) as m: + if port_file: + Path(port_file).write_text(str(m.port)) + + m.set_params( + { + "checkpoint-dir": str(ckpt_dir), + "extra-task-output-size-mb": [0.0, 0.0], + "extra-task-sleep-time": [0.0, 0.0], + "libcores": libcores, + "output-dir": str(out_dir), + "task-group": task_group, + "task-priority-mode": priority, + "wait-for-workers": 1, + } + ) + results = m.run( + wf, + targets=targets, + hoisting_modules=[sys.modules[__name__]], + env_files={"./vine_graph_workflow_examples.py": "vine_graph_workflow_examples.py"}, + ) or {} + if graph == "corner-cases": + assert results[corner_target] == "corner-cases-ok" + target_path = results[wf._corner_file_target] + assert Path(target_path).read_text() == "left" + return {str(i): value for i, value in enumerate(results.values())} + + +def run_graph( + graph, + n=None, + task_group=0, + port=0, + port_file=None, + work_root=None, + tag="run", + timeout_s=120.0, + priority="random", + manager_name=None, + libcores=4, +): + root = work_root or Path(tempfile.mkdtemp(prefix="vine_graph-run-")) + delete_root = work_root is None + logs = root / "logs" + out_d = root / "out" / tag + ckpt = root / "ckpt" / tag + for d in (logs, out_d, ckpt): + d.mkdir(parents=True, exist_ok=True) + + def on_alarm(signum, frame): + raise TimeoutError(timeout_s) + + try: + old = signal.signal(signal.SIGALRM, on_alarm) + signal.setitimer(signal.ITIMER_REAL, timeout_s) + try: + os.chdir(TEST_DIR) + return _run_vine_graph( + graph, + n, + task_group, + port, + port_file, + logs, + tag, + out_d, + ckpt, + priority, + manager_name, + libcores, + ) + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + signal.signal(signal.SIGALRM, old) + finally: + if delete_root: + shutil.rmtree(root, ignore_errors=True) + + +def main(): + check_rejected_corner_cases() + p = argparse.ArgumentParser() + p.add_argument("port_file", nargs="?") + p.add_argument("-G", "--graph", nargs="+") + p.add_argument("--case", action="append", dest="cases") + p.add_argument("--task-group", type=int, default=0) + p.add_argument("--task-priority-mode", default="random") + p.add_argument("--port", type=int, default=0) + p.add_argument("--manager-name") + p.add_argument("--libcores", type=int, default=4) + p.add_argument("--result-file") + p.add_argument("--timeout", type=float, default=120.0) + p.add_argument("--no-print-results", action="store_true") + args = p.parse_args() + + if args.cases: + cases = parse_cases(args.cases) + elif args.graph: + if len(args.graph) > 2: + p.error("-G takes GRAPH [N]") + n = int(args.graph[1]) if len(args.graph) == 2 else None + cases = [(args.graph[0], n)] + else: + p.error("need -G or --case") + + root = Path(tempfile.mkdtemp(prefix="vine_graph-run-")) if len(cases) > 1 else None + rc = 0 + results = {} + try: + for i, (g, n) in enumerate(cases): + try: + res = run_graph( + g, + n, + task_group=args.task_group, + port=args.port if args.port == 0 else args.port + i, + port_file=args.port_file, + work_root=root, + tag=f"{i:02d}-{g}-{n or 'na'}", + timeout_s=args.timeout, + priority=args.task_priority_mode, + manager_name=args.manager_name, + libcores=args.libcores, + ) + except Exception as e: + rc = 1 + print(g, n, "fail:", e) + continue + results[f"{g}:{'' if n is None else n}"] = res + if not args.no_print_results: + print(g, n, res) + finally: + if root: + shutil.rmtree(root, ignore_errors=True) + if args.result_file: + Path(args.result_file).write_text(json.dumps(results, sort_keys=True)) + return rc + + +if __name__ == "__main__": + raise SystemExit(main())