Fix missing defaults for GPU target arch'es (#293)

Move defaults for these up to ci_build python
We also add some error catching logic in the
python script to clean up weirdness coming
from the build scripts.

Also clean up some formatting and typos

(cherry picked from commit ef1d561)

Author: mrodden

build/rocm/ci_build

@@ -25,6 +25,10 @@ import logging
 import os
 import subprocess
 import sys
+from typing import List
+
+
+DEFAULT_GPU_DEVICE_TARGETS = "gfx906,gfx908,gfx90a,gfx942,gfx1030,gfx1100,gfx1101,gfx1200,gfx1201"
 
 
 LOG = logging.getLogger("ci_build")
@@ -37,10 +41,11 @@ def image_by_name(name):
     return image_id
 
 
-def create_manylinux_build_image(rocm_version, rocm_build_job, rocm_build_num):
+def create_manylinux_build_image(rocm_version, rocm_build_job, rocm_build_num, gpu_device_targets: List[str]) -> str:
     image_name = "jax-build-manylinux_2_28_x86_64_rocm%s" % rocm_version.replace(
         ".", ""
     )
+
     cmd = [
         "docker",
         "build",
@@ -50,6 +55,7 @@ def create_manylinux_build_image(rocm_version, rocm_build_job, rocm_build_num):
         "--build-arg=ROCM_BUILD_JOB=%s" % rocm_build_job,
         "--build-arg=ROCM_BUILD_NUM=%s" % rocm_build_num,
         "--tag=%s" % image_name,
+        "--build-arg=GPU_DEVICE_TARGETS=%s" % " ".join(gpu_device_targets),
         ".",
     ]
 
@@ -65,11 +71,15 @@ def dist_wheels(
     rocm_build_job="",
     rocm_build_num="",
     compiler="gcc",
+    gpu_device_targets : List[str] = None,
 ):
+    if not gpu_device_targets:
+        gpu_device_targets = DEFAULT_GPU_DEVICE_TARGETS.split(",")
+
     # We want to make sure the wheels we build are manylinux compliant. We'll
     # do the build in a container. Build the image for this.
     image_name = create_manylinux_build_image(
-        rocm_version, rocm_build_job, rocm_build_num
+        rocm_version, rocm_build_job, rocm_build_num, gpu_device_targets
     )
 
     if xla_path:
@@ -91,6 +101,8 @@ def dist_wheels(
         pyver_string,
         "--compiler",
         compiler,
+        "--gpu-device-targets",
+        ",".join(gpu_device_targets),
     ]
 
     if xla_path:
@@ -127,31 +139,6 @@ def dist_wheels(
         ]
     )
 
-    # Add command for unit tests
-    cmd.extend(
-        [
-            "&&",
-            "bazel",
-            "test",
-            "-k",
-            "--jobs=4",
-            "--test_verbose_timeout_warnings=true",
-            "--test_output=all",
-            "--test_summary=detailed",
-            "--local_test_jobs=1",
-            "--test_env=JAX_ACCELERATOR_COUNT=%i" % 4, 
-            "--test_env=JAX_SKIP_SLOW_TESTS=0",
-            "--verbose_failures=true",
-            "--config=rocm",
-            "--action_env=ROCM_PATH=/opt/rocm",
-            "--action_env=TF_ROCM_AMDGPU_TARGETS=%s" % "gfx90a",
-            "--test_tag_filters=-multiaccelerator",
-            "--test_env=XLA_PYTHON_CLIENT_ALLOCATOR=platform",
-            "--test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow",
-            "//tests:gpu_tests",
-        ]
-    )
-
     LOG.info("Running: %s", cmd)
     _ = subprocess.run(cmd, check=True)
 
@@ -196,10 +183,14 @@ def dist_docker(
     tag="rocm/jax-dev",
     dockerfile=None,
     keep_image=True,
+    gpu_device_targets : List[str] = None,
 ):
     if not dockerfile:
         dockerfile = "build/rocm/Dockerfile.ms"
 
+    if not gpu_device_targets:
+        gpu_device_targets = DEFAULT_GPU_DEVICE_TARGETS.split(",")
+
     python_version = python_versions[0]
 
     md = _fetch_jax_metadata(xla_path)
