Add variable to know if a workflow call was made, enable all package builds, change gsutil path

Author: nitins17

.github/workflows/build_artifacts.yml

@@ -51,6 +51,11 @@ on:
         required: true
         default: '${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
         type: string
+      is_workflow_call:
+        description: "Metadata variable to know whether a workflow call was made"
+        type: string
+        required: true
+        default: "1"
 
 jobs:
   determine_matrix:
@@ -70,7 +75,8 @@ jobs:
       - id: set-matrix
         run: |
             artifacts=()
-            if [[ ${{ github.event_name }}  == "pull_request" ]]; then
+            # Build every package if not a workflow call
+            if [[ ${${{ inputs.is_workflow_call }}:-0}  == "0" ]]; then
               artifacts=("'jax'" ", 'jaxlib'" ", 'jax-cuda-pjrt'" ", 'jax-cuda-plugin'")
             else
               if [[ ${${{ inputs.build_jax }}:-0} == "1" ]]; then
@@ -102,7 +108,7 @@ jobs:
       matrix:
         runner: ["windows-x86-n2-64", "linux-x86-n2-16", "linux-arm64-t2a-16"]
         artifact: ${{ fromJSON(needs.determine_matrix.outputs.artifact_matrix) }}
-        python: ["3.10"] #, "3.11", "3.12"]
+        python: ["3.10", "3.11", "3.12"]
         # jax-cuda-pjrt and jax are pure Python packages so they do not need to be built for each
         # Python version.
         exclude:
@@ -151,5 +157,5 @@ jobs:
       - name: Upload artifacts to GCS bucket
         # Upload if requested and one of the artifacts was built
         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 }}"/$PLATFORM
+        run:  gsutil -m cp -r $(pwd)/dist gs://general-ml-ci-transient/jax-github-actions/"${{ inputs.upload_destination }}"/$PLATFORM
 

.github/workflows/pytest_cpu_reuse.yml

@@ -60,7 +60,7 @@ jobs:
       - 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 }}/$PLATFORM $(pwd)/dist
+          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 }}