Skip to content

feat: Add GPU type support and update GPU resource limit … #2402

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

feat: Add GPU type support and update GPU resource limit … #2402

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fc623e0
feat(kubernetes): Add GPU type support and update GPU resource limit …
May 8, 2025
1c25404
feat(argo): 添加 GPU 类型配置
May 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions metaflow/metaflow_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -364,6 +364,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 GPU type to use by K8S jobs created by Metaflow
KUBERNETES_GPU_TYPE = from_conf("KUBERNETES_GPU_TYPE", "gpu")
# Default container image for K8S
KUBERNETES_CONTAINER_IMAGE = from_conf(
"KUBERNETES_CONTAINER_IMAGE", DEFAULT_CONTAINER_IMAGE
Expand Down
4 changes: 2 additions & 2 deletions metaflow/plugins/airflow/airflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -440,8 +440,8 @@ def _to_job(self, node):
limits={
**qos_limits,
**{
"%s.com/gpu".lower()
% k8s_deco.attributes["gpu_vendor"]: str(k8s_deco.attributes["gpu"])
("%s.com/%s".lower()
% (k8s_deco.attributes["gpu_vendor"], k8s_deco.attributes["gpu_type"])): str(k8s_deco.attributes["gpu"])
for k in [0]
# Don't set GPU limits if gpu isn't specified.
if k8s_deco.attributes["gpu"] is not None
Expand Down
12 changes: 8 additions & 4 deletions metaflow/plugins/argo/argo_workflows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2033,6 +2033,7 @@ def _container_templates(self):
disk=str(resources["disk"]),
gpu=resources["gpu"],
gpu_vendor=str(resources["gpu_vendor"]),
gpu_type=str(resources["gpu_type"]),
tolerations=resources["tolerations"],
use_tmpfs=use_tmpfs,
tmpfs_tempdir=tmpfs_tempdir,
Expand Down Expand Up @@ -2248,10 +2249,13 @@ def _container_templates(self):
limits={
**qos_limits,
**{
"%s.com/gpu".lower()
% resources["gpu_vendor"]: str(
resources["gpu"]
)
(
"%s.com/%s".lower()
% (
resources["gpu_vendor"],
resources["gpu_type"],
)
): str(resources["gpu"])
for k in [0]
if resources["gpu"] is not None
},
Expand Down
6 changes: 6 additions & 0 deletions metaflow/plugins/kubernetes/kubernetes_decorator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
KUBERNETES_DISK,
KUBERNETES_FETCH_EC2_METADATA,
KUBERNETES_GPU_VENDOR,
KUBERNETES_GPU_TYPE,
KUBERNETES_IMAGE_PULL_POLICY,
KUBERNETES_MEMORY,
KUBERNETES_LABELS,
Expand Down Expand Up @@ -90,6 +91,8 @@ 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.
gpu_type : str , optional, default KUBERNETES_GPU_TYPE
The type of the GPUs 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.
Expand Down Expand Up @@ -145,6 +148,7 @@ class KubernetesDecorator(StepDecorator):
"namespace": None,
"gpu": None, # value of 0 implies that the scheduled node should not have GPUs
"gpu_vendor": None,
"gpu_type": 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"}
Expand Down Expand Up @@ -181,6 +185,8 @@ def init(self):
self.attributes["gpu_vendor"] = KUBERNETES_GPU_VENDOR
if not self.attributes["node_selector"] and KUBERNETES_NODE_SELECTOR:
self.attributes["node_selector"] = KUBERNETES_NODE_SELECTOR
if not self.attributes["gpu_type"]:
self.attributes["gpu_type"] = KUBERNETES_GPU_TYPE
if not self.attributes["tolerations"] and KUBERNETES_TOLERATIONS:
self.attributes["tolerations"] = json.loads(KUBERNETES_TOLERATIONS)
if (
Expand Down
6 changes: 2 additions & 4 deletions metaflow/plugins/kubernetes/kubernetes_job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,10 +171,8 @@ def create_job_spec(self):
limits={
**qos_limits,
**{
"%s.com/gpu".lower()
% self._kwargs["gpu_vendor"]: str(
self._kwargs["gpu"]
)
("%s.com/%s".lower()
% (self._kwargs["gpu_vendor"], self._kwargs["gpu_type"])): str(self._kwargs["gpu"])
for k in [0]
# Don't set GPU limits if gpu isn't specified.
if self._kwargs["gpu"] is not None
Expand Down
7 changes: 3 additions & 4 deletions metaflow/plugins/kubernetes/kubernetes_jobsets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -670,10 +670,9 @@ def dump(self):
limits={
**qos_limits,
**{
"%s.com/gpu".lower()
% self._kwargs["gpu_vendor"]: str(
self._kwargs["gpu"]
)
("%s.com/%s".lower()
% (self._kwargs["gpu_vendor"], self._kwargs["gpu_type"])): str(
self._kwargs["gpu"])
for k in [0]
# Don't set GPU limits if gpu isn't specified.
if self._kwargs["gpu"] is not None
Expand Down