build_image.yml

name: Build and Push Orbax-checkpoint IT runner Docker Images
on:
  schedule:
  # Run the job daily at 12AM UTC
    - cron: '0 0 * * *'
  push:
    branches:
      - 'test_*'
permissions:
  contents: read
jobs:
  build_and_push:
    runs-on: linux-x86-n2-16-buildkit
    container: google/cloud-sdk:524.0.0
    steps:
      - name: Checkout Orbax-checkpoint
        uses: actions/checkout@v5
      - name: Mark git repositories as safe
        run: git config --global --add safe.directory '*'
      - name: Configure Docker
        run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
      - name: Set up Docker BuildX
        uses: docker/setup-buildx-action@v3.11.1
        with:
          driver: remote
          endpoint: tcp://localhost:1234
      - name: Build and push Docker image
        uses: docker/build-push-action@v6
        with:
          push: true
          context: .
          file: ./checkpoint/orbax/checkpoint/_src/testing/oss/Dockerfile
          tags: gcr.io/orbax-checkpoint/orbax-benchmarks-integration_tests-runner:latest
          cache-from: type=gha
          outputs: type=image,compression=zstd,force-compression=true
          build-args: |
            DEVICE=tpu
            JAX_VERSION=newest
            BRANCH=main
            GITHUB_RUNNER=true
      - name: Add tags to Docker images
        shell: bash
        run: |
          SOURCE_IMAGE="gcr.io/orbax-checkpoint/orbax-benchmarks-integration_tests-runner"
          # Add Orbax-checkpoint tag
          orbax_hash=$(git rev-parse --short HEAD)
          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:orbax_${orbax_hash}" --quiet