%load_ext autoreload
%autoreload 2

!pip install -q -U kfp

"""
Explanation: Compile and deploy the TFX pipeline to Kubeflow Pipelines
This notebook is the second of two notebooks that guide you through automating the Real-time Item-to-item Recommendation with BigQuery ML Matrix Factorization and ScaNN solution with a pipeline.
Use this notebook to compile the TFX pipeline to a Kubeflow Pipelines (KFP) package. This process creates an Argo YAML file in a .tar.gz package, and is accomplished through the following steps:

Build a custom container image that includes the solution modules.
Compile the TFX Pipeline using the TFX command-line interface (CLI).
Deploy the compiled pipeline to KFP.

The pipeline workflow is implemented in the pipeline.py module. The runner.py module reads the configuration settings from the config.py module, defines the runtime parameters of the pipeline,  and creates a KFP format that is executable on AI Platform pipelines. 
Before starting this notebook, you must run the tfx01_interactive notebook to create the TFX pipeline.
Install required libraries
End of explanation
"""


import os

os.environ['PROJECT_ID'] = 'yourProject' # Set your project.
os.environ['BUCKET'] = 'yourBucket' # Set your bucket.
os.environ['GKE_CLUSTER_NAME'] = 'yourCluster' # Set your GKE cluster name.
os.environ['GKE_CLUSTER_ZONE'] = 'yourClusterZone' # Set your GKE cluster zone.

os.environ['IMAGE_NAME'] = 'tfx-ml'
os.environ['TAG'] = 'tfx0.25.0'
os.environ['ML_IMAGE_URI']=f'gcr.io/{os.environ.get("PROJECT_ID")}/{os.environ.get("IMAGE_NAME")}:{os.environ.get("TAG")}'

os.environ['NAMESPACE'] = 'kubeflow-pipelines'
os.environ['ARTIFACT_STORE_URI'] = f'gs://{os.environ.get("BUCKET")}/tfx_artifact_store'
os.environ['GCS_STAGING_PATH'] = f'{os.environ.get("ARTIFACT_STORE_URI")}/staging'

os.environ['RUNTIME_VERSION'] = '2.2'
os.environ['PYTHON_VERSION'] = '3.7'
os.environ['BEAM_RUNNER'] = 'DirectRunner'
os.environ['MODEL_REGISTRY_URI'] = f'{os.environ.get("ARTIFACT_STORE_URI")}/model_registry'

os.environ['PIPELINE_NAME'] = 'tfx_bqml_scann'

from tfx_pipeline import config

for key, value in config.__dict__.items():
  if key.isupper(): print(f'{key}: {value}')

"""
Explanation: Set environment variables
Update the following variables to reflect the values for your GCP environment:

PROJECT_ID: The ID of the Google Cloud project you are using to implement this solution.
BUCKET: The name of the Cloud Storage bucket you created to use with this solution. The BUCKET value should be just the bucket name, so myBucket rather than gs://myBucket.
GKE_CLUSTER_NAME: The name of the Kubernetes Engine cluster used by the AI Platform pipeline. You can find this by looking at the Cluster column of the kubeflow-pipelines pipeline instance on the AI Platform Pipelines page.
GKE_CLUSTER_ZONE: The zone of the Kubernetes Engine cluster used by the AI Platform pipeline. You can find this by looking at the Zone column of the kubeflow-pipelines pipeline instance on the AI Platform Pipelines page.
End of explanation
"""


import kfp
import tfx
from tfx.orchestration.beam.beam_dag_runner import BeamDagRunner
from tfx_pipeline import pipeline as pipeline_module
import tensorflow as tf
import ml_metadata as mlmd
from ml_metadata.proto import metadata_store_pb2
import logging

logging.getLogger().setLevel(logging.INFO)

print("TFX Version:", tfx.__version__)

pipeline_root = f'{config.ARTIFACT_STORE_URI}/{config.PIPELINE_NAME}_beamrunner'
model_regisrty_uri = f'{config.MODEL_REGISTRY_URI}_beamrunner'
local_mlmd_sqllite = 'mlmd/mlmd.sqllite'

