from google.colab import auth

auth.authenticate_user()# @title
from fastdot.core import *

tfx_components = [
    "FileListGen",
    "BatchPredictionGen",
    "PerformanceEvaluator",
    "SpanPreparator",
    "PipelineTrigger",
]
block = "TFX Component Workflow"

g = graph_items(seq_cluster(tfx_components, block))
g_file_list_gen_module_file = "file_list_gen.py"%%writefile {_file_list_gen_module_file}

import tfx
from tfx.dsl.component.experimental.decorators import component
from tfx.dsl.component.experimental.annotations import Parameter
from tfx.dsl.component.experimental.annotations import OutputArtifact
from tfx.types.standard_artifacts import String
from google.cloud import storage
from absl import logging


@component
def FileListGen(
    outpath: OutputArtifact[String],
    project: Parameter[str],
    gcs_source_bucket: Parameter[str],
    gcs_source_prefix: Parameter[str] = "",
    output_filename: Parameter[str] = "test-images.txt",
):
    logging.info("FileListGen started")

    client = storage.Client(project=project)
    bucket = client.get_bucket(gcs_source_bucket)
    blobs = bucket.list_blobs(prefix=gcs_source_prefix)
    logging.info("Successfully retrieve the file(jpg) list from GCS path")

    f = open(output_filename, "w")
    for blob in blobs:
        if blob.name.split(".")[-1] == "jpg":
            prefix = ""
            if gcs_source_prefix != "":
                prefix = f"/{gcs_source_prefix}"
            line = f"gs://{gcs_source_bucket}{prefix}/{blob.name}\n"
            f.write(line)
    f.close()
    logging.info(
        f"Successfully created the file list file({output_filename}) in local storage"
    )

    prefix = ""
    if gcs_source_prefix != "":
        prefix = f"{gcs_source_prefix}/"
    blob = bucket.blob(f"{prefix}{output_filename}")
    blob.upload_from_filename(output_filename)
    logging.info(f"Successfully uploaded the file list ({prefix}{output_filename})")

    outpath.value = gcs_source_bucket + "/" + prefix + output_filename_batch_pred_module_file = 'batch_prediction_vertex.py'%%writefile {_batch_pred_module_file}

from google.cloud import storage

from tfx.dsl.component.experimental.annotations import Parameter, InputArtifact
from tfx.dsl.component.experimental.decorators import component
from tfx.types import artifact_utils
from tfx.types.standard_artifacts import String
import google.cloud.aiplatform as vertex_ai

from typing import Union, Sequence
from absl import logging


@component
def BatchPredictionGen(
    gcs_source: InputArtifact[String],
    project: Parameter[str],
    location: Parameter[str],
    model_resource_name: Parameter[str],
    job_display_name: Parameter[str],
    gcs_destination: Parameter[str],
    instances_format: Parameter[str] = "file-list",
    machine_type: Parameter[str] = "n1-standard-2",
    accelerator_count: Parameter[int] = 0,
    accelerator_type: Parameter[str] = None,
    starting_replica_count: Parameter[int] = 1,
    max_replica_count: Parameter[int] = 1,
):
    storage_client = storage.Client()

    # Read GCS Source (gcs_source contains the full path of GCS object)
    # 1-1. get bucketname from gcs_source
    gcs_source_uri = gcs_source.uri.split("//")[1:][0].split("/")
    bucketname = gcs_source_uri[0]
    bucket = storage_client.get_bucket(bucketname)
    logging.info(f"bucketname: {bucketname}")

    # 1-2. get object path without the bucketname
    objectpath = "/".join(gcs_source_uri[1:])

    # 1-3. read the object to get value set by OutputArtifact from FileListGen
    blob = bucket.blob(objectpath)
    logging.info(f"objectpath: {objectpath}")

    gcs_source = f"gs://{blob.download_as_text()}"

    # Get Model
    vertex_ai.init(project=project, location=location)

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

    # Batch Predictions
    logging.info("Starting batch prediction job.")
    logging.info(f"GCS path where file list is: {gcs_source}")

    batch_prediction_job = model.batch_predict(
        job_display_name=job_display_name,
        instances_format=instances_format,
        gcs_source=gcs_source,
        gcs_destination_prefix=gcs_destination,
        machine_type=machine_type,
        accelerator_count=accelerator_count,
        accelerator_type=accelerator_type,
        starting_replica_count=starting_replica_count,
        max_replica_count=max_replica_count,
        sync=True,
    )

    logging.info(batch_prediction_job.display_name)
    logging.info(batch_prediction_job.resource_name)
    logging.info(batch_prediction_job.state)_evaluator_module_file = 'batch_pred_evaluator.py'%%writefile {_evaluator_module_file}

