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