Home | History | Annotate | Download | only in gcs_test 
      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 # In-container wrapper for GCS smoke test.
     18 #
     19 # This script invokes gcs_smoke.py and performs tear down afterwards.
     20 #
     21 # Usage:
     22 #   gcs_smoke_wrapper.sh <GCS_BUCKET_URL>
     23 
     24 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
     25 
     26 # Helper function: Exit on failure.
     27 die () {
     28   echo $@
     29   exit 1
     30 }
     31 
     32 print_usage() {
     33   echo "Usage: gcs_smoke_wrapper.sh <GCS_BUCKET_URL>"
     34   echo ""
     35 }
     36 
     37 # Sanity check on command-line arguments.
     38 GCS_BUCKET_URL=$1
     39 if [[ -z "${GCS_BUCKET_URL}" ]]; then
     40   print_usage
     41   die "ERROR: Command-line argument GCS_BUCKET_URL is not supplied"
     42 fi
     43 
     44 # Check that gcloud and gsutil binaries are available.
     45 GCLOUD_BIN="/var/gcloud/google-cloud-sdk/bin/gcloud"
     46 if [[ ! -f "${GCLOUD_BIN}" ]]; then
     47   die "ERROR: Unable to find gcloud at path ${GCLOUD_BIN}"
     48 fi
     49 
     50 GSUTIL_BIN="/var/gcloud/google-cloud-sdk/bin/gsutil"
     51 if [[ ! -f "${GSUTIL_BIN}" ]]; then
     52   die "ERROR: Unable to find gsutil at path ${GSUTIL_BIN}"
     53 fi
     54 
     55 # Check environment variable for gcloud credentials
     56 if [[ -z "${GOOGLE_APPLICATION_CREDENTIALS}" ]]; then
     57   die "ERROR: Required gcloud environment variable "\
     58 "${GOOGLE_APPLICATION_CREDENTIALS} is not set."
     59 fi
     60 
     61 # Locate main Python file
     62 GCS_SMOKE_PY="${SCRIPT_DIR}/python/gcs_smoke.py"
     63 if [[ ! -f "${GCS_SMOKE_PY}" ]]; then
     64   die "ERROR: Unable to find Python file at ${GCS_SMOKE_PY}"
     65 fi
     66 
     67 
     68 LOG_FILE="/tmp/tf-gcs-test.log"
     69 rm -rf ${LOG_FILE} || \
     70     die "ERROR: Failed to remove existing log file ${LOG_FILE}"
     71 
     72 # Invoke main Python file
     73 python "${GCS_SMOKE_PY}" --gcs_bucket_url="${GCS_BUCKET_URL}" \
     74     > "${LOG_FILE}" 2>&1
     75 
     76 if [[ $? != "0" ]]; then
     77   cat ${LOG_FILE}
     78   die "FAIL: End-to-end test of GCS access from TensorFlow failed."
     79 fi
     80 
     81 cat ${LOG_FILE}
     82 echo ""
     83 
     84 # Clean up the newly created tfrecord file in GCS bucket.
     85 # First, activate gcloud service account
     86 "${GCLOUD_BIN}" auth activate-service-account \
     87     --key-file "${GOOGLE_APPLICATION_CREDENTIALS}" || \
     88     die "ERROR: Failed to activate gcloud service account with JSON key file"
     89 
     90 NEW_TFREC_URL=$(grep "Using input path" "${LOG_FILE}" | \
     91                 awk '{print $NF}')
     92 if [[ -z ${NEW_TFREC_URL} ]]; then
     93   die "FAIL: Unable to determine the URL to the new tfrecord file in GCS"
     94 fi
     95 if "${GSUTIL_BIN}" rm "${NEW_TFREC_URL}"
     96 then
     97   echo "Cleaned up new tfrecord file in GCS: ${NEW_TFREC_URL}"
     98 else
     99   die "FAIL: Unable to clean up new tfrecord file in GCS: ${NEW_TFREC_URL}"
    100 fi
    101