import os

# The Vertex AI Workbench Notebook product has specific requirements
IS_WORKBENCH_NOTEBOOK = os.getenv("DL_ANACONDA_HOME")
IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(
    "/opt/deeplearning/metadata/env_version"
)

# Vertex AI Notebook requires dependencies to be installed with '--user'
USER_FLAG = ""
if IS_WORKBENCH_NOTEBOOK:
    USER_FLAG = "--user"
    USER_FLAG = ""

! pip3 install --upgrade google-cloud-aiplatform $USER_FLAG -q
! pip3 install tensorflow-datasets $USER_FLAG -q

"""
Explanation: E2E ML on GCP: MLOps stage 2 : experiments: get started with prebuilt TFHub models
<table align="left">
  <td>
    <a href="https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage2/get_started_with_tfhub_models.ipynb">
      <img src="https://cloud.google.com/ml-engine/images/github-logo-32px.png" alt="GitHub logo">
      View on GitHub
    </a>
  </td>
  <td>
        <a href="https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage2/get_started_with_tfhub_models.ipynb">
        <img src="https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt="Colab logo"> Run in Colab
        </a>
  </td>
  <td>
    <a href="https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage2/get_started_with_tfhub_models.ipynb">
      <img src="https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32" alt="Vertex AI logo">
      Open in Vertex AI Workbench
    </a>
  </td>
</table>
<br/><br/><br/>
Overview
This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 2 : experimentation: get started with prebuilt TensorFlow Hub (TFHub) models.
Dataset
The dataset used for this tutorial is the CIFAR10 dataset from TensorFlow Datasets. The version of the dataset you will use is built into TensorFlow. The trained model predicts which type of class an image is from ten classes: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck.
Objective
In this tutorial, you learn how to use Vertex AI Training with prebuilt models from TensorFlow Hub.
This tutorial uses the following Google Cloud ML services:

Vertex AI Training
Vertex AI Model resource

The steps performed include:

Download a TensorFlow Hub prebuilt model.
Add the task component as a classifier for the CIFAR-10 dataset.
Fine tune locally the model with transfer learning training.
Construct a custom training script:
Get training data from TensorFlow Datasets
Get model architecture from TensorFlow Hub
Train then model
Save model artifacts and upload as Vertex AI Model resource.



Costs
This tutorial uses billable components of Google Cloud:

Vertex AI
Cloud Storage

Learn about Vertex AI
pricing and Cloud Storage
pricing, and use the Pricing
Calculator
to generate a cost estimate based on your projected usage.
Installation
Install the following packages to execute this notebook.
End of explanation
"""


! pip3 install -U google-cloud-storage $USER_FLAG

"""
Explanation: Install the latest GA version of google-cloud-storage library as well.
End of explanation
"""


import os

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

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

"""
Explanation: Restart the kernel
Once you've installed the additional packages, you need to restart the notebook kernel so it can find the packages.
End of explanation
"""


PROJECT_ID = "[your-project-id]"  # @param {type:"string"}

if PROJECT_ID == "" or PROJECT_ID is None or PROJECT_ID == "[your-project-id]":
    # Get your GCP project id from gcloud
    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null
    PROJECT_ID = shell_output[0]
    print("Project ID:", PROJECT_ID)

! gcloud config set project $PROJECT_ID

"""
Explanation: Before you begin
GPU runtime
Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select Runtime > Change Runtime Type > GPU
Set up your Google Cloud project
The following steps are required, regardless of your notebook environment.


Select or create a Google Cloud project. When you first create an account, you get a $300 free credit towards your compute/storage costs.


Make sure that billing is enabled for your project.


Enable the following APIs: Vertex AI APIs, Compute Engine APIs, and Cloud Storage.


If you are running this notebook locally, you will need to install the Cloud SDK.


Enter your project ID in the cell below. Then run the  cell to make sure the
Cloud SDK uses the right project for all the commands in this notebook.


Note: Jupyter runs lines prefixed with ! as shell commands, and it interpolates Python variables prefixed with $.
Set your project ID
If you don't know your project ID, you may be able to get your project ID using gcloud.
End of explanation
"""


REGION = "[your-region]"  # @param {type: "string"}

