From d180514ce68e9abce37180e2a49109e630e04f23 Mon Sep 17 00:00:00 2001
From: Will Cromar <wcromar@google.com>
Date: Thu, 13 Jun 2024 23:18:43 +0000
Subject: [PATCH] Create `_setup.yml` to wrap common setup steps.

---
 .github/workflows/_build_torch_xla.yml | 26 ++++----------------
 .github/workflows/_setup.yml           | 34 ++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 21 deletions(-)
 create mode 100644 .github/workflows/_setup.yml

diff --git a/.github/workflows/_build_torch_xla.yml b/.github/workflows/_build_torch_xla.yml
index 56e6b5408c3..83f6035f0f1 100644
--- a/.github/workflows/_build_torch_xla.yml
+++ b/.github/workflows/_build_torch_xla.yml
@@ -27,31 +27,15 @@ jobs:
       image: ${{ inputs.dev-image }}
     env:
       GCLOUD_SERVICE_KEY: ${{ secrets.gcloud-service-key }}
-      GOOGLE_APPLICATION_CREDENTIALS: /tmp/default_credentials.json
-      BAZEL_JOBS: 16
       BAZEL_REMOTE_CACHE: ${{ github.event_name == 'push' || github.event.pull_request.head.repo.full_name == github.repository }}
+      BAZEL_JOBS: 16
       BUILD_CPP_TESTS: 1
     steps:
-      # See https://github.com/actions/checkout/issues/1014#issuecomment-1906802802
-      - name: Clean up workspace
-        run: |
-          ls -la
-          rm -rvf ${GITHUB_WORKSPACE}/*
-      - name: Setup gcloud
-        shell: bash
-        run: |
-          echo "${GCLOUD_SERVICE_KEY}" > $GOOGLE_APPLICATION_CREDENTIALS
-      - name: Checkout PyTorch Repo
-        uses: actions/checkout@v4
-        with:
-          repository: pytorch/pytorch
-          path: pytorch
-          ref: ${{ inputs.torch-commit }}
-          submodules: recursive
-      - name: Checkout PyTorch/XLA Repo
-        uses: actions/checkout@v4
+      - name: Setup
+        uses: ./.github/workflows/_setup.yml
         with:
-          path: pytorch/xla
+          torch-commit: ${{ inputs.torch-commit }}
+          setup-remote-cache: ${{ env.BAZEL_REMOTE_CACHE }}
       - name: Build
         shell: bash
         run: |
diff --git a/.github/workflows/_setup.yml b/.github/workflows/_setup.yml
new file mode 100644
index 00000000000..13fed90829d
--- /dev/null
+++ b/.github/workflows/_setup.yml
@@ -0,0 +1,34 @@
+name: Set up PyTorch/XLA
+inputs:
+  torch-commit:
+    required: true
+    type: string
+  remote-cache:
+    required: true
+    type: boolean
+    description: Whether to set up Google Cloud credentials for Bazel cache.
+runs:
+  using: "composite"
+  steps:
+    # See https://github.com/actions/checkout/issues/1014#issuecomment-1906802802
+    - name: Clean up workspace
+      run: |
+        ls -la
+        rm -rvf ${GITHUB_WORKSPACE}/*
+    - name: Setup gcloud
+      shell: bash
+      run: |
+        echo "${GCLOUD_SERVICE_KEY}" > /tmp/default_credentials.json
+        echo "GOOGLE_APPLICATION_CREDENTIALS=/tmp/default_credentials.json" >> $GITHUB_ENV
+      if: ${{ inputs.remote-cache }}
+    - name: Checkout PyTorch Repo
+      uses: actions/checkout@v4
+      with:
+        repository: pytorch/pytorch
+        path: pytorch
+        ref: ${{ inputs.torch-commit }}
+        submodules: recursive
+    - name: Checkout PyTorch/XLA Repo
+      uses: actions/checkout@v4
+      with:
+        path: pytorch/xla