# Reference: https://bit.ly/vertex-batch

from tfx.dsl.component.experimental.annotations import Parameter
from tfx.dsl.component.experimental.annotations import OutputArtifact
from tfx.dsl.component.experimental.decorators import component
from tfx.types.experimental.simple_artifacts import Dataset

from absl import logging
import os
import json

@component
def PerformanceEvaluator(
    gcs_destination: Parameter[str],
    local_directory: Parameter[str],
    threshold: Parameter[float],
    trigger_pipeline: OutputArtifact[Dataset],
):
    full_gcs_results_dir = f"{gcs_destination}/{local_directory}"

    # Create missing directories.
    os.makedirs(local_directory, exist_ok=True)

    # Get the Cloud Storage paths for each result.
    os.system(f"gsutil -m cp -r {full_gcs_results_dir} {local_directory}")

    # Get most recently modified directory.
    latest_directory = max(
        [os.path.join(local_directory, d) for d in os.listdir(local_directory)],
        key=os.path.getmtime,
    )

    # Get downloaded results in directory.
    results_files = []
    for dirpath, subdirs, files in os.walk(latest_directory):
        for file in files:
            if file.startswith("prediction.results"):
                results_files.append(os.path.join(dirpath, file))

    # Consolidate all the results into a list.
    results = []
    for results_file in results_files:
        # Download each result.
        with open(results_file, "r") as file:
            results.extend([json.loads(line) for line in file.readlines()])

    # Calculate performance.
    num_correct = 0

    for result in results:
        label = os.path.basename(result["instance"]).split("_")[0]
        prediction = result["prediction"]["label"]

        if label == prediction:
            num_correct = num_correct + 1

    accuracy = num_correct / len(results)
    logging.info(f"Accuracy: {accuracy*100}%")
    trigger_pipeline.set_string_custom_property("result", str(accuracy >= threshold))_span_preparator_module_file = 'span_preparator.py'%%writefile {_span_preparator_module_file}

import tfx
from tfx.dsl.component.experimental.decorators import component
from tfx.dsl.component.experimental.annotations import Parameter
from tfx.dsl.component.experimental.annotations import OutputArtifact, InputArtifact
from tfx.types.experimental.simple_artifacts import Dataset
from google.cloud import storage
from absl import logging

from datetime import datetime
import tensorflow as tf
import random
import gzip
import os

# Label-mapping.
LABEL_DICT = {
    "airplane": 0,
    "automobile": 1,
    "bird": 2,
    "cat": 3,
    "deer": 4,
    "dog": 5,
    "frog": 6,
    "horse": 7,
    "ship": 8,
    "truck": 9,
}

# Images are byte-strings.
def _bytestring_feature(list_of_bytestrings):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=list_of_bytestrings))


# Classes would be integers.
def _int_feature(list_of_ints):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=list_of_ints))


# Function that prepares a record for the tfrecord file
# a record contains the image and its label.
def to_tfrecord(img_bytes, label):
    feature = {
        "image": _bytestring_feature([img_bytes]),
        "label": _int_feature([label]),
    }
    return tf.train.Example(features=tf.train.Features(feature=feature))


def write_tfrecords(filepaths, dest_gcs, tfrecord_filename, new_span, is_train):
    # For this project, we are serializing the images in one TFRecord only.
    # For more realistic purposes, this should be sharded.
    folder = "train" if is_train else "test"

    with tf.io.TFRecordWriter(tfrecord_filename) as writer:
        for path in filepaths:
            image_string = tf.io.read_file(path).numpy()
            class_name = path.split("/")[-1].split("_")[0]
            label = LABEL_DICT[class_name]

            example = to_tfrecord(image_string, label)
            writer.write(example.SerializeToString())

    # Copy over the zipped TFRecord file to the GCS Bucket and
    # remove the temporary files.
    logging.info(f"gsutil cp {tfrecord_filename} {dest_gcs}/span-{new_span}/{folder}/")
    os.system(f"gsutil cp {tfrecord_filename} {dest_gcs}/span-{new_span}/{folder}/")
    os.remove(tfrecord_filename)


