update-s3-dependencies.yml

name: Update S3 HTML dependencies for download.pytorch.org nightly and test

on:
  workflow_dispatch:
    inputs:
      mode:
        description: 'Operation mode'
        required: true
        type: choice
        default: update-packages
        options:
          - update-packages
          - create-target
      dryrun:
        description: 'Enable dry run mode'
        required: true
        type: choice
        default: enabled
        options:
          - enabled
          - disabled
      package:
        description: 'Package to update (only for update-packages mode)'
        required: false
        type: choice
        default: torch
        options:
          - all
          - torch
          - triton
          - torchtune
          - torch_xpu
          - torch_rocm
          - vllm
      target:
        description: 'Target to create (e.g., rocm7.2, cu130) - only for create-target mode'
        required: false
        type: string
      include-prod:
        description: 'Update prod (whl) dependencies only, via promote-env (skips nightly/test; update-packages mode only)'
        required: false
        type: choice
        default: disabled
        options:
          - disabled
          - enabled

permissions:
  id-token: write
  contents: read

jobs:
  # Nightly + test updates. Skipped when include-prod is enabled, so a prod
  # run only touches prod (the update-prod job below) and never nightly/test.
  update:
    if: ${{ github.event.inputs.include-prod != 'enabled' }}
    runs-on: ubuntu-22.04
    environment: pytorchbot-env
    container:
      image: continuumio/miniconda3:23.10.0-1
    steps:
      - name: Validate inputs for create-target mode
        if: ${{ github.event.inputs.mode == 'create-target' }}
        shell: bash
        env:
          TARGET: ${{ github.event.inputs.target }}
        run: |
          set -e
          if [[ -z "${TARGET}" ]]; then
            echo "ERROR: target is required for create-target mode"
            exit 1
          fi
          # Validate target matches expected patterns
          if [[ ! "${TARGET}" =~ ^(cu[0-9]+|rocm[0-9]+\.[0-9]+)$ ]]; then
            echo "ERROR: Invalid target format '${TARGET}'"
            echo "Valid formats:"
            echo "  - CUDA: cu118, cu121, cu126, cu128, cu129, cu130, etc."
            echo "  - ROCm: rocm5.7, rocm6.0, rocm6.4, rocm7.1, rocm7.2, etc."
            exit 1
          fi
          echo "Target '${TARGET}' is valid"

      - name: Configure AWS credentials
        id: aws_creds
        uses: aws-actions/configure-aws-credentials@50ac8dd1e1b10d09dac7b8727528b91bed831ac0 # v3.0.2
        with:
          role-to-assume: arn:aws:iam::749337293305:role/gha_workflow_s3_update
          aws-region: us-east-1

      - name: Checkout repository test-infra
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: pytorch/test-infra
          ref: ${{ github.ref }}

      - name: Install dependencies
        run: |
          pip install -r s3_management/requirements.txt

      - name: Update packages
        if: ${{ github.event.inputs.mode == 'update-packages' }}
        shell: bash
        env:
          DRYRUN: ${{ github.event.inputs.dryrun }}
          PACKAGE: ${{ github.event.inputs.package || 'torch' }}
          R2_ACCOUNT_ID: ${{ secrets.R2_PYTORCH_ACCOUNT_ID }}
          R2_ACCESS_KEY_ID: ${{ secrets.R2_PYTORCH_PUBLIC_ACCESS_KEY_ID }}
          R2_SECRET_ACCESS_KEY: ${{ secrets.R2_PYTORCH_PUBLIC_SECRET_ACCESS_KEY }}
          R2_BUCKET_NAME: "pytorch-downloads"
        run: |
          set -ex

          ARGS="--package ${PACKAGE}"

          if [[ "${DRYRUN}" == "enabled" ]]; then
            ARGS="${ARGS} --dry-run"
          fi

          # shellcheck disable=SC2086
          python3 s3_management/update_dependencies.py ${ARGS}

      - name: Create new target
        if: ${{ github.event.inputs.mode == 'create-target' }}
        shell: bash
        env:
          TARGET: ${{ github.event.inputs.target }}
          DRYRUN: ${{ github.event.inputs.dryrun }}
          R2_ACCOUNT_ID: ${{ secrets.R2_PYTORCH_ACCOUNT_ID }}
          R2_ACCESS_KEY_ID: ${{ secrets.R2_PYTORCH_PUBLIC_ACCESS_KEY_ID }}
          R2_SECRET_ACCESS_KEY: ${{ secrets.R2_PYTORCH_PUBLIC_SECRET_ACCESS_KEY }}
          R2_BUCKET_NAME: "pytorch-downloads"
        run: |
          set -ex

          # Create the target directory with torch dependencies
          ARGS="--create-target --target ${TARGET}"

          if [[ "${DRYRUN}" == "enabled" ]]; then
            ARGS="${ARGS} --dry-run"
          fi

          # shellcheck disable=SC2086
          python3 s3_management/update_dependencies.py ${ARGS}

          # Run update-packages for torch to ensure all dependencies are current
          ARGS="--package torch"

          if [[ "${DRYRUN}" == "enabled" ]]; then
            ARGS="${ARGS} --dry-run"
          fi

          # shellcheck disable=SC2086
          python3 s3_management/update_dependencies.py ${ARGS}

  # Prod (whl) dependency updates are gated behind the protected promote-env,
  # mirroring release-download-pytorch-org.yml. Opt-in via the include-prod
  # input; runs only in update-packages mode.
  update-prod:
    if: ${{ github.event.inputs.mode == 'update-packages' && github.event.inputs.include-prod == 'enabled' }}
    runs-on: ubuntu-22.04
    environment: promote-env
    container:
      image: continuumio/miniconda3:23.10.0-1
    steps:
      - name: Configure AWS credentials
        uses: aws-actions/configure-aws-credentials@50ac8dd1e1b10d09dac7b8727528b91bed831ac0 # v3.0.2
        with:
          role-to-assume: arn:aws:iam::749337293305:role/gha_workflow_promote_wheels
          aws-region: us-east-1

      - name: Checkout repository test-infra
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: pytorch/test-infra
          ref: ${{ github.ref }}

      - name: Install dependencies
        run: |
          pip install -r s3_management/requirements.txt

      - name: Update prod (stable) packages
        shell: bash
        env:
          DRYRUN: ${{ github.event.inputs.dryrun }}
          PACKAGE: ${{ github.event.inputs.package || 'torch' }}
          R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }}
          R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }}
          R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }}
          R2_BUCKET_NAME: "pytorch-downloads"
        run: |
          set -ex

          ARGS="--package ${PACKAGE} --stable-only"

          if [[ "${DRYRUN}" == "enabled" ]]; then
            ARGS="${ARGS} --dry-run"
          fi

          # shellcheck disable=SC2086
          python3 s3_management/update_dependencies.py ${ARGS}