Optuna Cloud Batch Template

This is a template repository for running Optuna optimization on Google Cloud Batch and GCS.

Requirements

• python>=3.11
• docker
• make
• yq
• gcloud

How to use

1. Create two service accounts, SA1 and SA2. SA1 is for optimizer and SA2 is for job batches. Add permissions roles/storage.objectUser, roles/batch.jobsEditor, roles/artifactregistry.reader and roles/logging.logWriter to a service account for SA1 and roles/storage.admin, roles/artifactregistry.reader, roles/batch.agentReporter and roles/logging.logWriter for SA2.
2. Edit optunabatch/custom.py to define your objective function and study settings.
3. Edit config.yaml to define your cloud batch settings. Service account for the job instances (SA2) should be specified in service_account field.
4. make build to build the Docker images.
5. make push to push the Docker images to the Artifact Resigtry.
6. Create a GCE instance. You can use make create-instance SERVICE_ACCOUNT={SA1} to create a VM. The service account should have permissions for Storage Object User, Batch Job Editor and Logs Writer.
7. make update-container to deploy the Docker images to the GCP.
8. Trials will be stored in the GCS bucket specified in config.yaml.

The command make run-optimizer is for debugging in the local environment. Note that it uses credentials stored in ~/.config/gcloud/.

References

• Run basic job with Cloud Batch
• Optuna