pipeline.py

#/!/usr/bin/env python
# Copyright 2016 Google Inc. 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.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.
"""Flowers Sample Cloud Runner."""

from __future__ import print_function
import argparse
import base64
import datetime
import errno
import io
import json
import multiprocessing
import os
import subprocess
import time
import uuid
import apache_beam as beam
from PIL import Image
from tensorflow.python.lib.io import file_io
from tensorflow.python.framework import errors

import trainer.preprocess as preprocess_lib

# Model variables
MODEL_NAME = 'flowers'
TRAINER_NAME = 'trainer-0.1.tar.gz'
METADATA_FILE_NAME = 'metadata.json'
EXPORT_SUBDIRECTORY = 'model'
CONFIG_FILE_NAME = 'config.yaml'
MODULE_NAME = 'trainer.task'
SAMPLE_IMAGE = \
  'gs://cloud-ml-data/img/flower_photos/tulips/4520577328_a94c11e806_n.jpg'

# Number of seconds to wait before sending next online prediction after
# an online prediction fails due to model deployment not being complete.
PREDICTION_WAIT_TIME = 30

def process_args():
  """Define arguments and assign default values to the ones that are not set.

  Returns:
    args: The parsed namespace with defaults assigned to the flags.
  """

  parser = argparse.ArgumentParser(
      description='Runs Flowers Sample E2E pipeline.')
  parser.add_argument(
      '--project',
      default=None,
      help='The project to which the job will be submitted.')
  parser.add_argument(
      '--cloud',
      action='store_true',
      help='Run preprocessing on the cloud.')
  parser.add_argument(
      '--train_input_path',
      default=None,
      help='Input specified as uri to CSV file for the train set')
  parser.add_argument(
      '--eval_input_path',
      default=None,
      help='Input specified as uri to CSV file for the eval set.')
  parser.add_argument(
      '--eval_set_size',
      default=50,
      help='The size of the eval dataset.')
  parser.add_argument(
      '--input_dict',
      default=None,
      help='Input dictionary. Specified as text file uri. '
      'Each line of the file stores one label.')
  parser.add_argument(
      '--deploy_model_name',
      default='flowerse2e',
      help=('If --cloud is used, the model is deployed with this '
            'name. The default is flowerse2e.'))
  parser.add_argument(
      '--dataflow_sdk_path',
      default=None,
      help=('Path to Dataflow SDK location. If None, Pip will '
            'be used to download the latest published version'))
  parser.add_argument(
      '--max_deploy_wait_time',
      default=600,
      help=('Maximum number of seconds to wait after a model is deployed.'))
  parser.add_argument(
      '--deploy_model_version',
      default='v' + uuid.uuid4().hex[:4],
      help=('If --cloud is used, the model is deployed with this '
            'version. The default is four random characters.'))
  parser.add_argument(
      '--preprocessed_train_set',
      default=None,
      help=('If specified, preprocessing steps will be skipped.'
            'The provided preprocessed dataset wil be used in this case.'
            'If specified, preprocessed_eval_set must also be provided.'))
  parser.add_argument(
      '--preprocessed_eval_set',
      default=None,
      help=('If specified, preprocessing steps will be skipped.'
            'The provided preprocessed dataset wil be used in this case.'
            'If specified, preprocessed_train_set must also be provided.'))
  parser.add_argument(
      '--pretrained_model_path',
      default=None,
      help=('If specified, preprocessing and training steps ares skipped.'
            'The pretrained model will be deployed in this case.'))
  parser.add_argument(
      '--sample_image_uri',
      default=SAMPLE_IMAGE,
      help=('URI for a single Jpeg image to be used for online prediction.'))
  parser.add_argument(
      '--gcs_bucket',
      default=None,
      help=('Google Cloud Storage bucket to be used for uploading intermediate '
            'data')),
  parser.add_argument(
      '--output_dir',
      default=None,
      help=('Google Cloud Storage or Local directory in which '
            'to place outputs.'))
  parser.add_argument(
      '--runtime_version',
      default=os.getenv('CLOUDSDK_ML_DEFAULT_RUNTIME_VERSION', '1.0'),
      help=('Tensorflow version for model training and prediction.'))

  args, _ = parser.parse_known_args()

  if args.cloud and not args.project:
    args.project = get_cloud_project()

  return args


