xref: /aosp_15_r20/external/grpc-grpc/src/python/grpcio/_parallel_compile_patch.py (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2) 
1# Copyright 2023 The gRPC Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14#
15# This file has been automatically generated from a template file.
16# Please make modifications to
17# `$REPO_ROOT/templates/src/python/grpcio/_parallel_compile_patch.py.template`
18# instead. This file can be regenerated from the template by running
19# `tools/buildgen/generate_projects.sh`.
20
21"""Patches the compile() to allow enable parallel compilation of C/C++.
22
23build_ext has lots of C/C++ files and normally them one by one.
24Enabling parallel build helps a lot.
25"""
26
27import os
28
29try:
30    BUILD_EXT_COMPILER_JOBS = int(
31        os.environ["GRPC_PYTHON_BUILD_EXT_COMPILER_JOBS"]
32    )
33except KeyError:
34    import multiprocessing
35
36    BUILD_EXT_COMPILER_JOBS = multiprocessing.cpu_count()
37except ValueError:
38    BUILD_EXT_COMPILER_JOBS = 1
39
40
41# monkey-patch for parallel compilation
42def _parallel_compile(
43    self,
44    sources,
45    output_dir=None,
46    macros=None,
47    include_dirs=None,
48    debug=0,
49    extra_preargs=None,
50    extra_postargs=None,
51    depends=None,
52):
53    # setup the same way as distutils.ccompiler.CCompiler
54    # https://github.com/python/cpython/blob/31368a4f0e531c19affe2a1becd25fc316bc7501/Lib/distutils/ccompiler.py#L564
55    macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
56        str(output_dir), macros, include_dirs, sources, depends, extra_postargs
57    )
58    cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
59
60    def _compile_single_file(obj):
61        try:
62            src, ext = build[obj]
63        except KeyError:
64            return
65        self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
66
67    # run compilation of individual files in parallel
68    import multiprocessing.pool
69
70    multiprocessing.pool.ThreadPool(BUILD_EXT_COMPILER_JOBS).map(
71        _compile_single_file, objects
72    )
73    return objects
74
75
76def monkeypatch_compile_maybe():
77    """
78    Monkeypatching is dumb, but the build speed gain is worth it.
79    After python 3.12, we won't find distutils if SETUPTOOLS_USE_DISTUTILS=stdlib.
80    """
81    use_distutils = os.environ.get("SETUPTOOLS_USE_DISTUTILS", "")
82    if BUILD_EXT_COMPILER_JOBS > 1 and use_distutils != "stdlib":
83        import distutils.ccompiler  # pylint: disable=wrong-import-position
84
85        distutils.ccompiler.CCompiler.compile = _parallel_compile
86