From 1fa6e91af6eb58b8cbc529d7dd9a76b2fabca6c0 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Thu, 9 Jan 2025 16:53:30 +0000 Subject: [PATCH 1/3] Add option to ci_build to run different tests --- build/rocm/ci_build | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/build/rocm/ci_build b/build/rocm/ci_build index 2556633482ff..86faa5e42fdd 100755 --- a/build/rocm/ci_build +++ b/build/rocm/ci_build @@ -202,7 +202,7 @@ def dist_docker( subprocess.check_call(cmd) -def test(image_name): +def test(image_name, test_cmd): """Run unit tests like CI would inside a JAX image.""" gpu_args = [ @@ -236,7 +236,7 @@ def test(image_name): cmd.extend(mounts) cmd.extend(gpu_args) - container_cmd = "cd /jax && ./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh" + container_cmd = "cd /jax && " + test_cmd cmd.append(image_name) cmd.extend( [ @@ -299,6 +299,7 @@ def parse_args(): testp = subp.add_parser("test") testp.add_argument("image_name") + testp.add_argument("--test-cmd", default="./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh") ddp = subp.add_parser("dist_docker") ddp.add_argument("--dockerfile", default="build/rocm/Dockerfile.ms") @@ -322,7 +323,7 @@ def main(): ) elif args.action == "test": - test(args.image_name) + test(args.image_name, args.test_cmd) elif args.action == "dist_docker": dist_wheels( From 5405752649d863dfe956630e2859cdbea981243a Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Thu, 9 Jan 2025 16:54:38 +0000 Subject: [PATCH 2/3] Only run core tests for CI --- .github/workflows/rocm-ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/rocm-ci.yml b/.github/workflows/rocm-ci.yml index 7f9cd760aee5..12e16cce6f24 100644 --- a/.github/workflows/rocm-ci.yml +++ b/.github/workflows/rocm-ci.yml @@ -59,5 +59,5 @@ jobs: - name: Run tests run: | cd $WORKSPACE_DIR - python3 build/rocm/ci_build test $TEST_IMAGE + python3 build/rocm/ci_build test $TEST_IMAGE --test-cmd pytest tests/core_test.py From bcc2417d3963082761c38ca367af54e9f2e53f66 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Thu, 9 Jan 2025 17:56:35 +0000 Subject: [PATCH 3/3] Quote test command in workflow file --- .github/workflows/rocm-ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/rocm-ci.yml b/.github/workflows/rocm-ci.yml index 12e16cce6f24..e910e8723b9b 100644 --- a/.github/workflows/rocm-ci.yml +++ b/.github/workflows/rocm-ci.yml @@ -59,5 +59,5 @@ jobs: - name: Run tests run: | cd $WORKSPACE_DIR - python3 build/rocm/ci_build test $TEST_IMAGE --test-cmd pytest tests/core_test.py + python3 build/rocm/ci_build test $TEST_IMAGE --test-cmd "pytest tests/core_test.py"