if REGION == "[your-region]":
    REGION = "us-central1"

"""
Explanation: Region
You can also change the REGION variable, which is used for operations
throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.

Americas: us-central1
Europe: europe-west4
Asia Pacific: asia-east1

You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.
Learn more about Vertex AI regions.
End of explanation
"""


from datetime import datetime

TIMESTAMP = datetime.now().strftime("%Y%m%d%H%M%S")

"""
Explanation: Timestamp
If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append the timestamp onto the name of resources you create in this tutorial.
End of explanation
"""


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

import os
import sys

# If on Vertex AI Workbench, then don't execute this code
IS_COLAB = False
if not os.path.exists("/opt/deeplearning/metadata/env_version") and not os.getenv(
    "DL_ANACONDA_HOME"
):
    if "google.colab" in sys.modules:
        IS_COLAB = True
        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 ''

"""
Explanation: Authenticate your Google Cloud account
If you are using Vertex AI Workbench Notebooks, your environment is already authenticated. Skip this step.
If you are using Colab, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.
Otherwise, follow these steps:
In the Cloud Console, go to the Create service account key page.
Click Create service account.
In the Service account name field, enter a name, and click Create.
In the Grant this service account access to project section, click the Role drop-down list. Type "Vertex" into the filter box, and select Vertex Administrator. Type "Storage Object Admin" into the filter box, and select Storage Object Admin.
Click Create. A JSON file that contains your key downloads to your local environment.
Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell.
End of explanation
"""


BUCKET_NAME = "[your-bucket-name]"  # @param {type:"string"}
BUCKET_URI = f"gs://{BUCKET_NAME}"

if BUCKET_NAME == "" or BUCKET_NAME is None or BUCKET_NAME == "[your-bucket-name]":
    BUCKET_URI = "gs://" + PROJECT_ID + "aip-" + TIMESTAMP

"""
Explanation: Create a Cloud Storage bucket
The following steps are required, regardless of your notebook environment.
When you submit a custom training job using the Vertex client library for Python, you upload a Python package
containing your training code to a Cloud Storage bucket. Vertex AI runs
the code from this package. In this tutorial, Vertex AI also saves the
trained model that results from your job in the same bucket. You can then
create an Endpoint resource based on this output in order to serve
online predictions.
Set the name of your Cloud Storage bucket below. Bucket names must be globally unique across all Google Cloud projects, including those outside of your organization.
End of explanation
"""


! gsutil mb -l $REGION $BUCKET_URI

"""
Explanation: Only if your bucket doesn't already exist: Run the following cell to create your Cloud Storage bucket.
End of explanation
"""


! gsutil ls -al $BUCKET_URI

"""
Explanation: Finally, validate access to your Cloud Storage bucket by examining its contents:
End of explanation
"""


import google.cloud.aiplatform as aip

"""
Explanation: Set up variables
Next, set up some variables used throughout the tutorial.
Import libraries and define constants
Import Vertex AI SDK
Import the Vertex AI SDK into your Python environment.
End of explanation
"""


aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)

"""
Explanation: Initialize Vertex AI SDK for Python
Initialize the Vertex AI SDK for Python for your project and corresponding bucket.
End of explanation
"""


if os.getenv("IS_TESTING_TRAIN_GPU"):
    TRAIN_GPU, TRAIN_NGPU = (
        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,
        int(os.getenv("IS_TESTING_TRAIN_GPU")),
    )
else:
    TRAIN_GPU, TRAIN_NGPU = (None, None)

if os.getenv("IS_TESTING_DEPLOY_GPU"):
    DEPLOY_GPU, DEPLOY_NGPU = (
        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,
        int(os.getenv("IS_TESTING_DEPLOY_GPU")),
    )
else:
    DEPLOY_GPU, DEPLOY_NGPU = (None, None)

"""
Explanation: Set hardware accelerators
You can set hardware accelerators for training and prediction.
Set the variables TRAIN_GPU/TRAIN_NGPU and DEPLOY_GPU/DEPLOY_NGPU to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:
(aip.AcceleratorType.NVIDIA_TESLA_K80, 4)

Otherwise specify (None, None) to use a container image to run on a CPU.
Learn more about hardware accelerator support for your region.
Note: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It is a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support.
End of explanation
"""


