xref: /aosp_15_r20/external/grpc-grpc/templates/src/python/_parallel_compile_patch.py.include (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2) 
1"""Patches the compile() to allow enable parallel compilation of C/C++.
2
3build_ext has lots of C/C++ files and normally them one by one.
4Enabling parallel build helps a lot.
5"""
6
7import os
8
9try:
10    BUILD_EXT_COMPILER_JOBS = int(
11        os.environ["GRPC_PYTHON_BUILD_EXT_COMPILER_JOBS"]
12    )
13except KeyError:
14    import multiprocessing
15
16    BUILD_EXT_COMPILER_JOBS = multiprocessing.cpu_count()
17except ValueError:
18    BUILD_EXT_COMPILER_JOBS = 1
19
20
21# monkey-patch for parallel compilation
22def _parallel_compile(
23    self,
24    sources,
25    output_dir=None,
26    macros=None,
27    include_dirs=None,
28    debug=0,
29    extra_preargs=None,
30    extra_postargs=None,
31    depends=None,
32):
33    # setup the same way as distutils.ccompiler.CCompiler
34    # https://github.com/python/cpython/blob/31368a4f0e531c19affe2a1becd25fc316bc7501/Lib/distutils/ccompiler.py#L564
35    macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
36        str(output_dir), macros, include_dirs, sources, depends, extra_postargs
37    )
38    cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
39
40    def _compile_single_file(obj):
41        try:
42            src, ext = build[obj]
43        except KeyError:
44            return
45        self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
46
47    # run compilation of individual files in parallel
48    import multiprocessing.pool
49
50    multiprocessing.pool.ThreadPool(BUILD_EXT_COMPILER_JOBS).map(
51        _compile_single_file, objects
52    )
53    return objects
54
55
56def monkeypatch_compile_maybe():
57    """
58    Monkeypatching is dumb, but the build speed gain is worth it.
59    After python 3.12, we won't find distutils if SETUPTOOLS_USE_DISTUTILS=stdlib.
60    """
61    use_distutils = os.environ.get("SETUPTOOLS_USE_DISTUTILS", "")
62    if BUILD_EXT_COMPILER_JOBS > 1 and use_distutils != "stdlib":
63        import distutils.ccompiler  # pylint: disable=wrong-import-position
64
65        distutils.ccompiler.CCompiler.compile = _parallel_compile
66