Dask compat fallback (#4059)

* auto fallback to compat DaskVine

* add test

Author: btovar

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

@@ -71,23 +71,16 @@ class DaskVineWarning(UserWarning):
     vd = Version(dask.__version__)
     vr = Version("2024.12.0")
     if vd < vr:
-        warnings.warn("ndcctools.taskvine.DaskVine only works with dask version >= 2024.12.0", DaskVineWarning)
+        raise ImportError
 
     from .dask_executor import DaskVine
     from .dask_dag import DaskVineDag
-except ImportError as e:
-    warnings.warn(f"DaskVine not available. Couldn't find module: {e.name}", DaskVineWarning)
+except (ImportError, ModuleNotFoundError):
+    warnings.warn("Dask >= 2024.12.0 not available, using DaskVine legacy task graph representation.", DaskVineWarning)
 
-    ##
-    # DaskVine compatibility class.
-    # See @ref dask_executor.DaskVine
-    class DaskVine:
-        exception = ImportError()
+    from .compat import DaskVine
+    from .compat import DaskVineDag
 
-        def __init__(*args, **kwargs):
-            raise DaskVine.exception
-
-    DaskVine.exception = e
 
 __all__ = [
     "Manager",

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

@@ -24,32 +24,23 @@
 # form ndcctools.taskvine.compat import DaskVine
 # @endcode
 #
-
-import dask
 import warnings
 from packaging.version import Version
-vd = Version(dask.__version__)
-vr = Version("2024.12.0")
-if vd >= vr:
-    warnings.warn("ndcctools.taskvine.compat only works with dask version < 2024.12.0")
-
 
 try:
-    from .dask_executor import DaskVine
-    from .dask_dag import DaskVineDag
-except ImportError as e:
-    print(f"DaskVine not available. Couldn't find module: {e.name}")
+    import dask
+
+    vd = Version(dask.__version__)
+    vr = Version("2024.12.0")
 
-    ##
-    # DaskVine compatibility class.
-    # See @ref dask_executor.DaskVine
-    class DaskVine:
-        exception = ImportError()
+    if vd >= vr:
+        warnings.warn("ndcctools.taskvine.compat only works with dask version < 2024.12.0")
+except (ImportError, ModuleNotFoundError):
+    pass
 
-        def __init__(*args, **kwargs):
-            raise DaskVine.exception
 
-    DaskVine.exception = e
+from .dask_executor import DaskVine
+from .dask_dag import DaskVineDag
 
 __all__ = [
     "DaskVine",

taskvine/test/TR_vine_task_graph_compat.sh

@@ -0,0 +1,52 @@
+#!/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=$(pwd)/../src/worker:$(pwd)/../../batch_job/src:$PATH
+
+STATUS_FILE=vine.status
+PORT_FILE=vine.port
+
+check_needed()
+{
+	[ -n "${CCTOOLS_PYTHON_TEST_EXEC}" ] || return 1
+	"${CCTOOLS_PYTHON_TEST_EXEC}" -c "import cloudpickle; import dask"  || return 1
+
+	return 0
+}
+
+prepare()
+{
+	rm -f $STATUS_FILE
+	rm -f $PORT_FILE
+
+	return 0
+}
+
+run()
+{
+	${CCTOOLS_PYTHON_TEST_EXEC} vine_task_graph_compat.py
+	echo $? > $STATUS_FILE
+
+	# retrieve taskvine exit status
+	status=$(cat $STATUS_FILE)
+	if [ $status -ne 0 ]
+	then
+		exit 1
+	fi
+
+	exit 0
+}
+
+clean()
+{
+	rm -rf vine-run-info
+	exit 0
+}
+
+dispatch "$@"

taskvine/test/vine_task_graph_compat.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python
+
+# Copyright (C) 2023- The University of Notre Dame
+# This software is distributed under the GNU General Public License.
+# See the file COPYING for details.
+
+# This example shows TaskVine executing a manually constructed dask graph.
+# See vine_example_dask_delayed.py for an example where the graph
+# is constructed by dask.
+
+import ndcctools.taskvine as vine
+from ndcctools.taskvine.compat import DaskVine
+import argparse
+import getpass
+import sys
+import traceback
+
+
+from operator import add  # use add function in the example graph
+dsk_graph = {
+    "x": 1,
+    "y": 2,
+    "z": (add, "x", "y"),
+    "w": (sum, ["x", "y", "z"]),
+    "v": [(sum, ["w", "z"]), 2],
+    "t": (sum, "v")
+}
+
+expected_result = 11
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="vine_example_dask_graph.py",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="""This example shows TaskVine executing a manually constructed dask graph.
+See vine_example_dask_delayed.py for an example where the graph
+is constructed by dask.""")
+    parser.add_argument(
+        "--name",
+        nargs="?",
+        type=str,
+        help="name to assign to the manager.",
+        default=f"vine-dask-graph-{getpass.getuser()}",
+    )
+    parser.add_argument(
+        "--port",
+        nargs="?",
+        type=int,
+        help="port for the manager to listen for connections. If 0, pick any available.",
+        default=9123,
+    )
+    parser.add_argument(
+        "--disable-peer-transfers",
+        action="store_true",
+        help="disable transfers among workers.",
+        default=False,
+    )
+
+    args = parser.parse_args()
+
+    m = DaskVine(port=args.port, ssl=True)
+    m.set_name(args.name)
+    print(f"Listening for workers at port: {m.port}")
+
+    if args.disable_peer_transfers:
+        m.disable_peer_transfers()
+
+    # checkpoint at even levels when nodes have at least one dependency
+    def checkpoint(dag, key):
+        if dag.depth_of(key) % 2 == 0 and len(dag.get_dependencies(key)) > 0:
+            print(f"checkpoint for {key}")
+            return True
+        return False
+
+    f = vine.Factory(manager=m)
+    f.cores = 4
+    f.disk = 2000
+    f.memory = 2000
+    f.max_workers = 1
+    f.min_workers = 1
+    with f:
+        desired_keys = ["t", "v"]
+        desired_keys = list(dsk_graph.keys())
+        print(f"dask graph example is:\n{dsk_graph}")
+        print(f"desired keys are {desired_keys}")
+
+        try:
+            results = m.get(dsk_graph, desired_keys, lazy_transfers=True, checkpoint_fn=checkpoint, resources={"cores": 1})
+            print({k: v for k, v in zip(desired_keys, results)})
+        except Exception:
+            traceback.print_exc()
+
+        print("Terminating workers...", end="")
+
+    print("done!")
+sys.exit(0)
+# vim: set sts=4 sw=4 ts=4 expandtab ft=python: