import os

# The Vertex AI Workbench Notebook product has specific requirements
IS_WORKBENCH_NOTEBOOK = os.getenv("DL_ANACONDA_HOME")
IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(
    "/opt/deeplearning/metadata/env_version"
)

# Vertex AI Notebook requires dependencies to be installed with '--user'
USER_FLAG = ""
if IS_WORKBENCH_NOTEBOOK:
    USER_FLAG = "--user"

! pip3 install -U tensorflow==2.5 $USER_FLAG -q
! pip3 install -U tensorflow-data-validation==1.2 $USER_FLAG -q
! pip3 install -U tensorflow-transform==1.2 $USER_FLAG -q
! pip3 install -U tensorflow-io==0.18 $USER_FLAG -q
! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG -q
! pip3 install --upgrade google-cloud-pipeline-components $USER_FLAG -q
! pip3 install --upgrade kfp $USER_FLAG -q

"""
Explanation: E2E ML on GCP: MLOps stage 3 : formalization: get started with Hyperparameter Tuning pipeline components
<table align="left">
    <td>
    <a href="https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage3/get_started_with_hpt_pipeline_components.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/main/notebooks/community/ml_ops/stage3/get_started_with_hpt_pipeline_components.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/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage3/get_started_with_hpt_pipeline_components.ipynb">
        <img src="https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32" alt="Vertex AI logo">
      Open in Google Cloud Notebooks
    </a>
  </td>
</table>
<br/><br/><br/>
Overview
This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 3 : formalization: get started with Hyperparameter Tuning pipeline components.
Dataset
The dataset used for this tutorial is the Horses or Humans from TensorFlow Datasets. The trained model predicts whether an image is a horse or human being.
Objective
In this tutorial, you learn how to use prebuilt Google Cloud Pipeline Components for Vertex AI Hyperparameter Tuning.
This tutorial uses the following Google Cloud ML services:

Google Cloud Pipeline Components
Vertex AI Dataset, Model and Endpoint resources
Vertex AI Hyperparameter Tuning

The steps performed include:

Construct a pipeline for:
Hyperparameter tune/train a custom model.
Retrieve the tuned hyperparameter values and metrics to optimize.
If the metrics exceed a specified threshold.
Get the location of the model artifacts for the best tuned model.
Upload the model artifacts to a Vertex AI Model resource.


Execute a Vertex AI pipeline.

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.
Installations
Install the required packages for executing the notebook.
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
"""


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


If you are running this notebook locally, you will need to install the Cloud SDK.


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 $ into these commands.
Set your project ID
If you don't know your project ID, you may be able to get your project ID using gcloud.
End of explanation
"""


REGION = "[your-region]"  # @param {type:"string"}

if REGION == "[your-region]":
    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 Vertex AI Workbench, then don't execute this code
IS_COLAB = False
if not os.path.exists("/opt/deeplearning/metadata/env_version") and not os.getenv(
    "DL_ANACONDA_HOME"
):
    if "google.colab" in sys.modules:
        IS_COLAB = True
        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 Vertex AI Workbench Notebooks, 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_URI = "gs://[your-bucket-name]"  # @param {type:"string"}

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 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 service account from gcloud
    if not IS_COLAB:
        shell_output = !gcloud auth list 2>/dev/null
        SERVICE_ACCOUNT = shell_output[2].replace("*", "").strip()

    if IS_COLAB:
        shell_output = ! gcloud projects describe  $PROJECT_ID
        project_number = shell_output[-1].split(":")[1].strip().replace("'", "")
        SERVICE_ACCOUNT = f"{project_number}-compute@developer.gserviceaccount.com"

    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

import json

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

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


import tensorflow as tf

"""
Explanation: Import TensorFlow
Import the TensorFlow package into your Python environment.
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
"""


import os

if os.getenv("IS_TESTING_TRAIN_GPU"):
    TRAIN_GPU, TRAIN_NGPU = (
        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,
        int(os.getenv("IS_TESTING_TRAIN_GPU")),
    )
else:
    TRAIN_GPU, TRAIN_NGPU = (None, None)

if os.getenv("IS_TESTING_DEPLOY_GPU"):
    DEPLOY_GPU, DEPLOY_NGPU = (
        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,
        int(os.getenv("IS_TESTING_DEPLOY_GPU")),
    )
else:
    DEPLOY_GPU, DEPLOY_NGPU = (None, None)

"""
Explanation: Set hardware accelerators
You can set hardware accelerators for training and prediction.
Set the variables TRAIN_GPU/TRAIN_NGPU and DEPLOY_GPU/DEPLOY_NGPU to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:
(aip.AcceleratorType.NVIDIA_TESLA_K80, 4)

