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: Pipelines introduction for KFP
<table align="left">
  <td>
    <a href="https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/official/pipelines/pipelines_intro_kfp.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/blob/master/notebooks/official/pipelines/pipelines_intro_kfp.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/blob/master/notebooks/official/pipelines/pipelines_intro_kfp.ipynb">
        <img src="https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32" alt="Vertex AI logo">
      Open in Vertex AI Workbench
    </a>
  </td>
</table>
<br/><br/><br/>
Overview
This notebook provides an introduction to using Vertex AI Pipelines with the Kubeflow Pipelines (KFP) SDK.
Objective
The steps performed include:

Define and compile a pipeline.
Schedule a recurring pipeline run.
Specify which service account to use for a pipeline run.

Costs
This tutorial uses billable components of Google Cloud:

Vertex AI
Cloud Storage
Cloud Functions
Cloud Scheduler

Learn about Vertex AI pricing,
Cloud Storage pricing,
Cloud Functions pricing, and
Clould Scheduler 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 AI 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 --upgrade

"""
Explanation: Install the latest GA version of KFP SDK 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__))"

"""
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 = "[your-region]"  # @param {type:"string"}

if REGION == "[your-region]":
    REGION = "us-central1"

"""
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 = "[your-bucket-name]"  # @param {type:"string"}
BUCKET_URI = f"gs://{BUCKET_NAME}"

if BUCKET_URI == "" or BUCKET_URI is None or BUCKET_URI == "gs://[your-bucket-name]":
    BUCKET_URI = "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_URI

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

"""
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].replace("*", "").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_URI

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

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


# API service endpoint
API_ENDPOINT = "{}-aiplatform.googleapis.com".format(REGION)

"""
Explanation: Vertex AI constants
Setup up the following constants for Vertex AI:

API_ENDPOINT: The Vertex AI API service endpoint for Dataset, Model, Job, Pipeline and Endpoint services.
End of explanation
"""


PIPELINE_ROOT = "{}/pipeline_root/intro".format(BUCKET_URI)

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


from typing import NamedTuple

from kfp import dsl
from kfp.v2 import compiler
from kfp.v2.dsl import component

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


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

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


@component(output_component_file="hw.yaml", base_image="python:3.9")
def hello_world(text: str) -> str:
    print(text)
    return text

"""
Explanation: Define Python function-based pipeline components
In this tutorial, you define a simple pipeline that has three steps, where each step is defined as a component.
Define hello_world component
First, define a component based on a very simple Python function. It takes a string input parameter and returns that value as output.
Note the use of the @component decorator, which compiles the function to a KFP component when evaluated.  For example purposes, this example specifies a base image to use for the component (python:3.9), and a component YAML file, hw.yaml. The compiled component specification is written to this file.  (The default base image is python:3.7, which would of course work just fine too).
End of explanation
"""


@component(packages_to_install=["google-cloud-storage"])
def two_outputs(
    text: str,
) -> NamedTuple(
    "Outputs",
    [
        ("output_one", str),  # Return parameters
        ("output_two", str),
    ],
):
    # the import is not actually used for this simple example, but the import
    # is successful, as it was included in the `packages_to_install` list.
    from google.cloud import storage  # noqa: F401

    o1 = f"output one from text: {text}"
    o2 = f"output two from text: {text}"
    print("output one: {}; output_two: {}".format(o1, o2))
    return (o1, o2)

"""
Explanation: As you'll see below, compilation of this component creates a task factory function—called hello_world— that you can use in defining a pipeline step.
While not shown here, if you want to share this component definition, or use it in another context, you could also load it from its yaml file like this:
hello_world_op = components.load_component_from_file('./hw.yaml').
You can also use the load_component_from_url method, if your component yaml file is stored online. (For GitHub URLs, load the 'raw' file.)
Define two_outputs component
The first component below, two_outputs, demonstrates installing a package -- in this case the google-cloud-storage package. Alternatively, you can specify a base image that includes the necessary installations.
Note: The component function won't actually use the package.
Alternatively, you can specify a base image that includes the necessary installations.
The two_outputs component returns two named outputs.
End of explanation
"""


