Added support for google TPU

metaflow/metaflow_config.py

@@ -359,6 +359,8 @@
 KUBERNETES_ANNOTATIONS = from_conf("KUBERNETES_ANNOTATIONS", "")
 # Default GPU vendor to use by K8S jobs created by Metaflow (supports nvidia, amd)
 KUBERNETES_GPU_VENDOR = from_conf("KUBERNETES_GPU_VENDOR", "nvidia")
+# Default TPU vendor to use by K8S jobs created by Metaflow (supports google)
+KUBERNETES_TPU_VENDOR = from_conf("KUBERNETES_TPU_VENDOR", "google")
 # Default container image for K8S
 KUBERNETES_CONTAINER_IMAGE = from_conf(
     "KUBERNETES_CONTAINER_IMAGE", DEFAULT_CONTAINER_IMAGE

metaflow/plugins/airflow/airflow.py

@@ -446,6 +446,13 @@ def _to_job(self, node):
                     # Don't set GPU limits if gpu isn't specified.
                     if k8s_deco.attributes["gpu"] is not None
                 },
+                **{
+                    "%s.com/tpu".lower()
+                    % k8s_deco.attributes["tpu_vendor"]: str(k8s_deco.attributes["tpu"])
+                    for k in [0]
+                    # Don't set TPU limits if gpu isn't specified.
+                    if k8s_deco.attributes["tpu"] is not None
+                },
             },
         )
 

metaflow/plugins/argo/argo_workflows.py

@@ -1940,6 +1940,8 @@ def _container_templates(self):
                     disk=str(resources["disk"]),
                     gpu=resources["gpu"],
                     gpu_vendor=str(resources["gpu_vendor"]),
+                    tpu=resources["tpu"],
+                    tpu_vendor=str(resources["tpu_vendor"]),
                     tolerations=resources["tolerations"],
                     use_tmpfs=use_tmpfs,
                     tmpfs_tempdir=tmpfs_tempdir,
@@ -2159,6 +2161,14 @@ def _container_templates(self):
                                             for k in [0]
                                             if resources["gpu"] is not None
                                         },
+                                        **{
+                                            "%s.com/tpu".lower()
+                                            % resources["tpu_vendor"]: str(
+                                                resources["tpu"]
+                                            )
+                                            for k in [0]
+                                            if resources["tpu"] is not None
+                                        },
                                     },
                                 ),
                                 # Configure secrets

metaflow/plugins/kubernetes/kubernetes.py

