update_model_artifacts.yml

# Copyright 2023 The OpenXLA Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# Workflow that updates model artifacts.

name: Update Model Artifacts

on:
  workflow_dispatch:
  pull_request:

concurrency:
  # A PR number if a pull request and otherwise the commit hash. This cancels
  # queued and in-progress runs for the same PR (presubmit) or commit
  # (postsubmit).
  group: ${{ github.workflow }}-${{ github.event.number || github.sha }}
  cancel-in-progress: true

jobs:
  setup:
    runs-on: ubuntu-22.04
    outputs:
      runner-group: ${{ steps.configure.outputs.runner-group }}
    steps:
      - name: "Checking out PR repository"
        uses: actions/checkout@e2f20e631ae6d7dd3b768f56a5d2af784dd54791  # v2.5.0
      - name: "Configuring CI options"
        id: configure
        env:
          RUNNER_GROUP: ${{ github.event_name == 'pull_request' && 'presubmit' || 'postsubmit' }}
        run: |
          # Just informative logging. There should only be two commits in the
          # history here, but limiting the depth helps when copying from a local
          # repo instead of using checkout, e.g. with
          # https://github.com/nektos/act where there will be more.
          git log --oneline --graph --max-count=3
          # Workflow jobs can't access `env` in `runs-on`, so we need to make
          # `runner-group` a job output variable.
          echo "runner-group=${RUNNER_GROUP}" > "${GITHUB_OUTPUT}"

  generate_artifacts:
    needs: setup
    runs-on:
      - self-hosted  # must come first
      - runner-group=${{ needs.setup.outputs.runner-group }}
      - environment=prod
      - machine-type=c2-standard-16
    steps:
      - name: "Checking out PR repository"
        uses: actions/checkout@e2f20e631ae6d7dd3b768f56a5d2af784dd54791  # v2.5.0
      - name: "Generate JAX model artifacts"
        run: |
          mkdir jax
          # Generate enabled models in comparative_benchmark/jax/benchmark_xla.sh
          docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
            --env "PYTHON=python3" \
            --env "WITH_CUDA=0" \
            --env "AUTO_UPLOAD=0" \
            --env "OUTPUT_DIR=jax" \
            "gcr.io/iree-oss/openxla-benchmark/base@sha256:1bf3e319465ec8fb465baae3f6ba9a5b09cb84a5349a675c671a552fc77f2251" \
            "common_benchmark_suite/openxla/benchmark/comparative_suite/jax/scripts/generate_model_artifacts.sh" --filter \
              "RESNET50_FP32_JAX_.+" \
              "BERT_LARGE_FP32_JAX_.+_BATCH(1|16|24|32|48|64|512)" \
              "T5_LARGE_FP32_JAX_.+_BATCH(1|16|24|32|48|64)" \
              "T5_4CG_LARGE_FP32_JAX_.+" \
              "GPT2LMHEAD_FP32_JAX_.+"
      - name: "Upload JAX model artifacts"
        run: |
          gcloud storage cp -r "jax/*"  "gs://iree-model-artifacts/jax"