Allow multi-GPU tests and enable MPS (#1153)

- CI ran in 15m50s with this change (more tests), vs. 18m before
NVIDIA · Nov 23, 2024 · b0e6753 · b0e6753
1 parent b4d8558
commit b0e6753
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/.github/container/test-jax.sh b/.github/container/test-jax.sh
@@ -113,8 +113,6 @@ for t in $*; do
 done
 
 TEST_TAG_FILTER_ARRAY=()
-TEST_TAG_FILTER_ARRAY+=('-multiaccelerator')
-
 COMMON_FLAGS=$(cat << EOF
 --@local_config_cuda//:enable_cuda
 --cache_test_results=${CACHE_TEST_RESULTS}
@@ -163,7 +161,10 @@ case "${BATTERY}" in
         ;;
 esac
 
-TEST_TAG_FILTERS=$(IFS=, ; echo "--test_tag_filters=${TEST_TAG_FILTER_ARRAY[*]}")
+TEST_TAG_FILTERS=""
+if [[ ${#TEST_TAG_FILTER_ARRAY[@]} > 0 ]]; then
+    TEST_TAG_FILTERS=$(IFS=, ; echo "--test_tag_filters=${TEST_TAG_FILTER_ARRAY[*]}")
+fi
 
 print_var NCPUS
 print_var NGPUS

diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
@@ -289,6 +289,7 @@ jobs:
         docker run -i --shm-size=1g --gpus all \
         ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }} \
         bash <<"EOF" |& tee tee test-gpu.log
+          nvidia-cuda-mps-control -d
           test-jax.sh -b gpu
         EOF
       STATISTICS_SCRIPT: |