if os.getenv("IS_TESTING_TF"):
    TF = os.getenv("IS_TESTING_TF")
else:
    TF = "2.1".replace(".", "-")

if TF[0] == "2":
    if TRAIN_GPU:
        TRAIN_VERSION = "tf-gpu.{}".format(TF)
    else:
        TRAIN_VERSION = "tf-cpu.{}".format(TF)
    if DEPLOY_GPU:
        DEPLOY_VERSION = "tf2-gpu.{}".format(TF)
    else:
        DEPLOY_VERSION = "tf2-cpu.{}".format(TF)
else:
    if TRAIN_GPU:
        TRAIN_VERSION = "tf-gpu.{}".format(TF)
    else:
        TRAIN_VERSION = "tf-cpu.{}".format(TF)
    if DEPLOY_GPU:
        DEPLOY_VERSION = "tf-gpu.{}".format(TF)
    else:
        DEPLOY_VERSION = "tf-cpu.{}".format(TF)

TRAIN_IMAGE = "{}-docker.pkg.dev/vertex-ai/training/{}:latest".format(
    REGION.split("-")[0], TRAIN_VERSION
)
DEPLOY_IMAGE = "{}-docker.pkg.dev/vertex-ai/prediction/{}:latest".format(
    REGION.split("-")[0], DEPLOY_VERSION
)

print("Training:", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)
print("Deployment:", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)

"""
Explanation: Set pre-built containers
Set the pre-built Docker container image for training and prediction.
For the latest list, see Pre-built containers for training.
For the latest list, see Pre-built containers for prediction.
End of explanation
"""


if os.getenv("IS_TESTING_TRAIN_MACHINE"):
    MACHINE_TYPE = os.getenv("IS_TESTING_TRAIN_MACHINE")
else:
    MACHINE_TYPE = "n1-standard"

VCPU = "4"
TRAIN_COMPUTE = MACHINE_TYPE + "-" + VCPU
print("Train machine type", TRAIN_COMPUTE)

"""
Explanation: Set machine type
Next, set the machine type to use for training.

Set the variable TRAIN_COMPUTE to configure  the compute resources for the VMs you will use for for training.
machine type
n1-standard: 3.75GB of memory per vCPU.
n1-highmem: 6.5GB of memory per vCPU
n1-highcpu: 0.9 GB of memory per vCPU


vCPUs: number of [2, 4, 8, 16, 32, 64, 96 ]

Note: The following is not supported for training:

standard: 2 vCPUs
highcpu: 2, 4 and 8 vCPUs

Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs.
End of explanation
"""


# Make folder for Python training script
! rm -rf custom
! mkdir custom

# Add package information
! touch custom/README.md

setup_cfg = "[egg_info]\n\ntag_build =\n\ntag_date = 0"
! echo "$setup_cfg" > custom/setup.cfg

setup_py = """import setuptools\n\nsetuptools.setup( \
    install_requires=[       
        'tensorflow==2.5.0',      
        'tensorflow_datasets==1.3.0',
        'tensorflow-hub'
    ],
    packages=setuptools.find_packages())"""
! echo "$setup_py" > custom/setup.py

pkg_info = "Metadata-Version: 1.0\n\nName: CIFAR-10 image classification\n\nVersion: 0.0.0\n\nSummary: Demostration training script\n\nHome-page: www.google.com\n\nAuthor: Google\n\nAuthor-email: aferlitsch@google.com\n\nLicense: Public\n\nDescription: Demo\n\nPlatform: Vertex"
! echo "$pkg_info" > custom/PKG-INFO

# Make the training subfolder
! mkdir custom/trainer
! touch custom/trainer/__init__.py

"""
Explanation: Examine the training package
Package layout
Before you start the training, you will look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.

PKG-INFO
README.md
setup.cfg
setup.py
trainer
__init__.py
task.py
data.py
model.py

The files setup.cfg and setup.py are the instructions for installing the package into the operating environment of the Docker image.
The file trainer/task.py is the Python script for executing the custom training job. Note, when we referred to it in the worker pool specification, we replace the directory slash with a dot (trainer.task) and dropped the file suffix (.py).
Package Assembly
In the following cells, you will assemble the training package.
End of explanation
"""


