From 62728530304951fd91193e4b1f9714babe63ccb8 Mon Sep 17 00:00:00 2001
From: "Adam J. Stewart" <ajstewart426@gmail.com>
Date: Fri, 17 Mar 2023 10:13:27 -0600
Subject: [PATCH] Bazel: limit parallelism (#36002)

* Bazel: limit parallelism

* Patch packages that don't directly invoke bazel

* Style fixes

* flag comes after build, not bazel

* flag comes after build, not bazel

* command is only attribute if specific package
---
 .../repos/builtin/packages/bazel/package.py      |  5 +++--
 .../repos/builtin/packages/distbench/package.py  |  3 +++
 .../repos/builtin/packages/py-jaxlib/package.py  | 16 ++++++++--------
 .../repos/builtin/packages/py-ray/package.py     |  8 ++++++++
 .../packages/py-tensorflow-metadata/package.py   | 15 +++++++++++++++
 5 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/var/spack/repos/builtin/packages/bazel/package.py b/var/spack/repos/builtin/packages/bazel/package.py
index 496f80d8c4..7e37b0574b 100644
--- a/var/spack/repos/builtin/packages/bazel/package.py
+++ b/var/spack/repos/builtin/packages/bazel/package.py
@@ -270,11 +270,12 @@ def install_test(self):
                 )
 
             # Spack's logs don't handle colored output well
-            bazel = Executable(self.prefix.bin.bazel)
+            bazel = Executable(self.spec["bazel"].command.path)
             bazel(
                 "--output_user_root=/tmp/spack/bazel/spack-test",
                 "build",
                 "--color=no",
+                f"--jobs={make_jobs}",
                 "//:bazel-test",
             )
 
@@ -282,7 +283,7 @@ def install_test(self):
             assert exe(output=str) == "Hi!\n"
 
     def setup_dependent_package(self, module, dependent_spec):
-        module.bazel = Executable("bazel")
+        module.bazel = Executable(self.spec["bazel"].command.path)
 
     @property
     def parallel(self):
diff --git a/var/spack/repos/builtin/packages/distbench/package.py b/var/spack/repos/builtin/packages/distbench/package.py
index 4d03eafc88..e2ee3e0cdc 100644
--- a/var/spack/repos/builtin/packages/distbench/package.py
+++ b/var/spack/repos/builtin/packages/distbench/package.py
@@ -16,3 +16,6 @@ class Distbench(MakefilePackage):
     version("1.0rc4", sha256="adc8da85890219800207d0d4cd7ffd63193d2c4007dba7c44cf545cc13675ff7")
 
     depends_on("bazel", type="build")
+
+    def patch(self):
+        filter_file("bazel build", f"bazel build --jobs={make_jobs}", "Makefile", string=True)
diff --git a/var/spack/repos/builtin/packages/py-jaxlib/package.py b/var/spack/repos/builtin/packages/py-jaxlib/package.py
index 9844b749b2..7d0c8b99da 100644
--- a/var/spack/repos/builtin/packages/py-jaxlib/package.py
+++ b/var/spack/repos/builtin/packages/py-jaxlib/package.py
@@ -50,19 +50,19 @@ class PyJaxlib(PythonPackage, CudaPackage):
     def patch(self):
         self.tmp_path = tempfile.mkdtemp(prefix="spack")
         self.buildtmp = tempfile.mkdtemp(prefix="spack")
-        # triple quotes necessary because of a variety
-        # of other embedded quote(s)
         filter_file(
-            """f"--output_path={output_path}",""",
-            """f"--output_path={output_path}","""
-            """f"--sources_path=%s","""
-            """f"--nohome_rc'","""
-            """f"--nosystem_rc'",""" % self.tmp_path,
+            'f"--output_path={output_path}",',
+            'f"--output_path={output_path}",'
+            f' "--sources_path={self.tmp_path}",'
+            ' "--nohome_rc",'
+            ' "--nosystem_rc",'
+            f' "--jobs={make_jobs}",',
             "build/build.py",
+            string=True,
         )
         filter_file(
             "args = parser.parse_args()",
-            "args,junk = parser.parse_known_args()",
+            "args, junk = parser.parse_known_args()",
             "build/build_wheel.py",
             string=True,
         )
diff --git a/var/spack/repos/builtin/packages/py-ray/package.py b/var/spack/repos/builtin/packages/py-ray/package.py
index 942198d80c..8cbfad66f9 100644
--- a/var/spack/repos/builtin/packages/py-ray/package.py
+++ b/var/spack/repos/builtin/packages/py-ray/package.py
@@ -78,6 +78,14 @@ class PyRay(PythonPackage):
 
     build_directory = "python"
 
+    def patch(self):
+        filter_file(
+            'bazel_flags = ["--verbose_failures"]',
+            f'bazel_flags = ["--verbose_failures", "--jobs={make_jobs}"]',
+            join_path("python", "setup.py"),
+            string=True,
+        )
+
     def setup_build_environment(self, env):
         env.set("SKIP_THIRDPARTY_INSTALL", "1")
 
diff --git a/var/spack/repos/builtin/packages/py-tensorflow-metadata/package.py b/var/spack/repos/builtin/packages/py-tensorflow-metadata/package.py
index 4ded5a7a43..3a2b44c66e 100644
--- a/var/spack/repos/builtin/packages/py-tensorflow-metadata/package.py
+++ b/var/spack/repos/builtin/packages/py-tensorflow-metadata/package.py
@@ -37,6 +37,21 @@ class PyTensorflowMetadata(PythonPackage):
     depends_on("py-googleapis-common-protos@1.52:1", type=("build", "run"))
     depends_on("py-protobuf@3.13:3", type=("build", "run"))
 
+    def patch(self):
+        filter_file(
+            "self._additional_build_options = ['--copt=-DWIN32_LEAN_AND_MEAN']",
+            "self._additional_build_options = ['--copt=-DWIN32_LEAN_AND_MEAN',"
+            f" '--jobs={make_jobs}']",
+            "setup.py",
+            string=True,
+        )
+        filter_file(
+            "self._additional_build_options = []",
+            f"self._additional_build_options = ['--jobs={make_jobs}']",
+            "setup.py",
+            string=True,
+        )
+
     def setup_build_environment(self, env):
         tmp_path = tempfile.mkdtemp(prefix="spack")
         env.set("TEST_TMPDIR", tmp_path)