import yaml

# Set `PATH` to include the directory containing TFX CLI and skaffold.
PATH=%env PATH
%env PATH=/home/jupyter/.local/bin:{PATH}

"""
Explanation: Continuous training with TFX and Google Cloud AI Platform
Learning Objectives

Use the TFX CLI to build a TFX pipeline.
Deploy a TFX pipeline version without tuning to a hosted AI Platform Pipelines instance.
Create and monitor a TFX pipeline run using the TFX CLI.
Deploy a new TFX pipeline version with tuning enabled to a hosted AI Platform Pipelines instance.
Create and monitor another TFX pipeline run directly in the KFP UI.

In this lab, you use utilize the following tools and services to deploy and run a TFX pipeline on Google Cloud that automates the development and deployment of a TensorFlow 2.3 WideDeep Classifer to predict forest cover from cartographic data:

The TFX CLI utility to build and deploy a TFX pipeline.
A hosted AI Platform Pipeline instance (Kubeflow Pipelines) for TFX pipeline orchestration.
Dataflow jobs for scalable, distributed data processing for TFX components.
A AI Platform Training job for model training and flock management for parallel tuning trials. 
AI Platform Prediction as a model server destination for blessed pipeline model versions.
CloudTuner and AI Platform Vizier for advanced model hyperparameter tuning using the Vizier algorithm.

You will then create and monitor pipeline runs using the TFX CLI as well as the KFP UI.
Setup
Update lab environment PATH to include TFX CLI and skaffold
End of explanation
"""


!python -c "import tfx; print('TFX version: {}'.format(tfx.__version__))"
!python -c "import kfp; print('KFP version: {}'.format(kfp.__version__))"

"""
Explanation: Validate lab package version installation
End of explanation
"""


%pip install --upgrade --user tfx==0.25.0
%pip install --upgrade --user kfp==1.0.4

"""
Explanation: Note: this lab was built and tested with the following package versions:
TFX version: 0.25.0
KFP version: 1.0.4
(Optional) If running the above command results in different package versions or you receive an import error, upgrade to the correct versions by running the cell below:
End of explanation
"""


%cd pipeline

!ls -la

"""
Explanation: Note: you may need to restart the kernel to pick up the correct package versions.
Validate creation of AI Platform Pipelines cluster
Navigate to AI Platform Pipelines page in the Google Cloud Console.
Note you may have already deployed an AI Pipelines instance during the Setup for the lab series. If so, you can proceed using that instance. If not:
1.  Create or select an existing Kubernetes cluster (GKE) and deploy AI Platform. Make sure to select "Allow access to the following Cloud APIs https://www.googleapis.com/auth/cloud-platform" to allow for programmatic access to your pipeline by the Kubeflow SDK for the rest of the lab. Also, provide an App instance name such as "tfx" or "mlops". 
Validate the deployment of your AI Platform Pipelines instance in the console before proceeding.
Review: example TFX pipeline design pattern for Google Cloud
The pipeline source code can be found in the pipeline folder.
End of explanation
"""


# Use the following command to identify the GCS bucket for metadata and pipeline storage.
!gsutil ls

"""
Explanation: The config.py module configures the default values for the environment specific settings and the default values for the pipeline runtime parameters. 
The default values can be overwritten at compile time by providing the updated values in a set of environment variables. You will set custom environment variables later on this lab.
The pipeline.py module contains the TFX DSL defining the workflow implemented by the pipeline.
The preprocessing.py module implements the data preprocessing logic  the Transform component.
The model.py module implements the training, tuning, and model building logic for the Trainer and Tuner components.
The runner.py module configures and executes KubeflowDagRunner. At compile time, the KubeflowDagRunner.run() method converts the TFX DSL into the pipeline package in the argo format for execution on your hosted AI Platform Pipelines instance.
The features.py module contains feature definitions common across preprocessing.py and model.py.
Exercise: build your pipeline with the TFX CLI
You will use TFX CLI to compile and deploy the pipeline. As explained in the previous section, the environment specific settings can be provided through a set of environment variables and embedded into the pipeline package at compile time.
Configure your environment resource settings
Update  the below constants  with the settings reflecting your lab environment. 

GCP_REGION - the compute region for AI Platform Training, Vizier, and Prediction.
ARTIFACT_STORE - An existing GCS bucket. You can use any bucket or use the GCS bucket created during installation of AI Platform Pipelines. The default bucket name will contain the kubeflowpipelines- prefix.
End of explanation
"""


#TODO: Set your environment resource settings here for GCP_REGION, ARTIFACT_STORE_URI, ENDPOINT, and CUSTOM_SERVICE_ACCOUNT.
GCP_REGION = ''
ARTIFACT_STORE_URI = ''
ENDPOINT = ''
CUSTOM_SERVICE_ACCOUNT = ''

PROJECT_ID = !(gcloud config get-value core/project)
PROJECT_ID = PROJECT_ID[0]

# Set your resource settings as environment variables. These override the default values in pipeline/config.py.
%env GCP_REGION={GCP_REGION}
%env ARTIFACT_STORE_URI={ARTIFACT_STORE_URI}
%env CUSTOM_SERVICE_ACCOUNT={CUSTOM_SERVICE_ACCOUNT}
%env PROJECT_ID={PROJECT_ID}

"""
Explanation: CUSTOM_SERVICE_ACCOUNT - In the gcp console Click on the Navigation Menu. Navigate to IAM &amp; Admin, then to Service Accounts and use the service account starting with prefix - 'tfx-tuner-caip-service-account'. This enables CloudTuner and the Google Cloud AI Platform extensions Tuner component to work together and allows for distributed and parallel tuning backed by AI Platform Vizier's hyperparameter search algorithm. Please see the lab setup README for setup instructions.


ENDPOINT - set the ENDPOINT constant to the endpoint to your AI Platform Pipelines instance. The endpoint to the AI Platform Pipelines instance can be found on the AI Platform Pipelines page in the Google Cloud Console. Open the SETTINGS for your instance and use the value of the host variable in the Connect to this Kubeflow Pipelines instance from a Python client via Kubeflow Pipelines SKD section of the SETTINGS window. The format is '...pipelines.googleusercontent.com'.
End of explanation
"""


PIPELINE_NAME = 'tfx_covertype_continuous_training'
MODEL_NAME = 'tfx_covertype_classifier'
DATA_ROOT_URI = 'gs://workshop-datasets/covertype/small'
CUSTOM_TFX_IMAGE = 'gcr.io/{}/{}'.format(PROJECT_ID, PIPELINE_NAME)
RUNTIME_VERSION = '2.3'
PYTHON_VERSION = '3.7'
USE_KFP_SA=False
ENABLE_TUNING=False

%env PIPELINE_NAME={PIPELINE_NAME}
%env MODEL_NAME={MODEL_NAME}
%env DATA_ROOT_URI={DATA_ROOT_URI}
%env KUBEFLOW_TFX_IMAGE={CUSTOM_TFX_IMAGE}
%env RUNTIME_VERSION={RUNTIME_VERSION}
%env PYTHON_VERIONS={PYTHON_VERSION}
%env USE_KFP_SA={USE_KFP_SA}
%env ENABLE_TUNING={ENABLE_TUNING}

"""
Explanation: Set the compile time settings to first create a pipeline version without hyperparameter tuning
Default pipeline runtime environment values are configured in the pipeline folder config.py. You will set their values directly below:


PIPELINE_NAME - the pipeline's globally unique name. For each pipeline update, each pipeline version uploaded to KFP will be reflected on the Pipelines tab in the Pipeline name &gt; Version name dropdown in the format PIPELINE_NAME_datetime.now().


MODEL_NAME - the pipeline's unique model output name for AI Platform Prediction. For multiple pipeline runs, each pushed blessed model will create a new version with the format 'v{}'.format(int(time.time())).


DATA_ROOT_URI - the URI for the raw lab dataset gs://workshop-datasets/covertype/small.


CUSTOM_TFX_IMAGE - the image name of your pipeline container build by skaffold and published by Cloud Build to Cloud Container Registry in the format 'gcr.io/{}/{}'.format(PROJECT_ID, PIPELINE_NAME).


RUNTIME_VERSION - the TensorFlow runtime version. This lab was built and tested using TensorFlow 2.3.


PYTHON_VERSION - the Python runtime version. This lab was built and tested using Python 3.7.


USE_KFP_SA - The pipeline can run using a security context of the GKE default node pool's service account or the service account defined in the user-gcp-sa secret of the Kubernetes namespace hosting Kubeflow Pipelines. If you want to use the user-gcp-sa service account you change the value of USE_KFP_SA to True. Note that the default AI Platform Pipelines configuration does not define the user-gcp-sa secret.


ENABLE_TUNING - boolean value indicating whether to add the Tuner component to the pipeline or use hyperparameter defaults. See the model.py and pipeline.py files for details on how this changes the pipeline topology across pipeline versions. You will create pipeline versions without and with tuning enabled in the subsequent lab exercises for comparison.
End of explanation
"""


