Allow GPU memory growth in CI

nengo · Feb 27, 2020 · a13fb9f · a13fb9f
1 parent 629373a
commit a13fb9f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/.templates/gpu.sh.template b/.templates/gpu.sh.template
@@ -13,9 +13,11 @@
         echo "Waiting for lock on GPU $GPU_NUM"
         (
             flock -x -w 540 200 || exit 1
-            CUDA_VISIBLE_DEVICES="$GPU_NUM" pytest $TEST_ARGS nengo_dl/tests/test_benchmarks.py::test_performance --performance -v --durations 20 --color=yes || exit 1
-            CUDA_VISIBLE_DEVICES="$GPU_NUM" pytest $TEST_ARGS nengo_dl -v --durations 20 --color=yes --cov=nengo_dl --cov-report=xml --cov-report=term-missing || exit 1
-            CUDA_VISIBLE_DEVICES="$GPU_NUM" pytest $TEST_ARGS --pyargs nengo -v --durations 20 --color=yes --cov=nengo_dl --cov-report=xml --cov-report=term-missing --cov-append || exit 1
+            export CUDA_VISIBLE_DEVICES="$GPU_NUM"
+            export TF_FORCE_GPU_ALLOW_GROWTH=true
+            pytest $TEST_ARGS nengo_dl/tests/test_benchmarks.py::test_performance --performance -v --durations 20 --color=yes || exit 1
+            pytest $TEST_ARGS nengo_dl -v -n 2 --durations 20 --color=yes --cov=nengo_dl --cov-report=xml --cov-report=term-missing || exit 1
+            pytest $TEST_ARGS --pyargs nengo -v -n 2 --durations 20 --color=yes --cov=nengo_dl --cov-report=xml --cov-report=term-missing --cov-append || exit 1
         ) 200>/var/lock/.travis-ci.exclusivelock."$GPU_NUM" || REMOTE_STATUS=1
 {% endblock %}
 

diff --git a/.templates/remote-script.sh.template b/.templates/remote-script.sh.template
@@ -23,7 +23,7 @@
         echo "Waiting for lock on GPU $GPU_NUM"
         (
             flock -x -w 540 200 || exit 1
-            CUDA_VISIBLE_DEVICES="$GPU_NUM" bash .ci/{{ remote_script }}.sh script || exit 1
+            CUDA_VISIBLE_DEVICES="$GPU_NUM" TF_FORCE_GPU_ALLOW_GROWTH=true bash .ci/{{ remote_script }}.sh script || exit 1
         ) 200>/var/lock/.travis-ci.exclusivelock."$GPU_NUM" || REMOTE_STATUS=1
 {% endblock %}