from google.colab import auth
auth.authenticate_user()import tensorflow as tf

print("TensorFlow version: {}".format(tf.__version__))
from tfx import v1 as tfx

print("TFX version: {}".format(tfx.__version__))
import kfp

print("KFP version: {}".format(kfp.__version__))

from google.cloud import aiplatform as vertex_ai
import osGOOGLE_CLOUD_PROJECT = "gcp-ml-172005"  # @param {type:"string"}
GOOGLE_CLOUD_REGION = "us-central1"  # @param {type:"string"}
GCS_BUCKET_NAME = "cifar10-experimental-csp2"  # @param {type:"string"}
DATA_ROOT = "gs://cifar10-csp-public2"  # @param {type:"string"}

if not (GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_REGION and GCS_BUCKET_NAME):
    from absl import logging

    logging.error("Please set all required parameters.")PIPELINE_NAME = "continuous-adaptation-for-data-changes"

# Path to various pipeline artifact.
PIPELINE_ROOT = "gs://{}/pipeline_root/{}".format(GCS_BUCKET_NAME, PIPELINE_NAME)

# Paths for users' Python module.
MODULE_ROOT = "gs://{}/pipeline_module/{}".format(GCS_BUCKET_NAME, PIPELINE_NAME)

# This is the path where your model will be pushed for serving.
SERVING_MODEL_DIR = "gs://{}/serving_model/{}".format(GCS_BUCKET_NAME, PIPELINE_NAME)

print("PIPELINE_ROOT: {}".format(PIPELINE_ROOT))_trainer_module_file = 'trainer.py'%%writefile {_trainer_module_file}

from typing import List
from absl import logging
from tensorflow import keras
from tfx import v1 as tfx
import tensorflow as tf

_IMAGE_FEATURES = {
    "image": tf.io.FixedLenFeature([], tf.string),
    "label": tf.io.FixedLenFeature([], tf.int64),
}

_CONCRETE_INPUT = "numpy_inputs"
_TRAIN_BATCH_SIZE = 64
_EVAL_BATCH_SIZE = 64
_INPUT_SHAPE = (32, 32, 3)
_EPOCHS = 2


def _parse_fn(example):
    example = tf.io.parse_single_example(example, _IMAGE_FEATURES)
    image = tf.image.decode_jpeg(example["image"], channels=3)
    class_label = tf.cast(example["label"], tf.int32)
    return image, class_label


def _input_fn(file_pattern: List[str], batch_size: int) -> tf.data.Dataset:
    print(f"Reading data from: {file_pattern}")
    tfrecord_filenames = tf.io.gfile.glob(file_pattern[0] + ".gz")
    print(tfrecord_filenames)
    dataset = tf.data.TFRecordDataset(tfrecord_filenames, compression_type="GZIP")
    dataset = dataset.map(_parse_fn).batch(batch_size)
    return dataset.repeat()


def _make_keras_model() -> tf.keras.Model:
    """Creates a ResNet50-based model for classifying flowers data.

  Returns:
  A Keras Model.
  """
    inputs = keras.Input(shape=_INPUT_SHAPE)
    base_model = keras.applications.ResNet50(
        include_top=False, input_shape=_INPUT_SHAPE, pooling="avg"
    )
    base_model.trainable = False
    x = tf.keras.applications.resnet.preprocess_input(inputs)
    x = base_model(
        x, training=False
    )  # Ensures BatchNorm runs in inference model in this model
    outputs = keras.layers.Dense(10, activation="softmax")(x)
    model = keras.Model(inputs, outputs)

    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )

    model.summary(print_fn=logging.info)
    return model


def _preprocess(bytes_input):
    decoded = tf.io.decode_jpeg(bytes_input, channels=3)
    resized = tf.image.resize(decoded, size=(32, 32))
    return resized


@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
def preprocess_fn(bytes_inputs):
    decoded_images = tf.map_fn(
        _preprocess, bytes_inputs, dtype=tf.float32, back_prop=False
    )
    return {_CONCRETE_INPUT: decoded_images}


