diff --git a/.github/workflows/ssh-runner.yml b/.github/workflows/ssh-runner.yml index 5ea83f2005..ab87b2d5fa 100644 --- a/.github/workflows/ssh-runner.yml +++ b/.github/workflows/ssh-runner.yml @@ -9,6 +9,9 @@ on: docker_image: description: 'Name of the Docker image' required: true + num_gpus: + description: 'Type of the number of gpus to use (`single` or `multi`)' + required: true env: IS_GITHUB_CI: "1" @@ -20,12 +23,13 @@ env: RUN_SLOW: yes # For gated repositories, we still need to agree to share information on the Hub repo. page in order to get access. # This token is created under the bot `hf-transformers-bot`. SIGOPT_API_TOKEN: ${{ secrets.SIGOPT_API_TOKEN }} TF_FORCE_GPU_ALLOW_GROWTH: true + CUDA_VISIBLE_DEVICES: 0,1 RUN_PT_TF_CROSS_TESTS: 1 jobs: ssh_runner: name: "SSH" - runs-on: [single-gpu, nvidia-gpu, "${{ github.event.inputs.runner_type }}", ci] + runs-on: ["${{ github.event.inputs.num_gpus }}-gpu", nvidia-gpu, "${{ github.event.inputs.runner_type }}", ci] container: image: ${{ github.event.inputs.docker_image }} options: --gpus all --privileged --ipc host -v /mnt/cache/.cache/huggingface:/mnt/cache/