Skip to content

Commit

Permalink
fix format and lint
Browse files Browse the repository at this point in the history
  • Loading branch information
nkwangleiGIT committed Nov 12, 2024
1 parent 262a4ce commit 0ff4387
Show file tree
Hide file tree
Showing 4 changed files with 18 additions and 10 deletions.
5 changes: 3 additions & 2 deletions sky/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -3161,8 +3161,9 @@ def _get_kubernetes_node_info_table(context: Optional[str]):

node_info_dict = kubernetes_utils.get_kubernetes_node_info(context)
for node_name, node_info in node_info_dict.items():
available = node_info.free[common_utils.get_gpu_resource_name()] if node_info.free[
common_utils.get_gpu_resource_name()] != -1 else no_permissions_str
available = node_info.free[common_utils.get_gpu_resource_name(
)] if node_info.free[common_utils.get_gpu_resource_name(
)] != -1 else no_permissions_str
node_table.add_row([
node_name, node_info.gpu_type,
node_info.total[common_utils.get_gpu_resource_name()], available
Expand Down
3 changes: 2 additions & 1 deletion sky/clouds/service_catalog/kubernetes_catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,8 @@ def list_accelerators_realtime(
continue

accelerator_count = int(
node.status.allocatable.get(common_utils.get_gpu_resource_name(), 0))
node.status.allocatable.get(
common_utils.get_gpu_resource_name(), 0))

# Generate the GPU quantities for the accelerators
if accelerator_name and accelerator_count > 0:
Expand Down
15 changes: 10 additions & 5 deletions sky/provision/kubernetes/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1833,8 +1833,9 @@ def get_kubernetes_node_info(
else:
accelerator_name = None

accelerator_count = int(node.status.allocatable.get(
common_utils.get_gpu_resource_name(), 0))
accelerator_count = int(
node.status.allocatable.get(common_utils.get_gpu_resource_name(),
0))

if pods is None:
accelerators_available = -1
Expand All @@ -1856,9 +1857,13 @@ def get_kubernetes_node_info(
node_info_dict[node.metadata.name] = KubernetesNodeInfo(
name=node.metadata.name,
gpu_type=accelerator_name,
total={common_utils.get_gpu_resource_name(): int(accelerator_count)},
free={common_utils.get_gpu_resource_name(): int(accelerators_available)})

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


Expand Down
5 changes: 3 additions & 2 deletions sky/utils/common_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -698,7 +698,8 @@ def truncate_long_string(s: str, max_length: int = 35) -> str:
prefix += s[len(prefix):max_length]
return prefix + '...'

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

def get_gpu_resource_name(name: str = 'nvidia.com/gpu'):
"""
Get the GPU resource name.
Expand All @@ -713,7 +714,7 @@ def get_gpu_resource_name(name="nvidia.com/gpu"):
"""
# 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")
custom_name = os.getenv('CUSTOM_GPU_RESOURCE_NAME')

# If the environment variable is not defined, return the default name
if custom_name is None:
Expand Down

0 comments on commit 0ff4387

Please sign in to comment.