name: Pytest GPU
on:
  workflow_call:
    inputs:
      container:
        required: true
        type: string
      mcloud-timeout:
        required: false
        type: number
        default: 1800
      name:
        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
    secrets:
      mcloud-api-key:
        required: true
jobs:
  pytest-gpu:
    timeout-minutes: 60 # ${{ inputs.gha-timeout }} for some reason not able to turn this into an input
    runs-on: ubuntu-latest
    env:
      MOSAICML_API_KEY: ${{ secrets.mcloud-api-key }}
    steps:
    - name: Checkout Repo
      uses: actions/checkout@v3
    - 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 init --mcloud
        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/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/mcp/mcp_pytest.py \
              --image '${{ inputs.container }}' \
              --pytest_markers '${{ inputs.pytest-markers }}' \
              --pytest_command '${{ inputs.pytest-command }}' \
              --timeout ${{ inputs.mcloud-timeout }} ${REF_ARGS}