!tfx pipeline compile --engine kubeflow --pipeline_path runner.py

"""
Explanation: Compile your pipeline code
You can build and upload the pipeline to the AI Platform Pipelines instance in one step, using the tfx pipeline create command. The tfx pipeline create goes through the following steps:
- (Optional) Builds the custom image to that provides a runtime environment for TFX components or uses the latest image of the installed TFX version 
- Compiles the pipeline code into a pipeline package 
- Uploads the pipeline package via the ENDPOINT to the hosted AI Platform instance.
As you debug the pipeline DSL, you may prefer to first use the tfx pipeline compile command, which only executes the compilation step. After the DSL compiles successfully you can use tfx pipeline create to go through all steps.
End of explanation
"""


# TODO: Your code here to use the TFX CLI to deploy your pipeline image to AI Platform Pipelines.


"""
Explanation: Note: you should see a {PIPELINE_NAME}.tar.gz file appear in your current pipeline directory.
Exercise: deploy your pipeline container to AI Platform Pipelines with TFX CLI
After the pipeline code compiles without any errors you can use the tfx pipeline create command to perform the full build and deploy the pipeline. You will deploy your compiled pipeline container hosted on Google Container Registry e.g. gcr.io/[PROJECT_ID]/tfx_covertype_continuous_training to run on AI Platform Pipelines with the TFX CLI.
End of explanation
"""


# TODO: your code here to trigger a pipeline run with the TFX CLI


"""
Explanation: Hint: review the TFX CLI documentation on the "pipeline group" to create your pipeline. You will need to specify the --pipeline_path to point at the pipeline DSL and runner defined locally in runner.py, --endpoint, and --build_target_image arguments using the environment variables specified above.
Note: you should see a build.yaml file in your pipeline folder created by skaffold. The TFX CLI compile triggers a custom container to be built with skaffold using the instructions in the Dockerfile.
If you need to redeploy the pipeline you can first delete the previous version using tfx pipeline delete or you can update the pipeline in-place using tfx pipeline update.
To delete the pipeline:
tfx pipeline delete --pipeline_name {PIPELINE_NAME} --endpoint {ENDPOINT}
To update the pipeline:
tfx pipeline update --pipeline_path runner.py --endpoint {ENDPOINT}
Create and monitor a pipeline run with the TFX CLI
After the pipeline has been deployed, you can trigger and monitor pipeline runs using TFX CLI.
Hint: review the TFX CLI documentation on the "run group".
End of explanation
"""


!tfx run list --pipeline_name {PIPELINE_NAME} --endpoint {ENDPOINT}

"""
Explanation: To view the status of existing pipeline runs:
End of explanation
"""


RUN_ID='[YOUR RUN ID]'

!tfx run status --pipeline_name {PIPELINE_NAME} --run_id {RUN_ID} --endpoint {ENDPOINT}

"""
Explanation: To retrieve the status of a given run:
End of explanation
"""


ENABLE_TUNING=True

%env ENABLE_TUNING={ENABLE_TUNING}

"""
Explanation: Important
A full pipeline run without tuning enabled will take about 40 minutes to complete. You can view the run's progress using the TFX CLI commands above or in the
Exercise: deploy a pipeline version with tuning enabled
Incorporating automatic model hyperparameter tuning into a continuous training TFX pipeline workflow enables faster experimentation, development, and deployment of a top performing model.
The previous pipeline version read from hyperparameter default values in the search space defined in _get_hyperparameters() in model.py and used these values to build a TensorFlow WideDeep Classifier model.
Let's now deploy a new pipeline version with the Tuner component added to the pipeline that calls out to the AI Platform Vizier service for distributed and parallelized hyperparameter tuning. The Tuner component "best_hyperparameters" artifact will be passed directly to your Trainer component to deploy the top performing model. Review pipeline.py to see how this environment variable changes the pipeline topology. Also, review the tuning function in model.py for configuring CloudTuner.
Note that you might not want to tune the hyperparameters every time you retrain your model due to the computational cost. Once you have used Tuner determine a good set of hyperparameters, you can remove Tuner from your pipeline and use model hyperparameters defined in your model code or use a ImporterNode to import the Tuner "best_hyperparameters"artifact from a previous Tuner run to your model Trainer.
End of explanation
"""


!tfx pipeline compile --engine kubeflow --pipeline_path runner.py

"""
Explanation: Compile your pipeline code
End of explanation
"""


#TODO: your code to update your pipeline 


"""
Explanation: Deploy your pipeline container to AI Platform Pipelines with the TFX CLI
End of explanation
"""