Otherwise specify (None, None) to use a container image to run on a CPU.
Learn more about hardware accelerator support for your region.
Note: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support.
End of explanation
"""


if os.getenv("IS_TESTING_TF"):
    TF = os.getenv("IS_TESTING_TF")
else:
    TF = "2.5".replace(".", "-")

if TF[0] == "2":
    if DEPLOY_GPU:
        DEPLOY_VERSION = "tf2-gpu.{}".format(TF)
    else:
        DEPLOY_VERSION = "tf2-cpu.{}".format(TF)
else:
    if DEPLOY_GPU:
        DEPLOY_VERSION = "tf-gpu.{}".format(TF)
    else:
        DEPLOY_VERSION = "tf-cpu.{}".format(TF)

DEPLOY_IMAGE = "{}-docker.pkg.dev/vertex-ai/prediction/{}:latest".format(
    REGION.split("-")[0], DEPLOY_VERSION
)

print("Deployment:", DEPLOY_IMAGE, DEPLOY_GPU)

"""
Explanation: Set pre-built containers
Set the pre-built Docker container image for prediction.

Set the variable TF to the TensorFlow version of the container image. For example, 2-1 would be version 2.1, and 1-15 would be version 1.15. The following list shows some of the pre-built images available:

For the latest list, see Pre-built containers for prediction.
End of explanation
"""


if os.getenv("IS_TESTING_TRAIN_MACHINE"):
    MACHINE_TYPE = os.getenv("IS_TESTING_TRAIN_MACHINE")
else:
    MACHINE_TYPE = "n1-standard"

VCPU = "4"
TRAIN_COMPUTE = MACHINE_TYPE + "-" + VCPU
print("Train machine type", TRAIN_COMPUTE)

if os.getenv("IS_TESTING_DEPLOY_MACHINE"):
    MACHINE_TYPE = os.getenv("IS_TESTING_DEPLOY_MACHINE")
else:
    MACHINE_TYPE = "n1-standard"

VCPU = "4"
DEPLOY_COMPUTE = MACHINE_TYPE + "-" + VCPU
print("Deploy machine type", DEPLOY_COMPUTE)

"""
Explanation: Set machine type
Next, set the machine type to use for training and prediction.

Set the variables TRAIN_COMPUTE and DEPLOY_COMPUTE to configure  the compute resources for the VMs you will use for for training and prediction.
machine type
n1-standard: 3.75GB of memory per vCPU.
n1-highmem: 6.5GB of memory per vCPU
n1-highcpu: 0.9 GB of memory per vCPU


vCPUs: number of [2, 4, 8, 16, 32, 64, 96 ]

Note: The following is not supported for training:

standard: 2 vCPUs
highcpu: 2, 4 and 8 vCPUs

Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs.
End of explanation
"""


# Make folder for Python tuning script
! rm -rf custom
! mkdir custom

# Add package information
! touch custom/README.md

setup_cfg = "[egg_info]\n\ntag_build =\n\ntag_date = 0"
! echo "$setup_cfg" > custom/setup.cfg

setup_py = "import setuptools\n\nsetuptools.setup(\n\n    install_requires=[\n\n        'tensorflow==2.5.0',\n\n        'tensorflow_datasets==1.3.0',\n\n    ],\n\n    packages=setuptools.find_packages())"
! echo "$setup_py" > custom/setup.py

pkg_info = "Metadata-Version: 1.0\n\nName: Horses or Humans image classification\n\nVersion: 0.0.0\n\nSummary: Demostration tuning script\n\nHome-page: www.google.com\n\nAuthor: Google\n\nAuthor-email: aferlitsch@google.com\n\nLicense: Public\n\nDescription: Demo\n\nPlatform: Vertex"
! echo "$pkg_info" > custom/PKG-INFO

# Make the training subfolder
! mkdir custom/trainer
! touch custom/trainer/__init__.py

"""
Explanation: Examine the tuning package
Package layout
Before you start the tuning, you will look at how a Python package is assembled for a custom tuning job. When unarchived, the package contains the following directory/file layout.