@component
def consumer(text1: str, text2: str, text3: str):
    print(f"text1: {text1}; text2: {text2}; text3: {text3}")

"""
Explanation: Define the consumer component
The third component, consumer, takes three string inputs and prints them out.
End of explanation
"""


@dsl.pipeline(
    name="hello-world-v2",
    description="A simple intro pipeline",
    pipeline_root=PIPELINE_ROOT,
)
def pipeline(text: str = "hi there"):
    hw_task = hello_world(text)
    two_outputs_task = two_outputs(text)
    consumer_task = consumer(  # noqa: F841
        hw_task.output,
        two_outputs_task.outputs["output_one"],
        two_outputs_task.outputs["output_two"],
    )

"""
Explanation: Define a pipeline that uses the components
Next, define a pipeline that uses these three components.
By evaluating the component definitions above, you've created task factory functions that are used in the pipeline definition to create the pipeline steps.
The pipeline takes an input parameter, and passes that parameter as an argument to the first two pipeline steps (hw_task and two_outputs_task).
Then, the third pipeline step (consumer_task) consumes the outputs of the first and second steps.  Because the hello_world component definition just returns one unnamed output, you refer to it as hw_task.output.  The two_outputs task returns two named outputs, which you access as two_outputs_task.outputs["&lt;output_name&gt;"].
Note: In the @dsl.pipeline decorator, you're defining the PIPELINE_ROOT Cloud Storage path to use.  If you had not included that info here, it would be required to specify it when creating the pipeline run, as you'll see below.
End of explanation
"""


from kfp.v2 import compiler  # noqa: F811

compiler.Compiler().compile(pipeline_func=pipeline, package_path="intro_pipeline.json")

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


DISPLAY_NAME = "intro_" + TIMESTAMP

job = aip.PipelineJob(
    display_name=DISPLAY_NAME,
    template_path="intro_pipeline.json",
    pipeline_root=PIPELINE_ROOT,
)

job.run()

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


job.delete()

"""
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/intro_pipeline.png" target="_blank"><img src="https://storage.googleapis.com/amy-jo/images/mp/intro_pipeline.png" width="60%"/></a>
Delete the pipeline job
You can delete the pipeline job with the method delete().
End of explanation
"""


if not os.getenv("IS_TESTING"):
    from kfp.v2.google.client import AIPlatformClient  # noqa: F811

    api_client = AIPlatformClient(project_id=PROJECT_ID, region=REGION)

    # adjust time zone and cron schedule as necessary
    response = api_client.create_schedule_from_job_spec(
        job_spec_path="intro_pipeline.json",
        schedule="2 * * * *",
        time_zone="America/Los_Angeles",  # change this as necessary
        parameter_values={"text": "Hello world!"},
        # pipeline_root=PIPELINE_ROOT  # this argument is necessary if you did not specify PIPELINE_ROOT as part of the pipeline definition.
    )

"""
Explanation: Recurring pipeline runs: create a scheduled pipeline job
This section shows how to create a scheduled pipeline job.  You do this using the pipeline you already defined.
Under the hood, the scheduled jobs are supported by the Cloud Scheduler and a Cloud Functions function.  Check first that the APIs for both of these services are enabled.
You will need to first enable the enable the Cloud Scheduler API and the Cloud Functions and Cloud Build APIs if you have not already done so.
Note:you need to create an App Engine app for your project if one does not already exist.
See the Cloud Scheduler documentation for more on the cron syntax.
Create a scheduled pipeline job, passing as an argument the job specification file that you compiled above.
Note: You can pass a parameter_values dict that specifies the pipeline input parameters you want to use.
End of explanation
"""


if not os.getenv("IS_TESTING"):
    response = api_client.create_run_from_job_spec(
        job_spec_path="intro_pipeline.json",
        pipeline_root=PIPELINE_ROOT,
        service_account=SERVICE_ACCOUNT,  # <-- CHANGE to use non-default service account
    )

