• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env bash
2# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ==============================================================================
16#
17# Build and test TensorFlow docker images.
18# The tests include Python unit tests-on-install and tutorial tests.
19#
20# Usage: docker_test.sh <IMAGE_TYPE> <TAG> <WHL_PATH>
21# Arguments:
22#   IMAGE_TYPE : Type of the image: (CPU|GPU|ROCM)
23#   TAG        : Docker image tag
24#   WHL_PATH   : Path to the whl file to be installed inside the docker image
25#
26#   e.g.: docker_test.sh CPU someone/tensorflow:0.8.0 pip_test/whl/tensorflow-0.8.0-cp27-none-linux_x86_64.whl
27#
28
29# Helper functions
30# Exit after a failure
31die() {
32  echo $@
33  exit 1
34}
35
36# Convert to lower case
37to_lower () {
38  echo "$1" | tr '[:upper:]' '[:lower:]'
39}
40
41
42# Helper function to traverse directories up until given file is found.
43function upsearch () {
44  test / == "$PWD" && return || \
45      test -e "$1" && echo "$PWD" && return || \
46      cd .. && upsearch "$1"
47}
48
49
50# Verify command line argument
51if [[ $# != "3" ]]; then
52  die "Usage: $(basename $0) <IMAGE_TYPE> <TAG> <WHL_PATH>"
53fi
54IMAGE_TYPE=$(to_lower "$1")
55DOCKER_IMG_TAG=$2
56WHL_PATH=$3
57
58# Verify image type
59if [[ "${IMAGE_TYPE}" == "cpu" ]]; then
60  DOCKERFILE="tensorflow/tools/docker/Dockerfile"
61elif [[ "${IMAGE_TYPE}" == "gpu" ]]; then
62  DOCKERFILE="tensorflow/tools/docker/Dockerfile.gpu"
63elif [[ "${IMAGE_TYPE}" == "rocm" ]]; then
64  DOCKERFILE="tensorflow/tools/docker/Dockerfile.rocm"
65else
66  die "Unrecognized image type: $1"
67fi
68
69# Verify docker binary existence
70if [[ -z $(which docker) ]]; then
71  die "FAILED: docker binary unavailable"
72fi
73
74# Locate the base directory
75BASE_DIR=$(upsearch "${DOCKERFILE}")
76if [[ -z "${BASE_DIR}" ]]; then
77  die "FAILED: Unable to find the base directory where the dockerfile "\
78"${DOCKERFILE} resides"
79fi
80echo "Base directory: ${BASE_DIR}"
81
82pushd ${BASE_DIR} > /dev/null
83
84# Build docker image
85DOCKERFILE_PATH="${BASE_DIR}/${DOCKERFILE}"
86DOCKERFILE_DIR="$(dirname ${DOCKERFILE_PATH})"
87
88# Check to make sure that the whl file exists
89test -f ${WHL_PATH} || \
90    die "whl file does not exist: ${WHL_PATH}"
91
92TMP_WHL_DIR="${DOCKERFILE_DIR}/whl"
93mkdir -p "${TMP_WHL_DIR}"
94cp "${WHL_PATH}" "${TMP_WHL_DIR}/" || \
95    die "FAILED to copy whl file from ${WHL_PATH} to ${TMP_WHL_DIR}/"
96
97docker build -t "${DOCKER_IMG_TAG}" -f "${DOCKERFILE_PATH}" \
98"${DOCKERFILE_DIR}" || \
99    die "FAILED to build docker image from Dockerfile ${DOCKERFILE_PATH}"
100
101# Clean up
102rm -rf "${TMP_WHL_DIR}" || \
103    die "Failed to remove temporary directory ${TMP_WHL_DIR}"
104
105
106# Add extra params for cuda devices and libraries for GPU container.
107if [ "${IMAGE_TYPE}" == "gpu" ]; then
108  devices=$(\ls /dev/nvidia* | xargs -I{} echo '--device {}:{}')
109  libs=$(\ls /usr/lib/x86_64-linux-gnu/libcuda.* | xargs -I{} echo '-v {}:{}')
110  GPU_EXTRA_PARAMS="${devices} ${libs}"
111elif [ "${IMAGE_TYPE}" == "rocm" ]; then
112  ROCM_EXTRA_PARAMS="--device=/dev/kfd --device=/dev/dri --group-add video \
113  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size 16G"
114else
115  GPU_EXTRA_PARAMS=""
116  ROCM_EXTRA_PARAMS=""
117fi
118
119# Run docker image with source directory mapped
120docker run -v ${BASE_DIR}:/tensorflow-src -w /tensorflow-src \
121${GPU_EXTRA_PARAMS} ${ROCM_EXTRA_PARAMS} \
122"${DOCKER_IMG_TAG}" \
123/bin/bash -c "tensorflow/tools/ci_build/builds/run_pip_tests.sh && "\
124"tensorflow/tools/ci_build/builds/test_tutorials.sh && "\
125"tensorflow/tools/ci_build/builds/integration_tests.sh"
126
127RESULT=$?
128
129popd > /dev/null
130if [[ ${RESULT} == 0 ]]; then
131  echo "SUCCESS: Built and tested docker image: ${DOCKER_IMG_TAG}"
132else
133  die "FAILED to build and test docker image: ${DOCKER_IMG_TAG}"
134fi
135