[llvm-branch-commits] [llvm] 9946b16 - [lit] Use os.cpu_count() to cleanup TODO

[Julian Lettner via llvm-branch-commits]

Author: Julian Lettner
Date: 2021-01-25T11:44:18-08:00
New Revision: 9946b169c379daee603436a4753acfef8be373dd

URL: https://github.com/llvm/llvm-project/commit/9946b169c379daee603436a4753acfef8be373dd
DIFF: https://github.com/llvm/llvm-project/commit/9946b169c379daee603436a4753acfef8be373dd.diff

LOG: [lit] Use os.cpu_count() to cleanup TODO

We can now use Python3.  Let's use `os.cpu_count()` to cleanup this
helper.

Differential Revision: https://reviews.llvm.org/D94734

Added: 
    

Modified: 
    llvm/utils/lit/lit/cl_arguments.py
    llvm/utils/lit/lit/run.py
    llvm/utils/lit/lit/util.py

Removed: 
    


################################################################################
diff  --git a/llvm/utils/lit/lit/cl_arguments.py b/llvm/utils/lit/lit/cl_arguments.py
index ce49c3c48a97..ddb4d25cbba1 100644
--- a/llvm/utils/lit/lit/cl_arguments.py
+++ b/llvm/utils/lit/lit/cl_arguments.py
@@ -23,7 +23,7 @@ def parse_args():
             metavar="N",
             help="Number of workers used for testing",
             type=_positive_int,
-            default=lit.util.detectCPUs())
+            default=lit.util.usable_core_count())
     parser.add_argument("--config-prefix",
             dest="configPrefix",
             metavar="NAME",

diff  --git a/llvm/utils/lit/lit/run.py b/llvm/utils/lit/lit/run.py
index 5aef77e5f605..255e63329432 100644
--- a/llvm/utils/lit/lit/run.py
+++ b/llvm/utils/lit/lit/run.py
@@ -110,7 +110,7 @@ def _update_test(self, local_test, remote_test):
     # threads counts toward the current process limit. Try to raise the (soft)
     # process limit so that tests don't fail due to resource exhaustion.
     def _increase_process_limit(self):
-        ncpus = lit.util.detectCPUs()
+        ncpus = lit.util.usable_core_count()
         desired_limit = self.workers * ncpus * 2 # the 2 is a safety factor
 
         # Importing the resource module will likely fail on Windows.

diff  --git a/llvm/utils/lit/lit/util.py b/llvm/utils/lit/lit/util.py
index d7afbdabcff9..0d8f0040f41c 100644
--- a/llvm/utils/lit/lit/util.py
+++ b/llvm/utils/lit/lit/util.py
@@ -109,32 +109,23 @@ def to_unicode(s):
     return s
 
 
-# TODO(yln): multiprocessing.cpu_count()
-# TODO(python3): len(os.sched_getaffinity(0)) and os.cpu_count()
-def detectCPUs():
-    """Detects the number of CPUs on a system.
-
-    Cribbed from pp.
+def usable_core_count():
+    """Return the number of cores the current process can use, if supported.
+    Otherwise, return the total number of cores (like `os.cpu_count()`).
+    Default to 1 if undetermined.
 
     """
-    # Linux, Unix and MacOS:
-    if hasattr(os, 'sysconf'):
-        if 'SC_NPROCESSORS_ONLN' in os.sysconf_names:
-            # Linux & Unix:
-            ncpus = os.sysconf('SC_NPROCESSORS_ONLN')
-            if isinstance(ncpus, int) and ncpus > 0:
-                return ncpus
-        else:  # OSX:
-            return int(subprocess.check_output(['sysctl', '-n', 'hw.ncpu'],
-                                               stderr=subprocess.STDOUT))
-    # Windows:
-    if 'NUMBER_OF_PROCESSORS' in os.environ:
-        ncpus = int(os.environ['NUMBER_OF_PROCESSORS'])
-        if ncpus > 0:
-            # With more than 32 processes, process creation often fails with
-            # "Too many open files".  FIXME: Check if there's a better fix.
-            return min(ncpus, 32)
-    return 1  # Default
+    try:
+        n = len(os.sched_getaffinity(0))
+    except AttributeError:
+        n = os.cpu_count() or 1
+
+    # On Windows, with more than 32 processes, process creation often fails with
+    # "Too many open files".  FIXME: Check if there's a better fix.
+    if platform.system() == 'Windows':
+        return min(n, 32)
+
+    return n
 
 
 def mkdir(path):