"""
Explanation: Once the scheduled job is created, you can see it listed in the Cloud Scheduler panel in the Console.
<a href="https://storage.googleapis.com/amy-jo/images/kf-pls/pipelines_scheduler.png" target="_blank"><img src="https://storage.googleapis.com/amy-jo/images/kf-pls/pipelines_scheduler.png" width="95%"/></a>
You can test the setup from the Cloud Scheduler panel by clicking 'RUN NOW'.

Note: The implementation is using a Cloud Functions function, which you can see listed in the Cloud Functions panel in the console as templated_http_request-v1.
Don't delete this function, as it will prevent the Cloud Scheduler jobs from actually kicking off the pipeline run.  If you do delete it, create a new scheduled job in order to recreate the function.

When you're done experimenting, you probably want to PAUSE your scheduled job from the Cloud Scheduler panel, so that the recurrent jobs do not keep running.
Specifying a service account to use for a pipeline run
By default, the service account used for your pipeline run is your default compute engine service account.
However, you might want to run pipelines with permissions to access different roles than those configured for your default SA (e.g. perhaps using a more restricted set of permissions).
If you want to execute your pipeline using a different service account, this is straightforward to do.  You just need to give the new service account the correct permissions.
Create a service account
Once your service account is configured, you pass it as an argument to the create_run_from_job_spec method. The pipeline job runs with the permissions of the given service account.
Learn about creating and configuring a service account to work with Vertex AI Pipelines.
End of explanation
"""


job = aip.PipelineJob(
    display_name="intro_" + TIMESTAMP,
    template_path="intro_pipeline.json",
    enable_caching=False,
)

job.run()

"""
Explanation: Pipeline step caching
By default, pipeline step caching is enabled. This means that the results of previous step executions are reused when possible.
if you want to disable caching for a pipeline run, you can pass the enable_caching=False argument to the PipelineJob constructor when you submit the pipeline job, as shown below.  Try submitting the example pipeline job again, first with and then without caching enabled.
End of explanation
"""


job.delete()

"""
Explanation: Delete the pipeline job
You can delete the pipeline job with the method delete().
End of explanation
"""


! curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)"   -H "Content-Type: application/json"   https://{API_ENDPOINT}/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/pipelineJobs

"""
Explanation: Using the Pipelines REST API
At times you may want to use the REST API instead of the Python KFP SDK.  Below are examples of how to do that.
Where a command requires a pipeline ID, you can get that info from the "Run" column in the pipelines list— as shown below— as well as from the 'details' page for a given pipeline.  You can also see that info when you list pipeline job(s) via the API.
<a href="https://storage.googleapis.com/amy-jo/images/mp/pipeline_id.png" target="_blank"><img src="https://storage.googleapis.com/amy-jo/images/mp/pipeline_id.png" width="80%"/></a>
List pipeline jobs
Note: This request may generate a large response if you have many pipeline runs.
End of explanation
"""


output = ! curl -X POST  -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json; charset=utf-8"   https://{API_ENDPOINT}/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/pipelineJobs  --data "@intro_pipeline.json"

PIPELINE_RUN_ID = output[5].split("/")[-1].split('"')[0]
print(output)

"""
Explanation: Create a pipeline job
For this REST request, you submit a compiled pipeline job specification.  We're using the one you generated previously.
End of explanation
"""


! curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)"   -H "Content-Type: application/json"   https://{API_ENDPOINT}/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/pipelineJobs/{PIPELINE_RUN_ID}

"""
Explanation: Get a pipeline job from its ID
End of explanation
"""


! curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)"   -H "Content-Type: application/json"   https://{API_ENDPOINT}/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/pipelineJobs/{PIPELINE_RUN_ID}:cancel

"""
Explanation: Cancel a pipeline job given its ID
End of explanation
"""


! curl -X DELETE -H "Authorization: Bearer $(gcloud auth print-access-token)"   -H "Content-Type: application/json"   https://{API_ENDPOINT}/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/pipelineJobs/{PIPELINE_RUN_ID}

"""
Explanation: Delete a pipeline job given its ID
End of explanation
"""


delete_pipeline = True
delete_bucket = True

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_URI" in globals():
    ! gsutil rm -r $BUCKET_URI

"""
Explanation: 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.
End of explanation
"""


