! pip uninstall -y kfp
! pip install kfp

import kfp
import json
import os
from kfp.onprem import use_k8s_secret
from kfp import components
from kfp.components import load_component_from_file, load_component_from_url
from kfp import dsl
from kfp import compiler

import numpy as np
import logging


kfp.__version__

"""
Explanation: KubeFlow Pipelines :  Pytorch Cifar10 Image classification
This notebook shows PyTorch CIFAR10 end-to-end  classification example using Kubeflow Pipelines. 
An example notebook that demonstrates how to:

Get different tasks needed for the pipeline
Create a Kubeflow pipeline
Include Pytorch KFP components to preprocess, train, visualize and deploy the model in the pipeline
Submit a job for execution
Query(prediction and explain) the final deployed model
Interpretation of the model using the Captum Insights

import the necessary packages
End of explanation
"""


INGRESS_GATEWAY='http://istio-ingressgateway.istio-system.svc.cluster.local'
AUTH="<enter your auth token>"
NAMESPACE="kubeflow-user-example-com"
COOKIE="authservice_session="+AUTH
EXPERIMENT="Default"

"""
Explanation: Enter your gateway and the auth token
Use this extension on chrome to get token

Update values for the ingress gateway and auth session
End of explanation
"""


MINIO_ENDPOINT="http://minio-service.kubeflow:9000"
LOG_BUCKET="mlpipeline"
TENSORBOARD_IMAGE="public.ecr.aws/pytorch-samples/tboard:latest"

"""
Explanation: Set  the Log bucket and Tensorboard Image
End of explanation
"""


client = kfp.Client(host=INGRESS_GATEWAY+"/pipeline", cookies=COOKIE)

client.create_experiment(EXPERIMENT)
experiments = client.list_experiments(namespace=NAMESPACE)
my_experiment = experiments.experiments[0]
my_experiment

"""
Explanation: Set the client and create the experiment
End of explanation
"""


DEPLOY_NAME="torchserve"
MODEL_NAME="cifar10"
ISVC_NAME=DEPLOY_NAME+"."+NAMESPACE+"."+"example.com"
INPUT_REQUEST="https://raw.githubusercontent.com/kubeflow/pipelines/master/samples/contrib/pytorch-samples/cifar10/input.json"

"""
Explanation: Set  the Inference parameters
End of explanation
"""


! python utils/generate_templates.py cifar10/template_mapping.json

prepare_tensorboard_op = load_component_from_file("yaml/tensorboard_component.yaml")

prep_op = components.load_component_from_file(
    "yaml/preprocess_component.yaml"
)

train_op = components.load_component_from_file(
    "yaml/train_component.yaml"
)
deploy_op = load_component_from_file("yaml/deploy_component.yaml")

pred_op = load_component_from_file("yaml/prediction_component.yaml")

minio_op = components.load_component_from_file(
    "yaml/minio_component.yaml"
)


"""
Explanation: Load the the components yaml files for setting up the components
End of explanation
"""