%%writefile custom/trainer/data.py

import tensorflow as tf
import tensorflow_datasets as tfds

def get_data(buffer=10000):
  
    # Scaling CIFAR10 data from (0, 255] to (0., 1.]
    def scale(image, label):
        image = tf.cast(image, tf.float32)
        image /= 255.0
        return image, label


    datasets, info = tfds.load(name='cifar10',
                            with_info=True,
                            as_supervised=True
    )
    
    return datasets['train'].map(scale).cache().shuffle(buffer).repeat()

"""
Explanation: Introduction to TensorFlow Hub Models
TensorFlow Hub is a repository of pretrained models. For each pretrained large model, one can choose from:

A complete fully trained large model (i.e., Cloud model) -- with the task component.
A fully trained large model -- without the task component.

In the later case, the output of the model is the low-dimensional embedding, which is referred to as the feature vector in TensorFlow Hub. In the case of image models, the input shape can be overridden as well.
These feature vector models can be downloaded in a format that appears as a single TF.Keras layer using the method hub.KerasLayer(), from which a new model can be constructed.
Learn more about TensorFlow Hub.
Get the training data
Next, you write the Python module data.py for getting the training data.
End of explanation
"""


os.chdir("custom/trainer")
from data import get_data

train_dataset = get_data().batch(64)

os.chdir("../..")

"""
Explanation: Test the get_data() module locally
Next, you test the module you just wrote by importing the module and calling the get_data() function. This will return a TFDS generator for the CIFAR-10 training data.
End of explanation
"""


%%writefile custom/trainer/model.py

import tensorflow as tf
import tensorflow_hub as hub

def get_model(tfhub_url, input_shape, n_classes, lr=0.001):

    tfhub_model = tf.keras.Sequential([
        hub.KerasLayer(tfhub_url,
                       input_shape=input_shape)
    ])
    
    tfhub_model.trainable = False
    
    model = tf.keras.Sequential([
        tfhub_model,
        tf.keras.layers.Dense(n_classes, activation='softmax')
    ])
    
    model.compile(
        loss=tf.keras.losses.sparse_categorical_crossentropy,
        optimizer=tf.keras.optimizers.SGD(learning_rate=lr),
        metrics=['accuracy']
    )
    
    return model

"""
Explanation: Get pretrained model from TFHub
Next, you write the module model.py for getting the model architecture. This module does the following:

Takes as parameters:
tfhub_url: The URL to the TensorFlow Hub model.
input_shape: The input shape for the model.
n_classes: The number of output classes for classification.
lr: The learning rate for compiling the model.


Downloads the TensorFlow Hub model.
Freezes the weights of the layers of TensorFlow Hub model -- does not get retrained.
Adds the task component for classification -- these layers will get trained.
Compiles the model.
End of explanation
"""


os.chdir("custom/trainer")

from model import get_model

model = get_model(
    tfhub_url="https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/4",
    input_shape=(32, 32, 3),
    n_classes=10,
    lr=0.001,
)

os.chdir("../..")

"""
Explanation: Test the get_model() module locally
Next, you test the module you just wrote by importing the module and calling the get_model() function. This will return a TF.Keras model configured for a CIFAR-10 dataset.
End of explanation
"""


%%writefile custom/trainer/train.py

def train_model(model,
                train_dataset,
                epochs,
                steps_per_epoch
               ):

    history = model.fit(train_dataset, epochs=epochs, steps_per_epoch=steps_per_epoch)
    return model

"""
Explanation: Train the model
Next, you write the module train.py to train the model.
End of explanation
"""


os.chdir("custom/trainer")

from train import train_model

model = train_model(model, train_dataset, 10, 200)

os.chdir("../..")

"""
Explanation: Test the train_model() module locally
Next, you test the module you just wrote by importing the module and calling the train_model() function. This function will return the trained model.
End of explanation
"""


%%writefile custom/trainer/task.py
# Single, Mirror and Multi-Machine Distributed Training for CIFAR-10

import tensorflow_datasets as tfds
import tensorflow as tf
import argparse
import os
import sys
import logging
tfds.disable_progress_bar()

