Bug 1903620 - Calculate jobs using CPU affinity in mach. r=firefox-build-system-reviewers,glandium

Differential Revision: https://phabricator.services.mozilla.com/D214289

dom/bindings/mozwebidlcodegen/__init__.py

@@ -18,14 +18,14 @@ import mozpack.path as mozpath
 from mach.mixin.logging import LoggingMixin
 from mozbuild.makeutil import Makefile
 from mozbuild.pythonutil import iter_modules_in_path
-from mozbuild.util import FileAvoidWrite
+from mozbuild.util import FileAvoidWrite, cpu_count
 
 # There are various imports in this file in functions to avoid adding
 # dependencies to config.status. See bug 949875.
 
 # Limit the count on Windows, because of bug 1889842 and also the
 # inefficiency of fork on Windows.
-DEFAULT_PROCESS_COUNT = 4 if sys.platform == "win32" else os.cpu_count()
+DEFAULT_PROCESS_COUNT = 4 if sys.platform == "win32" else cpu_count()
 
 
 class WebIDLPool:

intl/icu_sources_data.py

@@ -11,13 +11,13 @@
 # copy of ICU has been updated.
 
 import glob
-import multiprocessing
 import os
 import shutil
 import subprocess
 import sys
 import tempfile
 
+from mozbuild.util import cpu_count
 from mozpack import path as mozpath
 
 # The following files have been determined to be dead/unused by a
@@ -277,7 +277,7 @@ def update_data_file(topsrcdir):
     print("Running ICU make...")
     if not try_run(
         "icu-make",
-        ["make", "--jobs=%d" % multiprocessing.cpu_count(), "--output-sync"],
+        ["make", "--jobs=%d" % cpu_count(), "--output-sync"],
         cwd=objdir,
     ):
         return False

layout/tools/reftest/runreftest.py