print(f'Pipeline artifacts root: {pipeline_root}')    
print(f'Model registry location: {model_regisrty_uri}')  

if tf.io.gfile.exists(pipeline_root):
  print("Removing previous artifacts...")
  tf.io.gfile.rmtree(pipeline_root)
if tf.io.gfile.exists('mlmd'):
  print("Removing local mlmd SQLite...")
  tf.io.gfile.rmtree('mlmd')
print("Creating mlmd directory...")
tf.io.gfile.mkdir('mlmd')

metadata_connection_config = metadata_store_pb2.ConnectionConfig()
metadata_connection_config.sqlite.filename_uri = local_mlmd_sqllite
metadata_connection_config.sqlite.connection_mode = 3
print("ML metadata store is ready.")

beam_pipeline_args = [
  f'--runner=DirectRunner',
  f'--project={config.PROJECT_ID}',
  f'--temp_location={config.ARTIFACT_STORE_URI}/beam/tmp'
]

pipeline_module.SCHEMA_DIR = 'tfx_pipeline/schema'
pipeline_module.LOOKUP_CREATOR_MODULE = 'tfx_pipeline/lookup_creator.py'
pipeline_module.SCANN_INDEXER_MODULE = 'tfx_pipeline/scann_indexer.py'

runner = BeamDagRunner()

pipeline = pipeline_module.create_pipeline(
  pipeline_name=config.PIPELINE_NAME,
  pipeline_root=pipeline_root,
  project_id=config.PROJECT_ID,
  bq_dataset_name=config.BQ_DATASET_NAME,
  min_item_frequency=15,
  max_group_size=10,
  dimensions=50,
  num_leaves=500,
  eval_min_recall=0.8,
  eval_max_latency=0.001,
  ai_platform_training_args=None,
  beam_pipeline_args=beam_pipeline_args,
  model_regisrty_uri=model_regisrty_uri,
  metadata_connection_config=metadata_connection_config,
  enable_cache=True
)

runner.run(pipeline)

"""
Explanation: Run the Pipeline locally by using the Beam runner
End of explanation
"""


!gcloud builds submit --tag $ML_IMAGE_URI tfx_pipeline

"""
Explanation: Build the container image
The pipeline uses a custom container image, which is a derivative of the tensorflow/tfx:0.25.0 image, as a runtime execution environment for the pipeline's components. The container image is defined in a Dockerfile.
The container image installs the required libraries and copies over the modules from the solution's tfx_pipeline directory, where the custom components are implemented. The container image is also used by AI Platform Training for executing the training jobs. 
Build the container image using Cloud Build and then store it in Cloud Container Registry:
End of explanation
"""


!rm ${PIPELINE_NAME}.tar.gz
!tfx pipeline compile \
    --engine=kubeflow \
    --pipeline_path=tfx_pipeline/runner.py 

"""
Explanation: Compile the TFX pipeline using the TFX CLI
Use the TFX CLI to compile the TFX pipeline to the KFP format, which allows the pipeline to be deployed and executed on AI Platform Pipelines. The output is a .tar.gz package containing an Argo definition of your pipeline.
End of explanation
"""


%%bash

gcloud container clusters get-credentials ${GKE_CLUSTER_NAME} --zone ${GKE_CLUSTER_ZONE}
export KFP_ENDPOINT=$(kubectl describe configmap inverse-proxy-config -n ${NAMESPACE} | grep "googleusercontent.com")

kfp --namespace=${NAMESPACE} --endpoint=${KFP_ENDPOINT} \
    pipeline upload \
    --pipeline-name=${PIPELINE_NAME} \
    ${PIPELINE_NAME}.tar.gz

"""
Explanation: Deploy the compiled pipeline to KFP
Use the KFP CLI to deploy the pipeline to a hosted instance of KFP on AI Platform Pipelines:
End of explanation
"""


