Merge pull request #242 from satra/enh-dask

enh: add support for dask

Author: djarecka

.travis.yml

@@ -33,6 +33,10 @@ env:
 # Useful for testing un-released upstream fixes
 matrix:
   include:
+  - python: 3.7
+    env:
+      - INSTALL_TYPE="develop"
+      - CHECK_TYPE="test_dask"
   - os: osx
     osx_image: xcode11.2
     language: generic

ci/none.sh

@@ -40,14 +40,20 @@ function travis_before_script {
             # extras don't seem possible with setup.py install, so switch to pip
             pip install ".[test]"
         fi
+    elif [ "$CHECK_TYPE" = "test_dask" ]; then
+        if [ "$INSTALL_TYPE" = "develop" ]; then
+            pip install -e ".[dask]"
+        fi
     elif [ "$CHECK_TYPE" = "style" ]; then
         pip install black==19.3b0
     fi
 }
 
 function travis_script {
     if [ "$CHECK_TYPE" = "test" ]; then
-        pytest -vs -n auto --cov pydra --cov-config .coveragerc --cov-report xml:cov.xml --doctest-modules pydra tutorial
+        pytest -vs -n auto --cov pydra --cov-config .coveragerc --cov-report xml:cov.xml --doctest-modules pydra
+    elif [ "$CHECK_TYPE" = "test_dask" ]; then
+        pytest -vs -n auto --cov pydra --cov-config .coveragerc --cov-report xml:cov.xml --doctest-modules --dask pydra/engine
     elif [ "$CHECK_TYPE" = "style" ]; then
         black --check pydra tools setup.py
     fi

pydra/conftest.py

@@ -0,0 +1,25 @@
+import shutil
+
+
+def pytest_addoption(parser):
+    parser.addoption("--dask", action="store_true", help="run all combinations")
+
+
+def pytest_generate_tests(metafunc):
+    if "plugin_dask_opt" in metafunc.fixturenames:
+        if bool(shutil.which("sbatch")):
+            Plugins = ["cf", "slurm"]
+        else:
+            Plugins = ["cf"]
+        if metafunc.config.getoption("dask"):
+            Plugins.append("dask")
+        metafunc.parametrize("plugin_dask_opt", Plugins)
+
+    if "plugin" in metafunc.fixturenames:
+        if metafunc.config.getoption("dask"):
+            Plugins = []
+        elif bool(shutil.which("sbatch")):
+            Plugins = ["cf", "slurm"]
+        else:
+            Plugins = ["cf"]
+        metafunc.parametrize("plugin", Plugins)

pydra/engine/helpers.py

@@ -338,6 +338,8 @@ def execute(cmd, strip=False):
     strip : :obj:`bool`
         TODO
 
+    """
+    rc, stdout, stderr = read_and_display(*cmd, strip=strip)
     """
     loop = get_open_loop()
     if loop.is_running():
@@ -346,6 +348,7 @@ def execute(cmd, strip=False):
         rc, stdout, stderr = loop.run_until_complete(
             read_and_display_async(*cmd, strip=strip)
         )
+    """
     return rc, stdout, stderr
 
 

pydra/engine/submitter.py

@@ -1,6 +1,6 @@
 """Handle execution backends."""
 import asyncio
-from .workers import SerialWorker, ConcurrentFuturesWorker, SlurmWorker
+from .workers import SerialWorker, ConcurrentFuturesWorker, SlurmWorker, DaskWorker
 from .core import is_workflow
 from .helpers import get_open_loop
 
@@ -33,6 +33,8 @@ def __init__(self, plugin="cf", **kwargs):
             self.worker = ConcurrentFuturesWorker(**kwargs)
         elif self.plugin == "slurm":
             self.worker = SlurmWorker(**kwargs)
+        elif self.plugin == "dask":
+            self.worker = DaskWorker(**kwargs)
         else:
             raise Exception("plugin {} not available".format(self.plugin))
         self.worker.loop = self.loop
@@ -176,9 +178,9 @@ def close(self):
         Do not close previously running loop.
 
         """
+        self.worker.close()
         if self._own_loop:
             self.loop.close()
-        self.worker.close()
 
 
 def get_runnable_tasks(graph):

pydra/engine/tests/test_boutiques.py

@@ -11,26 +11,21 @@
 
 need_bosh_docker = pytest.mark.skipif(
     shutil.which("docker") is None
-    or sp.call(["docker", "info"] or sp.call(["bosh", "version"])),
+    or sp.call(["docker", "info"])
+    or sp.call(["which", "bosh"]),
     reason="requires docker and bosh",
 )
 
-if bool(shutil.which("sbatch")):
-    Plugins = ["cf", "slurm"]
-else:
-    Plugins = ["cf"]
-
 Infile = Path(__file__).resolve().parent / "data_tests" / "test.nii.gz"
 
 
 @no_win
 @need_bosh_docker
-@pytest.mark.flaky(reruns=2)
+@pytest.mark.flaky(reruns=2)  # need for travis
 @pytest.mark.parametrize(
     "maskfile", ["test_brain.nii.gz", "test_brain", "test_brain.nii"]
 )
 @pytest.mark.parametrize("results_function", [result_no_submitter, result_submitter])
-@pytest.mark.parametrize("plugin", Plugins)
 def test_boutiques_1(maskfile, plugin, results_function):
     """ simple task to run fsl.bet using BoshTask"""
     btask = BoshTask(name="NA", zenodo_id="1482743")
@@ -99,7 +94,6 @@ def test_boutiques_spec_2():
 @pytest.mark.parametrize(
     "maskfile", ["test_brain.nii.gz", "test_brain", "test_brain.nii"]
 )
-@pytest.mark.parametrize("plugin", Plugins)
 def test_boutiques_wf_1(maskfile, plugin):
     """ wf with one task that runs fsl.bet using BoshTask"""
     wf = Workflow(name="wf", input_spec=["maskfile", "infile"])
@@ -130,7 +124,6 @@ def test_boutiques_wf_1(maskfile, plugin):
 @pytest.mark.parametrize(
     "maskfile", ["test_brain.nii.gz", "test_brain", "test_brain.nii"]
 )
-@pytest.mark.parametrize("plugin", Plugins)
 def test_boutiques_wf_2(maskfile, plugin):
     """ wf with two BoshTasks (fsl.bet and fsl.stats) and one ShellTask"""
     wf = Workflow(name="wf", input_spec=["maskfile", "infile"])