from . import data
from . import model as model_
from . import train

parser = argparse.ArgumentParser()
parser.add_argument('--model-dir', dest='model_dir',
                    default=os.getenv("AIP_MODEL_DIR"), type=str, help='Model dir.')
parser.add_argument('--lr', dest='lr',
                    default=0.01, type=float,
                    help='Learning rate.')
parser.add_argument('--epochs', dest='epochs',
                    default=10, type=int,
                    help='Number of epochs.')
parser.add_argument('--steps', dest='steps',
                    default=200, type=int,
                    help='Number of steps per epoch.')
parser.add_argument('--batch_size', dest='batch_size',
                    default=64, type=int,
                    help='Size of mini-batch.')
parser.add_argument('--distribute', dest='distribute', type=str, default='single',
                    help='distributed training strategy')
args = parser.parse_args()

logging.getLogger().setLevel(logging.INFO)

# Single Machine, single compute device
if args.distribute == 'single':
    if tf.test.is_gpu_available():
        strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
    else:
        strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0")
    logging.info("Single device training")
# Single Machine, multiple compute device
elif args.distribute == 'mirrored':
    strategy = tf.distribute.MirroredStrategy()
    logging.info("Mirrored Strategy distributed training")
# Multi Machine, multiple compute device
elif args.distribute == 'multiworker':
    strategy = tf.distribute.MultiWorkerMirroredStrategy()
    logging.info("Multi-worker Strategy distributed training")
    logging.info('TF_CONFIG = {}'.format(os.environ.get('TF_CONFIG', 'Not found')))
    
# Multi-worker configuration
logging.info('num_replicas_in_sync = {}'.format(strategy.num_replicas_in_sync))
NUM_WORKERS = strategy.num_replicas_in_sync
GLOBAL_BATCH_SIZE = args.batch_size * NUM_WORKERS


def _is_chief(task_type, task_id):
    ''' Check for primary if multiworker training
    '''
    return (task_type == 'chief') or (task_type == 'worker' and task_id == 0) or task_type is None

# Here the batch size scales up by number of workers since
# `tf.data.Dataset.batch` expects the global batch size
BUFFER_SIZE = 10000
train_data = data.get_data(buffer=BUFFER_SIZE).batch(GLOBAL_BATCH_SIZE)

with strategy.scope():
    # Creation of dataset, and model building/compiling need to be within
    # `strategy.scope()`.
    model = model_.get_model(
        tfhub_url="https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/4",
        input_shape=(32, 32, 3),
        n_classes=10,
        lr=args.lr
    )

model = train.train_model(model, train_data, epochs=args.epochs, steps_per_epoch=args.steps)

if args.distribute == 'multiworker':
    task_type, task_id = (strategy.cluster_resolver.task_type,
                              strategy.cluster_resolver.task_id)
else:
    task_type, task_id = None, None
    
# single, mirrored or primary for multiworker
if _is_chief(task_type, task_id):
     model.save(args.model_dir)
# non-primary workers for multi-workers
else:
    # each worker saves their model instance to a unique temp location
    worker_dir = args.model_dir + '/workertemp_' + str(task_id)
    tf.io.gfile.makedirs(worker_dir)
    model.save(worker_dir)

"""
Explanation: Task.py contents
In the next cell, you write the contents of the training script task.py. 


Parse the command line arguments for training the model

--model-dir: Get the directory where to save the model artifacts from the command line (--model_dir), and if not specified, then from the environment variable AIP_MODEL_DIR.
--lr: The learning rate.
--epochs: The number of epochs.
--steps: The number of steps per epoch.
--batch_size: The batch size per worker.



get_data(): Download and preprocess the CIFAR-10 dataset.

get_model(): Download the model architecture from TFHub and configure the model for CIFAR-10 training.
train_model(): Train the model.
Save the model artifacts to the specified Cloud Storage location.
End of explanation
"""


! rm -f custom.tar custom.tar.gz
! tar cvf custom.tar custom
! gzip custom.tar
! gsutil cp custom.tar.gz $BUCKET_URI/trainer_cifar10.tar.gz

"""
Explanation: Store training script on your Cloud Storage bucket
Next, you package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket.
End of explanation
"""


