Add Google Cloud/GCP Scheduler Support + TPUs

[d4l3k]

It would be nice to have GCP + TPU support in addition to our existing schedulers. Currently you can run on GCP via Kubernetes + the Kubernetes scheduler but would be handy to have direct training platform support.

Example scheduler: AWS Batch https://github.com/pytorch/torchx/blob/main/torchx/schedulers/aws_batch_scheduler.py

Scheduler documentation: https://pytorch.org/torchx/main/schedulers

GCP Docs:

• https://cloud.google.com/ai-platform/training/docs/training-jobs
• https://cloud.google.com/ai-platform/training/docs/reference/rest/v1/projects.jobs

Stretch Goal:

• add TPU integrations
  • https://cloud.google.com/ai-platform/training/docs/using-tpus
  • https://cloud.google.com/tpu/docs/training-on-tpu-pods

[d4l3k]

resources

v2-8

model name	: Intel(R) Xeon(R) CPU @ 2.00GHz
96 cores

$ free -h
              total        used        free      shared  buff/cache   available
Mem:          334Gi       1.5Gi       331Gi       1.0Mi       1.3Gi       331Gi
Swap:            0B          0B          0B

/dev/root        97G   15G   83G  16% /

v3-8

model name	: Intel(R) Xeon(R) CPU @ 2.00GHz
96 cores

tristanr@t1v-n-c76b9a61-w-0:~$ free -h
              total        used        free      shared  buff/cache   available
Mem:          334Gi       1.5Gi       331Gi       1.0Mi       1.5Gi       331Gi
Swap:            0B          0B          0B

tristanr@t1v-n-c76b9a61-w-0:~$ df -h /
Filesystem      Size  Used Avail Use% Mounted on
/dev/root        97G   15G   83G  15% /

[vwxyzjn]

Also see https://cloud.google.com/blog/products/compute/new-batch-service-processes-batch-jobs-on-google-cloud. GCP equivalent of AWS Batch recently announced.

[d4l3k]

@vwxyzjn thanks for sharing that! I hadn't seen it. Do you know if it supports TPUs? I didn't see anything listed on the announcement post

[priyaramani]

@vwxyzjn Hey there, we just added support for GCP batch (initial version w/ support for basic components - #621). Please try it out and let us know. Would appreciate early feedback. Thanks!

[priyaramani]

GCP Batch support has been added in TorchX in Prototype phase: See https://github.com/pytorch/torchx/releases/tag/v0.4.0