def _model_exporter(model: tf.keras.Model):
    m_call = tf.function(model.call).get_concrete_function(
        [tf.TensorSpec(shape=[None, 32, 32, 3], dtype=tf.float32, name=_CONCRETE_INPUT)]
    )

    @tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
    def serving_fn(bytes_inputs):
        # This function comes from the Computer Vision book from O'Reilly.
        labels = tf.constant(
            [
                "airplane",
                "automobile",
                "bird",
                "cat",
                "deer",
                "dog",
                "frog",
                "horse",
                "ship",
                "truck",
            ],
            dtype=tf.string,
        )
        images = preprocess_fn(bytes_inputs)

        probs = m_call(**images)
        indices = tf.argmax(probs, axis=1)
        pred_source = tf.gather(params=labels, indices=indices)
        pred_confidence = tf.reduce_max(probs, axis=1)
        return {"label": pred_source, "confidence": pred_confidence}

    return serving_fn


def run_fn(fn_args: tfx.components.FnArgs):
    print(fn_args)

    train_dataset = _input_fn(fn_args.train_files, batch_size=_TRAIN_BATCH_SIZE)
    eval_dataset = _input_fn(fn_args.eval_files, batch_size=_EVAL_BATCH_SIZE)

    model = _make_keras_model()
    model.fit(
        train_dataset,
        steps_per_epoch=fn_args.train_steps,
        validation_data=eval_dataset,
        validation_steps=fn_args.eval_steps,
        epochs=_EPOCHS,
    )

    _, acc = model.evaluate(eval_dataset, steps=fn_args.eval_steps)
    logging.info(f"Validation accuracy: {round(acc * 100, 2)}%")
    # The result of the training should be saved in `fn_args.serving_model_dir`
    # directory.
    tf.saved_model.save(
        model,
        fn_args.serving_model_dir,
        signatures={"serving_default": _model_exporter(model)},
    )os.path.join(MODULE_ROOT, _trainer_module_file)_vertex_uploader_module_file = "vertex_uploader.py"
_vertex_deployer_module_file = "vertex_deployer.py"%%writefile {_vertex_uploader_module_file}

import os
import tensorflow as tf

from tfx.dsl.component.experimental.decorators import component
from tfx.dsl.component.experimental.annotations import Parameter
from tfx.types.standard_artifacts import String
from google.cloud import aiplatform as vertex_ai
from tfx import v1 as tfx
from absl import logging


@component
def VertexUploader(
    project: Parameter[str],
    region: Parameter[str],
    model_display_name: Parameter[str],
    pushed_model_location: Parameter[str],
    serving_image_uri: Parameter[str],
    uploaded_model: tfx.dsl.components.OutputArtifact[String],
):

    vertex_ai.init(project=project, location=region)

    pushed_model_dir = os.path.join(
        pushed_model_location, tf.io.gfile.listdir(pushed_model_location)[-1]
    )

    logging.info(f"Model registry location: {pushed_model_dir}")

    vertex_model = vertex_ai.Model.upload(
        display_name=model_display_name,
        artifact_uri=pushed_model_dir,
        serving_container_image_uri=serving_image_uri,
        parameters_schema_uri=None,
        instance_schema_uri=None,
        explanation_metadata=None,
        explanation_parameters=None,
    )

    uploaded_model.set_string_custom_property(
        "model_resource_name", str(vertex_model.resource_name)
    )
    logging.info(f"Model resource: {str(vertex_model.resource_name)}")%%writefile {_vertex_deployer_module_file}

from tfx.dsl.component.experimental.decorators import component
from tfx.dsl.component.experimental.annotations import Parameter
from tfx.types.standard_artifacts import String
from google.cloud import aiplatform as vertex_ai
from tfx import v1 as tfx
from absl import logging


