import os

# Google Cloud Notebook
if os.path.exists("/opt/deeplearning/metadata/env_version"):
    USER_FLAG = "--user"
else:
    USER_FLAG = ""

! pip3 install --upgrade google-cloud-aiplatform $USER_FLAG

"""
Explanation: Vertex AI Pipelines:  model train, upload, and deploy using google-cloud-pipeline-components
<table align="left">
  <td>
    <a href="https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/official/pipelines/google_cloud_pipeline_components_model_train_upload_deploy.ipynb">
      <img src="https://cloud.google.com/ml-engine/images/colab-logo-32px.png" alt="Colab logo"> Run in Colab
    </a>
  </td>
  <td>
    <a href="https://github.com/GoogleCloudPlatform/vertex-ai-samples/notebooks/blob/master/official/pipelines/google_cloud_pipeline_components_model_train_upload_deploy.ipynb">
      <img src="https://cloud.google.com/ml-engine/images/github-logo-32px.png" alt="GitHub logo">
      View on GitHub
    </a>
  </td>
  <td>
    <a href="https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://github.com/GoogleCloudPlatform/vertex-ai-samples/notebooks/blob/master/official/pipelines/google_cloud_pipeline_components_model_train_upload_deploy.ipynb">
      Open in Vertex AI Workbench
    </a>
  </td>
</table>
<br/><br/><br/>
Overview
This notebook shows how to use the components defined in google_cloud_pipeline_components in conjunction with an experimental run_as_aiplatform_custom_job method, to build a Vertex AI Pipelines workflow that trains a custom model, uploads the model as a Model resource, creates an Endpoint resource, and deploys the Model resource to the Endpoint resource.
Dataset
The dataset used for this tutorial is Cloud Public Dataset Program London Bikes Rental combined with NOAA weather data 
The dataset predicts the duration of the bike rental.
Objective
In this tutorial, you create an custom model using a pipeline with components from google_cloud_pipeline_components and a custom pipeline component you build.
In addition, you'll use the kfp.v2.google.experimental.run_as_aiplatform_custom_job method to train a custom model.
The steps performed include:

Train a custom model.
Uploads the trained model as a Model resource.
Creates an Endpoint resource.
Deploys the Model resource to the Endpoint resource.

The components are documented here.
(From that page, see also the CustomPythonPackageTrainingJobRunOp and CustomContainerTrainingJobRunOp components, which similarly run 'custom' training, but as with the related google.cloud.aiplatform.CustomContainerTrainingJob and google.cloud.aiplatform.CustomPythonPackageTrainingJob methods from the Vertex AI SDK, also upload the trained model).
Costs
This tutorial uses billable components of Google Cloud:

Vertex AI
Cloud Storage

Learn about Vertex AI
pricing and Cloud Storage
pricing, and use the Pricing
Calculator
to generate a cost estimate based on your projected usage.
Set up your local development environment
If you are using Colab or Google Cloud Notebook, your environment already meets all the requirements to run this notebook. You can skip this step.
Otherwise, make sure your environment meets this notebook's requirements. You need the following:

The Cloud Storage SDK
Git
Python 3
virtualenv
Jupyter notebook running in a virtual environment with Python 3

The Cloud Storage guide to Setting up a Python development environment and the Jupyter installation guide provide detailed instructions for meeting these requirements. The following steps provide a condensed set of instructions:


Install and initialize the SDK.


Install Python 3.


Install virtualenv and create a virtual environment that uses Python 3.


Activate that environment and run pip3 install Jupyter in a terminal shell to install Jupyter.


Run jupyter notebook on the command line in a terminal shell to launch Jupyter.


Open this notebook in the Jupyter Notebook Dashboard.


Installation
Install the latest version of Vertex SDK for Python.
End of explanation
"""


! pip3 install -U google-cloud-storage $USER_FLAG

"""
Explanation: Install the latest GA version of google-cloud-storage library as well.
End of explanation
"""


! pip3 install $USER kfp google-cloud-pipeline-components --upgrade

"""
Explanation: Install the latest GA version of google-cloud-pipeline-components library as well.
End of explanation
"""


import os

if not os.getenv("IS_TESTING"):
    # Automatically restart kernel after installs
    import IPython

    app = IPython.Application.instance()
    app.kernel.do_shutdown(True)

"""
Explanation: Restart the kernel
Once you've installed the additional packages, you need to restart the notebook kernel so it can find the packages.
End of explanation
"""


! python3 -c "import kfp; print('KFP SDK version: {}'.format(kfp.__version__))"
! python3 -c "import google_cloud_pipeline_components; print('google_cloud_pipeline_components version: {}'.format(google_cloud_pipeline_components.__version__))"

