Merge pull request #13 from djarecka/satra-enh-dask

allowing for limitted dask testing

Author: satra

.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,6 +40,10 @@ 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
@@ -48,6 +52,8 @@ function travis_before_script {
 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
+    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/boutiques.py

@@ -0,0 +1,215 @@
+import typing as ty
+import json
+import attr
+from urllib.request import urlretrieve
+from pathlib import Path
+from functools import reduce
+
+from ..utils.messenger import AuditFlag
+from ..engine import ShellCommandTask
+from ..engine.specs import SpecInfo, ShellSpec, ShellOutSpec, File, attr_fields
+from .helpers_file import is_local_file
+
+
+class BoshTask(ShellCommandTask):
+    """Shell Command Task based on the Boutiques descriptor"""
+
+    def __init__(
+        self,
+        zenodo_id=None,
+        bosh_file=None,
+        audit_flags: AuditFlag = AuditFlag.NONE,
+        cache_dir=None,
+        input_spec_names: ty.Optional[ty.List] = None,
+        messenger_args=None,
+        messengers=None,
+        name=None,
+        output_spec_names: ty.Optional[ty.List] = None,
+        rerun=False,
+        strip=False,
+        **kwargs,
+    ):
+        """
+        Initialize this task.
+
+        Parameters
+        ----------
+        zenodo_id: :obj: str
+            Zenodo ID
+        bosh_file : : str
+            json file with the boutiques descriptors
+        audit_flags : :obj:`pydra.utils.messenger.AuditFlag`
+            Auditing configuration
+        cache_dir : :obj:`os.pathlike`
+            Cache directory
+        input_spec_names : :obj: list
+            Input names for input_spec.
+        messenger_args :
+            TODO
+        messengers :
+            TODO
+        name : :obj:`str`
+            Name of this task.
+        output_spec_names : :obj: list
+            Output names for output_spec.
+        strip : :obj:`bool`
+            TODO
+
+        """
+        self.cache_dir = cache_dir
+        if (bosh_file and zenodo_id) or not (bosh_file or zenodo_id):
+            raise Exception("either bosh or zenodo_id has to be specified")
+        elif zenodo_id:
+            self.bosh_file = self._download_spec(zenodo_id)
+        else:  # bosh_file
+            self.bosh_file = bosh_file
+
+        with self.bosh_file.open() as f:
+            self.bosh_spec = json.load(f)
+
+        self.input_spec = self._prepare_input_spec(names_subset=input_spec_names)
+        self.output_spec = self._prepare_output_spec(names_subset=output_spec_names)
+        self.bindings = ["-v", f"{self.bosh_file.parent}:{self.bosh_file.parent}:ro"]
+
+        super(BoshTask, self).__init__(
+            name=name,
+            input_spec=self.input_spec,
+            output_spec=self.output_spec,
+            executable=["bosh", "exec", "launch"],
+            args=["-s"],
+            audit_flags=audit_flags,
+            messengers=messengers,
+            messenger_args=messenger_args,
+            cache_dir=self.cache_dir,
+            strip=strip,
+            rerun=rerun,
+            **kwargs,
+        )
+        self.strip = strip
+
+    def _download_spec(self, zenodo_id):
+        """
+        usind boutiques Searcher to find url of zenodo file for a specific id,
+        and download the file to self.cache_dir
+        """
+        from boutiques.searcher import Searcher
+
+        searcher = Searcher(zenodo_id, exact_match=True)
+        hits = searcher.zenodo_search().json()["hits"]["hits"]
+        if len(hits) == 0:
+            raise Exception(f"can't find zenodo spec for {zenodo_id}")
+        elif len(hits) > 1:
+            raise Exception(f"too many hits for {zenodo_id}")
+        else:
+            zenodo_url = hits[0]["files"][0]["links"]["self"]
+            zenodo_file = self.cache_dir / f"zenodo.{zenodo_id}.json"
+            urlretrieve(zenodo_url, zenodo_file)
+            return zenodo_file
+
+    def _prepare_input_spec(self, names_subset=None):
+        """ creating input spec from the zenodo file
+            if name_subset provided, only names from the subset will be used in the spec
+        """
+        binputs = self.bosh_spec["inputs"]
+        self._input_spec_keys = {}
+        fields = []
+        for input in binputs:
+            name = input["id"]
+            if names_subset is None:
+                pass
+            elif name not in names_subset:
+                continue
+            else:
+                names_subset.remove(name)
+            if input["type"] == "File":
+                tp = File
+            elif input["type"] == "String":
+                tp = str
+            elif input["type"] == "Number":
+                tp = float
+            elif input["type"] == "Flag":
+                tp = bool
+            else:
+                tp = None
+            # adding list
+            if tp and "list" in input and input["list"]:
+                tp = ty.List[tp]
+
+            mdata = {
+                "help_string": input.get("description", None) or input["name"],
+                "mandatory": not input["optional"],
+                "argstr": input.get("command-line-flag", None),
+            }
+            fields.append((name, tp, mdata))
+            self._input_spec_keys[input["value-key"]] = "{" + f"{name}" + "}"
+        if names_subset:
+            raise RuntimeError(f"{names_subset} are not in the zenodo input spec")
+        spec = SpecInfo(name="Inputs", fields=fields, bases=(ShellSpec,))
+        return spec
+
+    def _prepare_output_spec(self, names_subset=None):
+        """ creating output spec from the zenodo file
+            if name_subset provided, only names from the subset will be used in the spec
+        """
+        boutputs = self.bosh_spec["output-files"]
+        fields = []
+        for output in boutputs:
+            name = output["id"]
+            if names_subset is None:
+                pass
+            elif name not in names_subset:
+                continue
+            else:
+                names_subset.remove(name)
+            path_template = reduce(
+                lambda s, r: s.replace(*r),
+                self._input_spec_keys.items(),
+                output["path-template"],
+            )
+            mdata = {
+                "help_string": output.get("description", None) or output["name"],
+                "mandatory": not output["optional"],
+                "output_file_template": path_template,
+            }
+            fields.append((name, attr.ib(type=File, metadata=mdata)))
+
+        if names_subset:
+            raise RuntimeError(f"{names_subset} are not in the zenodo output spec")
+        spec = SpecInfo(name="Outputs", fields=fields, bases=(ShellOutSpec,))
+        return spec
+
+    def _command_args_single(self, state_ind, ind=None):
+        """Get command line arguments for a single state"""
+        input_filepath = self._bosh_invocation_file(state_ind=state_ind, ind=ind)
+        cmd_list = (
+            self.inputs.executable
+            + [str(self.bosh_file), input_filepath]
+            + self.inputs.args
+            + self.bindings
+        )
+        return cmd_list
+
+    def _bosh_invocation_file(self, state_ind, ind=None):
+        """creating bosh invocation file - json file with inputs values"""
+        input_json = {}
+        for f in attr_fields(self.inputs):
+            if f.name in ["executable", "args"]:
+                continue
+            if self.state and f"{self.name}.{f.name}" in state_ind:
+                value = getattr(self.inputs, f.name)[state_ind[f"{self.name}.{f.name}"]]
+            else:
+                value = getattr(self.inputs, f.name)
+            # adding to the json file if specified by the user
+            if value is not attr.NOTHING and value != "NOTHING":
+                if is_local_file(f):
+                    value = Path(value)
+                    self.bindings.extend(["-v", f"{value.parent}:{value.parent}:ro"])
+                    value = str(value)
+
+                input_json[f.name] = value
+
+        filename = self.cache_dir / f"{self.name}-{ind}.json"
+        with open(filename, "w") as jsonfile:
+            json.dump(input_json, jsonfile)
+
+        return str(filename)

pydra/engine/core.py

@@ -532,29 +532,37 @@ def done(self):
                 return True
         return False
 
-    def _combined_output(self):
+    def _combined_output(self, return_inputs=False):
         combined_results = []
         for (gr, ind_l) in self.state.final_combined_ind_mapping.items():
-            combined_results.append([])
+            combined_results_gr = []
             for ind in ind_l:
                 result = load_result(self.checksum_states(ind), self.cache_locations)
                 if result is None:
                     return None
-                combined_results[gr].append(result)
+                if return_inputs is True or return_inputs == "val":
+                    result = (self.state.states_val[ind], result)
+                elif return_inputs == "ind":
+                    result = (self.state.states_ind[ind], result)
+                combined_results_gr.append(result)
+            combined_results.append(combined_results_gr)
         if len(combined_results) == 1 and self.state.splitter_rpn_final == []:
             # in case it's full combiner, removing the nested structure
             return combined_results[0]
         else:
             return combined_results
 
-    def result(self, state_index=None):
+    def result(self, state_index=None, return_inputs=False):
         """
         Retrieve the outcomes of this particular task.
 
         Parameters
         ----------
-        state_index :
-            TODO
+        state_index : :obj: `int`
+            index of the element for task with splitter and multiple states
+        return_inputs : :obj: `bool`, :obj:`str`
+            if True or "val" result is returned together with values of the input fields,
+            if "ind" result is returned together with indices of the input fields
 
         Returns
         -------
@@ -567,28 +575,50 @@ def result(self, state_index=None):
             if state_index is None:
                 # if state_index=None, collecting all results
                 if self.state.combiner:
-                    return self._combined_output()
+                    return self._combined_output(return_inputs=return_inputs)
                 else:
                     results = []
                     for checksum in self.checksum_states():
                         result = load_result(checksum, self.cache_locations)
                         if result is None:
                             return None
                         results.append(result)
-                    return results
+                    if return_inputs is True or return_inputs == "val":
+                        return list(zip(self.state.states_val, results))
+                    elif return_inputs == "ind":
+                        return list(zip(self.state.states_ind, results))
+                    else:
+                        return results
             else:  # state_index is not None
                 if self.state.combiner:
-                    return self._combined_output()[state_index]
+                    return self._combined_output(return_inputs=return_inputs)[
+                        state_index
+                    ]
                 result = load_result(
                     self.checksum_states(state_index), self.cache_locations
                 )
-                return result
+                if return_inputs is True or return_inputs == "val":
+                    return (self.state.states_val[state_index], result)
+                elif return_inputs == "ind":
+                    return (self.state.states_ind[state_index], result)
+                else:
+                    return result
         else:
             if state_index is not None:
                 raise ValueError("Task does not have a state")
             checksum = self.checksum
             result = load_result(checksum, self.cache_locations)
-            return result
+            if return_inputs is True or return_inputs == "val":
+                inputs_val = {
+                    f"{self.name}.{inp}": getattr(self.inputs, inp)
+                    for inp in self.input_names
+                }
+                return (inputs_val, result)
+            elif return_inputs == "ind":
+                inputs_ind = {f"{self.name}.{inp}": None for inp in self.input_names}
+                return (inputs_ind, result)
+            else:
+                return result
 
     def _reset(self):
         """Reset the connections between inputs and LazyFields."""