Merge branch 'master' into enh/bosh

Author: djarecka

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."""

pydra/engine/helpers.py

@@ -513,3 +513,31 @@ def output_from_inputfields(output_spec, inputs):
                 (field_name, attr.ib(type=File, metadata={"value": value}))
             )
     return output_spec
+
+
+def get_available_cpus():
+    """
+    Return the number of CPUs available to the current process or, if that is not
+    available, the total number of CPUs on the system.
+
+    Returns
+    -------
+    n_proc : :obj:`int`
+        The number of available CPUs.
+    """
+    # Will not work on some systems or if psutil is not installed.
+    # See https://psutil.readthedocs.io/en/latest/#psutil.Process.cpu_affinity
+    try:
+        import psutil
+
+        return len(psutil.Process().cpu_affinity())
+    except (AttributeError, ImportError, NotImplementedError):
+        pass
+
+    # Not available on all systems, including macOS.
+    # See https://docs.python.org/3/library/os.html#os.sched_getaffinity
+    if hasattr(os, "sched_getaffinity"):
+        return len(os.sched_getaffinity(0))
+
+    # Last resort
+    return os.cpu_count()

pydra/engine/helpers_file.py

@@ -87,27 +87,66 @@ def hash_file(afile, chunk_len=8192, crypto=sha256, raise_notfound=True):
     return crypto_obj.hexdigest()
 
 
-def hash_dir(dirpath, raise_notfound=True):
+def hash_dir(
+    dirpath,
+    crypto=sha256,
+    ignore_hidden_files=False,
+    ignore_hidden_dirs=False,
+    raise_notfound=True,
+):
+    """Compute hash of directory contents.
+
+    This function computes the hash of every file in directory `dirpath` and then
+    computes the hash of that list of hashes to return a single hash value. The
+    directory is traversed recursively.
+
+    Parameters
+    ----------
+    dirpath : :obj:`str`
+        Path to directory.
+    crypto : :obj: `function`
+        cryptographic hash functions
+    ignore_hidden_files : :obj:`bool`
+        If `True`, ignore filenames that begin with `.`.
+    ignore_hidden_dirs : :obj:`bool`
+        If `True`, ignore files in directories that begin with `.`.
+    raise_notfound : :obj:`bool`
+        If `True` and `dirpath` does not exist, raise `FileNotFound` exception. If
+        `False` and `dirpath` does not exist, return `None`.
+
+    Returns
+    -------
+    hash : :obj:`str`
+        Hash of the directory contents.
+    """
     from .specs import LazyField
 
     if dirpath is None or isinstance(dirpath, LazyField) or isinstance(dirpath, list):
         return None
     if not Path(dirpath).is_dir():
         if raise_notfound:
-            raise RuntimeError(f"Directory {dirpath} not found.")
+            raise FileNotFoundError(f"Directory {dirpath} not found.")
         return None
 
-    def search_dir(path):
-        path = Path(path)
-        file_list = []
-        for el in path.iterdir():
-            if el.is_file():
-                file_list.append(hash_file(el))
-            else:
-                file_list.append(search_dir(path / el))
-        return file_list
+    file_hashes = []
+    for dpath, dirnames, filenames in os.walk(dirpath):
+        # Sort in-place to guarantee order.
+        dirnames.sort()
+        filenames.sort()
+        dpath = Path(dpath)
+        if ignore_hidden_dirs and dpath.name.startswith(".") and str(dpath) != dirpath:
+            continue
+        for filename in filenames:
+            if ignore_hidden_files and filename.startswith("."):
+                continue
+            this_hash = hash_file(dpath / filename)
+            file_hashes.append(this_hash)
 
-    return search_dir(dirpath)
+    crypto_obj = crypto()
+    for h in file_hashes:
+        crypto_obj.update(h.encode())
+
+    return crypto_obj.hexdigest()
 
 
 def _parse_mount_table(exit_code, output):

pydra/engine/tests/test_helpers.py

@@ -1,10 +1,20 @@
+import os
+import hashlib
 from pathlib import Path
+import random
+import platform
 
 import pytest
 import cloudpickle as cp
 
 from .utils import multiply
-from ..helpers import hash_value, hash_function, save, create_pyscript
+from ..helpers import (
+    hash_value,
+    hash_function,
+    get_available_cpus,
+    save,
+    create_pyscript,
+)
 from .. import helpers_file
 from ..specs import File, Directory
 
@@ -135,23 +145,70 @@ def test_hash_value_dir(tmpdir):
     with open(file_2, "w") as f:
         f.write("hi")
 
-    assert hash_value(tmpdir, tp=Directory) == hash_value([file_1, file_2], tp=File)
-    assert hash_value(tmpdir, tp=Directory) == helpers_file.hash_dir(tmpdir)
+    test_sha = hashlib.sha256()
+    for fx in [file_1, file_2]:
+        test_sha.update(helpers_file.hash_file(fx).encode())
+
+    bad_sha = hashlib.sha256()
+    for fx in [file_2, file_1]:
+        bad_sha.update(helpers_file.hash_file(fx).encode())
+
+    orig_hash = helpers_file.hash_dir(tmpdir)
+
+    assert orig_hash == test_sha.hexdigest()
+    assert orig_hash != bad_sha.hexdigest()
+    assert orig_hash == hash_value(tmpdir, tp=Directory)
 
 
 def test_hash_value_nested(tmpdir):
+    hidden = tmpdir.mkdir(".hidden")
     nested = tmpdir.mkdir("nested")
     file_1 = tmpdir.join("file_1.txt")
-    file_2 = nested.join("file_2.txt")
-    file_3 = nested.join("file_3.txt")
-    with open(file_1, "w") as f:
-        f.write("hello")
-    with open(file_2, "w") as f:
-        f.write("hi")
-    with open(file_3, "w") as f:
-        f.write("hola")
+    file_2 = hidden.join("file_2.txt")
+    file_3 = nested.join(".file_3.txt")
+    file_4 = nested.join("file_4.txt")
+
+    test_sha = hashlib.sha256()
+    for fx in [file_1, file_2, file_3, file_4]:
+        with open(fx, "w") as f:
+            f.write(str(random.randint(0, 1000)))
+        test_sha.update(helpers_file.hash_file(fx).encode())
+
+    orig_hash = helpers_file.hash_dir(tmpdir)
 
-    assert hash_value(tmpdir, tp=Directory) == hash_value(
-        [file_1, [file_2, file_3]], tp=File
+    assert orig_hash == test_sha.hexdigest()
+    assert orig_hash == hash_value(tmpdir, tp=Directory)
+
+    nohidden_hash = helpers_file.hash_dir(
+        tmpdir, ignore_hidden_dirs=True, ignore_hidden_files=True
     )
-    assert hash_value(tmpdir, tp=Directory) == helpers_file.hash_dir(tmpdir)
+    nohiddendirs_hash = helpers_file.hash_dir(tmpdir, ignore_hidden_dirs=True)
+    nohiddenfiles_hash = helpers_file.hash_dir(tmpdir, ignore_hidden_files=True)
+
+    assert orig_hash != nohidden_hash
+    assert orig_hash != nohiddendirs_hash
+    assert orig_hash != nohiddenfiles_hash
+
+    file_3.remove()
+    assert helpers_file.hash_dir(tmpdir) == nohiddenfiles_hash
+    hidden.remove()
+    assert helpers_file.hash_dir(tmpdir) == nohidden_hash
+
+
+def test_get_available_cpus():
+    assert get_available_cpus() > 0
+    try:
+        import psutil
+
+        has_psutil = True
+    except ImportError:
+        has_psutil = False
+
+    if hasattr(os, "sched_getaffinity"):
+        assert get_available_cpus() == len(os.sched_getaffinity(0))
+
+    if has_psutil and platform.system().lower() != "darwin":
+        assert get_available_cpus() == len(psutil.Process().cpu_affinity())
+
+    if platform.system().lower() == "darwin":
+        assert get_available_cpus() == os.cpu_count()