"""
Explanation: Check the versions of the packages you installed.  The KFP SDK version should be >=1.6.
End of explanation
"""


PROJECT_ID = "[your-project-id]"  # @param {type:"string"}

if PROJECT_ID == "" or PROJECT_ID is None or PROJECT_ID == "[your-project-id]":
    # Get your GCP project id from gcloud
    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null
    PROJECT_ID = shell_output[0]
    print("Project ID:", PROJECT_ID)

! gcloud config set project $PROJECT_ID

"""
Explanation: Before you begin
GPU runtime
This tutorial does not require a GPU runtime.
Set up your Google Cloud project
The following steps are required, regardless of your notebook environment.


Select or create a Google Cloud project. When you first create an account, you get a $300 free credit towards your compute/storage costs.


Make sure that billing is enabled for your project.


Enable the Vertex AI APIs, Compute Engine APIs, and Cloud Storage.


The Google Cloud SDK is already installed in Google Cloud Notebook.


Enter your project ID in the cell below. Then run the  cell to make sure the
Cloud SDK uses the right project for all the commands in this notebook.


Note: Jupyter runs lines prefixed with ! as shell commands, and it interpolates Python variables prefixed with $.
End of explanation
"""


REGION = "us-central1"  # @param {type: "string"}

"""
Explanation: Region
You can also change the REGION variable, which is used for operations
throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.

Americas: us-central1
Europe: europe-west4
Asia Pacific: asia-east1

You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.
Learn more about Vertex AI regions
End of explanation
"""


from datetime import datetime

TIMESTAMP = datetime.now().strftime("%Y%m%d%H%M%S")

"""
Explanation: Timestamp
If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append the timestamp onto the name of resources you create in this tutorial.
End of explanation
"""


# If you are running this notebook in Colab, run this cell and follow the
# instructions to authenticate your GCP account. This provides access to your
# Cloud Storage bucket and lets you submit training jobs and prediction
# requests.

import os
import sys

# If on Google Cloud Notebook, then don't execute this code
if not os.path.exists("/opt/deeplearning/metadata/env_version"):
    if "google.colab" in sys.modules:
        from google.colab import auth as google_auth

        google_auth.authenticate_user()

    # If you are running this notebook locally, replace the string below with the
    # path to your service account key and run this cell to authenticate your GCP
    # account.
    elif not os.getenv("IS_TESTING"):
        %env GOOGLE_APPLICATION_CREDENTIALS ''

"""
Explanation: Authenticate your Google Cloud account
If you are using Google Cloud Notebook, your environment is already authenticated. Skip this step.
If you are using Colab, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.
Otherwise, follow these steps:
In the Cloud Console, go to the Create service account key page.
Click Create service account.
In the Service account name field, enter a name, and click Create.
In the Grant this service account access to project section, click the Role drop-down list. Type "Vertex" into the filter box, and select Vertex Administrator. Type "Storage Object Admin" into the filter box, and select Storage Object Admin.
Click Create. A JSON file that contains your key downloads to your local environment.
Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell.
End of explanation
"""


BUCKET_NAME = "gs://[your-bucket-name]"  # @param {type:"string"}

if BUCKET_NAME == "" or BUCKET_NAME is None or BUCKET_NAME == "gs://[your-bucket-name]":
    BUCKET_NAME = "gs://" + PROJECT_ID + "aip-" + TIMESTAMP

"""
Explanation: Create a Cloud Storage bucket
The following steps are required, regardless of your notebook environment.
When you initialize the Vertex AI SDK for Python, you specify a Cloud Storage staging bucket. The staging bucket is where all the data associated with your dataset and model resources are retained across sessions.
Set the name of your Cloud Storage bucket below. Bucket names must be globally unique across all Google Cloud projects, including those outside of your organization.
End of explanation
"""


! gsutil mb -l $REGION $BUCKET_NAME

"""
Explanation: Only if your bucket doesn't already exist: Run the following cell to create your Cloud Storage bucket.
End of explanation
"""


! gsutil ls -al $BUCKET_NAME

"""
Explanation: Finally, validate access to your Cloud Storage bucket by examining its contents:
End of explanation
"""


SERVICE_ACCOUNT = "[your-service-account]"  # @param {type:"string"}

if (
    SERVICE_ACCOUNT == ""
    or SERVICE_ACCOUNT is None
    or SERVICE_ACCOUNT == "[your-service-account]"
):
    # Get your GCP project id from gcloud
    shell_output = !gcloud auth list 2>/dev/null
    SERVICE_ACCOUNT = shell_output[2].strip()
    print("Service Account:", SERVICE_ACCOUNT)

