PROJECT_ID = "YOUR PROJECT ID"
BUCKET_NAME = "gs://YOUR BUCKET NAME"
REGION = "YOUR REGION"
SERVICE_ACCOUNT = "YOUR SERVICE ACCOUNT"

content_name = "tf-keras-img-cls-dist-multi-worker-cpu-cust-cont"

"""
Explanation: TF-Keras Image Classification Distributed Multi-Worker Training on CPU using Vertex Training with Custom Container
<table align="left">
  <td>
    <a href="https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/multi_worker_vertex_training_on_cpu_with_custom_container.ipynb">
      <img src="https://cloud.google.com/ml-engine/images/github-logo-32px.png" alt="GitHub logo">
      View on GitHub
    </a>
  </td>
</table>

Setup
End of explanation
"""


hostname = "gcr.io"
image_name = content_name
tag = "latest"

custom_container_image_uri = f"{hostname}/{PROJECT_ID}/{image_name}:{tag}"

! cd trainer && docker build -t $custom_container_image_uri -f cpu.Dockerfile .

! docker run --rm $custom_container_image_uri --epochs 2 --local-mode

! docker push $custom_container_image_uri

! gcloud container images list --repository $hostname/$PROJECT_ID

"""
Explanation: Vertex Training using Vertex SDK and Custom Container
Build Custom Container
End of explanation
"""


! pip install -r requirements.txt

from google.cloud import aiplatform

aiplatform.init(
    project=PROJECT_ID,
    staging_bucket=BUCKET_NAME,
    location=REGION,
)

"""
Explanation: Initialize Vertex SDK
End of explanation
"""


tensorboard = aiplatform.Tensorboard.create(
    display_name=content_name,
)

"""
Explanation: Create a Vertex Tensorboard Instance
End of explanation
"""


display_name = content_name
gcs_output_uri_prefix = f"{BUCKET_NAME}/{display_name}"

replica_count = 4
machine_type = "n1-standard-4"

container_args = [
    "--epochs",
    "50",
    "--batch-size",
    "32",
]

custom_container_training_job = aiplatform.CustomContainerTrainingJob(
    display_name=display_name,
    container_uri=custom_container_image_uri,
)

custom_container_training_job.run(
    args=container_args,
    base_output_dir=gcs_output_uri_prefix,
    replica_count=replica_count,
    machine_type=machine_type,
    tensorboard=tensorboard.resource_name,
    service_account=SERVICE_ACCOUNT,
)

print(f"Custom Training Job Name: {custom_container_training_job.resource_name}")
print(f"GCS Output URI Prefix: {gcs_output_uri_prefix}")

"""
Explanation: Option: Use a Previously Created Vertex Tensorboard Instance
tensorboard_name = "Your Tensorboard Resource Name or Tensorboard ID"
tensorboard = aiplatform.Tensorboard(tensorboard_name=tensorboard_name)
Run a Vertex SDK CustomContainerTrainingJob
End of explanation
"""


! gsutil ls $gcs_output_uri_prefix

"""
Explanation: Training Output Artifact
End of explanation
"""


! gsutil rm -rf $gcs_output_uri_prefix

"""
Explanation: Clean Up Artifact
End of explanation
"""


