name: Retrieval PyTorch Tests for Target Determination

on:
  workflow_call:

permissions:
  id-token: write
  contents: read

jobs:
  get-label-type:
    name: get-label-type
    # Don't run on forked repos
    if: github.repository_owner == 'pytorch'
    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.7
    with:
      triggering_actor: ${{ github.triggering_actor }}
      issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
      curr_branch: ${{ github.head_ref || github.ref_name }}
      curr_ref_type: ${{ github.ref_type }}

  llm-retrieval:
    # Don't run on forked repos
    if: github.repository_owner == 'pytorch'
    runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge"
    continue-on-error: true
    needs: get-label-type
    steps:
      - name: Clone PyTorch
        uses: actions/checkout@v3
        with:
          repository: pytorch/pytorch
          fetch-depth: 0
          path: pytorch

      - name: Setup Linux
        uses: ./pytorch/.github/actions/setup-linux

      - name: Clone CodeLlama
        uses: actions/checkout@v3
        with:
          repository: osalpekar/codellama
          ref: main
          path: codellama

      - name: Clone Target Determination Code
        uses: actions/checkout@v3
        with:
          repository: osalpekar/llm-target-determinator
          ref: v0.0.2
          path: llm-target-determinator

      - name: Setup miniconda
        uses: pytorch/test-infra/.github/actions/setup-miniconda@release/2.7
        with:
          python-version: "3.9"

      - name: Install requirements
        shell: bash
        run: |
          set -euxo pipefail
          ${CONDA_RUN} pip install -r llm-target-determinator/requirements.txt
          cd "${GITHUB_WORKSPACE}/codellama"
          ${CONDA_RUN} pip install -e .

      - name: Fetch CodeLlama Checkpoint
        shell: bash
        run: |
          set -euxo pipefail
          cd "${GITHUB_WORKSPACE}/codellama"
          mkdir "CodeLlama-7b-Python"
          aws s3 cp "s3://target-determinator-assets/CodeLlama-7b-Python" "CodeLlama-7b-Python" --recursive --no-progress

      - name: Fetch indexes
        uses: nick-fields/retry@v3.0.0
        with:
          max_attempts: 3
          retry_wait_seconds: 10
          timeout_minutes: 5
          shell: bash
          command: |
            set -euxo pipefail
            ${CONDA_RUN} python -m pip install awscli==1.29.40
            cd "${GITHUB_WORKSPACE}"/llm-target-determinator/assets
            aws s3 cp "s3://target-determinator-assets/indexes/latest" . --recursive

            unzip -o indexer-files\*.zip
            rm indexer-files*.zip

      - name: Run Retriever
        id: run_retriever
        continue-on-error: true  # ghstack not currently supported due to problems getting git diff
        shell: bash
        run: |
          set -euxo pipefail
          cd "${GITHUB_WORKSPACE}"/llm-target-determinator
          ${CONDA_RUN} torchrun \
            --standalone \
            --nnodes=1 \
            --nproc-per-node=1 \
            retriever.py \
            --experiment-name indexer-files \
            --pr-parse-format GITDIFF
          cd assets
          zip -r mappings.zip mappings

      - name: Upload results to s3
        uses: seemethere/upload-artifact-s3@v5
        if: ${{ steps.run_retriever.outcome == 'success' }}
        with:
          name: llm_results
          retention-days: 14
          if-no-files-found: warn
          path: llm-target-determinator/assets/mappings.zip
        env:
          AWS_ACCESS_KEY_ID: ""
          AWS_SECRET_ACCESS_KEY: ""
          AWS_SESSION_TOKEN: ""
          AWS_DEFAULT_REGION: ""
          AWS_REGION: ""

      - name: Teardown Linux
        uses: pytorch/test-infra/.github/actions/teardown-linux@release/2.7
        if: always()
