diff --git a/.github/workflows/ssh-runner.yml b/.github/workflows/ssh-runner.yml index db649876f60492..02b022698b0c5e 100644 --- a/.github/workflows/ssh-runner.yml +++ b/.github/workflows/ssh-runner.yml @@ -26,9 +26,38 @@ env: RUN_PT_TF_CROSS_TESTS: 1 jobs: + get_runner: + name: "Get runner to use" + runs-on: ubuntu-22.04 + outputs: + RUNNER: ${{ steps.set_runner.outputs.RUNNER }} + steps: + - name: Get runner to use + shell: bash + run: | + if [[ "${{ github.event.inputs.num_gpus }}" == "single" && "${{ github.event.inputs.runner_type }}" == "t4" ]]; then + echo "RUNNER=aws-g4dn-2xlarge-cache" >> $GITHUB_ENV + elif [[ "${{ github.event.inputs.num_gpus }}" == "multi" && "${{ github.event.inputs.runner_type }}" == "t4" ]]; then + echo "RUNNER=aws-g4dn-12xlarge-cache" >> $GITHUB_ENV + elif [[ "${{ github.event.inputs.num_gpus }}" == "single" && "${{ github.event.inputs.runner_type }}" == "a10" ]]; then + echo "RUNNER=aws-g5-4xlarge-cache" >> $GITHUB_ENV + elif [[ "${{ github.event.inputs.num_gpus }}" == "multi" && "${{ github.event.inputs.runner_type }}" == "a10" ]]; then + echo "RUNNER=aws-g5-12xlarge-cache" >> $GITHUB_ENV + else + echo "RUNNER=" >> $GITHUB_ENV + fi + + - name: Set runner to use + id: set_runner + run: | + echo ${{ env.RUNNER }} + echo "RUNNER=${{ env.RUNNER }}" >> $GITHUB_OUTPUT + ssh_runner: name: "SSH" - runs-on: ["${{ github.event.inputs.num_gpus }}-gpu", nvidia-gpu, "${{ github.event.inputs.runner_type }}", ci] + needs: get_runner + runs-on: + group: ${{ needs.get_runner.outputs.RUNNER }} container: image: ${{ github.event.inputs.docker_image }} options: --gpus all --privileged --ipc host -v /mnt/cache/.cache/huggingface:/mnt/cache/