!pip3 install cloudml-hypertune

"""
Explanation: LAB 5a:  Training Keras model on Cloud AI Platform.
Learning Objectives

Setup up the environment
Create trainer module's task.py to hold hyperparameter argparsing code
Create trainer module's model.py to hold Keras model code
Run trainer module package locally
Submit training job to Cloud AI Platform
Submit hyperparameter tuning job to Cloud AI Platform

Introduction
After having testing our training pipeline both locally and in the cloud on a susbset of the data, we can submit another (much larger) training job to the cloud. It is also a good idea to run a hyperparameter tuning job to make sure we have optimized the hyperparameters of our model. 
In this notebook, we'll be training our Keras model at scale using Cloud AI Platform.
In this lab, we will set up the environment, create the trainer module's task.py to hold hyperparameter argparsing code, create the trainer module's model.py to hold Keras model code, run the trainer module package locally, submit a training job to Cloud AI Platform, and submit a hyperparameter tuning job to Cloud AI Platform.
Each learning objective will correspond to a #TODO in this student lab notebook -- try to complete this notebook first and then review the solution notebook.
Set up environment variables and load necessary libraries
First we will install the cloudml-hypertune package on our local machine. This is the package which we will use to report hyperparameter tuning metrics to Cloud AI Platform. Installing the package will allow us to test our trainer package locally.
End of explanation
"""


import os

"""
Explanation: Import necessary libraries.
End of explanation
"""


PROJECT = !gcloud config list --format 'value(core.project)'
PROJECT = PROJECT[0]
BUCKET = PROJECT
REGION = "us-central1"

os.environ["PROJECT"] = PROJECT
os.environ["BUCKET"] = BUCKET
os.environ["REGION"] = REGION
os.environ["TFVERSION"] = "2.3"
os.environ["PYTHONVERSION"] = "3.7"

%%bash
gcloud config set project ${PROJECT}
gcloud config set ai_platform/region ${REGION}

%%bash
if ! gsutil ls | grep -q gs://${BUCKET}; then
    gsutil mb -l ${REGION} gs://${BUCKET}
fi

"""
Explanation: Set environment variables.
Set environment variables so that we can use them throughout the entire lab. We will be using our project ID for our bucket.
End of explanation
"""


%%bash
gsutil ls gs://${BUCKET}/babyweight/data/*000000000000.csv

"""
Explanation: Check data exists
Verify that you previously created CSV files we'll be using for training and evaluation. If not, go back to lab 1b_prepare_data_babyweight to create them.
End of explanation
"""


%%bash
mkdir -p babyweight/trainer
touch babyweight/trainer/__init__.py

"""
Explanation: Now that we have the Keras wide-and-deep code working on a subset of the data, we can package the TensorFlow code up as a Python module and train it on Cloud AI Platform.
Train on Cloud AI Platform
Training on Cloud AI Platform requires:
* Making the code a Python package
* Using gcloud to submit the training code to Cloud AI Platform
Ensure that the AI Platform API is enabled by going to this link.
Move code into a Python package
A Python package is simply a collection of one or more .py files along with an __init__.py file to identify the containing directory as a package. The __init__.py sometimes contains initialization code but for our purposes an empty file suffices.
The bash command touch creates an empty file in the specified location, the directory babyweight should already exist.
End of explanation
"""


%%writefile babyweight/trainer/task.py
import argparse
import json
import os

from trainer import model

import tensorflow as tf

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--job-dir",
        help="this model ignores this field, but it is required by gcloud",
        default="junk"
    )
    parser.add_argument(
        "--train_data_path",
        help="GCS location of training data",
        required=True
    )
    parser.add_argument(
        "--eval_data_path",
        help="GCS location of evaluation data",
        required=True
    )
    parser.add_argument(
        "--output_dir",
        help="GCS location to write checkpoints and export models",
        required=True
    )
    parser.add_argument(
        "--batch_size",
        help="Number of examples to compute gradient over.",
        type=int,
        default=512
    )

    # TODO: Add nnsize argument

    # TODO: Add nembeds argument

    # TODO: Add num_epochs argument

    # TODO: Add train_examples argument

    # TODO: Add eval_steps argument

    # Parse all arguments
    args = parser.parse_args()
    arguments = args.__dict__

    # Unused args provided by service
    arguments.pop("job_dir", None)
    arguments.pop("job-dir", None)

    # Modify some arguments
    arguments["train_examples"] *= 1000

    # Append trial_id to path if we are doing hptuning
    # This code can be removed if you are not using hyperparameter tuning
    arguments["output_dir"] = os.path.join(
        arguments["output_dir"],
        json.loads(
            os.environ.get("TF_CONFIG", "{}")
        ).get("task", {}).get("trial", "")
    )

    # Run the training job
    model.train_and_evaluate(arguments)