@@ -212,6 +203,7 @@ def dist_docker(
         "--target",
         "rt_build",
         "--build-arg=ROCM_VERSION=%s" % rocm_version,
+        "--build-arg=GPU_DEVICE_TARGETS=%s" % " ".join(gpu_device_targets),
         "--build-arg=ROCM_BUILD_JOB=%s" % rocm_build_job,
         "--build-arg=ROCM_BUILD_NUM=%s" % rocm_build_num,
         "--build-arg=BASE_DOCKER=%s" % base_docker,
@@ -278,6 +270,37 @@ def test(image_name, test_cmd):
     subprocess.check_call(cmd)
 
 
+def parse_gpu_targets(targets_string):
+    # catch case where targets_string was empty.
+    # None should already be caught by argparse, but
+    # it doesn't hurt to check twice
+    if not targets_string:
+        targets_string = DEFAULT_GPU_DEVICE_TARGETS
+
+    if "," in targets_string:
+        targets = targets_string.split(",")
+    elif " " in targets_string:
+        targets = targets_string.split(" ")
+    else:
+        targets = targets_string
+
+    res = []
+    # cleanup and validation
+    for t in targets:
+        if not t:
+            continue
+
+        if not t.startswith("gfx"):
+            raise ValueError("Invalid GPU architecture target: %r" % t)
+
+        res.append(t.strip())
+
+    if not res:
+        raise ValueError("GPU_DEVICE_TARGETS cannot be empty")
+
+    return res
+
+
 def parse_args():
     p = argparse.ArgumentParser()
     p.add_argument(
@@ -289,7 +312,7 @@ def parse_args():
     p.add_argument(
         "--python-versions",
         type=lambda x: x.split(","),
-        default="3.12",
+        default=["3.12"],
         help="Comma separated list of CPython versions to build wheels for",
     )
 
@@ -321,6 +344,11 @@ def parse_args():
         choices=["gcc", "clang"],
         help="Compiler backend to use when compiling jax/jaxlib",
     )
+    p.add_argument(
+        "--gpu-device-targets",
+        default=DEFAULT_GPU_DEVICE_TARGETS,
+        help="List of AMDGPU device targets passed from job",
+    )
 
     subp = p.add_subparsers(dest="action", required=True)
 
@@ -344,6 +372,7 @@ def parse_args():
 def main():
     logging.basicConfig(level=logging.INFO)
     args = parse_args()
+    gpu_device_targets = parse_gpu_targets(args.gpu_device_targets)
 
     if args.action == "dist_wheels":
         dist_wheels(
@@ -353,6 +382,7 @@ def main():
             args.rocm_build_job,
             args.rocm_build_num,
             compiler=args.compiler,
+            gpu_device_targets=gpu_device_targets,
         )
 
     elif args.action == "test":
@@ -363,9 +393,11 @@ def main():
             args.rocm_version,
             args.python_versions,
             args.xla_source_dir,
+            gpu_device_targets,
             args.rocm_build_job,
             args.rocm_build_num,
             compiler=args.compiler,
+            gpu_device_targets=gpu_device_targets,
         )
         dist_docker(
             args.rocm_version,
@@ -377,6 +409,7 @@ def main():
             tag=args.image_tag,
             dockerfile=args.dockerfile,
             keep_image=args.keep_image,
+            gpu_device_targets=gpu_device_targets,
         )
 
 

build/rocm/ci_build.sh

@@ -51,6 +51,7 @@ ROCM_BUILD_NUM=""
 BASE_DOCKER="ubuntu:22.04"
 CUSTOM_INSTALL=""
 JAX_USE_CLANG=""
+GPU_DEVICE_TARGETS=""
 POSITIONAL_ARGS=()
 
 RUNTIME_FLAG=0
@@ -98,6 +99,18 @@ while [[ $# -gt 0 ]]; do
           JAX_USE_CLANG="$2"
           shift 2
           ;;
+        --gpu_device_targets)
+          if [[ "$2" == "--custom_install" ]]; then
+            GPU_DEVICE_TARGETS=""
+            shift 2
+          elif [[ -n "$2" ]]; then
+            GPU_DEVICE_TARGETS="$2"
+            shift 2
+          else
+            GPU_DEVICE_TARGETS=""
+            shift 1
+          fi
+          ;;
         *)
           POSITIONAL_ARGS+=("$1")
           shift
@@ -164,6 +177,7 @@ fi
     --rocm-build-job=$ROCM_BUILD_JOB \
     --rocm-build-num=$ROCM_BUILD_NUM \
     --compiler=$JAX_COMPILER \
+    --gpu-device-targets="${GPU_DEVICE_TARGETS}" \
     dist_docker \
     --dockerfile $DOCKERFILE_PATH \
     --image-tag $DOCKER_IMG_NAME

build/rocm/test_ci_build.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import importlib.util
+import importlib.machinery
+
+
+def load_ci_build():
+    spec = importlib.util.spec_from_loader(
+        "ci_build", importlib.machinery.SourceFileLoader("ci_build", "./ci_build")
+    )
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+ci_build = load_ci_build()
+
+
+class CIBuildTestCase(unittest.TestCase):
+    def test_parse_gpu_targets(self):
+        targets = ["gfx908", "gfx940", "gfx1201"]
+
+        r = ci_build.parse_gpu_targets(" ".join(targets))
+        self.assertEqual(r, targets)
+
+        r = ci_build.parse_gpu_targets(",".join(targets))
+        self.assertEqual(r, targets)
+
+    def test_parse_gpu_targets_empty_string(self):
+        expected = ci_build.DEFAULT_GPU_DEVICE_TARGETS.split(",")
+        r = ci_build.parse_gpu_targets("")
+        self.assertEqual(r, expected)
+
+        self.assertRaises(ValueError, ci_build.parse_gpu_targets, " ")
+
+    def test_parse_gpu_targets_invalid_arch(self):
+        targets = ["gfx908", "gfx940", "--oops", "/jax"]
+        self.assertRaises(ValueError, ci_build.parse_gpu_targets, " ".join(targets))
+
+
+if __name__ == "__main__":
+    unittest.main()