diff --git a/llvm/utils/lit/lit/cl_arguments.py b/llvm/utils/lit/lit/cl_arguments.py index ddb4d25cbba1..ce49c3c48a97 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.usable_core_count()) + default=lit.util.detectCPUs()) 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 255e63329432..5aef77e5f605 100644 --- a/llvm/utils/lit/lit/run.py +++ b/llvm/utils/lit/lit/run.py @@ -110,7 +110,7 @@ class Run(object): # 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.usable_core_count() + ncpus = lit.util.detectCPUs() 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 0d8f0040f41c..d7afbdabcff9 100644 --- a/llvm/utils/lit/lit/util.py +++ b/llvm/utils/lit/lit/util.py @@ -109,23 +109,32 @@ def to_unicode(s): return s -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. +# 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. """ - 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 + # 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 def mkdir(path):