Repository: joblib/loky

Issue:
Possibility to ignore hyperthreading in loky.cpu_count
In many situations it's detrimental to use simultaneous multi-threading (smt) (hyperthreading for intel). It would be interesting to add the option to ignore smt to cpu_count.
```
loky.cpu_count(include_smt=False)
```

Comments:

Patch:
diff --git a/CHANGES.md b/CHANGES.md
index 6e2c1dc..6b9ab23 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -11,6 +11,9 @@
 - Fix support for Python 3.9 and test against python-nightly from now on
   (#250).
 
+- Add a parameter to ``cpu_count``, ``only_physical_cores``, to return the
+  number of physical cores instead of the number of logical cores (#271).
+
 ### 2.7.0 - 2020-04-30
 
 - Increase the residual memory increase threshold  (100MB -> 300MB) used by
diff --git a/loky/backend/context.py b/loky/backend/context.py
index dfeb4ed..76f6520 100644
--- a/loky/backend/context.py
+++ b/loky/backend/context.py
@@ -14,6 +14,8 @@ from __future__ import division
 
 import os
 import sys
+import subprocess
+import traceback
 import warnings
 import multiprocessing as mp
 
@@ -23,6 +25,10 @@ from .process import LokyProcess, LokyInitMainProcess
 START_METHODS = ['loky', 'loky_init_main']
 _DEFAULT_START_METHOD = None
 
+# Cache for the number of physical cores to avoid repeating subprocess calls.
+# It should not change during the lifetime of the program.
+physical_cores_cache = None
+
 if sys.version_info[:2] >= (3, 4):
     from multiprocessing import get_context as mp_get_context
     from multiprocessing.context import assert_spawning, set_spawning_popen
@@ -101,7 +107,7 @@ def get_start_method():
     return _DEFAULT_START_METHOD
 
 
-def cpu_count():
+def cpu_count(only_physical_cores=False):
     """Return the number of CPUs the current process can use.
 
     The returned number of CPUs accounts for:
@@ -113,15 +119,57 @@ def cpu_count():
        set by docker and similar container orchestration systems);
      * the value of the LOKY_MAX_CPU_COUNT environment variable if defined.
     and is given as the minimum of these constraints.
+
+    If ``only_physical_cores`` is True, return the number of physical cores
+    instead of the number of logical cores (hyperthreading / SMT). Note that
+    this option is not enforced if the number of usable cores is controlled in
+    any other way such as: process affinity, restricting CFS scheduler policy
+    or the LOKY_MAX_CPU_COUNT environment variable. If the number of physical
+    cores is not found, return the number of logical cores.
+ 
     It is also always larger or equal to 1.
     """
-    import math
-
+    # TODO: use os.cpu_count when dropping python 2 support
     try:
         cpu_count_mp = mp.cpu_count()
     except NotImplementedError:
         cpu_count_mp = 1
 
+    cpu_count_user = _cpu_count_user(cpu_count_mp)
+    aggregate_cpu_count = min(cpu_count_mp, cpu_count_user)
+
+    if only_physical_cores:
+        cpu_count_physical, exception = _count_physical_cores()
+        if cpu_count_user < cpu_count_mp:
+            # Respect user setting
+            cpu_count = max(cpu_count_user, 1)
+        elif cpu_count_physical == "not found":
+            # Fallback to default behavior
+            if exception is not None:
+                # warns only the first time
+                warnings.warn(
+                    "Could not find the number of physical cores for the "
+                    "following reason:\n" + str(exception) + "\n"
+                    "Returning the number of logical cores instead. You can "
+                    "silence this warning by setting LOKY_MAX_CPU_COUNT to "
+                    "the number of cores you want to use.")
+                if sys.version_info >= (3, 5):
+                    # TODO remove the version check when dropping py2 support
+                    traceback.print_tb(exception.__traceback__)
+
+            cpu_count = max(aggregate_cpu_count, 1)
+        else:
+            return cpu_count_physical
+    else:
+        cpu_count = max(aggregate_cpu_count, 1)
+
+    return cpu_count
+
+
+def _cpu_count_user(cpu_count_mp):
+    """Number of user defined available CPUs"""
+    import math
+
     # Number of available CPUs given affinity settings
     cpu_count_affinity = cpu_count_mp
     if hasattr(os, 'sched_getaffinity'):
@@ -146,11 +194,65 @@ def cpu_count():
             # float in python2.7. (See issue #165)
             cpu_count_cfs = int(math.ceil(cfs_quota_us / cfs_period_us))
 
-    # User defined soft-limit passed as an loky specific environment variable.
+    # User defined soft-limit passed as a loky specific environment variable.
     cpu_count_loky = int(os.environ.get('LOKY_MAX_CPU_COUNT', cpu_count_mp))
-    aggregate_cpu_count = min(cpu_count_mp, cpu_count_affinity, cpu_count_cfs,
-                              cpu_count_loky)
-    return max(aggregate_cpu_count, 1)
+
+    return min(cpu_count_affinity, cpu_count_cfs, cpu_count_loky)
+
+
+def _count_physical_cores():
+    """Return a tuple (number of physical cores, exception)
+
+    If the number of physical cores is found, exception is set to None.
+    If it has not been found, return ("not found", exception).
+
+    The number of physical cores is cached to avoid repeating subprocess calls.
+    """
+    exception = None
+
+    # First check if the value is cached
+    global physical_cores_cache
+    if physical_cores_cache is not None:
+        return physical_cores_cache, exception
+
+    # Not cached yet, find it
+    try:
+        if sys.platform == "linux":
+            cpu_info = subprocess.run(
+                "lscpu --parse=core".split(" "), capture_output=True)
+            cpu_info = cpu_info.stdout.decode("utf-8").splitlines()
+            cpu_info = {line for line in cpu_info if not line.startswith("#")}
+            cpu_count_physical = len(cpu_info)
+        elif sys.platform == "win32":
+            cpu_info = subprocess.run(
+                "wmic CPU Get NumberOfCores /Format:csv".split(" "),
+                capture_output=True)
+            cpu_info = cpu_info.stdout.decode('utf-8').splitlines()
+            cpu_info = [l.split(",")[1] for l in cpu_info
+                        if (l and l != "Node,NumberOfCores")]
+            cpu_count_physical = sum(map(int, cpu_info))
+        elif sys.platform == "darwin":
+            cpu_info = subprocess.run(
+                "sysctl -n hw.physicalcpu".split(" "), capture_output=True)
+            cpu_info = cpu_info.stdout.decode('utf-8')
+            cpu_count_physical = int(cpu_info)
+        else:
+            raise NotImplementedError(
+                "unsupported platform: {}".format(sys.platform))
+
+        # if cpu_count_physical < 1, we did not find a valid value
+        if cpu_count_physical < 1:
+            raise ValueError(
+                "found {} physical cores < 1".format(cpu_count_physical))
+        
+    except Exception as e:
+        exception = e
+        cpu_count_physical = "not found"
+
+    # Put the result in cache
+    physical_cores_cache = cpu_count_physical
+    
+    return cpu_count_physical, exception
 
 
 class LokyContext(BaseContext):


Test patch:
diff --git a/tests/test_loky_module.py b/tests/test_loky_module.py
index 382301a..076e6e4 100644
--- a/tests/test_loky_module.py
+++ b/tests/test_loky_module.py
@@ -1,12 +1,16 @@
+import multiprocessing as mp
 import os
 import sys
 import shutil
+import tempfile
 from subprocess import check_output
+import subprocess
 
 import pytest
 
 import loky
 from loky import cpu_count
+from loky.backend.context import _cpu_count_user
 
 
 def test_version():
@@ -19,9 +23,18 @@ def test_cpu_count():
     assert type(cpus) is int
     assert cpus >= 1
 
+    cpus_physical = cpu_count(only_physical_cores=True)
+    assert type(cpus_physical) is int
+    assert 1 <= cpus_physical <= cpus
+
+    # again to check that it's correctly cached
+    cpus_physical = cpu_count(only_physical_cores=True)
+    assert type(cpus_physical) is int
+    assert 1 <= cpus_physical <= cpus
+
 
 cpu_count_cmd = ("from loky.backend.context import cpu_count;"
-                 "print(cpu_count())")
+                 "print(cpu_count({args}))")
 
 
 def test_cpu_count_affinity():
@@ -40,9 +53,14 @@ def test_cpu_count_affinity():
         pytest.skip()
 
     res = check_output([taskset_bin, '-c', '0',
-                        python_bin, '-c', cpu_count_cmd])
+                        python_bin, '-c', cpu_count_cmd.format(args='')])
+    
+    res_physical = check_output([
+        taskset_bin, '-c', '0', python_bin, '-c',
+        cpu_count_cmd.format(args='only_physical_cores=True')])
 
     assert res.strip().decode('utf-8') == '1'