@component
def SpanPreparator(
    is_retrain: InputArtifact[Dataset],
    gcs_source_bucket: Parameter[str],
    gcs_destination_bucket: Parameter[str],
    latest_span_id: OutputArtifact[Dataset],
    gcs_source_prefix: Parameter[str] = "",
):
    if is_retrain.get_string_custom_property("result") == "False":
        last_span_str = tf.io.gfile.glob(f"{gcs_destination_bucket}/span-*")[-1]
        last_span = int(last_span_str.split("-")[-1])
        new_span = last_span + 1

        timestamp = datetime.utcnow().strftime("%y%m%d-%H%M%S")

        image_paths = tf.io.gfile.glob(f"gs://{gcs_source_bucket}/*.jpg")
        logging.info(image_paths)
        random.shuffle(image_paths)

        val_split = 0.2

        split_index = int(len(image_paths) * (1 - val_split))
        training_paths = image_paths[:split_index]
        validation_paths = image_paths[split_index:]

        write_tfrecords(
            training_paths,
            gcs_destination_bucket,
            tfrecord_filename=f"new_training_data_{timestamp}.tfrecord",
            new_span=new_span,
            is_train=True,
        )
        write_tfrecords(
            validation_paths,
            gcs_destination_bucket,
            tfrecord_filename=f"new_validation_data_{timestamp}.tfrecord",
            new_span=new_span,
            is_train=False,
        )

        logging.info("Removing images from batch prediction bucket.")
        os.system(
            f"gsutil mv gs://{gcs_source_bucket}/{gcs_source_prefix} gs://{gcs_source_bucket}/{gcs_source_prefix}_old"
        )
        # os.system(f"gsutil rm -rf gs://{gcs_source_bucket}/*")

        latest_span_id.set_string_custom_property("latest_span", str(new_span))_pipeline_trigger_module_file = 'training_pipeline_trigger.py'%%writefile {_pipeline_trigger_module_file}

import json

from google.cloud import storage

from kfp.v2.google.client import AIPlatformClient
from tfx.dsl.component.experimental.annotations import Parameter, InputArtifact
from tfx.dsl.component.experimental.decorators import component
from tfx.types.experimental.simple_artifacts import Dataset

from absl import logging

@component
def PipelineTrigger(
	is_retrain: InputArtifact[Dataset],
  latest_span_id: InputArtifact[Dataset],
	pipeline_spec_path: Parameter[str],
	project_id: Parameter[str],
	region: Parameter[str],
):
	if is_retrain.get_string_custom_property('result') == 'False':
		# Check if the pipeline spec exists.
		storage_client = storage.Client()

		path_parts = pipeline_spec_path.replace("gs://", "").split("/")
		bucket_name = path_parts[0]
		blob_name = "/".join(path_parts[1:])

		bucket = storage_client.bucket(bucket_name)
		blob = storage.Blob(bucket=bucket, name=blob_name)

		if not blob.exists(storage_client):
				raise ValueError(f"{pipeline_spec_path} does not exist.")

		# Initialize Vertex AI API client and submit for pipeline execution.
		api_client = AIPlatformClient(project_id=project_id, region=region)

		# Fetch the latest span.
		latest_span = latest_span_id.get_string_custom_property('latest_span')

		# Create a training job from pipeline spec.
		response = api_client.create_run_from_job_spec(pipeline_spec_path, 
																						enable_caching=False,
																						parameter_values={
																							'input-config': json.dumps({
																								'splits': [
																									{'name': 'train', 'pattern': f'span-[{int(latest_span)-1}{latest_span}]/train/*.tfrecord'},
																									{'name': 'val', 'pattern': f'span-[{int(latest_span)-1}{latest_span}]/test/*.tfrecord'}
																								]
																							}),
																							'output-config': json.dumps({})
																						})

		logging.info(response)# This bucket will be responsible for storing the pipeline related artifacts.
GOOGLE_CLOUD_PROJECT = "gcp-ml-172005"  # @param {type:"string"}
GOOGLE_CLOUD_REGION = "us-central1"

GCS_BUCKET_NAME = "cifar10-experimental-csp2"  # @param {type:"string"}

MODEL_RESOURCE_NAME = "resnet_cifar_latest"  # @param {type: "string"}

TEST_FILENAME = "test-images.txt"  # @param {type:"string"}
TEST_GCS_BUCKET = "batch-prediction-collection-3"  # @param {type:"string"}
TEST_GCS_PREFIX = ""  # @param {type: "string"}

TRAINING_PIPELINE_SPEC = "gs://cifar10-experimental-csp2/pipeline_root/continuous-adaptation-for-data-changes/continuous-adaptation-for-data-changes_pipeline.json"  # @param {type: "string"}
TRAINING_DATA_PATH = "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-batch'

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

print('PIPELINE_ROOT: {}'.format(PIPELINE_ROOT))DISPLAY_NAME = "batch-predictions-pipeline"
VERSION = "tfx-1-2-0-34"
TFX_IMAGE_URI = f"gcr.io/{GOOGLE_CLOUD_PROJECT}/{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 google-cloud-storage kfp==1.6.1from datetime import datetime

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

from tfx import v1 as tfx

