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