@dsl.pipeline(
    name="Training Cifar10 pipeline", description="Cifar 10 dataset pipeline"
)
def pytorch_cifar10( # pylint: disable=too-many-arguments
    minio_endpoint=MINIO_ENDPOINT,
    log_bucket=LOG_BUCKET,
    log_dir=f"tensorboard/logs/{dsl.RUN_ID_PLACEHOLDER}",
    mar_path=f"mar/{dsl.RUN_ID_PLACEHOLDER}/model-store",
    config_prop_path=f"mar/{dsl.RUN_ID_PLACEHOLDER}/config",
    model_uri=f"s3://mlpipeline/mar/{dsl.RUN_ID_PLACEHOLDER}",
    tf_image=TENSORBOARD_IMAGE,
    deploy=DEPLOY_NAME,
    isvc_name=ISVC_NAME,
    model=MODEL_NAME,
    namespace=NAMESPACE,
    confusion_matrix_log_dir=f"confusion_matrix/{dsl.RUN_ID_PLACEHOLDER}/",
    checkpoint_dir="checkpoint_dir/cifar10",
    input_req=INPUT_REQUEST,
    cookie=COOKIE,
    ingress_gateway=INGRESS_GATEWAY,
):

    def sleep_op(seconds):
        """Sleep for a while."""
        return dsl.ContainerOp(
            name="Sleep " + str(seconds) + " seconds",
            image="python:alpine3.6",
            command=["sh", "-c"],
            arguments=[
                'python -c "import time; time.sleep($0)"',
                str(seconds)
            ],
        )

    """This method defines the pipeline tasks and operations"""
    pod_template_spec = json.dumps({
        "spec": {
            "containers": [{
                "env": [
                    {
                        "name": "AWS_ACCESS_KEY_ID",
                        "valueFrom": {
                            "secretKeyRef": {
                                "name": "mlpipeline-minio-artifact",
                                "key": "accesskey",
                            }
                        },
                    },
                    {
                        "name": "AWS_SECRET_ACCESS_KEY",
                        "valueFrom": {
                            "secretKeyRef": {
                                "name": "mlpipeline-minio-artifact",
                                "key": "secretkey",
                            }
                        },
                    },
                    {
                        "name": "AWS_REGION",
                        "value": "minio"
                    },
                    {
                        "name": "S3_ENDPOINT",
                        "value": f"{minio_endpoint}",
                    },
                    {
                        "name": "S3_USE_HTTPS",
                        "value": "0"
                    },
                    {
                        "name": "S3_VERIFY_SSL",
                        "value": "0"
                    },
                ]
            }]
        }
    })

    prepare_tb_task = prepare_tensorboard_op(
        log_dir_uri=f"s3://{log_bucket}/{log_dir}",
        image=tf_image,
        pod_template_spec=pod_template_spec,
    ).set_display_name("Visualization")

    prep_task = (
        prep_op().after(prepare_tb_task
                       ).set_display_name("Preprocess & Transform")
    )
    confusion_matrix_url = f"minio://{log_bucket}/{confusion_matrix_log_dir}"
    script_args = f"model_name=resnet.pth," \
                  f"confusion_matrix_url={confusion_matrix_url}"
    # For GPU, set number of gpus and accelerator type
    ptl_args = f"max_epochs=1, gpus=0, accelerator=None, profiler=pytorch"
    train_task = (
        train_op(
            input_data=prep_task.outputs["output_data"],
            script_args=script_args,
            ptl_arguments=ptl_args
        ).after(prep_task).set_display_name("Training")
    )
    # For GPU uncomment below line and set GPU limit and node selector
    # ).set_gpu_limit(1).add_node_selector_constraint
    # ('cloud.google.com/gke-accelerator','nvidia-tesla-p4')

    (
        minio_op(
            bucket_name="mlpipeline",
            folder_name=log_dir,
            input_path=train_task.outputs["tensorboard_root"],
            filename="",
        ).after(train_task).set_display_name("Tensorboard Events Pusher")
    )

    (
        minio_op(
            bucket_name="mlpipeline",
            folder_name=checkpoint_dir,
            input_path=train_task.outputs["checkpoint_dir"],
            filename="",
        ).after(train_task).set_display_name("checkpoint_dir Pusher")
    )

    minio_mar_upload = (
        minio_op(
            bucket_name="mlpipeline",
            folder_name=mar_path,
            input_path=train_task.outputs["checkpoint_dir"],
            filename="cifar10_test.mar",
        ).after(train_task).set_display_name("Mar Pusher")
    )

    (
        minio_op(
            bucket_name="mlpipeline",
            folder_name=config_prop_path,
            input_path=train_task.outputs["checkpoint_dir"],
            filename="config.properties",
        ).after(train_task).set_display_name("Conifg Pusher")
    )

    model_uri = str(model_uri)
    # pylint: disable=unused-variable
    isvc_yaml = """
    apiVersion: "serving.kubeflow.org/v1beta1"
    kind: "InferenceService"
    metadata:
      name: {}
      namespace: {}
    spec:
      predictor:
        serviceAccountName: sa
        pytorch:
          storageUri: {}
          resources:
            requests: 
              cpu: 4
              memory: 16Gi
            limits:
              cpu: 4
              memory: 16Gi
    """.format(
        deploy, namespace, model_uri
    )
    
    # For GPU inference use below yaml with gpu count and accelerator
    gpu_count = "1"
    accelerator = "nvidia-tesla-p4"
    isvc_gpu_yaml = """# pylint: disable=unused-variable
    apiVersion: "serving.kubeflow.org/v1beta1"
    kind: "InferenceService"
    metadata:
      name: {}
      namespace: {}
    spec:
      predictor:
        serviceAccountName: sa
        pytorch:
          storageUri: {}
          resources:
            requests: 
              cpu: 4
              memory: 16Gi
            limits:
              cpu: 4
              memory: 16Gi
              nvidia.com/gpu: {}
          nodeSelector:
            cloud.google.com/gke-accelerator: {}
""".format(deploy, namespace, model_uri, gpu_count, accelerator)
    # Update inferenceservice_yaml for GPU inference
    deploy_task = (
        deploy_op(action="apply", inferenceservice_yaml=isvc_yaml
                 ).after(minio_mar_upload).set_display_name("Deployer")
    )
    # Wait here for model to be loaded in torchserve for inference
    sleep_task = sleep_op(5).after(deploy_task).set_display_name("Sleep")
    # Make Inference request
    pred_task = (
        pred_op(
            host_name=isvc_name,
            input_request=input_req,
            cookie=cookie,
            url=ingress_gateway,
            model=model,
            inference_type="predict",
        ).after(sleep_task).set_display_name("Prediction")
    )
    (
        pred_op(
            host_name=isvc_name,
            input_request=input_req,
            cookie=cookie,
            url=ingress_gateway,
            model=model,
            inference_type="explain",
        ).after(pred_task).set_display_name("Explanation")
    )

    dsl.get_pipeline_conf().add_op_transformer(
        use_k8s_secret(
            secret_name="mlpipeline-minio-artifact",
            k8s_secret_key_to_env={
                "secretkey": "MINIO_SECRET_KEY",
                "accesskey": "MINIO_ACCESS_KEY",
            },
        )
    )