from tfx.orchestration.pipeline import Pipeline
from custom_components.file_list_gen import FileListGen
from custom_components.batch_prediction_vertex import BatchPredictionGen
from custom_components.batch_pred_evaluator import PerformanceEvaluator
from custom_components.span_preparator import SpanPreparator
from custom_components.training_pipeline_trigger import PipelineTrigger


def _create_pipeline(
    pipeline_name: str,
    pipeline_root: str,
    data_gcs_bucket: str,
    data_gcs_prefix: data_types.RuntimeParameter,
    batch_job_gcs: str,
    job_display_name: str,
    model_resource_name: str,
    project_id: str,
    region: str,
    threshold: float,
    data_gcs_destination: str,
    training_pipeline_spec: str,
) -> Pipeline:

    # Generate a file list for batch preditions.
    # More details on the structure of this file here:
    # https://bit.ly/3BzfHVu.
    filelist_gen = FileListGen(
        project=project_id,
        gcs_source_bucket=data_gcs_bucket,
        gcs_source_prefix=data_gcs_prefix,
    ).with_id("filelist_gen")

    # Submit a batch prediction job.
    batch_pred_component = BatchPredictionGen(
        project=project_id,
        location=region,
        job_display_name=job_display_name,
        model_resource_name=model_resource_name,
        gcs_source=filelist_gen.outputs["outpath"],
        gcs_destination=f"gs://{batch_job_gcs}/results/",
        accelerator_count=0,
        accelerator_type=None,
    ).with_id("bulk_inferer_vertex")
    batch_pred_component.add_upstream_node(filelist_gen)

    # Evaluate the performance of the predictions.
    # In a real-world project, this evaluation takes place
    # separately, typically with the help of domain experts.
    final_gcs_destination = f"gs://{batch_job_gcs}/results/"
    evaluator = PerformanceEvaluator(
        gcs_destination=f'gs://{final_gcs_destination.split("/")[2]}',
        local_directory=final_gcs_destination.split("/")[-2],
        threshold=threshold,
    ).with_id("batch_prediction_evaluator")
    evaluator.add_upstream_node(batch_pred_component)

    span_preparator = SpanPreparator(
        is_retrain=evaluator.outputs["trigger_pipeline"],
        gcs_source_bucket=data_gcs_bucket,
        gcs_source_prefix=data_gcs_prefix,
        gcs_destination_bucket=data_gcs_destination,
    ).with_id("span_preparator")
    span_preparator.add_upstream_node(evaluator)

    trigger = PipelineTrigger(
        is_retrain=evaluator.outputs["trigger_pipeline"],
        latest_span_id=span_preparator.outputs["latest_span_id"],
        pipeline_spec_path=training_pipeline_spec,
        project_id=project_id,
        region=region,
    ).with_id("training_pipeline_trigger")
    trigger.add_upstream_node(span_preparator)

    components = [
        filelist_gen,
        batch_pred_component,
        evaluator,
        span_preparator,
        trigger,
    ]

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

import tfx
from tfx.orchestration import data_types
from tfx.orchestration.kubeflow.v2.kubeflow_v2_dag_runner import KubeflowV2DagRunner
from tfx.orchestration.kubeflow.v2.kubeflow_v2_dag_runner import (
    KubeflowV2DagRunnerConfig,
)

PIPELINE_DEFINITION_FILE = PIPELINE_NAME + "_pipeline.json"
THRESHOLD = 0.9

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

_ = runner.run(
    _create_pipeline(
        pipeline_name=PIPELINE_NAME,
        pipeline_root=PIPELINE_ROOT,
        data_gcs_bucket=TEST_GCS_BUCKET,
        data_gcs_prefix=data_types.RuntimeParameter(
            name="data_gcs_prefix", default="", ptype=str
        ),
        batch_job_gcs=GCS_BUCKET_NAME,
        job_display_name=f"{MODEL_RESOURCE_NAME}_{TIMESTAMP}",
        project_id=GOOGLE_CLOUD_PROJECT,
        region=GOOGLE_CLOUD_REGION,
        model_resource_name=MODEL_RESOURCE_NAME,
        threshold=THRESHOLD,
        data_gcs_destination=TRAINING_DATA_PATH,
        training_pipeline_spec=TRAINING_PIPELINE_SPEC,
    )
)from kfp.v2.google import client

pipelines_client = client.AIPlatformClient(
    project_id=GOOGLE_CLOUD_PROJECT, region=GOOGLE_CLOUD_REGION,
)

_ = pipelines_client.create_run_from_job_spec(
    PIPELINE_DEFINITION_FILE,
    enable_caching=False,
    parameter_values={"data_gcs_prefix": "2021-10"},
)