comma seperated gpus

mosaicml · May 21, 2024 · 2f0f83d · 2f0f83d
1 parent 38fc62c
commit 2f0f83d
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 13 deletions.
diff --git a/.github/mcli/mcli_pytest.py b/.github/mcli/mcli_pytest.py
@@ -67,11 +67,19 @@
     pip install --upgrade --user .{args.pip_deps}
 
     export COMMON_ARGS="-v --durations=20 -m '{args.pytest_markers}' {clear_tmp_path_flag}"
+    '''
 
-    make test PYTEST='{args.pytest_command}' EXTRA_ARGS="$COMMON_ARGS --codeblocks"
-
-    make test-dist PYTEST='{args.pytest_command}' EXTRA_ARGS="$COMMON_ARGS" WORLD_SIZE=2
+    if args.gpu_num == 1:
+        command += f'''
+        make test PYTEST='{args.pytest_command}' EXTRA_ARGS="$COMMON_ARGS --codeblocks"
+        '''
+    else:
+        world_size = args.gpu_num
+        command += f'''
+        make test-dist PYTEST='{args.pytest_command}' EXTRA_ARGS="$COMMON_ARGS" WORLD_SIZE={world_size}
+        '''
 
+    command += '''
     python -m coverage combine
 
     python -m coverage report

diff --git a/.github/workflows/pytest-gpu.yaml b/.github/workflows/pytest-gpu.yaml
@@ -39,6 +39,9 @@ on:
         required: false
         type: boolean
         default: false
+      gpu_configurations:
+        required: true
+        type: string
     secrets:
       mcloud-api-key:
         required: true
@@ -56,16 +59,8 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        gpu_num: [1, 2, 4]  # Define the GPU configurations to test
-        job_name: ['pytest-gpu-1', 'pytest-gpu-2', 'pytest-gpu-4']
-      include:
-        - gpu_num: 1
-          job_name: pytest-gpu-1
-        - gpu_num: 2
-          job_name: pytest-gpu-2
-        - gpu_num: 4
-          job_name: pytest-gpu-4
-    name: ${{ matrix.job_name }}
+        gpu_num: [${{ fromJSON(inputs.gpu_configurations) }}]
+    name: Pytest GPU ${{ matrix.gpu_num }}
     env:
       MOSAICML_API_KEY: ${{ secrets.mcloud-api-key }}
     steps: