Index PyTorch Tests for Target Determination #18
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # Owner(s): ["module: ci"] | |
| name: Index PyTorch Tests for Target Determination | |
| on: | |
| workflow_dispatch: | |
| schedule: | |
| - cron: '0 0 * * *' | |
| permissions: | |
| id-token: write | |
| contents: read | |
| jobs: | |
| get-label-type: | |
| if: github.repository_owner == 'pytorch' | |
| name: get-label-type | |
| 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 }} | |
| check_experiments: "lf" | |
| index: | |
| needs: get-label-type | |
| # OSDC A10G GPU runner (equivalent of linux.g5.4xlarge.nvidia.gpu). | |
| runs-on: "${{ needs.get-label-type.outputs.label-type }}l-x86aavx2-29-113-a10g" | |
| environment: target-determinator-env | |
| # Run inside the CUDA CI image on OSDC (ARC has no host docker daemon). The | |
| # public GHCR mirror needs no ECR creds; --gpus all exposes the A10G. | |
| container: | |
| image: ghcr.io/pytorch/ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-${{ needs.get-label-type.outputs.ci-docker-hash }} | |
| options: "--gpus all" | |
| steps: | |
| - name: Setup Linux | |
| uses: pytorch/pytorch/.github/actions/setup-linux@main | |
| with: | |
| use-arc: true | |
| - name: Clone CodeLlama | |
| uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 | |
| with: | |
| repository: osalpekar/codellama | |
| ref: 1ec50e0cfc0fadc3b6ceb146617e2119ab26eb34 | |
| 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: Set up sibling workspace for the indexer | |
| shell: bash | |
| run: | | |
| set -eux | |
| # llm-target-determinator@v0.0.2 (config.py) resolves its own realpath | |
| # and then reads the corpus from `<parent>/pytorch` and the model from | |
| # `<parent>/codellama` — i.e. it expects pytorch / codellama / itself to | |
| # be sibling dirs. setup-linux checks pytorch out at GITHUB_WORKSPACE and | |
| # the two helper repos as subdirs, so the realpath-based lookup misses | |
| # the corpus. Move the helper repos into a clean workspace and point | |
| # `pytorch` at the checkout via a symlink. Because `.resolve()` follows | |
| # symlinks, the helper repos must be *real* dirs here (not symlinks). | |
| # Moving them out of the pytorch tree also keeps them out of the | |
| # indexer's os.walk corpus. | |
| WS="${RUNNER_TEMP}/td-workspace" | |
| mkdir -p "${WS}" | |
| mv "${GITHUB_WORKSPACE}/codellama" "${WS}/codellama" | |
| mv "${GITHUB_WORKSPACE}/llm-target-determinator" "${WS}/llm-target-determinator" | |
| ln -sfn "${GITHUB_WORKSPACE}" "${WS}/pytorch" | |
| echo "TD_WORKSPACE=${WS}" >> "${GITHUB_ENV}" | |
| - name: Configure AWS credentials | |
| uses: aws-actions/configure-aws-credentials@ececac1a45f3b08a01d2dd070d28d111c5fe6722 # v4.1.0 | |
| with: | |
| role-to-assume: arn:aws:iam::308535385114:role/gha_target_determinator_s3_read_write | |
| aws-region: us-east-1 | |
| - name: Download checkpoint | |
| shell: bash | |
| env: | |
| AWS_DEFAULT_REGION: us-east-1 | |
| run: | | |
| # Do this outside of docker so I don't have to put env vars in | |
| pip3 install -q awscli==1.29.40 | |
| cd "${TD_WORKSPACE}/codellama" | |
| mkdir "CodeLlama-7b-Python" | |
| # --quiet: the checkpoint is many files; without it the per-object | |
| # "download:" lines flood the log. | |
| aws s3 cp \ | |
| "s3://target-determinator-assets/CodeLlama-7b-Python" \ | |
| "CodeLlama-7b-Python" \ | |
| --recursive \ | |
| --quiet | |
| - name: Run indexer | |
| shell: bash -l {0} | |
| env: | |
| AWS_DEFAULT_REGION: us-east-1 | |
| run: | | |
| # Sibling layout (pytorch / codellama / llm-target-determinator) was | |
| # set up in "Set up sibling workspace for the indexer". `pytorch` is a | |
| # symlink to the checkout; the script cd's between the siblings. | |
| cd "${TD_WORKSPACE}" | |
| bash pytorch/.github/scripts/td_llm_indexer.sh | |
| - name: Upload to s3 | |
| shell: bash -l {0} | |
| env: | |
| AWS_DEFAULT_REGION: us-east-1 | |
| run: | | |
| cd "${TD_WORKSPACE}/llm-target-determinator/assets" | |
| TIMESTAMP=$(date -Iseconds) | |
| ZIP_NAME="indexer-files-${TIMESTAMP}.zip" | |
| # Create a zipfile with all the generated indices | |
| zip -r "${ZIP_NAME}" indexer-files | |
| # Note that because the below 2 operations are not atomic, there will | |
| # be a period of a few seconds between these where there is no index | |
| # present in the latest/ folder. To account for this, the retriever | |
| # should have some retry logic with backoff to ensure fetching the | |
| # index doesn't fail. | |
| # Move the old index into the archived/ folder | |
| aws s3 mv \ | |
| "s3://target-determinator-assets/indexes/latest" \ | |
| "s3://target-determinator-assets/indexes/archived" \ | |
| --recursive | |
| # Move the new index into the latestl/ folder | |
| aws s3 cp \ | |
| "${ZIP_NAME}" \ | |
| "s3://target-determinator-assets/indexes/latest/${ZIP_NAME}" | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }} | |
| cancel-in-progress: true |