vine: dagvine implementation

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)
 

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) 
 

taskvine/src/Makefile

@@ -1,12 +1,13 @@
 include ../../config.mk
 include ../../rules.mk
 
-TARGETS=manager worker tools bindings examples
+TARGETS=manager worker tools bindings examples graph
 
 all: $(TARGETS)
 
 worker: manager
-bindings: manager
+graph: manager
+bindings: manager graph
 tools: manager
 examples: manager worker tools bindings
 

taskvine/src/bindings/python3/Makefile

@@ -34,4 +34,4 @@ 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
-	cp taskvine.py $(CCTOOLS_PYTHON3_PATH)/
+	cp taskvine.py $(CCTOOLS_PYTHON3_PATH)/

taskvine/src/bindings/python3/ndcctools/taskvine/compat/dask_executor.py

@@ -138,8 +138,10 @@ def get(self, dsk, keys, *,
             hoisting_modules=None,  # Deprecated, use lib_modules
             import_modules=None,    # Deprecated, use lib_modules
             lazy_transfers=True,    # Deprecated, use worker_tranfers
+            extra_serialize_time_sec=0,
             ):
         try:
+            self.extra_serialize_time_sec = extra_serialize_time_sec
             self.set_property("framework", "dask")
             if retries and retries < 1:
                 raise ValueError("retries should be larger than 0")
@@ -213,10 +215,13 @@ def __call__(self, *args, **kwargs):
         return self.get(*args, **kwargs)
 
     def _dask_execute(self, dsk, keys):
+
         indices = {k: inds for (k, inds) in find_dask_keys(keys)}
         keys_flatten = indices.keys()
 
+        time_start = time.time()
         dag = DaskVineDag(dsk, low_memory_mode=self.low_memory_mode, prune_depth=self.prune_depth)
+        print(f"Time taken to enqueue tasks: {time.time() - time_start:.6f} seconds")
         tag = f"dag-{id(dag)}"
 
         # create Library if using 'function-calls' task mode.
@@ -437,7 +442,8 @@ def _enqueue_dask_calls(self, dag, tag, rs, retries, enqueued_calls):
                                      extra_files=self.extra_files,
                                      retries=retries,
                                      worker_transfers=lazy,
-                                     wrapper=self.wrapper)
+                                     wrapper=self.wrapper,
+                                     extra_serialize_time_sec=self.extra_serialize_time_sec)
 
                 t.set_priority(priority)
                 t.set_tag(tag)  # tag that identifies this dag
@@ -557,7 +563,9 @@ def __init__(self, m,
                  env_vars=None,
                  retries=5,
                  worker_transfers=False,
-                 wrapper=None):
+                 wrapper=None,
+                 extra_serialize_time_sec=0):
+        time.sleep(extra_serialize_time_sec)
         self._key = key
         self._sexpr = sexpr
 
@@ -658,7 +666,9 @@ def __init__(self, m,
                  extra_files=None,
                  retries=5,
                  worker_transfers=False,
-                 wrapper=None):
+                 wrapper=None,
+                 extra_serialize_time_sec=0):
+        time.sleep(extra_serialize_time_sec)
 
         self._key = key
         self.resources = resources

taskvine/src/bindings/python3/ndcctools/taskvine/manager.py

@@ -946,8 +946,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
@@ -969,7 +970,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"
@@ -1017,7 +1019,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)

taskvine/src/bindings/python3/taskvine.i

@@ -1,5 +1,5 @@
 /* taskvine.i */
-%module cvine
+%module(package="ndcctools.taskvine") cvine
 
 %include carrays.i
 %array_functions(struct rmsummary *, rmsummayArray);

taskvine/src/graph/Makefile

@@ -0,0 +1,14 @@
+include ../../../config.mk
+include ../../../rules.mk
+
+SUBDIRS = dagvine
+
+all clean install test lint format: $(SUBDIRS)
+
+$(SUBDIRS): %:
+	$(MAKE) -C $@ $(MAKECMDGOALS)
+
+.PHONY: all clean install test lint format $(SUBDIRS) example
+
+example:
+	PYTHONPATH=../../../ python -m taskvine.src.graph.example_blueprint

taskvine/src/graph/dagvine/.gitignore

@@ -0,0 +1,7 @@
+*.a
+*.so
+*.o
+*_wrap.c
+*_wrap.0
+build/
+context_graph/

taskvine/src/graph/dagvine/Makefile

@@ -0,0 +1,36 @@
+include ../../../../config.mk
+include ../../../../rules.mk
+
+MODULE_DIR := $(CCTOOLS_PYTHON3_PATH)/ndcctools/taskvine/dagvine
+SUBDIRS    := vine_graph blueprint_graph
+
+.PHONY: all install clean lint format $(SUBDIRS)
+
+all: $(SUBDIRS)
+
+$(SUBDIRS):
+	$(MAKE) -C $@ all
+
+install: all
+	mkdir -p $(MODULE_DIR)
+	cp dagvine.py $(MODULE_DIR)
+	cp __init__.py $(MODULE_DIR)
+	@for dir in $(SUBDIRS); do \
+		$(MAKE) -C $$dir install; \
+	done
+
+clean:
+	@for dir in $(SUBDIRS); do \
+		$(MAKE) -C $$dir clean; \
+	done
+	rm -rf build
+
+lint:
+	@for dir in $(SUBDIRS); do \
+		$(MAKE) -C $$dir lint; \
+	done
+
+format:
+	@for dir in $(SUBDIRS); do \
+		$(MAKE) -C $$dir format; \
+	done

taskvine/src/graph/dagvine/__init__.py

@@ -0,0 +1,7 @@
+# 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 .dagvine import DAGVine
+
+__all__ = ["DAGVine"]

taskvine/src/graph/dagvine/blueprint_graph/.gitignore

@@ -0,0 +1,2 @@
+adaptor_test.py
+__pycache__/

taskvine/src/graph/dagvine/blueprint_graph/Makefile

@@ -0,0 +1,31 @@
+include ../../../../../config.mk
+include ../../../../../rules.mk
+
+PROJECT_NAME = dagvine
+
+SOURCE_DIR = $(CCTOOLS_HOME)/taskvine/src/graph/$(PROJECT_NAME)/blueprint_graph
+MODULE_ROOT = $(CCTOOLS_PYTHON3_PATH)/ndcctools/taskvine/$(PROJECT_NAME)
+MODULE_DIR = $(MODULE_ROOT)/blueprint_graph
+
+PY_SOURCES = $(wildcard $(SOURCE_DIR)/*.py)
+
+.PHONY: all install clean lint format
+
+all:
+	@true
+
+install:
+	mkdir -p $(MODULE_DIR)
+	cp $(PY_SOURCES) $(MODULE_DIR)
+
+clean:
+	rm -rf __pycache__
+
+lint:
+	flake8 --ignore=$(CCTOOLS_FLAKE8_IGNORE_ERRORS) \
+		--exclude=$(CCTOOLS_FLAKE8_IGNORE_FILES) \
+		$(SOURCE_DIR)/
+
+format:
+	@true
+

taskvine/src/graph/dagvine/blueprint_graph/__init__.py

@@ -0,0 +1,16 @@
+# 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 .blueprint_graph import BlueprintGraph, TaskOutputRef, TaskOutputWrapper
+from .proxy_functions import compute_single_key
+from .adaptor import Adaptor
+
+__all__ = [
+    "BlueprintGraph",
+    "TaskOutputRef",
+    "TaskOutputWrapper",
+    "compute_single_key",
+    "Adaptor",
+]