"""
Explanation: Define the pipeline
End of explanation
"""


compiler.Compiler().compile(pytorch_cifar10, 'pytorch.tar.gz', type_check=True)

"""
Explanation: Compile the pipeline
End of explanation
"""


run = client.run_pipeline(my_experiment.id, 'pytorch-cifar10', 'pytorch.tar.gz')

"""
Explanation: Execute the pipeline
End of explanation
"""


!kubectl get isvc $DEPLOY

"""
Explanation: Wait for inference service below to go to READY True state
End of explanation
"""


INFERENCE_SERVICE_LIST = ! kubectl get isvc {DEPLOY_NAME} -n {NAMESPACE} -o json | python3 -c "import sys, json; print(json.load(sys.stdin)['status']['url'])"| tr -d '"' | cut -d "/" -f 3
INFERENCE_SERVICE_NAME = INFERENCE_SERVICE_LIST[0]
INFERENCE_SERVICE_NAME

"""
Explanation: Get  the Inference service name
End of explanation
"""


!curl -v -H "Host: $INFERENCE_SERVICE_NAME" -H "Cookie: $COOKIE" "$INGRESS_GATEWAY/v1/models/$MODEL_NAME:predict" -d @./cifar10/input.json > cifar10_prediction_output.json

! cat cifar10_prediction_output.json

"""
Explanation: Use the deployed model for prediction request and save the output into a json
End of explanation
"""


!curl -v -H "Host: $INFERENCE_SERVICE_NAME" -H "Cookie: $COOKIE" "$INGRESS_GATEWAY/v1/models/$MODEL_NAME:explain" -d @./cifar10/input.json > cifar10_explanation_output.json

"""
Explanation: Use the deployed model for explain request and save the output into a json
End of explanation
"""


!./install-dependencies.sh

"""
Explanation: Model Interpretation using Captum Vis and Insights
Install dependencies for Captum Insights
End of explanation
"""


from PIL import Image
import numpy as np 
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import torchvision.transforms as transforms
import torch
import torch.nn.functional as F
import json
import captum
from captum.attr import LayerAttribution
from captum.attr import visualization as viz
import base64
import os
import io

"""
Explanation: import the necessary packages
End of explanation
"""


prediction_json = json.loads(open("./cifar10_prediction_output.json", "r").read())

explainations_json = json.loads(open("./cifar10_explanation_output.json", "r").read())

labels_path = './cifar10/class_mapping.json'
with open(labels_path) as json_data:
    idx_to_labels = json.load(json_data) 

count = 0
for i in range(0, len(explainations_json["explanations"])):

    image = base64.b64decode(explainations_json["explanations"][i]["b64"])       
    fileName = 'captum_kitten_{}.jpeg'.format(count)

    imagePath = ( os.getcwd() +"/" + fileName)
    img = Image.open(io.BytesIO(image))
    img = img.convert('RGB')
    img.save(imagePath, 'jpeg', quality=100)
    print("Saving ", imagePath)
    count += 1
    

