added separate option in decorator for TPU

Netflix · Jan 15, 2025 · 814693f · 814693f
1 parent 7d53b75
commit 814693f
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 0 deletions.
diff --git a/metaflow/plugins/kubernetes/kubernetes.py b/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,

diff --git a/metaflow/plugins/kubernetes/kubernetes_cli.py b/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,

diff --git a/metaflow/plugins/kubernetes/kubernetes_jobsets.py b/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=(