Merge pull request #2804 from varungandhi-apple/vg-python3-is-lit

Cherry-pick Python3 changes to apple/stable/20210107

Author: varungandhi-apple

llvm/utils/lit/lit.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 
 from lit.main import main
 

llvm/utils/lit/lit/cl_arguments.py

@@ -31,7 +31,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",

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.

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):