from IPython.display import Image
Image(filename='captum_kitten_0.jpeg') 

Image(filename='captum_kitten_1.jpeg') 

Image(filename='captum_kitten_2.jpeg') 

"""
Explanation: Read the prediction, explanation, and the class mapping file which saved during the prediction and expalain requests.
End of explanation
"""


from minio import Minio
from kubernetes import client, config
import base64

config.load_incluster_config()
v1 = client.CoreV1Api()
sec = v1.read_namespaced_secret("mlpipeline-minio-artifact", NAMESPACE).data
minio_accesskey = base64.b64decode(sec["accesskey"]).decode('UTF-8')
minio_secretkey = base64.b64decode(sec["secretkey"]).decode('UTF-8')


minio_config = {
            "HOST": "minio-service.kubeflow:9000",
            "ACCESS_KEY": minio_accesskey,
            "SECRET_KEY": minio_secretkey,
            "BUCKET": "mlpipeline",
            "FOLDER": "checkpoint_dir/cifar10"}

def _initiate_minio_client(minio_config):
        minio_host = minio_config["HOST"]
        access_key = minio_config["ACCESS_KEY"]
        secret_key = minio_config["SECRET_KEY"]
        client = Minio(minio_host, access_key=access_key, secret_key=secret_key, secure=False)
        return client

client= _initiate_minio_client(minio_config)
client

def download_artifact_from_minio(folder: str, artifact: str):
        artifact_name = artifact.split("/")[-1]
        result = client.fget_object(
            minio_config["BUCKET"],
            os.path.join(folder, artifact_name),
            artifact,
        )

download_artifact_from_minio(minio_config["FOLDER"],"resnet.pth")
print("[INFO] Downloaded the Model Pth File.....")
download_artifact_from_minio(minio_config["FOLDER"],"cifar10_train.py")
print("[INFO] Downloaded the Model Classifier File.....")

"""
Explanation: Captum Insights can also be used for visualization
Define the minio client for downloading the artifactes from minio storage ( model pth file and training file)
End of explanation
"""


from cifar10_train import CIFAR10Classifier
model = CIFAR10Classifier()

model_pt_path ="./resnet.pth"
model.load_state_dict(torch.load(model_pt_path,map_location=torch.device('cpu')))
model.eval()

#Lets read two test images and make the prediction and use these images for captum Insights. 

imgs = ['./cifar10/kitten.png',"./cifar10/horse.png"]

for img in imgs:
    img = Image.open(img)
    transformed_img = transform(img)
    input_img = transform_normalize(transformed_img)
    input_img = input_img.unsqueeze(0) # the model requires a dummy batch dimension

    output = model(input_img)
    output = F.softmax(output, dim=1)
    prediction_score, pred_label_idx = torch.topk(output, 1)
    pred_label_idx.squeeze_()
    predicted_label = idx_to_labels[str(pred_label_idx.squeeze_().item())]
    print('Predicted:', predicted_label, '/', pred_label_idx.item(), ' (', prediction_score.squeeze().item(), ')')

from captum.insights import AttributionVisualizer, Batch
from captum.insights.attr_vis.features import ImageFeature

# Baseline is all-zeros input - this may differ depending on your data
def baseline_func(input):
    return input * 0

# merging our image transforms from above
def full_img_transform(input):
    i = Image.open(input)
    i = transform(i)
    i = transform_normalize(i)
    i = i.unsqueeze(0)
    i.requires_grad = True
    return i


input_imgs = torch.cat(list(map(lambda i: full_img_transform(i), imgs)), 0)

visualizer = AttributionVisualizer(
    models=[model],
    score_func=lambda o: torch.nn.functional.softmax(o, 1),
    classes=list(map(lambda k: idx_to_labels[k], idx_to_labels.keys())),
    features=[
        ImageFeature(
            "Photo",
            baseline_transforms=[baseline_func],
            input_transforms=[],
        )
    ],
    dataset=[Batch(input_imgs, labels=[3,7])]
)

visualizer.serve(debug=True,port=6080)

"""
Explanation: Load the downloaded model pth file and classifer
End of explanation
"""


! kubectl delete --all isvc -n $NAMESPACE

! kubectl delete pod --field-selector=status.phase==Succeeded -n $NAMESPACE

"""
Explanation: Captum Insights output image

Clean up
Delete Viewers, Inference Services and Completed pods
End of explanation
"""