"""
Explanation: Service Account
If you don't know your service account, try to get your service account using gcloud command by executing the second cell below.
End of explanation
"""


! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_NAME

! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_NAME

"""
Explanation: Set service account access for Vertex AI Pipelines
Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step -- you only need to run these once per service account.
End of explanation
"""


import google.cloud.aiplatform as aip

"""
Explanation: Set up variables
Next, set up some variables used throughout the tutorial.
Import libraries and define constants
End of explanation
"""


PIPELINE_ROOT = "{}/pipeline_root/bikes_weather".format(BUCKET_NAME)

"""
Explanation: Vertex AI Pipelines constants
Setup up the following constants for Vertex AI Pipelines:
End of explanation
"""


import kfp
from kfp.v2.dsl import component

"""
Explanation: Additional imports.
End of explanation
"""


aip.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)

"""
Explanation: Initialize Vertex AI SDK for Python
Initialize the Vertex AI SDK for Python for your project and corresponding bucket.
End of explanation
"""


hp_dict: str = '{"num_hidden_layers": 3, "hidden_size": 32, "learning_rate": 0.01, "epochs": 1, "steps_per_epoch": -1}'
data_dir: str = "gs://aju-dev-demos-codelabs/bikes_weather/"
TRAINER_ARGS = ["--data-dir", data_dir, "--hptune-dict", hp_dict]

# create working dir to pass to job spec
WORKING_DIR = f"{PIPELINE_ROOT}/{TIMESTAMP}"

MODEL_DISPLAY_NAME = f"train_deploy{TIMESTAMP}"
print(TRAINER_ARGS, WORKING_DIR, MODEL_DISPLAY_NAME)


@kfp.dsl.pipeline(name="train-endpoint-deploy" + TIMESTAMP)
def pipeline(
    project: str = PROJECT_ID,
    model_display_name: str = MODEL_DISPLAY_NAME,
    serving_container_image_uri: str = "us-docker.pkg.dev/cloud-aiplatform/prediction/tf2-cpu.2-3:latest",
):
    from google_cloud_pipeline_components.types import artifact_types
    from google_cloud_pipeline_components.v1.custom_job import \
        CustomTrainingJobOp
    from google_cloud_pipeline_components.v1.endpoint import (EndpointCreateOp,
                                                              ModelDeployOp)
    from google_cloud_pipeline_components.v1.model import ModelUploadOp
    from kfp.v2.components import importer_node

    custom_job_task = CustomTrainingJobOp(
        project=project,
        display_name="model-training",
        worker_pool_specs=[
            {
                "containerSpec": {
                    "args": TRAINER_ARGS,
                    "env": [{"name": "AIP_MODEL_DIR", "value": WORKING_DIR}],
                    "imageUri": "gcr.io/google-samples/bw-cc-train:latest",
                },
                "replicaCount": "1",
                "machineSpec": {
                    "machineType": "n1-standard-16",
                    "accelerator_type": aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,
                    "accelerator_count": 2,
                },
            }
        ],
    )

    import_unmanaged_model_task = importer_node.importer(
        artifact_uri=WORKING_DIR,
        artifact_class=artifact_types.UnmanagedContainerModel,
        metadata={
            "containerSpec": {
                "imageUri": "us-docker.pkg.dev/cloud-aiplatform/prediction/tf2-cpu.2-3:latest",
            },
        },
    ).after(custom_job_task)

    model_upload_op = ModelUploadOp(
        project=project,
        display_name=model_display_name,
        unmanaged_container_model=import_unmanaged_model_task.outputs["artifact"],
    )
    model_upload_op.after(import_unmanaged_model_task)

    endpoint_create_op = EndpointCreateOp(
        project=project,
        display_name="pipelines-created-endpoint",
    )

    ModelDeployOp(
        endpoint=endpoint_create_op.outputs["endpoint"],
        model=model_upload_op.outputs["model"],
        deployed_model_display_name=model_display_name,
        dedicated_resources_machine_type="n1-standard-16",
        dedicated_resources_min_replica_count=1,
        dedicated_resources_max_replica_count=1,
    )

"""
Explanation: Define custom model pipeline that uses components from google_cloud_pipeline_components
Next, you define the pipeline.
The experimental.run_as_aiplatform_custom_job method takes as arguments the previously defined component, and the list of worker_pool_specs— in this case one— with which the custom training job is configured.
Then, google_cloud_pipeline_components components are used to define the rest of the pipeline: upload the model, create an endpoint, and deploy the model to the endpoint.
Note: While not shown in this example, the model deploy will create an endpoint if one is not provided.
End of explanation
"""


from kfp.v2 import compiler  # noqa: F811