DISPLAY_NAME = "cifar10_" + TIMESTAMP

job = aip.CustomPythonPackageTrainingJob(
    display_name=DISPLAY_NAME,
    python_package_gcs_uri=f"{BUCKET_URI}/trainer_cifar10.tar.gz",
    python_module_name="trainer.task",
    container_uri=TRAIN_IMAGE,
    model_serving_container_image_uri=DEPLOY_IMAGE,
    project=PROJECT_ID,
)

"""
Explanation: Create and run custom training job
To train a custom model, you perform two steps: 1) create a custom training job, and 2) run the job.
Create custom training job
A custom training job is created with the CustomTrainingJob class, with the following parameters:

display_name: The human readable name for the custom training job.

container_uri: The training container image.


python_package_gcs_uri: The location of the Python training package as a tarball.

python_module_name: The relative path to the training script in the Python package.
model_serving_container_uri: The container image for deploying the model.

Note: There is no requirements parameter. You specify any requirements in the setup.py script in your Python package.
End of explanation
"""


MODEL_DIR = "{}/{}".format(BUCKET_URI, TIMESTAMP)

EPOCHS = 10
STEPS = 200

CMDARGS = [
    "--epochs=" + str(EPOCHS),
    "--steps=" + str(STEPS),
]

"""
Explanation: Prepare your command-line arguments
Now define the command-line arguments for your custom training container:

args: The command-line arguments to pass to the executable that is set as the entry point into the container.
--model-dir : For our demonstrations, we use this command-line argument to specify where to store the model artifacts.
direct: You pass the Cloud Storage location as a command line argument to your training script (set variable DIRECT = True), or
indirect: The service passes the Cloud Storage location as the environment variable AIP_MODEL_DIR to your training script (set variable DIRECT = False). In this case, you tell the service the model artifact location in the job specification.


"--epochs=" + EPOCHS: The number of epochs for training.
"--steps=" + STEPS: The number of steps per epoch.
End of explanation
"""


if TRAIN_GPU:
    model = job.run(
        model_display_name="cifar10_" + TIMESTAMP,
        args=CMDARGS,
        replica_count=1,
        machine_type=TRAIN_COMPUTE,
        accelerator_type=TRAIN_GPU.name,
        accelerator_count=TRAIN_NGPU,
        base_output_dir=MODEL_DIR,
        sync=False,
    )
else:
    model = job.run(
        model_display_name="cifar10_" + TIMESTAMP,
        args=CMDARGS,
        replica_count=1,
        machine_type=TRAIN_COMPUTE,
        base_output_dir=MODEL_DIR,
        sync=False,
    )

"""
Explanation: Run the custom training job
Next, you run the custom job to start the training job by invoking the method run, with the following parameters:

model_display_name: The human readable name for the Model resource.
args: The command-line arguments to pass to the training script.
replica_count: The number of compute instances for training (replica_count = 1 is single node training).
machine_type: The machine type for the compute instances.
accelerator_type: The hardware accelerator type.
accelerator_count: The number of accelerators to attach to a worker replica.
base_output_dir: The Cloud Storage location to write the model artifacts to.
sync: Whether to block until completion of the job.
End of explanation
"""


_job = job.list(filter=f"display_name={DISPLAY_NAME}")
print(_job)

"""
Explanation: List a custom training job
End of explanation
"""


model.wait()

"""
Explanation: Wait for completion of custom training job
Next, wait for the custom training job to complete. Alternatively, one can set the parameter sync to True in the run() method to block until the custom training job is completed.
End of explanation
"""


model.delete()

"""
Explanation: Delete the model
You can delete the model using the method delete().
End of explanation
"""


job.delete()

! rm -rf custom custom.tar.gz

"""
Explanation: Delete the custom job
You can delete the custom job using the method delete().
End of explanation
"""


delete_bucket = False

if delete_bucket and "BUCKET_URI" in globals():
    ! gsutil rm -r $BUCKET_URI

"""
Explanation: Cleaning up
To clean up all Google Cloud resources used in this project, you can delete the Google Cloud
project you used for the tutorial.
Otherwise, you can delete the individual resources you created in this tutorial.
End of explanation
"""


