[k8s] support to use custom gpu resource name if it's not nvidia.com/gpu #4337

nkwangleiGIT · 2024-11-12T15:01:42Z

We're using customized gpu resource name when using the device plugin, so we have the following GPUs in the node capability:

nvidia.com/gpu-h100
nvidia.com/gpu-h20
nvidia.com/gpu-4090

So we have to use the resource names above when launch to local K8S, such as

sky launch --image-id skypilot:20240613 --cpus 8 --memory 32 --gpus gpu-3090:2 -c my-sky-cluster --cloud kubernetes

So this PR will support to use CUSTOM_GPU_RESOURCE_NAME from environment variable to overwrite the default nvidia.com/gpu in the resources.

Tested (run the relevant ones):

Code formatting: bash format.sh
Any manual or new tests for this PR (please specify below)
All smoke tests: pytest tests/test_smoke.py
Relevant individual smoke tests: pytest tests/test_smoke.py::test_fill_in_the_name
Backward compatibility tests: conda deactivate; bash -i tests/backward_compatibility_tests.sh

romilbhardwaj

Thanks and welcome to SkyPilot @nkwangleiGIT!

romilbhardwaj · 2024-11-12T20:21:38Z

sky/utils/common_utils.py

+def get_gpu_resource_name(name: str = 'nvidia.com/gpu'):
+    """Get the GPU resource name.
+
+    The function first checks for an environment variable.
+    If defined, it uses its value; otherwise, it returns the default value.
+
+    Args:
+        name (str): Default GPU resource name, default is "nvidia.com/gpu".
+
+    Returns:
+        str: The selected GPU resource name.
+    """
+    # retrieve GPU resource name from environment variable,
+    # can be nvidia.com/gpu-h100 if it's customized by the device plugin
+    custom_name = os.getenv('CUSTOM_GPU_RESOURCE_NAME')
+
+    # If the environment variable is not defined, return the default name
+    if custom_name is None:
+        return name
+
+    return custom_name


Can we move this to sky.provision.kubernetes.utils since only k8s relies on this method?

moved to sky.provision.kubernetes.utils

romilbhardwaj · 2024-11-12T20:26:50Z

sky/utils/common_utils.py

+    # retrieve GPU resource name from environment variable,
+    # can be nvidia.com/gpu-h100 if it's customized by the device plugin


Nit

Suggested change

# retrieve GPU resource name from environment variable,

# can be nvidia.com/gpu-h100 if it's customized by the device plugin

# Retrieve GPU resource name from environment variable, if set.

# E.g., can be nvidia.com/gpu-h100, amd.com/gpu etc.

romilbhardwaj · 2024-11-12T20:28:00Z

sky/utils/common_utils.py

+
+
+def get_gpu_resource_name(name: str = 'nvidia.com/gpu'):
+    """Get the GPU resource name.


Suggested change

"""Get the GPU resource name.

"""Get the GPU resource name to use in kubernetes.

romilbhardwaj · 2024-11-12T20:29:20Z

sky/utils/common_utils.py

@@ -697,3 +697,26 @@ def truncate_long_string(s: str, max_length: int = 35) -> str:
    if len(prefix) < max_length:
        prefix += s[len(prefix):max_length]
    return prefix + '...'
+
+
+def get_gpu_resource_name(name: str = 'nvidia.com/gpu'):


Instead of having name as a arg, we probably want to move the hardcoded string to the top of kubernetes_utils.py as DEFAULT_GPU_RESOURCE = 'nvidia.com/gpu' and use that in the method if the envvar does not exist.

Suggested change

def get_gpu_resource_name(name: str = 'nvidia.com/gpu'):

def get_gpu_resource_name():

romilbhardwaj · 2024-11-12T20:32:17Z

sky/provision/kubernetes/utils.py

+            total={
+                common_utils.get_gpu_resource_name(): int(accelerator_count)
+            },
+            free={
+                common_utils.get_gpu_resource_name():
+                    int(accelerators_available)
+            })


Cleaner to invoke once:

Suggested change

total={

common_utils.get_gpu_resource_name(): int(accelerator_count)

},

free={

common_utils.get_gpu_resource_name():

int(accelerators_available)

})

gpu_resource = common_utils.get_gpu_resource_name()

total={gpu_resource: int(accelerator_count)},

free={gpu_resource: int(accelerators_available)

})

romilbhardwaj · 2024-11-14T07:23:45Z

Hi @nkwangleiGIT - this PR was close to being merged. Would you like to reopen it?

romilbhardwaj

Thanks @nkwangleiGIT! Mostly looks good, left some minor comments. Should be good to go soon!

romilbhardwaj · 2024-11-15T17:36:22Z

sky/provision/kubernetes/utils.py

@@ -2233,10 +2233,11 @@ def get_node_accelerator_count(attribute_dict: dict) -> int:
        Number of accelerators allocated or available from the node. If no
            resource is found, it returns 0.
    """
