metaflow
metaflow copied to clipboard
Added support for google TPU
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"
)
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 could you please review again, I've added TPU parameter to k8s decorator, but didn't tested it with Argo and Airflow.