google_ai_platform_helpers.py

import glob
import logging
import os
import time

from google.cloud import storage as cloud_storage
from googleapiclient import discovery, errors


SAVED_MODEL_PATH = '/path/to/local/model'
CLOUD_STORAGE_BUCKET = 'bucket-name'
PROJECT_NAME = 'project-id-found-in-json-credentials'
MODEL_NAME = 'existing_model_name' 
PROJECT_ID = f'projects/{PROJECT_NAME}'
MODEL_ID = f'models/{MODEL_NAME}'
# set env vr that google expects
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = '/path/to/your/json/credentials.json'
ml = discovery.build('ml', 'v1')


def upload_model_to_gcs(version_name: str) -> None:
    gcs = cloud_storage.Client()
    bucket = gcs.get_bucket(CLOUD_STORAGE_BUCKET)
    local_path = f'{SAVED_MODEL_PATH}/{version_name}'
    copy_local_directory_to_gcs(local_path, bucket, version_name)
    logging.info(f'finished uploading model version: {version_name}')


def copy_local_directory_to_gcs(directory: str, bucket: cloud_storage.bucket.Bucket, gcs_path: str) -> None:
    """
    recursively copy a directory of directories to google cloud storage.
    directory should not have a trailing slash.
    adapted from:
    https://stackoverflow.com/questions/48514933/how-to-copy-a-directory-to-google-cloud-storage-using-google-cloud-python-api
    """
    assert os.path.isdir(directory)
    for storage_location in glob.glob(directory + '/**'):
        storage_location_name = storage_location[1 + len(directory):]
        if os.path.isfile(storage_location):
            remote_path = os.path.join(gcs_path, storage_location_name)
            blob = bucket.blob(remote_path)
            blob.upload_from_filename(storage_location)
            logging.info(f'uploaded {storage_location} to {remote_path}')
        elif os.path.isdir(storage_location):
            copy_local_directory_to_gcs(storage_location, bucket, f'{gcs_path}/{storage_location_name}')


def set_model_default(version_name: str, backoff: int = 1, tries: int = 0) -> None:
    MAX_TRIES = 3
    version_id = f'versions/{version_name}'
    request = ml.projects().models().versions().setDefault(name=f'{PROJECT_ID}/{MODEL_ID}/{version_id}')

    try:
        request.execute()
        logging.info(f'set new ai platform model default to version: {version_name}')
    except errors.HttpError as e:
        logging.info(f'there was an error setting the ai platform model default: {e._get_reason()}')
        if tries < MAX_TRIES:
            logging.info(f'sleeping for {backoff} minutes and retrying set_model_default...')
            time.sleep(backoff * 60)
            set_model_default(version_name, backoff * 2, tries + 1)


def create_model_version(version_name: str) -> None:
    request_dict = {
        "name": version_name,
        "deploymentUri": f"gs://{CLOUD_STORAGE_BUCKET}/{version_name}",
        "runtimeVersion": "1.15",
        "framework": "tensorflow",
        "pythonVersion": "3.7",
    }
    request = ml.projects().models().versions().create(parent=f'{PROJECT_ID}/{MODEL_ID}', body=request_dict)

    try:
        request.execute()
        logging.info(f'created new ai platform model version: {version_name}')
    except errors.HttpError as e:
        logging.info(f'there was an error creating the ai platform model version: {e._get_reason()}')