Repository: nipype/pydra

Issue:
Change default `n_procs` arg in `ConcurrentFuturesWorker`
What would you like changed/added and why?
- Change the `n_procs` default to the number of CPUs the current process has access to, not the number of CPUs on the machine.

What would be the benefit? Does the change make something easier to use?
- The benefit would be that the CPUs are not oversubscribed.

By default, `n_procs` is set to the number of processors on the machine (that is the default behavior of [concurrent.futures.ProcessPoolExecutor](https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ProcessPoolExecutor), but this could oversubscribe the CPUs. Instead, if `n_procs` is `None`, I propose that `n_procs` should be translated to `len(os.sched_getaffinity(0))`, which returns the number of _usable_ CPUs ([documentation](https://docs.python.org/3.8/library/os.html#os.cpu_count)). This would help with things like SLURM, where jobs can be started with access to a subset of CPUs on a machine.

Comments:

Patch:
diff --git a/pydra/engine/helpers.py b/pydra/engine/helpers.py
index 2762111..e10335e 100644
--- a/pydra/engine/helpers.py
+++ b/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()
diff --git a/pydra/engine/workers.py b/pydra/engine/workers.py
index ff418b8..615ec52 100644
--- a/pydra/engine/workers.py
+++ b/pydra/engine/workers.py
@@ -7,7 +7,7 @@ from pathlib import Path
 
 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()


Test patch:
diff --git a/pydra/engine/tests/test_helpers.py b/pydra/engine/tests/test_helpers.py
index 97a30fb..b0e298d 100644
--- a/pydra/engine/tests/test_helpers.py
+++ b/pydra/engine/tests/test_helpers.py
@@ -1,10 +1,18 @@
+import os
 from pathlib import Path
+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
 
@@ -155,3 +163,22 @@ 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
+    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()