_parallel_compile_patch.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. # Copyright 2018 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. """Patches the compile() to allow enable parallel compilation of C/C++.
  15. build_ext has lots of C/C++ files and normally them one by one.
  16. Enabling parallel build helps a lot.
  17. """
  18. import distutils.ccompiler
  19. import os
  20. try:
  21. BUILD_EXT_COMPILER_JOBS = int(
  22. os.environ['GRPC_PYTHON_BUILD_EXT_COMPILER_JOBS'])
  23. except KeyError:
  24. import multiprocessing
  25. BUILD_EXT_COMPILER_JOBS = multiprocessing.cpu_count()
  26. # monkey-patch for parallel compilation
  27. def _parallel_compile(self,
  28. sources,
  29. output_dir=None,
  30. macros=None,
  31. include_dirs=None,
  32. debug=0,
  33. extra_preargs=None,
  34. extra_postargs=None,
  35. depends=None):
  36. # setup the same way as distutils.ccompiler.CCompiler
  37. # https://github.com/python/cpython/blob/31368a4f0e531c19affe2a1becd25fc316bc7501/Lib/distutils/ccompiler.py#L564
  38. macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
  39. output_dir, macros, include_dirs, sources, depends, extra_postargs)
  40. cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
  41. def _compile_single_file(obj):
  42. try:
  43. src, ext = build[obj]
  44. except KeyError:
  45. return
  46. self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
  47. # run compilation of individual files in parallel
  48. import multiprocessing.pool
  49. multiprocessing.pool.ThreadPool(BUILD_EXT_COMPILER_JOBS).map(
  50. _compile_single_file, objects)
  51. return objects
  52. def monkeypatch_compile_maybe():
  53. """Monkeypatching is dumb, but the build speed gain is worth it."""
  54. if BUILD_EXT_COMPILER_JOBS > 1:
  55. distutils.ccompiler.CCompiler.compile = _parallel_compile