Allow multi-GPU tests and enable MPS

NVIDIA · Nov 13, 2024 · 10688b6 · 10688b6
1 parent cc2922c
commit 10688b6
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 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}

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: |