split build and test jobs

Author: nitins17

.github/workflows/pytest_gpu.yml

@@ -16,7 +16,35 @@ on:
         - 'no'
 
 jobs:
-  run_pytest_gpu:
+  build_artifacts:
+    strategy:
+      matrix:
+        python: ["3.10"]
+
+    runs-on: "linux-x86-n2-16"
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
+    
+    name: "Pytest GPU (Build wheels on CUDA 12.3)"
+    env:
+      JAXCI_CLONE_MAIN_XLA: 1
+      JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
+
+    steps:
+      - uses: actions/checkout@v3
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Build jaxlib
+        run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jaxlib.env"
+      - name: Build jax-cuda-plugin
+        run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jax-cuda-plugin.env"
+      - name: Build jax-cuda-pjrt
+        run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jax-cuda-pjrt.env"
+  
+  run_tests:
+    needs: build_artifacts
     strategy:
       matrix:
         test_env: [
@@ -31,9 +59,8 @@ jobs:
     container:
       image: ${{ matrix.test_env.image }}
 
-    name: "Pytest GPU (Build wheels on CUDA 12.3 and test on CUDA ${{ matrix.test_env.cuda_version }})"
+    name: "Pytest GPU (Test on CUDA ${{ matrix.test_env.cuda_version }})"
     env:
-      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
 
     steps:
@@ -43,12 +70,6 @@ jobs:
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: Build jaxlib
-        run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jaxlib.env"
-      - name: Build jax-cuda-plugin
-        run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jax-cuda-plugin.env"
-      - name: Build jax-cuda-pjrt
-        run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jax-cuda-pjrt.env"
       - name: Install pytest
         env:
           JAXCI_PYTHON: python${{ matrix.python }}