Refine if conditions and update the upload destination

Author: nitins17

.github/workflows/build_artifacts.yml

@@ -53,7 +53,7 @@ on:
         type: string
 
 jobs:
-  build:
+  build_artifacts:
     continue-on-error: true
     defaults:
       run:
@@ -104,19 +104,20 @@ jobs:
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: Build jax
-        if: inputs.build_jax && matrix.artifact == 'jax'
-        run: ./ci/build_artifacts.sh "jax"
-      - name: Build jaxlib
-        if: inputs.build_jaxlib && matrix.artifact == 'jaxlib'
-        run: ./ci/build_artifacts.sh "jaxlib"
-      - name: Build jax-cuda-plugin
-        if: inputs.build_jax_cuda_plugin && matrix.artifact == 'jax-cuda-plugin'
-        run: ./ci/build_artifacts.sh "jax-cuda-plugin"
-      - name: Build jax-cuda-pjrt
-        if: inputs.build_jax_cuda_pjrt && matrix.artifact == 'jax-cuda-pjrt'
-        run: ./ci/build_artifacts.sh "jax-cuda-pjrt"
+      - name: Build ${{ matrix.artifact }}
+        if: >-
+          (inputs.build_jax && matrix.artifact == 'jax') ||
+          (inputs.build_jaxlib && matrix.artifact == 'jaxlib') ||
+          (inputs.build_jax_cuda_plugin && matrix.artifact == 'jax-cuda-plugin') ||
+          (inputs.build_jax_cuda_pjrt && matrix.artifact == 'jax-cuda-pjrt')
+        run: ./ci/build_artifacts.sh "${{ matrix.artifact }}"
+      - name: Set Platform
+        run: |
+          echo "PLATFORM=$(uname)_$(uname -m)" >> $GITHUB_ENV
       - name: Upload artifacts to GCS bucket
-        if: inputs.upload_artifacts
-        run:  ~/usr/local/bin/google-cloud-sdk/bin/gsutil -m cp -r $(pwd)/dist gs://general-ml-ci-transient/jax-github-actions/"${{ inputs.upload_destination }}"
+        # Upload if requested and one of the artifacts was built
+        if: >-
+          inputs.upload_artifacts &&
+          (inputs.build_jax || inputs.build_jaxlib ||  inputs.build_jax_cuda_plugin || inputs.build_jax_cuda_pjrt)
+        run:  ~/usr/local/bin/google-cloud-sdk/bin/gsutil -m cp -r $(pwd)/dist gs://general-ml-ci-transient/jax-github-actions/"${{ inputs.upload_destination }}"/$PLATFORM
 

.github/workflows/pytest_cpu_reuse.yml

@@ -54,8 +54,13 @@ jobs:
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Set Platform
+        run: |
+          echo "PLATFORM=$(uname)_$(uname -m)" >> $GITHUB_ENV
       - name: Download the artifacts built in the "build_artifacts" job
-        run: mkdir -p $(pwd)/dist && ~/usr/local/bin/google-cloud-sdk/bin/gsutil -m cp -r gs://general-ml-ci-transient/jax-github-actions/"${{ github.workflow }}"/${{ github.run_number }}/${{ github.run_attempt }} $(pwd)/dist
+        run: >-
+          mkdir -p $(pwd)/dist &&
+          ~/usr/local/bin/google-cloud-sdk/bin/gsutil -m cp -r gs://general-ml-ci-transient/jax-github-actions/"${{ github.workflow }}"/${{ github.run_number }}/${{ github.run_attempt }}/$PLATFORM $(pwd)/dist
       - name: Install pytest
         env:
           JAXCI_PYTHON: python${{ matrix.python }}