pytest-gpu.yaml

name: Pytest GPU
on:
  workflow_call:
    inputs:
      container:
        required: true
        type: string
      gha-timeout:
        required: false
        type: number
        default: 60
      git_repo:
        required: true
        type: string
      name:
        required: true
        type: string
      pip_deps:
        required: true
        type: string
      pytest-command:
        required: true
        type: string
      pytest-markers:
        required: true
        type: string
      python-version:
        required: false
        type: string
        default: 3.9
      mcloud-timeout:
        required: false
        type: number
        default: 2700
      composer_package_name:
        required: false
        type: string
      git-ssh-clone:
        required: false
        type: boolean
        default: false
      gpu_configurations:
        required: true
        type: string
    secrets:
      mcloud-api-key:
        required: true
      slack-notifications-bot-token:
        required: false
      code-eval-device:
        required: false
      code-eval-url:
        required: false
      code-eval-apikey:
        required: false
jobs:
  pytest-gpu:
    timeout-minutes: ${{ inputs.gha-timeout }}
    runs-on: ubuntu-latest
    strategy:
      matrix:
        gpu_num: [${{ fromJSON(inputs.gpu_configurations) }}]
    name: Pytest GPU ${{ matrix.gpu_num }}
    env:
      MOSAICML_API_KEY: ${{ secrets.mcloud-api-key }}
    steps:
    - name: Checkout Repo
      uses: actions/checkout@v3
      with:
        repository: "mosaicml/ci-testing"
    - name: Setup Python
      uses: actions/setup-python@v4
      with:
        python-version: ${{ inputs.python-version }}
    - name: Cache pip
      uses: actions/cache@v3
      with:
        # This path is specific to Ubuntu
        path: ~/.cache/pip
        # Look to see if there is a cache hit for the corresponding requirements file
        key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
        restore-keys: |
          ${{ runner.os }}-pip-
          ${{ runner.os }}-
    - name: Setup MCLI
      run: |
        set -ex
        python -m pip install mosaicml-cli
        mcli version
    - name: Submit Run
      id: tests
      run: |
        set -ex

        PR_NUMBER="$(jq --raw-output .pull_request.number "$GITHUB_EVENT_PATH")"
        REF_ARGS=""

        # Use the PR number if it exists, commit SHA for protected branches and the branch name otherwise
        if [ -z "$PR_NUMBER" ] || [ "$PR_NUMBER" = "null" ]; then
          if [[ "$GITHUB_REF" =~ "refs/heads/dev" || "$GITHUB_REF" =~ "refs/heads/main" || \
            "$GITHUB_REF" =~ "refs/heads/release" ]]; then
            REF_ARGS="--git_commit $GITHUB_SHA"
          else
            REF_ARGS="--git_branch $GITHUB_REF_NAME"
          fi
        else
          REF_ARGS="--pr_number $PR_NUMBER"
        fi

        python .github/mcli/mcli_pytest.py \
          --image '${{ inputs.container }}' \
          --git_repo '${{ inputs.git_repo }}' \
          --pip_deps '${{ inputs.pip_deps }}' \
          --pip_package_name '${{ inputs.composer_package_name }}' \
          --pytest_markers '${{ inputs.pytest-markers }}' \
          --pytest_command '${{ inputs.pytest-command }}' \
          --timeout ${{ inputs.mcloud-timeout }} \
          --gpu_num ${{ matrix.gpu_num }} \
          --git_ssh_clone ${{ inputs.git-ssh-clone }} \
          ${REF_ARGS}