[k8s] support to use custom gpu resource name if it's not nvidia.com/gpu #4337
+47
−18
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
sky/utils/common_utils.py
Outdated
Comment on lines
702
to
722
| 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 |
There was a problem hiding this comment.
Can we move this to sky.provision.kubernetes.utils since only k8s relies on this method?
There was a problem hiding this comment.
moved to sky.provision.kubernetes.utils
sky/utils/common_utils.py
Outdated
Comment on lines
714
to
715
| # retrieve GPU resource name from environment variable, | ||
| # can be nvidia.com/gpu-h100 if it's customized by the device plugin |
There was a problem hiding this comment.
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. |
sky/utils/common_utils.py
Outdated
|
|
||
|
|
||
| def get_gpu_resource_name(name: str = 'nvidia.com/gpu'): | ||
| """Get the GPU resource name. |
There was a problem hiding this comment.
Suggested change
| """Get the GPU resource name. | |
| """Get the GPU resource name to use in kubernetes. |
sky/utils/common_utils.py
Outdated
| @@ -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'): | |||
There was a problem hiding this comment.
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(): |
sky/provision/kubernetes/utils.py
Outdated
Comment on lines
1860
to
1866
| total={ | ||
| common_utils.get_gpu_resource_name(): int(accelerator_count) | ||
| }, | ||
| free={ | ||
| common_utils.get_gpu_resource_name(): | ||
| int(accelerators_available) | ||
| }) |
There was a problem hiding this comment.
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) | |
| }) |
nkwangleiGIT
force-pushed
the
master
branch
3 times, most recently
from
18f7aac
to
8538f07
Compare
Signed-off-by: nkwangleiGIT <nkwanglei@126.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
We're using customized gpu resource name when using the device plugin, so we have the following GPUs in the node capability:
So we have to use the resource names above when launch to local K8S, such as
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):
bash format.shpytest tests/test_smoke.pypytest tests/test_smoke.py::test_fill_in_the_nameconda deactivate; bash -i tests/backward_compatibility_tests.sh