class FlowersE2E(object):
  """The end-2-end pipeline for Flowers Sample."""

  def  __init__(self, args=None):
    if not args:
      self.args = process_args()
    else:
      self.args = args

  def preprocess(self):
    """Runs the pre-processing pipeline.

    It tiggers two Dataflow pipelines in parallel for train and eval.
    Returns:
      train_output_prefix: Path prefix for the preprocessed train dataset.
      eval_output_prefix: Path prefix for the preprocessed eval dataset.
    """

    train_dataset_name = 'train'
    eval_dataset_name = 'eval'

    # Prepare the environment to run the Dataflow pipeline for preprocessing.
    if self.args.dataflow_sdk_path:
      dataflow_sdk = self.args.dataflow_sdk_path
      if dataflow_sdk.startswith('gs://'):
        subprocess.check_call(
            ['gsutil', 'cp', self.args.dataflow_sdk_path, '.'])
        dataflow_sdk = self.args.dataflow_sdk_path.split('/')[-1]
    else:
      dataflow_sdk = None

    subprocess.check_call(['python', 'setup.py', 'sdist', '--format=gztar'])

    trainer_uri = os.path.join(self.args.output_dir, TRAINER_NAME)
    subprocess.check_call(
        ['gsutil', '-q', 'cp', os.path.join('dist', TRAINER_NAME), trainer_uri])

    thread_pool = multiprocessing.pool.ThreadPool(2)

    train_output_prefix = os.path.join(self.args.output_dir, 'preprocessed',
                                       train_dataset_name)
    eval_output_prefix = os.path.join(self.args.output_dir, 'preprocessed',
                                      eval_dataset_name)

    train_args = (train_dataset_name, self.args.train_input_path,
                  train_output_prefix, dataflow_sdk, trainer_uri)
    eval_args = (eval_dataset_name, self.args.eval_input_path,
                 eval_output_prefix, dataflow_sdk, trainer_uri)

    # make a pool to run two pipelines in parallel.
    pipeline_pool = [thread_pool.apply_async(self.run_pipeline, train_args),
                     thread_pool.apply_async(self.run_pipeline, eval_args)]
    _ = [res.get() for res in pipeline_pool]
    return train_output_prefix, eval_output_prefix

  def run_pipeline(self, dataset_name, input_csv, output_prefix,
                   dataflow_sdk_location, trainer_uri):
    """Runs a Dataflow pipeline to preprocess the given dataset.

    Args:
      dataset_name: The name of the dataset ('eval' or 'train').
      input_csv: Path to the input CSV file which contains an image-URI with
                 its labels in each line.
      output_prefix:  Output prefix to write results to.
      dataflow_sdk_location: path to Dataflow SDK package.
      trainer_uri: Path to the Flower's trainer package.
    """
    job_name = ('cloud-ml-sample-flowers-' +
                datetime.datetime.now().strftime('%Y%m%d%H%M%S')  +
                '-' + dataset_name)

    options = {
        'staging_location':
            os.path.join(self.args.output_dir, 'tmp', dataset_name, 'staging'),
        'temp_location':
            os.path.join(self.args.output_dir, 'tmp', dataset_name),
        'project':
            self.args.project,
        'job_name': job_name,
        'extra_packages': [trainer_uri],
        'save_main_session':
            True,
    }
    if dataflow_sdk_location:
      options['sdk_location'] = dataflow_sdk_location

    pipeline_name = 'DataflowRunner' if self.args.cloud else 'DirectRunner'

    opts = beam.pipeline.PipelineOptions(flags=[], **options)
    args = argparse.Namespace(**vars(self.args))
    vars(args)['input_path'] = input_csv
    vars(args)['input_dict'] = self.args.input_dict
    vars(args)['output_path'] = output_prefix
    # execute the pipeline
    with beam.Pipeline(pipeline_name, options=opts) as pipeline:
      preprocess_lib.configure_pipeline(pipeline, args)

  def train(self, train_file_path, eval_file_path):
    """Train a model using the eval and train datasets.

    Args:
      train_file_path: Path to the train dataset.
      eval_file_path: Path to the eval dataset.
    """
    trainer_args = [
        '--output_path', self.args.output_dir,
        '--eval_data_paths', eval_file_path,
        '--eval_set_size', str(self.args.eval_set_size),
        '--train_data_paths', train_file_path
    ]

    if self.args.cloud:
      job_name = 'flowers_model' + datetime.datetime.now().strftime(
          '_%y%m%d_%H%M%S')
      command = [
          'gcloud', 'ml-engine', 'jobs', 'submit', 'training', job_name,
          '--stream-logs',
          '--module-name', MODULE_NAME,
          '--staging-bucket', self.args.gcs_bucket,
          '--region', 'us-central1',
          '--project', self.args.project,
          '--package-path', 'trainer',
          '--runtime-version', self.args.runtime_version,
          '--'
      ] + trainer_args
    else:
      command = [
          'gcloud', 'ml-engine', 'local', 'train',
          '--module-name', MODULE_NAME,
          '--package-path', 'trainer',
          '--',
      ] + trainer_args
    subprocess.check_call(command)

  def deploy_model(self, model_path):
    """Deploys the trained model.

    Args:
      model_path: Path to the trained model.
    """

    create_model_cmd = [
        'gcloud', 'ml-engine', 'models', 'create', self.args.deploy_model_name,
        '--regions', 'us-central1',
        '--project', self.args.project,
    ]

    print(create_model_cmd)
    subprocess.check_call(create_model_cmd)

    submit = [
        'gcloud', 'ml-engine', 'versions', 'create',
        self.args.deploy_model_version,
        '--model', self.args.deploy_model_name,
        '--origin', model_path,
        '--project', self.args.project,
        '--runtime-version', self.args.runtime_version,
    ]
    if not model_path.startswith('gs://'):
      submit.extend(['--staging-bucket', self.args.gcs_bucket])
    print(submit)
    subprocess.check_call(submit)

    self.adaptive_wait()

    print('Deployed %s version: %s' % (self.args.deploy_model_name,
                                       self.args.deploy_model_version))

  def adaptive_wait(self):
    """Waits for a model to be fully deployed.

       It keeps sending online prediction requests until a prediction is
       successful or maximum wait time is reached. It sleeps between requests.
    """
    start_time = datetime.datetime.utcnow()
    elapsed_time = 0
    while elapsed_time < self.args.max_deploy_wait_time:
      try:
        self.predict(self.args.sample_image_uri)
        return
      except Exception as e:
        time.sleep(PREDICTION_WAIT_TIME)
        elapsed_time = (datetime.datetime.utcnow() - start_time).total_seconds()
        continue

  def predict(self, image_uri):
    """Sends a predict request for the deployed model for the given image.

    Args:
      image_uri: The input image URI.
    """
    output_json = 'request.json'
    self.make_request_json(image_uri, output_json)
    cmd = [
        'gcloud', 'ml-engine', 'predict',
        '--model', self.args.deploy_model_name,
        '--version', self.args.deploy_model_version,
        '--json-instances', 'request.json',
        '--project', self.args.project
    ]
    subprocess.check_call(cmd)

  def make_request_json(self, uri, output_json):
    """Produces a JSON request suitable to send to CloudML Prediction API.

    Args:
      uri: The input image URI.
      output_json: File handle of the output json where request will be written.
    """
    def _open_file_read_binary(uri):
      try:
        return file_io.FileIO(uri, mode='rb')
      except errors.InvalidArgumentError:
        return file_io.FileIO(uri, mode='r')

    with open(output_json, 'w') as outf:
      with _open_file_read_binary(uri) as f:
        image_bytes = f.read()
        image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
        image = image.resize((299, 299), Image.BILINEAR)
        resized_image = io.BytesIO()
        image.save(resized_image, format='JPEG')
        encoded_image = base64.b64encode(resized_image.getvalue())
        row = json.dumps({'key': uri, 'image_bytes': {'b64': encoded_image}})
        outf.write(row)
        outf.write('\n')

  def run(self):
    """Runs the pipeline."""
    model_path = self.args.pretrained_model_path
    if not model_path:
      train_prefix, eval_prefix = (self.args.preprocessed_train_set,
                                   self.args.preprocessed_eval_set)

      if not train_prefix or not eval_prefix:
        train_prefix, eval_prefix = self.preprocess()
      self.train(train_prefix + '*', eval_prefix + '*')
      model_path = os.path.join(self.args.output_dir, EXPORT_SUBDIRECTORY)
    self.deploy_model(model_path)


def get_cloud_project():
  cmd = [
      'gcloud', '-q', 'config', 'list', 'project',
      '--format=value(core.project)'
  ]
  with open(os.devnull, 'w') as dev_null:
    try:
      res = subprocess.check_output(cmd, stderr=dev_null).strip()
      if not res:
        raise Exception('--cloud specified but no Google Cloud Platform '
                        'project found.\n'
                        'Please specify your project name with the --project '
                        'flag or set a default project: '
                        'gcloud config set project YOUR_PROJECT_NAME')
      return res
    except OSError as e:
      if e.errno == errno.ENOENT:
        raise Exception('gcloud is not installed. The Google Cloud SDK is '
                        'necessary to communicate with the Cloud ML service. '
                        'Please install and set up gcloud.')
      raise


def main():
  pipeline = FlowersE2E()
  pipeline.run()

if __name__ == '__main__':
  main()