metaflow icon indicating copy to clipboard operation
metaflow copied to clipboard

Added support for google TPU

Open kulikovv opened this issue 1 year ago • 2 comments

set vendor as google and ready to and proper nodeSelector

@kubernetes(
       image="image_name",
       node_selector={
            "cloud.google.com/gke-tpu-accelerator": "tpu-v6e-slice",
            "cloud.google.com/gke-tpu-topology": "1x1"
       },
       tpu=1,
       tpu_vendor="google"
   )

kulikovv avatar Nov 13 '24 14:11 kulikovv

thanks for the PR! it might make more sense to introduce tpu as an arg rather than overriding gpu - since one can access gpus in gcp too. also, this approach would mimic our support for trainium.

savingoyal avatar Nov 14 '24 17:11 savingoyal

@savingoyal could you please review again, I've added TPU parameter to k8s decorator, but didn't tested it with Argo and Airflow.

kulikovv avatar Jan 15 '25 14:01 kulikovv