"""
Explanation: We then use the %%writefile magic to write the contents of the cell below to a file called task.py in the babyweight/trainer folder.
Lab Task #1: Create trainer module's task.py to hold hyperparameter argparsing code.
The cell below writes the file babyweight/trainer/task.py which sets up our training job. Here is where we determine which parameters of our model to pass as flags during training using the parser module. Look at how batch_size is passed to the model in the code below. Use this as an example to parse arguements for the following variables
- nnsize which represents the hidden layer sizes to use for DNN feature columns
- nembeds which represents the embedding size of a cross of n key real-valued parameters
- train_examples which represents the number of examples (in thousands) to run the training job
- eval_steps which represents the positive number of steps for which to evaluate model
Be sure to include a default value for the parsed arguments above and specfy the type if necessary.
End of explanation
"""


%%writefile babyweight/trainer/model.py
import datetime
import os
import shutil
import numpy as np
import tensorflow as tf

# Determine CSV, label, and key columns
# TODO: Add CSV_COLUMNS and LABEL_COLUMN

# Set default values for each CSV column.
# Treat is_male and plurality as strings.
# TODO: Add DEFAULTS


def features_and_labels(row_data):
    # TODO: Add your code here
    pass


def load_dataset(pattern, batch_size=1, mode=tf.estimator.ModeKeys.EVAL):
    # TODO: Add your code here
    pass


def create_input_layers():
    # TODO: Add your code here
    pass


def categorical_fc(name, values):
    # TODO: Add your code here
    pass


def create_feature_columns(nembeds):
    # TODO: Add your code here
    pass


def get_model_outputs(wide_inputs, deep_inputs, dnn_hidden_units):
    # TODO: Add your code here
    pass


def rmse(y_true, y_pred):
    # TODO: Add your code here
    pass


def build_wide_deep_model(dnn_hidden_units=[64, 32], nembeds=3):
    # TODO: Add your code here
    pass


