aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/tools
diff options
context:
space:
mode:
authorGravatar Todd Wang <toddw@google.com>2018-10-03 15:15:23 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-10-03 15:24:02 -0700
commitefbee1ab2cac59f511cc0850d84414e711bbda3b (patch)
tree5e63a827df27b2fdac2de9c946612616fb53bbce /tensorflow/tools
parent148bc62dba0a0b9d26945ce48b6dcd903613de14 (diff)
Fix ci_parameterized_build to pass environment variables to tests.
This is particularly important when using --run_under with parallel_gpu_execute, since the envvars control the execution. PiperOrigin-RevId: 215637931
Diffstat (limited to 'tensorflow/tools')
-rwxr-xr-xtensorflow/tools/ci_build/ci_parameterized_build.sh18
1 files changed, 12 insertions, 6 deletions
diff --git a/tensorflow/tools/ci_build/ci_parameterized_build.sh b/tensorflow/tools/ci_build/ci_parameterized_build.sh
index 49a9048c03..99bdedf7b4 100755
--- a/tensorflow/tools/ci_build/ci_parameterized_build.sh
+++ b/tensorflow/tools/ci_build/ci_parameterized_build.sh
@@ -65,8 +65,6 @@
# TF_GPU_COUNT:
# Run this many parallel tests for serial builds.
# For now, only can be edited for PIP builds.
-# TODO(gunan): Find a way to pass this environment variable
-# to the script bazel runs (using --run_under).
# TF_BUILD_TEST_TUTORIALS:
# If set to any non-empty and non-0 value, will perform
# tutorials tests (Applicable only if TF_BUILD_IS_PIP is
@@ -150,6 +148,13 @@ ANDROID_FULL_CMD="${CI_BUILD_DIR}/builds/android_full.sh"
TF_GPU_COUNT=${TF_GPU_COUNT:-4}
PARALLEL_GPU_TEST_CMD='//tensorflow/tools/ci_build/gpu_build:parallel_gpu_execute'
+# Environment variables to set when running bazel tests. These are especially
+# important when using --run_under with parallel_gpu_execute.
+BAZEL_TEST_ENV=""\
+"--test_env=TF_GPU_COUNT=${TF_GPU_COUNT} "\
+"--test_env=TF_TESTS_PER_GPU=${TF_TESTS_PER_GPU} "\
+"--test_env=TF_PER_DEVICE_MEMORY_LIMIT_MB=${TF_PER_DEVICE_MEMORY_LIMIT_MB} "
+
BENCHMARK_CMD="${CI_BUILD_DIR}/builds/benchmark.sh"
EXTRA_PARAMS=""
@@ -410,13 +415,14 @@ if [[ ${TF_BUILD_IS_PIP} == "no_pip" ]] ||
if [[ ${CTYPE} == cpu* ]] || \
[[ ${CTYPE} == "debian.jessie.cpu" ]]; then
# CPU only command, fully parallel.
- NO_PIP_MAIN_CMD="${MAIN_CMD} ${BAZEL_CMD} ${OPT_FLAG} ${EXTRA_ARGS} -- "\
-"${BAZEL_TARGET}"
+ NO_PIP_MAIN_CMD="${MAIN_CMD} ${BAZEL_CMD} ${BAZEL_TEST_ENV} ${OPT_FLAG} "\
+ "${EXTRA_ARGS} -- ${BAZEL_TARGET}"
elif [[ ${CTYPE} == gpu* ]]; then
# GPU only command, run as many jobs as the GPU count only.
- NO_PIP_MAIN_CMD="${BAZEL_CMD} ${OPT_FLAG} "\
+ NO_PIP_MAIN_CMD="${BAZEL_CMD} ${BAZEL_TEST_ENV} ${OPT_FLAG} "\
"--local_test_jobs=${TF_GPU_COUNT} "\
-"--run_under=${PARALLEL_GPU_TEST_CMD} ${EXTRA_ARGS} -- ${BAZEL_TARGET}"
+"--run_under=${PARALLEL_GPU_TEST_CMD} "\
+"${EXTRA_ARGS} -- ${BAZEL_TARGET}"
elif [[ ${CTYPE} == "android" ]]; then
# Run android specific script for android build.
NO_PIP_MAIN_CMD="${ANDROID_CMD} ${OPT_FLAG} "