Rename GPU_DEVICE_LABEL to NVIDIA_GPU_DEVICE_LABEL

Signed-off-by: Andrey Velichkevich <[email protected]>
kubeflow · Nov 29, 2024 · e9c22b2 · e9c22b2
1 parent 9bea4cf
commit e9c22b2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/sdk_v2/kubeflow/training/constants/constants.py b/sdk_v2/kubeflow/training/constants/constants.py
@@ -54,8 +54,8 @@
 # The default type for CPU.
 CPU_DEVICE_TYPE = "cpu"
 
-# The label for GPU device in the container resources.
-GPU_DEVICE_LABEL = "nvidia.com/gpu"
+# The label for NVIDIA GPU device in the container resources.
+NVIDIA_GPU_DEVICE_LABEL = "nvidia.com/gpu"
 
 # The default type for GPU
 GPU_DEVICE_TYPE = "gpu"

diff --git a/sdk_v2/kubeflow/training/utils/utils.py b/sdk_v2/kubeflow/training/utils/utils.py
@@ -69,9 +69,9 @@ def get_container_devices(
         return device, device_count
 
     # TODO (andreyvelich): Support other resource labels (e.g. NPUs).
-    if constants.GPU_DEVICE_LABEL in resources.limits:
+    if constants.NVIDIA_GPU_DEVICE_LABEL in resources.limits:
         device = constants.GPU_DEVICE_TYPE
-        device_count = resources.limits[constants.GPU_DEVICE_LABEL]
+        device_count = resources.limits[constants.NVIDIA_GPU_DEVICE_LABEL]
     elif constants.TPU_DEVICE_LABEL in resources.limits:
         device = constants.TPU_DEVICE_TYPE
         device_count = resources.limits[constants.TPU_DEVICE_LABEL]