def train_and_evaluate(args):
    model = build_wide_deep_model(args["nnsize"], args["nembeds"])
    print("Here is our Wide-and-Deep architecture so far:\n")
    print(model.summary())

    trainds = load_dataset(
        args["train_data_path"],
        args["batch_size"],
        tf.estimator.ModeKeys.TRAIN)

    evalds = load_dataset(
        args["eval_data_path"], 1000, tf.estimator.ModeKeys.EVAL)
    if args["eval_steps"]:
        evalds = evalds.take(count=args["eval_steps"])

    num_batches = args["batch_size"] * args["num_epochs"]
    steps_per_epoch = args["train_examples"] // num_batches

    checkpoint_path = os.path.join(args["output_dir"], "checkpoints/babyweight")
    cp_callback = tf.keras.callbacks.ModelCheckpoint(
        filepath=checkpoint_path, verbose=1, save_weights_only=True)

    history = model.fit(
        trainds,
        validation_data=evalds,
        epochs=args["num_epochs"],
        steps_per_epoch=steps_per_epoch,
        verbose=2,  # 0=silent, 1=progress bar, 2=one line per epoch
        callbacks=[cp_callback])

    EXPORT_PATH = os.path.join(
        args["output_dir"], datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
    tf.saved_model.save(
        obj=model, export_dir=EXPORT_PATH)  # with default serving function
    print("Exported trained model to {}".format(EXPORT_PATH))

"""
Explanation: In the same way we can write to the file model.py the model that we developed in the previous notebooks. 
Lab Task #2: Create trainer module's model.py to hold Keras model code.
Complete the TODOs in the code cell below to create our model.py. We'll use the code we wrote for the Wide & Deep model. Look back at your 9_keras_wide_and_deep_babyweight notebook and copy/paste the necessary code from that notebook into its place in the cell below.
End of explanation
"""


%%bash
OUTDIR=babyweight_trained
rm -rf ${OUTDIR}
export PYTHONPATH=${PYTHONPATH}:${PWD}/babyweight
python3 -m trainer.task \
    --job-dir=./tmp \
    --train_data_path=gs://${BUCKET}/babyweight/data/train*.csv \
    --eval_data_path=gs://${BUCKET}/babyweight/data/eval*.csv \
    --output_dir=${OUTDIR} \
    --batch_size=# TODO: Add batch size
    --num_epochs=# TODO: Add the number of epochs to train for
    --train_examples=# TODO: Add the number of examples to train each epoch for
    --eval_steps=# TODO: Add the number of evaluation batches to run

"""
Explanation: Train locally
After moving the code to a package, make sure it works as a standalone. Note, we incorporated the --train_examples flag so that we don't try to train on the entire dataset while we are developing our pipeline. Once we are sure that everything is working on a subset, we can change it so that we can train on all the data. Even for this subset, this takes about 3 minutes in which you won't see any output ...
Lab Task #3: Run trainer module package locally.
Fill in the missing code in the TODOs below so that we can run a very small training job over a single file with a small batch size, 1 epoch, 1 train example, and 1 eval step.
End of explanation
"""


%%bash

OUTDIR=gs://${BUCKET}/babyweight/trained_model
JOBID=babyweight_$(date -u +%y%m%d_%H%M%S)

gcloud ai-platform jobs submit training ${JOBID} \
    --region=${REGION} \
    --module-name=trainer.task \
    --package-path=$(pwd)/babyweight/trainer \
    --job-dir=${OUTDIR} \
    --staging-bucket=gs://${BUCKET} \
    --master-machine-type=n1-standard-8 \
    --scale-tier=CUSTOM \
    --runtime-version=${TFVERSION} \
    --python-version=${PYTHONVERSION} \
    -- \
    --train_data_path=gs://${BUCKET}/babyweight/data/train*.csv \
    --eval_data_path=gs://${BUCKET}/babyweight/data/eval*.csv \
    --output_dir=${OUTDIR} \
    --num_epochs=10 \
    --train_examples=10000 \
    --eval_steps=100 \
    --batch_size=32 \
    --nembeds=8

"""
Explanation: Training on Cloud AI Platform
Now that we see everything is working locally, it's time to train on the cloud! 
To submit to the Cloud we use gcloud ai-platform jobs submit training [jobname] and simply specify some additional parameters for AI Platform Training Service:
- jobname: A unique identifier for the Cloud job. We usually append system time to ensure uniqueness
- job-dir: A GCS location to upload the Python package to
- runtime-version: Version of TF to use.
- python-version: Version of Python to use.
- region: Cloud region to train in. See here for supported AI Platform Training Service regions
Below the -- \ we add in the arguments for our task.py file.
End of explanation
"""


%%writefile hyperparam.yaml
trainingInput:
    scaleTier: STANDARD_1
    hyperparameters:
        hyperparameterMetricTag: # TODO: Add metric we want to optimize
        goal: # TODO: MAXIMIZE or MINIMIZE?
        maxTrials: 20
        maxParallelTrials: 5
        enableTrialEarlyStopping: True
        params:
        - parameterName: batch_size
          type: # TODO: What datatype?
          minValue: # TODO: Choose a min value
          maxValue: # TODO: Choose a max value
          scaleType: # TODO: UNIT_LINEAR_SCALE or UNIT_LOG_SCALE?
        - parameterName: nembeds
          type: # TODO: What datatype?
          minValue: # TODO: Choose a min value
          maxValue: # TODO: Choose a max value
          scaleType: # TODO: UNIT_LINEAR_SCALE or UNIT_LOG_SCALE?

%%bash

OUTDIR=gs://${BUCKET}/babyweight/hp_tuning
JOBID=babyweight_$(date -u +%y%m%d_%H%M%S)

gcloud ai-platform jobs submit training ${JOBID} \
    --region=${REGION} \
    --module-name=trainer.task \
    --package-path=$(pwd)/babyweight/trainer \
    --job-dir=${OUTDIR} \
    --staging-bucket=gs://${BUCKET} \
    --master-machine-type=n1-standard-8 \
    --scale-tier=CUSTOM \
    --config=hyperparam.yaml \
    --runtime-version=${TFVERSION} \
    --python-version=${PYTHONVERSION} \
    -- \
    --train_data_path=gs://${BUCKET}/babyweight/data/train*.csv \
    --eval_data_path=gs://${BUCKET}/babyweight/data/eval*.csv \
    --output_dir=${OUTDIR} \
    --num_epochs=10 \
    --train_examples=5000 \
    --eval_steps=100 \
    --batch_size=32 \
    --nembeds=8

"""
Explanation: The training job should complete within 10 to 15 minutes. You do not need to wait for this training job to finish before moving forward in the notebook, but will need a trained model to complete our next lab.
Lab Task #4: Hyperparameter tuning.
All of these are command-line parameters to my program.  To do hyperparameter tuning, create hyperparam.yaml and pass it as --config hyperparam.yaml.
This step will take <b>up to 2 hours</b> -- you can increase maxParallelTrials or reduce maxTrials to get it done faster.  Since maxParallelTrials is the number of initial seeds to start searching from, you don't want it to be too large; otherwise, all you have is a random search. Complete __#TODO__s in yaml file and gcloud training job bash command so that we can run hyperparameter tuning.
End of explanation
"""


%%bash
OUTDIR=gs://${BUCKET}/babyweight/trained_model_tuned
JOBID=babyweight_$(date -u +%y%m%d_%H%M%S)
echo ${OUTDIR} ${REGION} ${JOBNAME}
gsutil -m rm -rf ${OUTDIR}

gcloud ai-platform jobs submit training ${JOBID} \
    --region=${REGION} \
    --module-name=trainer.task \
    --package-path=$(pwd)/babyweight/trainer \
    --job-dir=${OUTDIR} \
    --staging-bucket=gs://${BUCKET} \
    --master-machine-type=n1-standard-8 \
    --scale-tier=CUSTOM \
    --runtime-version=${TFVERSION} \
    --python-version=${PYTHONVERSION} \
    -- \
    --train_data_path=gs://${BUCKET}/babyweight/data/train*.csv \
    --eval_data_path=gs://${BUCKET}/babyweight/data/eval*.csv \
    --output_dir=${OUTDIR} \
    --num_epochs=10 \
    --train_examples=20000 \
    --eval_steps=100 \
    --batch_size=32 \
    --nembeds=8

"""
Explanation: Repeat training
This time with tuned parameters for batch_size and nembeds. Note that your best results may differ from below. So be sure to fill yours in!
End of explanation
"""