compiler.Compiler().compile(
    pipeline_func=pipeline,
    package_path="tabular regression_pipeline.json".replace(" ", "_"),
)

"""
Explanation: Compile the pipeline
Next, compile the pipeline.
End of explanation
"""


DISPLAY_NAME = "bikes_weather_" + TIMESTAMP

job = aip.PipelineJob(
    display_name=DISPLAY_NAME,
    template_path="tabular regression_pipeline.json".replace(" ", "_"),
    pipeline_root=PIPELINE_ROOT,
    enable_caching=False,
)

job.run()

! rm tabular_regression_pipeline.json

"""
Explanation: Run the pipeline
Next, run the pipeline.
End of explanation
"""


delete_dataset = True
delete_pipeline = True
delete_model = True
delete_endpoint = True
delete_batchjob = True
delete_customjob = True
delete_hptjob = True
delete_bucket = True

try:
    if delete_model and "DISPLAY_NAME" in globals():
        models = aip.Model.list(
            filter=f"display_name={DISPLAY_NAME}", order_by="create_time"
        )
        model = models[0]
        aip.Model.delete(model)
        print("Deleted model:", model)
except Exception as e:
    print(e)

try:
    if delete_endpoint and "DISPLAY_NAME" in globals():
        endpoints = aip.Endpoint.list(
            filter=f"display_name={DISPLAY_NAME}_endpoint", order_by="create_time"
        )
        endpoint = endpoints[0]
        endpoint.undeploy_all()
        aip.Endpoint.delete(endpoint.resource_name)
        print("Deleted endpoint:", endpoint)
except Exception as e:
    print(e)

if delete_dataset and "DISPLAY_NAME" in globals():
    if "tabular" == "tabular":
        try:
            datasets = aip.TabularDataset.list(
                filter=f"display_name={DISPLAY_NAME}", order_by="create_time"
            )
            dataset = datasets[0]
            aip.TabularDataset.delete(dataset.resource_name)
            print("Deleted dataset:", dataset)
        except Exception as e:
            print(e)

    if "tabular" == "image":
        try:
            datasets = aip.ImageDataset.list(
                filter=f"display_name={DISPLAY_NAME}", order_by="create_time"
            )
            dataset = datasets[0]
            aip.ImageDataset.delete(dataset.resource_name)
            print("Deleted dataset:", dataset)
        except Exception as e:
            print(e)

    if "tabular" == "text":
        try:
            datasets = aip.TextDataset.list(
                filter=f"display_name={DISPLAY_NAME}", order_by="create_time"
            )
            dataset = datasets[0]
            aip.TextDataset.delete(dataset.resource_name)
            print("Deleted dataset:", dataset)
        except Exception as e:
            print(e)

    if "tabular" == "video":
        try:
            datasets = aip.VideoDataset.list(
                filter=f"display_name={DISPLAY_NAME}", order_by="create_time"
            )
            dataset = datasets[0]
            aip.VideoDataset.delete(dataset.resource_name)
            print("Deleted dataset:", dataset)
        except Exception as e:
            print(e)

try:
    if delete_pipeline and "DISPLAY_NAME" in globals():
        pipelines = aip.PipelineJob.list(
            filter=f"display_name={DISPLAY_NAME}", order_by="create_time"
        )
        pipeline = pipelines[0]
        aip.PipelineJob.delete(pipeline.resource_name)
        print("Deleted pipeline:", pipeline)
except Exception as e:
    print(e)

if delete_bucket and "BUCKET_NAME" in globals():
    ! gsutil rm -r $BUCKET_NAME

"""
Explanation: Click on the generated link to see your run in the Cloud Console.
<!-- It should look something like this as it is running:

<a href="https://storage.googleapis.com/amy-jo/images/mp/automl_tabular_classif.png" target="_blank"><img src="https://storage.googleapis.com/amy-jo/images/mp/automl_tabular_classif.png" width="40%"/></a> -->

In the UI, many of the pipeline DAG nodes will expand or collapse when you click on them. Here is a partially-expanded view of the DAG (click image to see larger version).
<a href="https://storage.googleapis.com/amy-jo/images/mp/train_endpoint_deploy.png" target="_blank"><img src="https://storage.googleapis.com/amy-jo/images/mp/train_endpoint_deploy.png" width="75%"/></a>
Cleaning up
To clean up all Google Cloud resources used in this project, you can delete the Google Cloud
project you used for the tutorial.
Otherwise, you can delete the individual resources you created in this tutorial -- Note: this is auto-generated and not all resources may be applicable for this tutorial:

Dataset
Pipeline
Model
Endpoint
Batch Job
Custom Job
Hyperparameter Tuning Job
Cloud Storage Bucket
End of explanation
"""