PKG-INFO
README.md
setup.cfg
setup.py
trainer
__init__.py
task.py

The files setup.cfg and setup.py are the instructions for installing the package into the operating environment of the Docker image.
The file trainer/task.py is the Python script for executing the custom tuning job. Note, when we referred to it in the worker pool specification, we replace the directory slash with a dot (trainer.task) and dropped the file suffix (.py).
Package Assembly
In the following cells, you will assemble the training package.
End of explanation
"""


%%writefile custom/trainer/task.py
import os
os.system('pip install cloudml-hypertune')  # alternaterly, this can be added to the Dockerfile

import tensorflow as tf
import tensorflow_datasets as tfds
import argparse
import hypertune


def get_args():
  '''Parses args. Must include all hyperparameters you want to tune.'''

  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--epochs',
      required=True,
      type=int,
      help='number of epochs')
  parser.add_argument(
      '--learning_rate',
      required=True,
      type=float,
      help='learning rate')
  parser.add_argument(
      '--momentum',
      required=False,
      type=float,
      default=0.5,
      help='SGD momentum value')
  parser.add_argument(
      '--batch_size',
      required=True,
      type=int,
      help='the batch size')
  parser.add_argument(
      '--model-dir',
      dest='model_dir',
      default=os.getenv('AIP_MODEL_DIR'),
      type=str, help='Model dir.')
  args = parser.parse_args()
  return args

def preprocess_data(image, label):
  '''Resizes and scales images.'''

  image = tf.image.resize(image, (150,150))
  return tf.cast(image, tf.float32) / 255., label


def get_data():
  '''Loads Horses Or Humans dataset and preprocesses data.'''

  data, info = tfds.load(name='horses_or_humans', as_supervised=True, with_info=True)

  # Create train dataset
  train_data = data['train'].map(preprocess_data)
  train_data  = train_data.shuffle(1000)
  train_data  = train_data.batch(64)

  # Create validation dataset
  validation_data = data['test'].map(preprocess_data)
  validation_data  = validation_data.batch(64)

  return train_data, validation_data


def get_model(learning_rate, momentum):
  '''Defines and complies model.'''

  inputs = tf.keras.Input(shape=(150, 150, 3))
  x = tf.keras.layers.Conv2D(16, (3, 3), activation='relu')(inputs)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')(x)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Conv2D(64, (3, 3), activation='relu')(x)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Flatten()(x)
  outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)
  model = tf.keras.Model(inputs, outputs)
  model.compile(
      loss='binary_crossentropy',
      optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=momentum),
      metrics=['accuracy'])
  return model

def train_model(model, train_data, validation_data, epochs, batch_size):

  history = model.fit(train_data, epochs=epochs, batch_size=batch_size, validation_data=validation_data)

  # DEFINE METRIC
  hp_metric = history.history['val_accuracy'][-1]

  hpt = hypertune.HyperTune()
  hpt.report_hyperparameter_tuning_metric(
      hyperparameter_metric_tag='accuracy',
      metric_value=hp_metric,
      global_step=epochs
  )

  return model


def main():
  args = get_args()
  train_data, validation_data = get_data()

  model = get_model(args.learning_rate, args.momentum)

  model = train_model(model, train_data, validation_data, args.epochs, args.batch_size)

  model.save(args.model_dir)


if __name__ == "__main__":
    main()

"""
Explanation: Create the task script for the Python training package
Next, you create the task.py script for driving the training package. Some noteable steps include:


Command-line arguments:

model-dir: The location to save the trained model. When using Vertex AI custom training, the location will be specified in the environment variable: AIP_MODEL_DIR,
epochs: The number of epochs to train for.
learning_rate: Hyperparameter for learning rate.
batch_size: Hyperparameter for batch size.