@@ -178,6 +178,8 @@ def create_jobset(
         cpu=None,
         gpu=None,
         gpu_vendor=None,
+        tpu=None,
+        tpu_vendor=None,
         disk=None,
         memory=None,
         use_tmpfs=None,
@@ -210,6 +212,8 @@ def create_jobset(
                 disk=disk,
                 gpu=gpu,
                 gpu_vendor=gpu_vendor,
+                tpu=tpu,
+                tpu_vendor=tpu_vendor,
                 timeout_in_seconds=run_time_limit,
                 # Retries are handled by Metaflow runtime
                 retries=0,
@@ -472,6 +476,8 @@ def create_job_object(
         cpu=None,
         gpu=None,
         gpu_vendor=None,
+        tpu=None,
+        tpu_vendor=None,
         disk=None,
         memory=None,
         use_tmpfs=None,
@@ -515,6 +521,8 @@ def create_job_object(
                 disk=disk,
                 gpu=gpu,
                 gpu_vendor=gpu_vendor,
+                tpu=tpu,
+                tpu_vendor=tpu_vendor,
                 timeout_in_seconds=run_time_limit,
                 # Retries are handled by Metaflow runtime
                 retries=0,

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -82,6 +82,8 @@ def kubernetes():
 @click.option("--memory", help="Memory requirement for Kubernetes pod.")
 @click.option("--gpu", help="GPU requirement for Kubernetes pod.")
 @click.option("--gpu-vendor", help="GPU vendor requirement for Kubernetes pod.")
+@click.option("--tpu", help="TPU requirement for Kubernetes pod.")
+@click.option("--tpu-vendor", help="TPU vendor requirement for Kubernetes pod.")
 @click.option("--run-id", help="Passed to the top-level 'step'.")
 @click.option("--task-id", help="Passed to the top-level 'step'.")
 @click.option("--input-paths", help="Passed to the top-level 'step'.")
@@ -163,6 +165,8 @@ def step(
     memory=None,
     gpu=None,
     gpu_vendor=None,
+    tpu=None,
+    tpu_vendor=None,
     use_tmpfs=None,
     tmpfs_tempdir=None,
     tmpfs_size=None,
@@ -305,6 +309,8 @@ def _sync_metadata():
                 memory=memory,
                 gpu=gpu,
                 gpu_vendor=gpu_vendor,
+                tpu=tpu,
+                tpu_vendor=tpu_vendor,
                 use_tmpfs=use_tmpfs,
                 tmpfs_tempdir=tmpfs_tempdir,
                 tmpfs_size=tmpfs_size,

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -17,6 +17,7 @@
     KUBERNETES_DISK,
     KUBERNETES_FETCH_EC2_METADATA,
     KUBERNETES_GPU_VENDOR,
+    KUBERNETES_TPU_VENDOR,
     KUBERNETES_IMAGE_PULL_POLICY,
     KUBERNETES_MEMORY,
     KUBERNETES_LABELS,
@@ -89,6 +90,11 @@ class KubernetesDecorator(StepDecorator):
         the scheduled node should not have GPUs.
     gpu_vendor : str, default KUBERNETES_GPU_VENDOR
         The vendor of the GPUs to be used for this step.
+    tpu : int, optional, default None
+        Number of TPUs required for this step. A value of zero implies that
+        the scheduled node should not have TPUs.
+    tpu_vendor : str, default KUBERNETES_TPU_VENDOR
+        The vendor of the TPUs to be used for this step.
     tolerations : List[str], default []
         The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
         Kubernetes tolerations to use when launching pod in Kubernetes.
@@ -136,6 +142,8 @@ class KubernetesDecorator(StepDecorator):
         "namespace": None,
         "gpu": None,  # value of 0 implies that the scheduled node should not have GPUs
         "gpu_vendor": None,
+        "tpu": None,  # value of 0 implies that the scheduled node should not have TPUs
+        "tpu_vendor": None,
         "tolerations": None,  # e.g., [{"key": "arch", "operator": "Equal", "value": "amd"},
         #                              {"key": "foo", "operator": "Equal", "value": "bar"}]
         "labels": None,  # e.g. {"test-label": "value", "another-label":"value2"}
@@ -169,6 +177,8 @@ def init(self):
             self.attributes["service_account"] = KUBERNETES_SERVICE_ACCOUNT
         if not self.attributes["gpu_vendor"]:
             self.attributes["gpu_vendor"] = KUBERNETES_GPU_VENDOR
+        if not self.attributes["tpu_vendor"]:
+            self.attributes["tpu_vendor"] = KUBERNETES_TPU_VENDOR
         if not self.attributes["node_selector"] and KUBERNETES_NODE_SELECTOR:
             self.attributes["node_selector"] = KUBERNETES_NODE_SELECTOR
         if not self.attributes["tolerations"] and KUBERNETES_TOLERATIONS:
@@ -342,9 +352,9 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
         for deco in decos:
             if isinstance(deco, ResourcesDecorator):
                 for k, v in deco.attributes.items():
-                    # If GPU count is specified, explicitly set it in self.attributes.
-                    if k == "gpu" and v != None:
-                        self.attributes["gpu"] = v
+                    # If GPU/TPU count is specified, explicitly set it in self.attributes.
+                    if k in ("gpu", "tpu") and v != None:
+                        self.attributes[k] = v
 
                     if k in self.attributes:
                         if self.defaults[k] is None:
@@ -366,6 +376,14 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                 )
             )
 
+        # Check TPU vendor.
+        if self.attributes["tpu_vendor"].lower() not in ("google"):
+            raise KubernetesException(
+                "TPU vendor *{}* for step *{step}* is not currently supported.".format(
+                    self.attributes["tpu_vendor"], step=step
+                )
+            )
+
         # CPU, Disk, and Memory values should be greater than 0.
         for attr in ["cpu", "disk", "memory"]:
             if not (
@@ -378,15 +396,20 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
                     )
                 )
 
-        if self.attributes["gpu"] is not None and not (
-            isinstance(self.attributes["gpu"], (int, unicode, basestring))
-            and float(self.attributes["gpu"]).is_integer()
-        ):
-            raise KubernetesException(
-                "Invalid GPU value *{}* for step *{step}*; it should be an integer".format(
-                    self.attributes["gpu"], step=step
+        for accelerator_type in ("gpu", "tpu"):
+            if self.attributes[accelerator_type] is not None and not (
+                isinstance(
+                    self.attributes[accelerator_type], (int, unicode, basestring)
+                )
+                and float(self.attributes[accelerator_type]).is_integer()
+            ):
+                raise KubernetesException(
+                    "Invalid {accelerator_type} value *{number}* for step *{step}*; it should be an integer".format(
+                        accelerator_type=accelerator_type.toUpper(),
+                        number=self.attributes[accelerator_type],
+                        step=step,
+                    )
                 )
-            )
 
         if self.attributes["tmpfs_size"]:
             if not (

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -172,6 +172,15 @@ def create_job_spec(self):
                                         # Don't set GPU limits if gpu isn't specified.
                                         if self._kwargs["gpu"] is not None
                                     },
+                                    **{
+                                        "%s.com/tpu".lower()
+                                        % self._kwargs["tpu_vendor"]: str(
+                                            self._kwargs["tpu"]
+                                        )
+                                        for k in [0]
+                                        # Don't set GPU limits if gpu isn't specified.
+                                        if self._kwargs["tpu"] is not None
+                                    },
                                 },
                             ),
                             volume_mounts=(

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -670,6 +670,15 @@ def dump(self):
                                                     # Don't set GPU limits if gpu isn't specified.
                                                     if self._kwargs["gpu"] is not None
                                                 },
+                                                **{
+                                                    "%s.com/tpu".lower()
+                                                    % self._kwargs["tpu_vendor"]: str(
+                                                        self._kwargs["tpu"]
+                                                    )
+                                                    for k in [0]
+                                                    # Don't set GPU limits if gpu isn't specified.
+                                                    if self._kwargs["tpu"] is not None
+                                                },
                                             },
                                         ),
                                         volume_mounts=(