[CI] auto select 4090

.github/workflows/run_tests.yml

@@ -3,27 +3,42 @@ name: Unit Tests (GPU)
 on:
   repository_dispatch:
   workflow_dispatch:
-    inputs:
-      test_names:
-        description: 'Input Test(s) to Run (default all)'
-        required: false
-        default: ''
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
 
 env:
-  DAY_OF_WEEK: Monday
   CUDA_DEVICE_ORDER: PCI_BUS_ID
-  CUDA_VISIBLE_DEVICES: 0
 
 jobs:
   build:
     runs-on: self-hosted
+    strategy:
+      matrix:
+        test-group:
+          - gpu_group1
+          - gpu_group2
     container:
       image: pytorch/pytorch:2.3.1-cuda12.1-cudnn8-devel
 
     steps:
+      - name: Find suitable GPU
+        run: |
+         suitable_gpu=$(nvidia-smi -L | grep "RTX 4090" | awk -F': ' '{print $1}' | sed 's/GPU //g' | while read gpu_id
+          do
+            mem_total=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits -i $gpu_id)
+            mem_used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i $gpu_id)
+            mem_used_pct=$((100 * mem_used / mem_total))
+            if [ $mem_used_pct -lt 2 ]; then # 2 -> 98% free
+              echo $gpu_id
+              break
+            fi
+          done)
+          if [ -z "$suitable_gpu" ]; then
+            echo "No suitable GPU found. Exiting with error."
+            exit 1
+          else
+            echo "CUDA_VISIBLE_DEVICES=$suitable_gpu" >> $GITHUB_ENV
+            echo "CUDA_VISIBLE_DEVICES set to $suitable_gpu"
+          fi
+
       - name: Checkout code
         uses: actions/checkout@v4
 
@@ -32,53 +47,28 @@ jobs:
           pip install -U pytest ninja parameterized protobuf
           pip install -v --no-build-isolation .
 
-      - name: test_lm_head.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_lm_head') }}
-        run: pytest tests/test_lm_head.py
-
-      - name: test_q4_exallama.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_q4_exallama') }}
-        run: pytest tests/test_q4_exallama.py
-
-      - name: test_q4_exallama_v2.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_q4_exallama_v2') }}
-        run: pytest tests/test_q4_exallama_v2.py
-
-      - name: test_q4_marlin.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_q4_marlin') }}
-        run: pytest tests/test_q4_marlin.py
-
-      - name: test_q4_triton.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_q4_triton') }}
-        run: pytest tests/test_q4_triton.py
-
-      - name: test_repacking.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_repacking') }}
-        run: pytest tests/test_repacking.py
-
-      - name: test_serialization.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_serialization') }}
-        run: pytest tests/test_serialization.py
-
-      - name: test_sharded.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_sharded') }}
-        run: pytest tests/test_sharded.py
-
-      - name: test_triton.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_triton') }}
-        run: pytest tests/test_triton.py
-
-      - name: test_quant_formats.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_quant_formats') }}
-        run: pytest tests/test_quant_formats.py
-
-      - name: test_q4_cuda.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_q4_cuda') }}
-        run: pytest tests/test_q4_cuda.py
-
-      - name: test_perplexity.py
-        if: ${{ !github.event.inputs.test_names || contains(github.event.inputs.test_names, 'test_perplexity') }}
+      - name: Run tests
+        if: matrix.test-group == 'gpu_group2'
         run: pytest tests/test_perplexity.py
-
-
-
+      - name: Run tests
+        if: matrix.test-group == 'gpu_group1'
+        run: |
+          run_test() {
+              echo "\033[1;34mRunning $1...\033[0m"
+              echo "========================================"
+              pytest tests/$1
+              echo "\033[1;32mFinished $1\033[0m"
+              echo "========================================\n\n\n"
+          }
+
+          run_test test_lm_head.py
+          run_test test_q4_exallama.py
+          run_test test_q4_exallama_v2.py
+          run_test test_q4_marlin.py
+          run_test test_q4_triton.py
+          run_test test_repacking.py
+          run_test test_serialization.py
+          run_test test_sharded.py
+          run_test test_triton.py
+          run_test test_quant_formats.py
+          run_test test_q4_cuda.py