Data preprocessing (get_data())

Loads and preprocesses the dataset as a tf.data.Dataset generator.



Model architecture (get_model()):

Builds the corresponding model architecture.



Training (train_model()):

Trains the model



Model artifact saving

Saves the model artifacts where the Cloud Storage location is specified.
End of explanation
"""


! rm -f custom.tar custom.tar.gz
! tar cvf custom.tar custom
! gzip custom.tar
! gsutil cp custom.tar.gz $BUCKET_URI/trainer_horses_or_humans.tar.gz

"""
Explanation: Store tuning script on your Cloud Storage bucket
Next, you package the tuning folder into a compressed tar ball, and then store it in your Cloud Storage bucket.
End of explanation
"""


%%writefile custom/Dockerfile

FROM gcr.io/deeplearning-platform-release/tf2-cpu.2-3
WORKDIR /root

WORKDIR /

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

"""
Explanation: Create a Docker file
To use your own custom training container, you build a Docker file and embed into the container your training scripts.
Write the Docker file contents
Your first step in containerizing your code is to create a Docker file. In your Docker you’ll include all the commands needed to run your container image. It’ll install all the libraries you’re using and set up the entry point for your training code.

Install a pre-defined container image from TensorFlow repository for deep learning images.
Copies in the Python training code, to be shown subsequently.
Sets the entry into the Python training script as trainer/task.py. Note, the .py is dropped in the ENTRYPOINT command, as it is implied.
End of explanation
"""


TRAIN_IMAGE = "gcr.io/" + PROJECT_ID + "/horses_or_humans:v1"

"""
Explanation: Build the container locally
Next, you will provide a name for your customer container that you will use when you submit it to the Google Container Registry.
End of explanation
"""


! docker build custom -t $TRAIN_IMAGE

"""
Explanation: Next, build the container.
End of explanation
"""


! docker push $TRAIN_IMAGE

"""
Explanation: Register the custom container
When you’ve finished running the container locally, push it to Google Container Registry.
End of explanation
"""


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


@component(packages_to_install=["google-cloud-aiplatform"])
def model_dir(base_output_directory: str, best_trial: str) -> str:
    from google.cloud.aiplatform_v1.types import study

    trial_proto = study.Trial.from_json(best_trial)
    model_id = trial_proto.id
    return f"{base_output_directory}/{model_id}/model"


@dsl.pipeline(
    name="hp-tuning", description="Custom image classification hyperparameter tuning"
)
def pipeline(
    display_name: str,
    worker_pool_specs: list,
    study_spec_metrics: list,
    study_spec_parameters: list,
    threshold: float,
    deploy_image: str,
    max_trial_count: int = 5,
    parallel_trial_count: int = 1,
    base_output_directory: str = PIPELINE_ROOT,
    labels: dict = {},
    project: str = PROJECT_ID,
    region: str = REGION,
):

    from google_cloud_pipeline_components.experimental import \
        hyperparameter_tuning_job
    from google_cloud_pipeline_components.types import artifact_types
    from google_cloud_pipeline_components.v1.hyperparameter_tuning_job import \
        HyperparameterTuningJobRunOp
    from google_cloud_pipeline_components.v1.model import ModelUploadOp
    from kfp.v2.components import importer_node

    tuning_op = HyperparameterTuningJobRunOp(
        display_name=display_name,
        project=project,
        location=region,
        worker_pool_specs=worker_pool_specs,
        study_spec_metrics=study_spec_metrics,
        study_spec_parameters=study_spec_parameters,
        max_trial_count=max_trial_count,
        parallel_trial_count=parallel_trial_count,
        base_output_directory=base_output_directory,
    )

    trials_op = hyperparameter_tuning_job.GetTrialsOp(
        gcp_resources=tuning_op.outputs["gcp_resources"]
    )

    best_trial_op = hyperparameter_tuning_job.GetBestTrialOp(
        trials=trials_op.output, study_spec_metrics=study_spec_metrics
    )

    threshold_op = hyperparameter_tuning_job.IsMetricBeyondThresholdOp(
        trial=best_trial_op.output,
        study_spec_metrics=study_spec_metrics,
        threshold=threshold,
    )

    with dsl.Condition(
        threshold_op.output == "true",
        name="deploy_decision",
    ):
        _ = hyperparameter_tuning_job.GetHyperparametersOp(trial=best_trial_op.output)

        model_dir_op = model_dir(base_output_directory, best_trial_op.output)

        import_unmanaged_model_op = importer_node.importer(
            artifact_uri=model_dir_op.output,
            artifact_class=artifact_types.UnmanagedContainerModel,
            metadata={
                "containerSpec": {
                    "imageUri": DEPLOY_IMAGE,
                },
            },
        ).after(model_dir_op)

        _ = ModelUploadOp(
            project=project,
            display_name=display_name,
            unmanaged_container_model=import_unmanaged_model_op.outputs["artifact"],
        ).after(import_unmanaged_model_op)

"""
Explanation: Construct hyperparameter tuning pipeline
Next, construct the pipeline with the following tasks:

Create/Execute a hyperparameter tuning job
Get all trial results.
Get the best trial results.
Determine if the best trial results exceed a threshold
Retrieve the hyperparameter values
Determine Cloud Storage location of the best model
Upload the best model as a Vertex AI Model resource.
End of explanation
"""


from google_cloud_pipeline_components.experimental import \
    hyperparameter_tuning_job

gpu = "ACCELERATOR_TYPE_UNSPECIFIED"
accelerator_count = 0

if TRAIN_GPU:
    gpu = TRAIN_GPU.name
    accelerator_count = 1

else:
    gpu = "ACCELERATOR_TYPE_UNSPECIFIED"
    accelerator_count = (
        0  # same problem with accelerator_count, if we keep is as "None" its not
    )

CMDARGS = [
    "--epochs=10",
]

# The spec of the worker pools including machine type and Docker image
worker_pool_specs = [
    {
        "machine_spec": {
            "machine_type": TRAIN_COMPUTE,
            "accelerator_type": gpu,
            "accelerator_count": accelerator_count,
        },
        "replica_count": 1,
        "container_spec": {"image_uri": TRAIN_IMAGE, "args": CMDARGS},
    }
]

# List serialized from the dictionary representing metrics to optimize.
# The dictionary key is the metric_id, which is reported by your training job,
# and the dictionary value is the optimization goal of the metric.
metric_spec = hyperparameter_tuning_job.serialize_metrics({"accuracy": "maximize"})

# List serialized from the parameter dictionary. The dictionary
# represents parameters to optimize. The dictionary key is the parameter_id,
# which is passed into your training job as a command line key word argument, and the
# dictionary value is the parameter specification of the metric.
parameter_spec = hyperparameter_tuning_job.serialize_parameters(
    {
        "learning_rate": aip.hyperparameter_tuning.DoubleParameterSpec(
            min=0.001, max=1, scale="log"
        ),
        "batch_size": aip.hyperparameter_tuning.DiscreteParameterSpec(
            values=[16, 32, 64], scale=None
        ),
    }
)

"""
Explanation: Create hyperparameter tuning specifications
Next, you construct the worker pool specification, and the study's metric and parameter specifications, as follows:
Worker pool specification
This specification describes the machine and container requirements, and scaling for executing the hyperparameter study. Since the training module is embedded in the docker image, you use the args field to specify any command-lime arguments, which are not part of the study, to the training module. In this example, you pass the number of epochs.
Parameter specification
This specification describes the hyperparameters to tune, and the range of values to tune then for. For each study, the trial values for these parameters are passed as command-line arguments to the training module, as --&lt;parameter_name&gt;=&lt;trial_value&gt;.
Metric specification
This specification describes the metric(s) to be evaluated in the study and wether to minize or maximize the metric.
End of explanation
"""


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

pipeline = aip.PipelineJob(
    display_name="hp_tuning",
    template_path="hp_tune_pipeline_job.json",
    pipeline_root=PIPELINE_ROOT,
    parameter_values={
        "display_name": "hp_tuning",
        "worker_pool_specs": worker_pool_specs,
        "study_spec_metrics": metric_spec,
        "study_spec_parameters": parameter_spec,
        "threshold": 0.7,
        "deploy_image": DEPLOY_IMAGE,
    },
    enable_caching=False,
)

pipeline.run()

! rm -rf hp_tune_pipeline_job.json custom custom.tar.gz

"""
Explanation: Compile and execute hyperparameter tuning pipeline
Next, you compile the pipeline and then exeute it. The pipeline takes the following parameters, which are passed as the dictionary parameter_values:

display_name: A human readable name for the pipeline job.
import_file: The Cloud Storage location to the dataset.
worker_pool_specs: The the machine and container, and auto-scaling requirements, as well as command line arguments.
study_spec_metrics: The metrics to optimize in the study trials.
study_spec_parameters: The parameters to tune.
End of explanation
"""


PROJECT_NUMBER = pipeline.gca_resource.name.split("/")[1]
print(PROJECT_NUMBER)


def print_pipeline_output(job, output_task_name):
    JOB_ID = job.name
    print(JOB_ID)
    for _ in range(len(job.gca_resource.job_detail.task_details)):
        TASK_ID = job.gca_resource.job_detail.task_details[_].task_id
        EXECUTE_OUTPUT = (
            PIPELINE_ROOT
            + "/"
            + PROJECT_NUMBER
            + "/"
            + JOB_ID
            + "/"
            + output_task_name
            + "_"
            + str(TASK_ID)
            + "/executor_output.json"
        )
        GCP_RESOURCES = (
            PIPELINE_ROOT
            + "/"
            + PROJECT_NUMBER
            + "/"
            + JOB_ID
            + "/"
            + output_task_name
            + "_"
            + str(TASK_ID)
            + "/gcp_resources"
        )
        EVAL_METRICS = (
            PIPELINE_ROOT
            + "/"
            + PROJECT_NUMBER
            + "/"
            + JOB_ID
            + "/"
            + output_task_name
            + "_"
            + str(TASK_ID)
            + "/evaluation_metrics"
        )
        if tf.io.gfile.exists(EXECUTE_OUTPUT):
            ! gsutil cat $EXECUTE_OUTPUT
            return EXECUTE_OUTPUT
        elif tf.io.gfile.exists(GCP_RESOURCES):
            ! gsutil cat $GCP_RESOURCES
            return GCP_RESOURCES
        elif tf.io.gfile.exists(EVAL_METRICS):
            ! gsutil cat $EVAL_METRICS
            return EVAL_METRICS

    return None


print("hyperparameter-tuning-job")
artifacts = print_pipeline_output(pipeline, "hyperparameter-tuning-job")
print("\n\n")
print("gettrialsop")
artifacts = print_pipeline_output(pipeline, "gettrialsop")
print("\n\n")
print("getbesttrialop")
artifacts = print_pipeline_output(pipeline, "getbesttrialop")
print("\n\n")
output = !gsutil cat $artifacts
output = json.loads(output[0])
best_trial = json.loads(output["parameters"]["Output"]["stringValue"])
model_id = best_trial["id"]
print("BEST MODEL", model_id)
parameters = best_trial["parameters"]
batch_size = parameters[0]["value"]
print("BATCH SIZE", batch_size)
learning_rate = parameters[1]["value"]
print("LR", learning_rate)
MODEL_DIR = f"{PIPELINE_ROOT}/{model_id}/model"

print("ismetricbeyondthresholdop")
artifacts = print_pipeline_output(pipeline, "ismetricbeyondthresholdop")
print("\n\n")
print("deploy-decision")
artifacts = print_pipeline_output(pipeline, "deploy-decision")
print("\n\n")
print("model-dir")
artifacts = print_pipeline_output(pipeline, "model-dir")
print("\n\n")
print("model-upload")
artifacts = print_pipeline_output(pipeline, "model-upload")
print("\n\n")

"""
Explanation: View the data pipeline execution results
End of explanation
"""


pipeline.delete()

"""
Explanation: Delete a pipeline job
After a pipeline job is completed, you can delete the pipeline job with the method delete().  Prior to completion, a pipeline job can be canceled with the method cancel().
End of explanation
"""


delete_bucket = False

if delete_bucket or os.getenv("IS_TESTING"):
    ! 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:

Cloud Storage Bucket
End of explanation
"""


