text: # Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License"). You
  # may not use this file except in compliance with the License. A copy of
  # the License is located at
  #
  #     http://aws.amazon.com/apache2.0/
  #
  # or in the "license" file accompanying this file. This file is
  # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
  # ANY KIND, either express or implied. See the License for the specific
  # language governing permissions and limitations under the License.
  
  import base64
  import boto3
  import docker
  import os
  import repo_helper
  import shutil
  
  import region_helper
  
  
  class BenchmarkScriptException(Exception):
      """An error building a benchmarking docker image."""
  
      def __init__(self, message):
          """Create a BenchmarkingScriptException."""
          super(BenchmarkScriptException, self).__init__(message)
  
  
  FROM_IMAGE = "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tensorflow:1.6.0-gpu-py2"
  
  
  class BenchmarkScript(object):
      """A script that performs benchmarking, built into a docker image."""
  
      def __init__(self, name, repository, script_name, tag, device):
          """Create a benchmarking script.
  
          Args:
              name (str): The name of the benchmarking script
              repository (str): The ECR repository to store the image in.
              script_name (str): The name of the script file name within ./docker/ to.
                  build.
              tag (str): A tag to apply to the image
              device (str): cpu or gpu, indicating whether the script uses a gpu device or not.
          """
          self.name = name
          self.script_name = script_name
          self.tag = tag
          self.repository = repository
          self.device = device
  
      @property
      def image(self):
          """The URI of the image containing the script."""
          return self.repository + ":" + self.tag
  
      def build(self, sdist_path):
          """Build the script into a docker image and upload to ECR.
  
          Args:
                sdist_path (str): The path to a sagemaker_tensorflow sdist .tar.gz that will be benchmarked.
          """
          # Copy in the sdist into a docker build directory
          docker_build_dir = ".docker-build-{}".format(self.name)
          if os.path.exists(docker_build_dir):
              shutil.rmtree(docker_build_dir)
  
          # Copy everything in the docker package data into the docker build dir
          docker_install_dir = '/' + '/'.join(list(__file__.split('/')[:-1]) + ['docker/'])
          shutil.copytree(docker_install_dir, docker_build_dir)
  
          # Copy sagemaker_tensorflow sdist
          sdist_name = os.path.basename(sdist_path)
          sdist_dest = "{}/{}".format(docker_build_dir, sdist_name)
          shutil.copyfile(sdist_path, sdist_dest)
  
          tf_version = sdist_name.split("-")[1][:3]
  
          client = docker.from_env()
          ecr_client = boto3.client('ecr', region_name=region_helper.region)
          token = ecr_client.get_authorization_token()
          username, password = base64.b64decode(token['authorizationData'][0]['authorizationToken']).decode().split(':')
          tag = "{}:{}".format(self.repository, self.tag)
          print "Pulling base image {}".format(FROM_IMAGE)
          client.images.pull(FROM_IMAGE, auth_config={'username': username, 'password': password})
          print "Building image {}".format(tag)
          client.images.build(
              path=docker_build_dir,
              tag=tag,
              buildargs={'sagemaker_tensorflow': sdist_name,
                         'device': self.device,
                         'tf_version': tf_version,
                         'script': self.script_name})
          print "Push image"
          client.images.push(tag,
                             auth_config={'username': username, 'password': password})
          print "Image pushed, cleaning up"
          shutil.rmtree(docker_build_dir)
  
  all_scripts = {script.name: script for script in [
      BenchmarkScript("InputOnly", repo_helper.repository(region=region_helper.region),
                      "input_only_script.py", "input-only", "cpu"),
      BenchmarkScript("GpuLoad", repo_helper.repository(region=region_helper.region),
                      "gpu_pipeline_script.py", "gpu-load", "gpu")
  ]}
  