build docker image in test

Author: tengyifei

.github/workflows/_torchprime_ci.yml

@@ -15,23 +15,77 @@ on:
     secrets:
       PERSONAL_ACCESS_TOKEN_FOR_TRIGGERING_TORCHPRIME:
         required: true
+      GCLOUD_SERVICE_KEY:
+        required: true
 jobs:
   torchprime-e2e-test:
     name: Run torchprime E2E tests
     timeout-minutes: ${{ inputs.timeout-minutes }}
     runs-on: ubuntu-22.04
     steps:
+      - name: Use Docker in rootless mode
+        if: inputs.has_code_changes == 'true'
+        uses: ScribeMD/[email protected]
+      - name: Add user to docker group
+        if: inputs.has_code_changes == 'true'
+        run: |
+          sudo usermod -aG docker $USER
+          newgrp docker
+        shell: bash
+      # Googlers: if this fails, follow http://shortn/_61iSj31q1b to debug.
+      - uses: google-github-actions/auth@v2
+        if: inputs.has_code_changes == 'true'
+        with:
+          credentials_json: '${{ secrets.GCLOUD_SERVICE_KEY }}'
+      - uses: google-github-actions/setup-gcloud@v2
+        if: inputs.has_code_changes == 'true'
+        with:
+          version: '>= 363.0.0'
+          install_components: 'beta,gke-gcloud-auth-plugin'
+      - name: Verify GCP setup
+        if: inputs.has_code_changes == 'true'
+        run: gcloud info
+        shell: bash
+      - name: Authenticate Docker
+        if: inputs.has_code_changes == 'true'
+        run: gcloud auth configure-docker --quiet
+        shell: bash
+      - name: Activate SA credentials
+        if: inputs.has_code_changes == 'true'
+        run: gcloud auth activate-service-account --key-file=$GOOGLE_APPLICATION_CREDENTIALS
+        shell: bash
       - name: Checkout actions
         if: inputs.has_code_changes == 'true'
         uses: actions/checkout@v4
         with:
           sparse-checkout: |
             .github/workflows/setup
           path: .actions
+      # Build a docker image for torchprime E2E test
+      # First download the torch-xla-wheels
+      - name: Fetch wheels
+        if: inputs.has_code_changes == 'true'
+        uses: actions/download-artifact@v4
+        with:
+          name: torch-xla-wheels
+          path: /tmp/wheels/
+      # Then run docker to install them and push a docker
+      - name: Build and push docker image
+        if: inputs.has_code_changes == 'true'
+        run: |
+          . ./infra/ansible/build_for_torchprime.sh
+        env:
+          XLA_WHEELS_PATH: /tmp/wheels/
+          DOCKER_IMAGE_NAME: torch-xla-torchprime-ci
+          DOCKER_IMAGE_TAG: latest
+          DOCKER_ORG: pytorch
+      # Trigger torchprime E2E test workflow
       - uses: convictional/[email protected]
+        if: inputs.has_code_changes == 'true'
         with:
           owner: AI-Hypercomputer
           repo: torchprime
           github_token: ${{ secrets.PERSONAL_ACCESS_TOKEN_FOR_TRIGGERING_TORCHPRIME }}
           workflow_file_name: e2e_test.yml
+          wait_interval: 60
           ref: yifeit/torchprime-ci

.github/workflows/build_and_test.yml

@@ -128,8 +128,7 @@ jobs:
   test-torchprime:
     name: "torchprime tests"
     uses: ./.github/workflows/_torchprime_ci.yml
-    # needs: [build-torch-xla, check_code_changes]
-    needs: [check_code_changes]
+    needs: [build-torch-xla, check_code_changes]
     with:
       timeout-minutes: 100
       has_code_changes: ${{ needs.check_code_changes.outputs.has_code_changes }}

infra/ansible/build_for_torchprime.Dockerfile

@@ -0,0 +1,19 @@
+FROM python:${python_version}-${debian_version} AS release
+
+WORKDIR /tmp/wheels
+COPY ./*.whl ./
+
+RUN echo "Installing the following wheels" && ls *.whl
+RUN pip install *.whl
+
+# Install the dependencies including libtpu.
+WORKDIR /ansible
+RUN pip install ansible
+COPY . /ansible
+
+ARG ansible_vars
+RUN ansible-playbook -vvv playbook.yaml -e "stage=release" -e "${ansible_vars}" --tags "install_deps"
+
+WORKDIR /
+
+RUN rm -rf /ansible /tmp/wheels

infra/ansible/build_for_torchprime.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+set -ex
+
+export IMAGE_NAME="gcr.io/tpu-pytorch/for-torchprime-ci:test"
+export DOCKERFILE_PATH="infra/ansible/build_for_torchprime.Dockerfile"
+export CONTEXT_PATH="/tmp/wheels"
+
+docker build -t "${IMAGE_NAME}" "${CONTEXT_PATH}" -f "${DOCKERFILE_PATH}" --build-arg ansible_vars='{"arch":"amd64","accelerator":"tpu","pytorch_git_rev":"main","xla_git_rev":"foobar","bundle_libtpu":"0","package_version":"2.8","nightly_release":true}'
+docker push "${IMAGE_NAME}"