@@ -6,7 +6,6 @@
 Runs the reftest test harness.
 """
 import json
-import multiprocessing
 import os
 import platform
 import posixpath
@@ -61,10 +60,12 @@ here = os.path.abspath(os.path.dirname(__file__))
 
 try:
     from mozbuild.base import MozbuildObject
+    from mozbuild.util import cpu_count
 
     build_obj = MozbuildObject.from_environment(cwd=here)
 except ImportError:
     build_obj = None
+    from multiprocessing import cpu_count
 
 
 def categoriesToRegex(categoryList):
@@ -698,7 +699,7 @@ class RefTest(object):
         if not getattr(options, "runTestsInParallel", False):
             return self.runSerialTests(manifests, options, cmdargs)
 
-        cpuCount = multiprocessing.cpu_count()
+        cpuCount = cpu_count()
 
         # We have the directive, technology, and machine to run multiple test instances.
         # Experimentation says that reftests are not overly CPU-intensive, so we can run

python/mach_commands.py

@@ -8,12 +8,12 @@ import os
 import subprocess
 import tempfile
 from concurrent.futures import ThreadPoolExecutor, as_completed, thread
-from multiprocessing import cpu_count
 
 import mozinfo
 from mach.decorators import Command, CommandArgument
 from manifestparser import TestManifest
 from manifestparser import filters as mpf
+from mozbuild.util import cpu_count
 from mozfile import which
 from tqdm import tqdm
 

python/mozbuild/mozbuild/backend/mach_commands.py

@@ -13,6 +13,7 @@ from mach.decorators import Command, CommandArgument
 from mozfile import which
 
 from mozbuild import build_commands
+from mozbuild.util import cpu_count
 
 
 @Command(
@@ -300,8 +301,6 @@ def setup_clangd_rust_in_vscode(command_context):
         if rc != 0:
             return rc
 
-    import multiprocessing
-
     from mozbuild.code_analysis.utils import ClangTidyConfig
 
     clang_tidy_cfg = ClangTidyConfig(command_context.topsrcdir)
@@ -316,7 +315,7 @@ def setup_clangd_rust_in_vscode(command_context):
         "cargo",
         "check",
         "-j",
-        str(multiprocessing.cpu_count() // 2),
+        str(cpu_count() // 2),
         "--all-crates",
         "--message-format-json",
     ]
@@ -344,7 +343,7 @@ def setup_clangd_rust_in_vscode(command_context):
         "clangd.path": clangd_path,
         "clangd.arguments": [
             "-j",
-            str(multiprocessing.cpu_count() // 2),
+            str(cpu_count() // 2),
             "--limit-results",
             "0",
             "--completion-style",

python/mozbuild/mozbuild/base.py

@@ -6,7 +6,6 @@ import errno
 import io
 import json
 import logging
-import multiprocessing
 import os
 import subprocess
 import sys
@@ -31,7 +30,7 @@ from .backend.configenvironment import ConfigEnvironment, ConfigStatusFailure
 from .configure import ConfigureSandbox
 from .controller.clobber import Clobberer
 from .mozconfig import MozconfigLoader, MozconfigLoadException
-from .util import memoize, memoized_property
+from .util import cpu_count, memoize, memoized_property
 
 try:
     import psutil
@@ -747,7 +746,7 @@ class MozbuildObject(ProcessExecutionMixin):
             if job_size == 0:
                 job_size = 2.0 if self.substs.get("CC_TYPE") == "gcc" else 1.0  # GiB
 
-            cpus = multiprocessing.cpu_count()
+            cpus = cpu_count()
             if not psutil or not job_size:
                 num_jobs = cpus
             else:

python/mozbuild/mozbuild/code_analysis/mach_commands.py

@@ -4,7 +4,6 @@
 import concurrent.futures
 import json
 import logging
-import multiprocessing
 import ntpath
 import os
 import pathlib
@@ -28,7 +27,7 @@ from six.moves import input
 from mozbuild import build_commands
 from mozbuild.controller.clobber import Clobberer
 from mozbuild.nodeutil import find_node_executable
-from mozbuild.util import memoize
+from mozbuild.util import cpu_count, memoize
 
 
 # Function used to run clang-format on a batch of files. It is a helper function
@@ -386,7 +385,7 @@ def check(
     ) as output_manager:
         import math
 
-        batch_size = int(math.ceil(float(len(source)) / multiprocessing.cpu_count()))
+        batch_size = int(math.ceil(float(len(source)) / cpu_count()))
         for i in range(0, len(source), batch_size):
             args = _get_clang_tidy_command(
                 command_context,
@@ -711,7 +710,7 @@ def autotest(
         )
         return TOOLS_UNSUPORTED_PLATFORM
 
-    max_workers = multiprocessing.cpu_count()
+    max_workers = cpu_count()
 
     command_context.log(
         logging.INFO,
@@ -1904,7 +1903,7 @@ def _run_clang_format_path(
     # Run clang-format in parallel trying to saturate all of the available cores.
     import math
 
-    max_workers = multiprocessing.cpu_count()
+    max_workers = cpu_count()
 
     # To maximize CPU usage when there are few items to handle,
     # underestimate the number of items per batch, then dispatch

python/mozbuild/mozbuild/frontend/reader.py

@@ -30,7 +30,6 @@ from collections import OrderedDict, defaultdict
 from concurrent.futures.process import ProcessPoolExecutor
 from io import StringIO
 from itertools import chain
-from multiprocessing import cpu_count
 
 import mozpack.path as mozpath
 from mozpack.files import FileFinder
@@ -41,6 +40,7 @@ from mozbuild.util import (
     EmptyValue,
     HierarchicalStringList,
     ReadOnlyDefaultDict,
+    cpu_count,
     memoize,
 )
 

python/mozbuild/mozbuild/util.py

@@ -22,6 +22,11 @@ import six
 
 from mozbuild.dirutils import ensureParentDir
 
+try:
+    import psutil
+except Exception:
+    psutil = None
+
 MOZBUILD_METRICS_PATH = os.path.abspath(
     os.path.join(__file__, "..", "..", "metrics.yaml")
 )
@@ -1227,3 +1232,22 @@ def hexdump(buf):
         line += "|\n"
         lines.append(line)
     return lines
+
+
+def cpu_count():
+    """
+    Returns the number of CPUs available to us. This may be different than
+    `os.cpu_count()` because of affinity.
+
+    See the Python documentation for `os.cpu_count()`.
+    """
+    try:
+        return len(os.sched_getaffinity(0))
+    except (AttributeError, OSError):
+        pass
+    if psutil:
+        try:
+            return len(psutil.Process().cpu_affinity())
+        except (AttributeError, OSError):
+            pass
+    return os.cpu_count()

python/mozbuild/mozpack/macpkg.py

@@ -16,6 +16,8 @@ import zlib
 from collections import deque, namedtuple
 from xml.etree.ElementTree import XML
 
+from mozbuild.util import cpu_count
+
 
 class ZlibFile(object):
     def __init__(self, fileobj):
@@ -108,7 +110,7 @@ class Pbzx(object):
         # check.
         chunk_size = fileobj.read(8)
         chunk_size = struct.unpack(">Q", chunk_size)[0]
-        executor = concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count())
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=cpu_count())
         self.chunk_getter = executor.map(self._uncompress_chunk, self._chunker(fileobj))
         self._init_one_chunk()
 

python/mozbuild/mozpack/pkg.py

@@ -4,7 +4,6 @@
 
 import concurrent.futures
 import lzma
-import os
 import plistlib
 import struct
 import subprocess
@@ -15,6 +14,8 @@ from urllib.parse import quote
 
 import mozfile
 
+from mozbuild.util import cpu_count
+
 TEMPLATE_DIRECTORY = Path(__file__).parent / "apple_pkg"
 PBZX_CHUNK_SIZE = 16 * 1024 * 1024  # 16MB chunks
 
@@ -132,7 +133,7 @@ def create_payload(destination: Path, root_path: Path, cpio_tool: str):
         with tmp_payload_path.open("rb") as f_in, destination.open(
             "wb"
         ) as f_out, concurrent.futures.ThreadPoolExecutor(
-            max_workers=os.cpu_count()
+            max_workers=cpu_count()
         ) as executor:
             f_out.write(b"pbzx")
             f_out.write(struct.pack(">Q", PBZX_CHUNK_SIZE))

python/mozlint/mozlint/roller.py

@@ -14,12 +14,13 @@ from concurrent.futures import ProcessPoolExecutor
 from concurrent.futures.process import _python_exit as futures_atexit
 from itertools import chain
 from math import ceil
-from multiprocessing import cpu_count, get_context
+from multiprocessing import get_context
 from multiprocessing.queues import Queue
 from subprocess import CalledProcessError
 from typing import Dict, Set
 
 import mozpack.path as mozpath
+from mozbuild.util import cpu_count
 from mozversioncontrol import (
     InvalidRepoPath,
     MissingUpstreamRepo,

testing/xpcshell/mach_commands.py

@@ -8,11 +8,11 @@ import errno
 import logging
 import os
 import sys
-from multiprocessing import cpu_count
 
 from mach.decorators import Command
 from mozbuild.base import BinaryNotFoundException, MozbuildObject
 from mozbuild.base import MachCommandConditions as conditions
+from mozbuild.util import cpu_count
 from mozlog import structured
 from xpcshellcommandline import parser_desktop, parser_remote
 

tools/moztreedocs/mach_commands.py

@@ -4,7 +4,6 @@
 
 import fnmatch
 import json
-import multiprocessing
 import os
 import re
 import subprocess
@@ -20,7 +19,7 @@ import sentry_sdk
 import yaml
 from mach.decorators import Command, CommandArgument, SubCommand
 from mach.registrar import Registrar
-from mozbuild.util import memoize
+from mozbuild.util import cpu_count, memoize
 from mozfile import load_source
 
 here = os.path.abspath(os.path.dirname(__file__))
@@ -81,7 +80,7 @@ BASE_LINK = "http://gecko-docs.mozilla.org-l1.s3-website.us-west-2.amazonaws.com
 @CommandArgument(
     "-j",
     "--jobs",
-    default=str(multiprocessing.cpu_count()),
+    default=str(cpu_count()),
     dest="jobs",
     help="Distribute the build over N processes in parallel.",
 )