# change these to try this notebook out
BUCKET = 'cloud-training-demos-ml'
PROJECT = 'cloud-training-demos'
REGION = 'us-central1'

import os
os.environ['BUCKET'] = BUCKET
os.environ['PROJECT'] = PROJECT
os.environ['REGION'] = REGION

"""
Explanation: Launcher notebook
This notebook shows how to launch the flights_model.ipynb notebook either in Deep Learning VM or on Kubeflow pipelines
End of explanation
"""


# Install papermill package in the current Jupyter kernel
import sys
!{sys.executable} -m pip install papermill

def run_flights_model(input_notebook: str, 
                      output_notebook: str, 
                      develop_mode: bool, 
                      bucket : str, 
                      project: str):
  # execute notebook using papermill
  import papermill as pm
  pm.execute_notebook(
    input_notebook,
    output_notebook,
    parameters = {'BUCKET': bucket, 'PROJECT': project, 'DEVELOP_MODE': develop_mode}
  )

%%bash
rm -rf kfp_output
mkdir kfp_output

run_flights_model('flights_model.ipynb',
                  os.path.join('kfp_output', 'flights_model.ipynb'),
                  develop_mode=True, bucket=BUCKET, project=PROJECT)

!ls kfp_output

"""
Explanation: Launch "locally" using Python
Make sure the notebook runs as intended when executed through papermill.
End of explanation
"""


%%writefile params.yaml
---
BUCKET: cloud-training-demos-ml
PROJECT: cloud-training-demos
DEVELOP_MODE: False

%%bash
gsutil -m rm -r gs://$BUCKET/flights/notebook
gsutil cp *.ipynb params.yaml gs://$BUCKET/flights/notebook

%%bash

set -x

GCS_FOLDER=gs://$BUCKET/flights/notebook

export IMAGE_FAMILY="tf-latest-cu100" # or put any required
export ZONE="us-west1-b"
export INSTANCE_NAME="notebookexecutor"
export INSTANCE_TYPE="n1-standard-4"
#export INPUT_NOTEBOOK="$GCS_FOLDER/flights_small.ipynb"
export INPUT_NOTEBOOK="$GCS_FOLDER/flights_model.ipynb"
export OUTPUT_NOTEBOOK_DIR=$GCS_FOLDER
export PARAMS="$GCS_FOLDER/params.yaml"
export LAUNCHER_SCRIPT=https://raw.githubusercontent.com/b0noI/gcp-notebook-executor/master/notebook_executor.sh

gcloud compute instances create $INSTANCE_NAME \
        --zone=$ZONE \
        --image-family=$IMAGE_FAMILY \
        --image-project=deeplearning-platform-release \
        --maintenance-policy=TERMINATE \
        --accelerator='type=nvidia-tesla-p100,count=1' \
        --machine-type=$INSTANCE_TYPE \
        --boot-disk-size=200GB \
        --scopes=https://www.googleapis.com/auth/cloud-platform \
        --metadata="parameters_file=$PARAMS,input_notebook=$INPUT_NOTEBOOK,output_notebook=$OUTPUT_NOTEBOOK_DIR,startup-script-url=$LAUNCHER_SCRIPT"


%%bash
gsutil mv gs://$BUCKET/flights/notebook/notebook.ipynb gs://$BUCKET/flights/notebook/flights_model_dlvm.ipynb
gsutil ls gs://$BUCKET/flights/notebook

"""
Explanation: Submit to Deep Learning VM to execute
This will launch a beefy virtual machine on the cloud and execute the notebook.
Change the params.yaml below to be your PROJECT and BUCKET
End of explanation
"""


import sys
!{sys.executable} -m pip install https://storage.googleapis.com/ml-pipeline/release/0.1.3-rc.2/kfp.tar.gz --upgrade

"""
Explanation: Submit to Kubeflow pipelines to run as a component
The submitnotebook docker image has the latest version of tensorflow and papermill installed. It invokes papermill on the supplied notebook.
End of explanation
"""


import kfp.components as comp
import kfp.dsl as dsl

# a single-op pipeline that runs the flights pipeline on the pod
@dsl.pipeline(
   name='FlightsPipeline',
   description='Trains, deploys flights model'
)
def flights_pipeline(
   inputnb=dsl.PipelineParam('inputnb'),
   outputnb=dsl.PipelineParam('outputnb'),
   params=dsl.PipelineParam('params')
):
    notebookop = dsl.ContainerOp(
      name='flightsmodel',
      # image needs to be a compile-time string
      image='gcr.io/cloud-training-demos/submitnotebook:latest',
      arguments=[
        inputnb,
        outputnb,
        params
      ]
    )

# compile the pipeline
pipeline_func = flights_pipeline
pipeline_filename = pipeline_func.__name__ + '.tar.gz'
import kfp.compiler as compiler
compiler.Compiler().compile(pipeline_func, pipeline_filename)
print(pipeline_filename)

!ls *.tar.gz

"""
Explanation: Create the pipeline
End of explanation
"""


%%writefile params.yaml
---
BUCKET: cloud-training-demos-ml
PROJECT: cloud-training-demos
DEVELOP_MODE: False

%%bash
gsutil -m rm -r gs://$BUCKET/flights/notebook
gsutil cp *.ipynb params.yaml gs://$BUCKET/flights/notebook

!gsutil ls gs://$BUCKET/flights/notebook

#Specify pipeline argument values
GCSDIR='gs://{}/flights/notebook'.format(BUCKET)
arguments = {
    'inputnb': '{}/flights_model.ipynb'.format(GCSDIR),
    'outputnb': '{}/flights_model_kfp.ipynb'.format(GCSDIR),
    'params': '{}/params.yaml'.format(GCSDIR),
}

#Get or create an experiment and submit a pipeline run
import kfp
client = kfp.Client()
list_experiments_response = client.list_experiments()
experiments = list_experiments_response.experiments
if not experiments:
    #The user does not have any experiments available. Creating a new one
    experiment = client.create_experiment('Flight pipeline experiment')
else:
    experiment = experiments[-1] #Using the last experiment

#Submit a pipeline run
from datetime import datetime
run_name = 'Flight pipeline {}'.format(datetime.now().strftime("%Y%m%d %H%M%S"))
run_result = client.run_pipeline(experiment.id, run_name, pipeline_filename, params=arguments)

!gsutil ls gs://$BUCKET/flights/notebook

"""
Explanation: Run the pipeline
End of explanation
"""


