make all SLURM job run on jumpbox runners

NVIDIA · Oct 15, 2024 · d0800ed · d0800ed
1 parent e8043a5
commit d0800ed
Show file tree

Hide file tree

Showing 7 changed files with 18 additions and 18 deletions.
diff --git a/.github/workflows/_runner_ondemand_slurm.yaml b/.github/workflows/_runner_ondemand_slurm.yaml
@@ -20,7 +20,7 @@ on:
 jobs:
 
   launch-slurm-runner:
-    runs-on: ubuntu-latest
+    runs-on: jumpbox
     steps:
       - name: Print environment variables
         run: env

diff --git a/.github/workflows/_test_maxtext.yaml b/.github/workflows/_test_maxtext.yaml
@@ -43,7 +43,7 @@ jobs:
         # - [1, 1, 1, 8] # PP, DP, FSDP, TP
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
 
     steps:
       - name: Print environment variables
@@ -194,7 +194,7 @@ jobs:
         - [1, 4, 2, 2]
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
 
     steps:
       - name: Print environment variables

diff --git a/.github/workflows/_test_pax_rosetta.yaml b/.github/workflows/_test_pax_rosetta.yaml
@@ -38,7 +38,7 @@ jobs:
         - [1, 1, 2, 4]
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
 
     env:
       BADGE_FILENAME_PREFIX: badge-rosetta-pax-single-process-multi-device-te
@@ -260,7 +260,7 @@ jobs:
             ADDITIONAL_ARGS: "--model-type LLaMA70BProxy --evaluate"
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
     env:
       BADGE_FILENAME_PREFIX: badge-rosetta-pax-multi-node-te
     steps:
@@ -458,7 +458,7 @@ jobs:
         - [4, 2, 1, 2]
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
     env:
       BADGE_FILENAME_PREFIX: badge-rosetta-pax-multi-node
     steps:
@@ -650,7 +650,7 @@ jobs:
         - [1, 8, 1, 1]
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
     env:
       BADGE_FILENAME_PREFIX: badge-rosetta-pax-single-node-dropout-te
     steps:
@@ -844,7 +844,7 @@ jobs:
         - [1, 8, 1, 1]
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
     env:
       BADGE_FILENAME_PREFIX: badge-rosetta-pax-single-process-evaluation-te
     steps:

diff --git a/.github/workflows/_test_slurm_pyxis.yaml b/.github/workflows/_test_slurm_pyxis.yaml
@@ -81,7 +81,7 @@ jobs:
 
   run-test:
     name: ${{ inputs.NAME }}
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
     outputs:
       SLURM_JOB_ID: ${{ steps.submit.outputs.SLURM_JOB_ID }}
       SLURM_STATE: ${{ steps.exit-info.outputs.SLURM_STATE }}

diff --git a/.github/workflows/_test_t5x_rosetta.yaml b/.github/workflows/_test_t5x_rosetta.yaml
@@ -48,7 +48,7 @@ jobs:
             EXTRA_GIN_ARGS: "--gin.train/utils.DatasetConfig.pack=False --gin.train_eval/utils.DatasetConfig.pack=False"
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
     env:
       BADGE_FILENAME_PREFIX: badge-rosetta-t5x-single-process-multi-device
     steps:
@@ -242,7 +242,7 @@ jobs:
             EXTRA_GIN_ARGS: ""
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
     env:
       BADGE_FILENAME_PREFIX: badge-rosetta-t5x-multi-gpu-multi-node
     steps:
@@ -428,7 +428,7 @@ jobs:
         N_GPU: [8]
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
     env:
       BADGE_FILENAME_PREFIX: badge-rosetta-t5x-vit-single-process-multi-device
     steps:
@@ -598,7 +598,7 @@ jobs:
         N_NODE: [1, 2]
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
     env:
       BADGE_FILENAME_PREFIX: badge-rosetta-t5x-vit-multi-gpu-multi-node
     steps:

diff --git a/.github/workflows/_test_upstream_pax.yaml b/.github/workflows/_test_upstream_pax.yaml
@@ -38,7 +38,7 @@ jobs:
         - [1, 1, 2, 4]
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
 
     steps:
       - name: Print environment variables
@@ -212,7 +212,7 @@ jobs:
             ADDITIONAL_ARGS: "--model-type LLaMA70BProxy --evaluate"
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
 
     steps:
       - name: Print environment variables
@@ -361,7 +361,7 @@ jobs:
         - [1, 8, 1, 1]
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
 
     steps:
       - name: Print environment variables

diff --git a/.github/workflows/_test_upstream_t5x.yaml b/.github/workflows/_test_upstream_t5x.yaml
@@ -46,7 +46,7 @@ jobs:
             ADDITIONAL_ARGS: "--enable-fmha 1"
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
 
     steps:
       - name: Print environment variables
@@ -191,7 +191,7 @@ jobs:
             ADDITIONAL_ARGS: "--enable-fmha 1"
       fail-fast: false
 
-    runs-on: ubuntu-22.04
+    runs-on: jumpbox
 
     steps:
       - name: Print environment variables