-    assert not (GPU_RESOURCE_KEY in attribute_dict and
+    gpu_resource_name = get_gpu_resource_key()


nit:

Suggested change

gpu_resource_name = get_gpu_resource_key()

gpu_resource_key = get_gpu_resource_key()

romilbhardwaj · 2024-11-15T17:37:57Z

sky/provision/kubernetes/utils.py

+    # Retrieve GPU resource name from environment variable, if set.
+    # E.g., can be nvidia.com/gpu-h100, amd.com/gpu etc.
+    custom_name = os.getenv('CUSTOM_GPU_RESOURCE_KEY')
+
+    # If the environment variable is not defined, return the default name
+    if custom_name is None:
+        return GPU_RESOURCE_KEY
+
+    return custom_name


Can simplify logic to:

Suggested change

# Retrieve GPU resource name from environment variable, if set.

# E.g., can be nvidia.com/gpu-h100, amd.com/gpu etc.

custom_name = os.getenv('CUSTOM_GPU_RESOURCE_KEY')

# If the environment variable is not defined, return the default name

if custom_name is None:

return GPU_RESOURCE_KEY

return custom_name

# Retrieve GPU resource name from environment variable, if set. Else use default.

# E.g., can be nvidia.com/gpu-h100, amd.com/gpu etc.

return os.getenv('CUSTOM_GPU_RESOURCE_KEY', default = GPU_RESOURCE_KEY)

romilbhardwaj · 2024-11-15T17:43:03Z

sky/provision/kubernetes/utils.py

@@ -2395,3 +2396,23 @@ def process_skypilot_pods(
        num_pods = len(cluster.pods)
        cluster.resources_str = f'{num_pods}x {cluster.resources}'
    return list(clusters.values()), jobs_controllers, serve_controllers
+
+
+def get_gpu_resource_key():


Can we also update L194 in sky/provision/kubernetes/instance.py to use get_gpu_resource_key() to make sure errors are handled correctly?

skypilot/sky/provision/kubernetes/instance.py

Line 194 in c1726ae

elif (('Insufficient nvidia.com/gpu'

Signed-off-by: nkwangleiGIT <[email protected]>

romilbhardwaj

thanks @nkwangleiGIT - lgtm. Should be good to merge after linter passes.

romilbhardwaj · 2024-11-22T18:25:20Z

Hey @nkwangleiGIT, quick follow up - you can run the linter with ./format.sh in the root of the repo. Let me know if you'd like me to run it for you and push directly to this branch :)

nkwangleiGIT · 2024-12-15T10:14:09Z

sorry for late, just fetch the latest code, fix format and lint issue, and push again now

Signed-off-by: nkwangleiGIT <[email protected]>

nkwangleiGIT force-pushed the master branch from 0ff4387 to aa94450 Compare November 12, 2024 15:39

Michaelvll requested a review from romilbhardwaj November 12, 2024 17:40

romilbhardwaj reviewed Nov 12, 2024

View reviewed changes

nkwangleiGIT force-pushed the master branch 3 times, most recently from 18f7aac to 8538f07 Compare November 13, 2024 02:22

nkwangleiGIT closed this Nov 14, 2024

nkwangleiGIT force-pushed the master branch from 8538f07 to 140125e Compare November 14, 2024 06:07

nkwangleiGIT reopened this Nov 14, 2024

nkwangleiGIT force-pushed the master branch from 4e138a3 to c1726ae Compare November 14, 2024 08:50

romilbhardwaj reviewed Nov 15, 2024

View reviewed changes

nkwangleiGIT force-pushed the master branch from 0a72839 to a03f546 Compare November 16, 2024 02:30

[k8s] support to use custom gpu resource name if it's not nvidia.com/gpu

a03f546

Signed-off-by: nkwangleiGIT <[email protected]>

romilbhardwaj approved these changes Nov 21, 2024

View reviewed changes

Merge remote-tracking branch 'upstream/master'

f40bcf8

nkwangleiGIT force-pushed the master branch 3 times, most recently from c070a45 to 2a59cf8 Compare December 15, 2024 10:13

fix format issue

2a59cf8

Signed-off-by: nkwangleiGIT <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[k8s] support to use custom gpu resource name if it's not nvidia.com/gpu #4337

[k8s] support to use custom gpu resource name if it's not nvidia.com/gpu #4337

nkwangleiGIT commented Nov 12, 2024 •

edited

Loading

romilbhardwaj left a comment

romilbhardwaj Nov 12, 2024

nkwangleiGIT Nov 13, 2024

romilbhardwaj Nov 12, 2024

romilbhardwaj Nov 12, 2024

romilbhardwaj Nov 12, 2024

nkwangleiGIT Nov 13, 2024

romilbhardwaj Nov 12, 2024

nkwangleiGIT Nov 13, 2024

romilbhardwaj commented Nov 14, 2024

romilbhardwaj left a comment

romilbhardwaj Nov 15, 2024

romilbhardwaj Nov 15, 2024

romilbhardwaj Nov 15, 2024

nkwangleiGIT Nov 16, 2024

romilbhardwaj left a comment

romilbhardwaj commented Nov 22, 2024

nkwangleiGIT commented Dec 15, 2024

		# retrieve GPU resource name from environment variable,
		# can be nvidia.com/gpu-h100 if it's customized by the device plugin



		def get_gpu_resource_name(name: str = 'nvidia.com/gpu'):
		"""Get the GPU resource name.

	"""Get the GPU resource name.
	"""Get the GPU resource name to use in kubernetes.

	def get_gpu_resource_name(name: str = 'nvidia.com/gpu'):
	def get_gpu_resource_name():

	gpu_resource_name = get_gpu_resource_key()
	gpu_resource_key = get_gpu_resource_key()

[k8s] support to use custom gpu resource name if it's not nvidia.com/gpu #4337

Are you sure you want to change the base?

[k8s] support to use custom gpu resource name if it's not nvidia.com/gpu #4337

Conversation

nkwangleiGIT commented Nov 12, 2024 • edited Loading

romilbhardwaj left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

romilbhardwaj commented Nov 14, 2024

romilbhardwaj left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

romilbhardwaj left a comment

Choose a reason for hiding this comment

romilbhardwaj commented Nov 22, 2024

nkwangleiGIT commented Dec 15, 2024

nkwangleiGIT commented Nov 12, 2024 •

edited

Loading