+    assert res_physical.strip().decode('utf-8') == '1'
 
 
 def test_cpu_count_cfs_limit():
@@ -64,6 +82,61 @@ def test_cpu_count_cfs_limit():
     res = check_output([docker_bin, 'run', '--rm', '--cpus', '0.5',
                         '-v', '%s:/loky' % loky_path,
                         'python:3.6',
-                        'python', '-c', cpu_count_cmd])
+                        'python', '-c', cpu_count_cmd.format(args='')])
 
     assert res.strip().decode('utf-8') == '1'
+
+
+def test_only_physical_cores_error():
+    # Check the warning issued by cpu_count(only_physical_cores=True) when
+    # unable to retrieve the number of physical cores.
+    if sys.platform != "linux":
+        pytest.skip()
+    
+    # if number of available cpus is already restricted, cpu_count will return
+    # that value and no warning is issued even if only_physical_cores == True.
+    # (tested in another test: test_only_physical_cores_with_user_limitation
+    cpu_count_mp = mp.cpu_count()
+    if _cpu_count_user(cpu_count_mp) < cpu_count_mp:
+        pytest.skip()
+
+    start_dir = os.path.abspath('.')
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        # Write bad lscpu program
+        lscpu_path = tmp_dir + '/lscpu'
+        with open(lscpu_path, 'w') as f:
+            f.write("#!/bin/sh\n"
+                    "exit(1)")
+        os.chmod(lscpu_path, 0o777)
+
+        try:
+            old_path = os.environ['PATH']
+            os.environ['PATH'] = tmp_dir + ":" + old_path
+
+            # clear the cache otherwise the warning is not triggered
+            import loky.backend.context
+            loky.backend.context.physical_cores_cache = None
+
+            with pytest.warns(UserWarning, match="Could not find the number of"
+                                                 " physical cores"):
+                cpu_count(only_physical_cores=True)
+
+            # Should not warn the second time
+            with pytest.warns(None) as record:
+                cpu_count(only_physical_cores=True)
+                assert not record
+
+        finally:
+            os.environ['PATH'] = old_path
+
+
+def test_only_physical_cores_with_user_limitation():
+    # Check that user limitation for the available number of cores is
+    # respected even if only_physical_cores == True
+    cpu_count_mp = mp.cpu_count()
+    cpu_count_user = _cpu_count_user(cpu_count_mp)
+
+    if cpu_count_user < cpu_count_mp:
+        assert cpu_count() == cpu_count_user
+        assert cpu_count(only_physical_cores=True) == cpu_count_user