diff options
author | Todd Wang <toddw@google.com> | 2018-10-03 15:15:23 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-10-03 15:24:02 -0700 |
commit | efbee1ab2cac59f511cc0850d84414e711bbda3b (patch) | |
tree | 5e63a827df27b2fdac2de9c946612616fb53bbce /tensorflow/tools | |
parent | 148bc62dba0a0b9d26945ce48b6dcd903613de14 (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-x | tensorflow/tools/ci_build/ci_parameterized_build.sh | 18 |
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} " |