build

huggingface · May 29, 2024 · 641b2e5 · 641b2e5
1 parent e126a8a
commit 641b2e5
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/.github/workflows/model_jobs.yml b/.github/workflows/model_jobs.yml
@@ -15,6 +15,9 @@ on:
       runner:
         required: true
         type: string
+      docker:
+        required: true
+        type: string
 
 env:
   HF_HOME: /mnt/cache
@@ -39,7 +42,7 @@ jobs:
         folders: ${{ fromJson(inputs.folder_slices)[inputs.slice_id] }}
     runs-on: ['${{ inputs.machine_type }}', nvidia-gpu, t4, '${{ inputs.runner }}']
     container:
-      image: huggingface/transformers-all-latest-gpu
+      image: ${{ inputs.docker }}
       options: --gpus all --shm-size "16gb" --ipc host -v /mnt/cache/.cache/huggingface:/mnt/cache/
     steps:
       - name: Echo input and matrix info

diff --git a/.github/workflows/self-nightly-caller.yml b/.github/workflows/self-nightly-caller.yml
@@ -17,6 +17,7 @@ jobs:
       job: run_models_gpu
       slack_report_channel: "#transformers-ci-past-future"
       runner: past-ci
+      docker: huggingface/transformers-all-latest-torch-nightly-gpu
     secrets: inherit
 
   deepspeed-ci:
@@ -26,4 +27,5 @@ jobs:
       job: run_torch_cuda_extensions_gpu
       slack_report_channel: "#transformers-ci-past-future"
       runner: past-ci
+      docker: huggingface/transformers-pytorch-deepspeed-nightly-gpu
     secrets: inherit
diff --git a/.github/workflows/self-scheduled.yml b/.github/workflows/self-scheduled.yml
@@ -18,6 +18,9 @@ on:
       runner:
         required: true
         type: string
+      docker:
+        required: true
+        type: string
 
 env:
   HF_HOME: /mnt/cache
@@ -100,6 +103,7 @@ jobs:
       machine_type: ${{ matrix.machine_type }}
       slice_id: ${{ matrix.slice_id }}
       runner: ${{ inputs.runner }}
+      docker: ${{ inputs.docker }}
     secrets: inherit
 
   run_pipelines_torch_gpu:
@@ -263,7 +267,7 @@ jobs:
         machine_type: [single-gpu, multi-gpu]
     runs-on: ['${{ matrix.machine_type }}', nvidia-gpu, t4, '${{ inputs.runner }}']
     container:
-      image: huggingface/transformers-pytorch-deepspeed-latest-gpu
+      image: ${{ inputs.docker }}
       options: --gpus all --shm-size "16gb" --ipc host -v /mnt/cache/.cache/huggingface:/mnt/cache/
     steps:
       - name: Update clone