add method pydra.helpers.get_available_cpus() to return n cpus

This method will attempt to return the number of cpus available to the current process. If those methods do not work, then `os.cpu_count` is used, which returns the number of CPUs on the machine (could be more than the cpus in the current process).

Author: kaczmarj

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/tests/test_helpers.py

@@ -1,10 +1,17 @@
+import os
 from pathlib import Path
 
 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
 
@@ -155,3 +162,18 @@ def test_hash_value_nested(tmpdir):
         [file_1, [file_2, file_3]], tp=File
     )
     assert hash_value(tmpdir, tp=Directory) == helpers_file.hash_dir(tmpdir)
+
+
+def test_get_available_cpus():
+    assert get_available_cpus() > 0
+    try:
+        import psutil
+
+        has_psutil = True
+        assert get_available_cpus() == len(psutil.Process().cpu_affinity())
+    except ImportError:
+        has_psutil = False
+    if hasattr(os, "sched_getaffinity"):
+        assert get_available_cpus() == len(os.sched_getaffinity(0))
+    elif not has_psutil:
+        assert get_available_cpus() == os.cpu_count()

pydra/engine/workers.py

@@ -7,7 +7,7 @@
 
 import concurrent.futures as cf
 
-from .helpers import create_pyscript, read_and_display_async, save
+from .helpers import create_pyscript, get_available_cpus, read_and_display_async, save
 
 import logging
 
@@ -164,7 +164,7 @@ class ConcurrentFuturesWorker(Worker):
     def __init__(self, n_procs=None):
         """Initialize Worker."""
         super(ConcurrentFuturesWorker, self).__init__()
-        self.n_procs = n_procs
+        self.n_procs = get_available_cpus() if n_procs is None else n_procs
         # added cpu_count to verify, remove once confident and let PPE handle
         self.pool = cf.ProcessPoolExecutor(self.n_procs)
         # self.loop = asyncio.get_event_loop()