from google.cloud import aiplatform as vertex_ai

"""
Explanation: Continuous training with TFX and Vertex
Learning Objectives

Containerize your TFX code into a pipeline package using Cloud Build.
Use the TFX CLI to compile a TFX pipeline.
Deploy a TFX pipeline version to run on Vertex Pipelines using the Vertex Python SDK.

Setup
End of explanation
"""


!python -c "import tensorflow as tf; print(f'TF version: {tf.__version__}')"
!python -c "import tfx; print(f'TFX version: {tfx.__version__}')"
!python -c "import kfp; print(f'KFP version: {kfp.__version__}')"
print(f"vertex_ai: {vertex_ai.__version__}")

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


%cd pipeline_vertex

!ls -la

"""
Explanation: Note: this lab was built and tested with the following package versions:
TF version: 2.6.2
TFX version: 1.4.0 
KFP version: 1.8.1
aiplatform: 1.7.1
Review: example TFX pipeline design pattern for Vertex
The pipeline source code can be found in the pipeline_vertex folder.
End of explanation
"""


# TODO: Set your environment resource settings here for GCP_REGION, ARTIFACT_STORE_URI, ENDPOINT, and CUSTOM_SERVICE_ACCOUNT.
REGION = "us-central1"
PROJECT_ID = !(gcloud config get-value core/project)
PROJECT_ID = PROJECT_ID[0]
ARTIFACT_STORE = f"gs://{PROJECT_ID}"

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

!gsutil ls | grep ^{ARTIFACT_STORE}/$ || gsutil mb -l {REGION} {ARTIFACT_STORE}

"""
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 TensorFlow model code and training logic for the Trainer component.
The runner.py module configures and executes KubeflowV2DagRunner. At compile time, the KubeflowDagRunner.run() method converts the TFX DSL into the pipeline package into a JSON format for execution on Vertex.
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. 

REGION - the compute region for AI Platform Training, Vizier, and Prediction.
ARTIFACT_STORE - An existing GCS bucket. You can use any bucket, but we will use here the bucket with the same name as the project.
End of explanation
"""


PIPELINE_NAME = "tfxcovertype"
DATA_ROOT_URI = f"gs://{PROJECT_ID}/data/tfxcovertype"
TFX_IMAGE_URI = f"gcr.io/{PROJECT_ID}/{PIPELINE_NAME}"
PIPELINE_JSON = f"{PIPELINE_NAME}.json"

TRAIN_STEPS = 10
EVAL_STEPS = 5

%env PIPELINE_NAME={PIPELINE_NAME}
%env DATA_ROOT_URI={DATA_ROOT_URI}
%env TFX_IMAGE_URI={TFX_IMAGE_URI}
%env PIPELINE_JSON={PIPELINE_JSON}
%env TRAIN_STEPS={TRAIN_STEPS}
%env EVAL_STEPS={EVAL_STEPS}

"""
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.


DATA_ROOT_URI - the URI for the raw lab dataset gs://{PROJECT_ID}/data/tfxcovertype.


TFX_IMAGE_URI - the image name of your pipeline container that will be used to execute each of your tfx components
End of explanation
"""


!gsutil cp  ../../../data/* $DATA_ROOT_URI/dataset.csv
!gsutil ls $DATA_ROOT_URI/*

"""
Explanation: Let us populate the data bucket at DATA_ROOT_URI:
End of explanation
"""


!gcloud builds submit --timeout 15m --tag $TFX_IMAGE_URI .

"""
Explanation: Let us build and push the TFX container image described in the Dockerfile:
End of explanation
"""


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

"""
Explanation: Compile your pipeline code
The following command will execute the KubeflowV2DagRunner that compiles the pipeline described in pipeline.py into a JSON representation consumable by Vertex:
End of explanation
"""


vertex_ai.init(project=PROJECT_ID, location=REGION)

pipeline = vertex_ai.PipelineJob(
    display_name="tfxcovertype4",
    template_path=PIPELINE_JSON,
    enable_caching=False,
)

pipeline.run()

"""
Explanation: Note: you should see a {PIPELINE_NAME}.json file appear in your current pipeline directory.
Exercise: deploy your pipeline on Vertex using the Vertex SDK
Once you have the {PIPELINE_NAME}.json available, you can run the tfx pipeline on Vertex by launching a pipeline job using the aiplatform handle:
End of explanation
"""