@component
def VertexDeployer(
    project: Parameter[str],
    region: Parameter[str],
    model_display_name: Parameter[str],
    deployed_model_display_name: Parameter[str],
):

    logging.info(f"Endpoint display: {deployed_model_display_name}")
    vertex_ai.init(project=project, location=region)

    endpoints = vertex_ai.Endpoint.list(
        filter=f"display_name={deployed_model_display_name}", order_by="update_time"
    )

    if len(endpoints) > 0:
        logging.info(f"Endpoint {deployed_model_display_name} already exists.")
        endpoint = endpoints[-1]
    else:
        endpoint = vertex_ai.Endpoint.create(deployed_model_display_name)

    model = vertex_ai.Model.list(
        filter=f"display_name={model_display_name}", order_by="update_time"
    )[-1]

    endpoint = vertex_ai.Endpoint.list(
        filter=f"display_name={deployed_model_display_name}", order_by="update_time"
    )[-1]

    deployed_model = endpoint.deploy(
        model=model,
        # Syntax from here: https://git.io/JBQDP
        traffic_split={"0": 100},
        machine_type="n1-standard-4",
        min_replica_count=1,
        max_replica_count=1,
    )

    logging.info(f"Model deployed to: {deployed_model}")DATASET_DISPLAY_NAME = "cifar10"
VERSION = "tfx-1-2-0"
TFX_IMAGE_URI = f"gcr.io/{GOOGLE_CLOUD_PROJECT}/{DATASET_DISPLAY_NAME}:{VERSION}"
print(f"URI of the custom image: {TFX_IMAGE_URI}")%%writefile Dockerfile

