import os

# The Google Cloud Notebook product has specific requirements
IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists("/opt/deeplearning/metadata/env_version")

# Google Cloud Notebook requires dependencies to be installed with '--user'
USER_FLAG = ""
if IS_GOOGLE_CLOUD_NOTEBOOK:
    USER_FLAG = "--user"# Automatically restart kernel after installs
import os

if not os.getenv("IS_TESTING"):
    # Automatically restart kernel after installs
    import IPython

    app = IPython.Application.instance()
    app.kernel.do_shutdown(True)import os

PROJECT_ID = "grounded-atrium-320207"

# Get your Google Cloud project ID from gcloud
if not os.getenv("IS_TESTING"):
    shell_output=!gcloud config list --format 'value(core.project)' 2>/dev/null
    PROJECT_ID = shell_output[0]
    print("Project ID: ", PROJECT_ID)import os
import sys

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

# The Google Cloud Notebook product has specific requirements
IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists("/opt/deeplearning/metadata/env_version")

# If on Google Cloud Notebooks, then don't execute this code
if not IS_GOOGLE_CLOUD_NOTEBOOK:
    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 ''BUCKET_NAME = "gs://vertexai_dual_example"
REGION      = "us-central1"  PATH=%env PATH
%env PATH={PATH}:/home/jupyter/.local/bin

USER = "chansung"
PIPELINE_ROOT = "{}/pipeline_root/{}".format(BUCKET_NAME, USER)
PIPELINE_ROOTimport kfp
from google.cloud import aiplatform
from google_cloud_pipeline_components import aiplatform as gcc_aip
from kfp.v2 import compiler
from kfp.v2.google.client import AIPlatformClient
from kfp.v2 import dsl
from kfp.v2.dsl import component@component(
    packages_to_install=["google-cloud-storage", "firebase-admin", "tensorflow"]
)
def push_to_firebase(
    credential_uri: str,
    model_bucket: str,
    firebase_dest_gcs_bucket: str,
    model_display_name: str,
    model_tag: str
):
    import firebase_admin
    from firebase_admin import ml
    from firebase_admin import storage
    from firebase_admin import credentials    
    from google.cloud import storage as gcs_storage
    
    gcs_client = gcs_storage.Client()
    
    # get credential for firebase  
    credential_gcs_bucket = credential_uri.split('//')[1].split('/')[0]
    credential_blob_path = '/'.join(credential_uri.split('//')[1].split('/')[1:])
    
    bucket = gcs_client.bucket(credential_gcs_bucket)
    blob = bucket.blob(credential_blob_path)
    blob.download_to_filename('credential.json')
    
    # get the latest model    
    tflite_blobs = gcs_client.get_bucket(model_bucket).list_blobs()
    tflite_blob = sorted(tflite_blobs, reverse=True, key=lambda blob: blob.name.split('/')[-2])[0]
    tflite_blob.download_to_filename('model.tflite')            
            
    firebase_admin.initialize_app(
        credentials.Certificate('credential.json'),
        options={
            'storageBucket': firebase_dest_gcs_bucket
        }
    )

    model_list = ml.list_models(list_filter=f'display_name={model_display_name}')
    # update
    if len(model_list.models) > 0:
        # get the first match model
        model = model_list.models[0]
        
        source = ml.TFLiteGCSModelSource.from_tflite_model_file('model.tflite')
        model.model_format = ml.TFLiteFormat(model_source=source)
        
        updated_model = ml.update_model(model)
        ml.publish_model(updated_model.model_id)
    # create
    else:    
        # Load a tflite file and upload it to Cloud Storage
        source = ml.TFLiteGCSModelSource.from_tflite_model_file('model.tflite')

        # Create the model object
        tflite_format = ml.TFLiteFormat(model_source=source)
        model = ml.Model(
            display_name=model_display_name,  # This is the name you use from your app to load the model.
            tags=[model_tag],             # Optional tags for easier management.
            model_format=tflite_format)

        # Add the model to your Firebase project and publish it
        new_model = ml.create_model(model)
        ml.publish_model(new_model.model_id)@kfp.dsl.pipeline(name="cloud-mobile-dual-deployment")
def pipeline(project: str = PROJECT_ID):
    ds_op = gcc_aip.ImageDatasetCreateOp(
       project=project,
       display_name="flowers-dataset",
       gcs_source="gs://dataset-meta-gde-csp/flowers_vertex.csv",
       import_schema_uri=aiplatform.schema.dataset.ioformat.image.multi_label_classification,
    )

    configs = [
       {
          "type": "CLOUD",
          "model_type": "CLOUD",
          "display_name": "train-cloud-model",
          "model_display_name": "cloud-model",
          "budget_milli_node_hours": 8000,
       },
       {
          "type": "MOBILE",
          "model_type": "MOBILE_TF_VERSATILE_1",
          "display_name": "train-mobile-model",
          "model_display_name": "mobile-model",
          "budget_milli_node_hours": 1000,
       }
    ]

    with kfp.dsl.ParallelFor(configs) as config:
        training_job_run_op = gcc_aip.AutoMLImageTrainingJobRunOp(
            project=project,
            display_name=config.display_name,
            prediction_type="classification",
            multi_label=True,
            model_type=config.model_type,
            base_model=None,
            dataset=ds_op.outputs["dataset"],
            model_display_name=config.model_display_name,
            budget_milli_node_hours=config.budget_milli_node_hours,
        )
        training_job_run_op.after(ds_op)

        with kfp.dsl.Condition(config.type=='CLOUD'):
            endpoint_op = gcc_aip.ModelDeployOp(
                project=project,
                model=training_job_run_op.outputs["model"]
            )
            endpoint_op.after(training_job_run_op)
            
        with kfp.dsl.Condition(config.type=='MOBILE'):
            export_op = gcc_aip.ModelExportOp( 
                project=project,
                model=training_job_run_op.outputs["model"],
                # tflite, edgetpu-tflite, tf-saved-model, tf-js, core-ml, custom-trained
                export_format_id="tflite",
                artifact_destination="gs://output-model-gde-csp/flower-models/"
            )
            export_op.after(training_job_run_op)

            credential_uri="gs://firebase-ml-bucket-gde-csp/grounded-atrium-320207-firebase-adminsdk-5n9sn-20dbda9947.json"
            model_bucket="output-model-gde-csp"
            firebase_bucket="grounded-atrium-320207.appspot.com"
            
            firebase_op = push_to_firebase(
                ins=export_op.outputs['exported_dataset'],
                credential_uri=credential_uri,
                model_bucket=model_bucket,
                firebase_dest_gcs_bucket=firebase_bucket,
                model_display_name="custom_model",
                model_tag="from_dual_deployment"
            )
            firebase_op.after(export_op)
                    from kfp.v2 import compiler

compiler.Compiler().compile(
    pipeline_func=pipeline, package_path="cloud-mobile-dual-deployment.json"
)from kfp.v2.google.client import AIPlatformClient

api_client = AIPlatformClient(project_id=PROJECT_ID, region=REGION)response = api_client.create_run_from_job_spec(
    "cloud-mobile-dual-deployment.json",
    pipeline_root=PIPELINE_ROOT,
    parameter_values={"project": PROJECT_ID},
)