pytorch/.github/workflows/llm_td_retrieval.yml at develop · ROCm/pytorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
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@main
    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@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: pytorch/pytorch
          fetch-depth: 0
          path: pytorch

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

      - name: Clone CodeLlama
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: osalpekar/codellama
          ref: main
          path: codellama

      - name: Clone Target Determination Code
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: osalpekar/llm-target-determinator
          ref: v0.0.2
          path: llm-target-determinator

      - name: Install requirements
        shell: bash
        run: |
          set -euxo pipefail
          python3 -m pip install -r llm-target-determinator/requirements.txt
          cd "${GITHUB_WORKSPACE}/codellama"
          python3 -m 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@7152eba30c6575329ac0576536151aca5a72780e # v3.0.0
        with:
          max_attempts: 3
          retry_wait_seconds: 10
          timeout_minutes: 5
          shell: bash
          command: |
            set -euxo pipefail
            python3 -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
          export PATH="$HOME/.local/bin:$PATH"
          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@baba72d0712b404f646cebe0730933554ebce96a # v5.1.0
        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@main
        if: always()