FROM gcr.io/tfx-oss-public/tfx:1.2.0
RUN mkdir -p custom_components
COPY custom_components/* ./custom_components/
RUN pip install --upgrade google-cloud-aiplatform# Specify training worker configurations. To minimize costs we can even specify two
# different configurations: a beefier machine for the Endpoint model and slightly less
# powerful machine for the mobile model.
TRAINING_JOB_SPEC = {
    "project": GOOGLE_CLOUD_PROJECT,
    "worker_pool_specs": [
        {
            "machine_spec": {
                "machine_type": "n1-standard-4",
                "accelerator_type": "NVIDIA_TESLA_K80",
                "accelerator_count": 1,
            },
            "replica_count": 1,
            "container_spec": {
                "image_uri": "gcr.io/tfx-oss-public/tfx:{}".format(tfx.__version__),
            },
        }
    ],
}SERVING_JOB_SPEC = {
    "endpoint_name": PIPELINE_NAME.replace("-", "_"),  # '-' is not allowed.
    "project_id": GOOGLE_CLOUD_PROJECT,
    "min_replica_count": 1,
    "max_replica_count": 1,
    "machine_type": "n1-standard-2",
}from datetime import datetime

TIMESTAMP = datetime.now().strftime("%Y%m%d%H%M%S")import tfxfrom tfx.orchestration import data_types

from tfx import v1 as tfx
from tfx.proto import example_gen_pb2, range_config_pb2
from tfx.components.example_gen import utils

from custom_components.vertex_uploader import VertexUploader
from custom_components.vertex_deployer import VertexDeployer


def _create_pipeline(
    input_config: data_types.RuntimeParameter,
    output_config: data_types.RuntimeParameter,
    pipeline_name: str,
    pipeline_root: str,
    data_root: str,
    serving_model_dir: str,
    trainer_module: str,
    project_id: str,
    region: str,
) -> tfx.dsl.Pipeline:
    """Creates a three component flowers pipeline with TFX."""
    example_gen = tfx.components.ImportExampleGen(
        input_base=data_root, input_config=input_config, output_config=output_config
    )

    # Trainer
    trainer = tfx.extensions.google_cloud_ai_platform.Trainer(
        module_file=trainer_module,
        examples=example_gen.outputs["examples"],
        train_args=tfx.proto.TrainArgs(splits=["train"], num_steps=50000 // 64),
        eval_args=tfx.proto.EvalArgs(splits=["val"], num_steps=10000 // 64),
        custom_config={
            tfx.extensions.google_cloud_ai_platform.ENABLE_VERTEX_KEY: True,
            tfx.extensions.google_cloud_ai_platform.VERTEX_REGION_KEY: region,
            tfx.extensions.google_cloud_ai_platform.TRAINING_ARGS_KEY: TRAINING_JOB_SPEC,
            "use_gpu": True,
        },
    ).with_id("trainer")

    # Pushes the model to a filesystem destination.
    pushed_model_location = os.path.join(serving_model_dir, "resnet50")
    resnet_pusher = tfx.components.Pusher(
        model=trainer.outputs["model"],
        push_destination=tfx.proto.PushDestination(
            filesystem=tfx.proto.PushDestination.Filesystem(
                base_directory=pushed_model_location
            )
        ),
    ).with_id("resnet_pusher")

    # Vertex AI upload.
    model_display_name = "resnet_cifar_latest"
    uploader = VertexUploader(
        project=project_id,
        region=region,
        model_display_name=model_display_name,
        pushed_model_location=pushed_model_location,
        serving_image_uri="us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-5:latest",
    ).with_id("vertex_uploader")
    uploader.add_upstream_node(resnet_pusher)

    # Create an endpoint.
    deployer = VertexDeployer(
        project=project_id,
        region=region,
        model_display_name=model_display_name,
        deployed_model_display_name=model_display_name + "_" + TIMESTAMP,
    ).with_id("vertex_deployer")
    deployer.add_upstream_node(uploader)

    components = [
        example_gen,
        trainer,
        resnet_pusher,
        uploader,
        deployer,
    ]

    return tfx.dsl.Pipeline(
        pipeline_name=pipeline_name,
        pipeline_root=pipeline_root,
        components=components,
        enable_cache=True,
    )import os

PIPELINE_DEFINITION_FILE = PIPELINE_NAME + "_pipeline.json"

# Important: We need to pass the custom Docker image URI to the
# `KubeflowV2DagRunnerConfig` to take effect.
runner = tfx.orchestration.experimental.KubeflowV2DagRunner(
    config=tfx.orchestration.experimental.KubeflowV2DagRunnerConfig(
        default_image=TFX_IMAGE_URI
    ),
    output_filename=PIPELINE_DEFINITION_FILE,
)

_ = runner.run(
    _create_pipeline(
        input_config=tfx.dsl.experimental.RuntimeParameter(
            name="input-config",
            default='{"input_config": {"splits": [{"name":"train", "pattern":"span-1/train/tfrecord"}, {"name":"val", "pattern":"span-1/test/tfrecord"}]}}',
            ptype=str,
        ),
        output_config=tfx.dsl.experimental.RuntimeParameter(
            name="output-config", default="{}", ptype=str,
        ),
        pipeline_name=PIPELINE_NAME,
        pipeline_root=PIPELINE_ROOT,
        data_root=DATA_ROOT,
        serving_model_dir=SERVING_MODEL_DIR,
        trainer_module=os.path.join(MODULE_ROOT, _trainer_module_file),
        project_id=GOOGLE_CLOUD_PROJECT,
        region=GOOGLE_CLOUD_REGION,
    )
)from kfp.v2.google import client

pipelines_client = client.AIPlatformClient(
    project_id=GOOGLE_CLOUD_PROJECT, region=GOOGLE_CLOUD_REGION,
)import json
from tfx.orchestration import data_types

_ = pipelines_client.create_run_from_job_spec(
    PIPELINE_DEFINITION_FILE,
    enable_caching=False,
    parameter_values={
        "input-config": json.dumps(
            {
                "splits": [
                    {"name": "train", "pattern": "span-[12]/train/*.tfrecord"},
                    {"name": "val", "pattern": "span-[12]/test/*.tfrecord"},
                ]
            }
        ),
        "output-config": json.dumps({}),
    },
)