Make GPU/TPU names case-insensitive (skypilot-org#463)

Author: WoosukKwon

examples/time_estimators.py

@@ -49,7 +49,7 @@ def _v100(num_v100s):
         return _v100(num_v100s)
 
     elif isinstance(resources.cloud, sky.GCP):
-        accelerators = resources.get_accelerators()
+        accelerators = resources.accelerators
         if accelerators is None:
             assert False, 'not supported'
 
@@ -131,8 +131,8 @@ def resnet50_infer_estimate_runtime(resources):
         # TODO: this ignores offline vs. online.  It's a huge batch.
         estimated_run_time_seconds = \
             flops_for_one_image * num_images / utilized_flops
-    elif resources.get_accelerators() is not None:
-        accs = resources.get_accelerators()
+    elif resources.accelerators is not None:
+        accs = resources.accelerators
         for acc, acc_count in accs.items():
             break
         assert acc == 'T4' and acc_count == 1, resources

sky/backends/cloud_vm_ray_backend.py

@@ -88,7 +88,7 @@ def _get_task_demands_dict(task: Task) -> Optional[Tuple[Optional[str], int]]:
         assert len(task.resources) == 1, task.resources
         resources = list(task.resources)[0]
     if resources is not None:
-        accelerator_dict = resources.get_accelerators()
+        accelerator_dict = resources.accelerators
     return accelerator_dict
 
 

sky/cli.py

@@ -218,7 +218,7 @@ def _default_interactive_node_name(node_type: str):
 
 def _infer_interactive_node_type(resources: sky.Resources):
     """Determine interactive node type from resources."""
-    accelerators = resources.get_accelerators()
+    accelerators = resources.accelerators
     cloud = resources.cloud
     if accelerators:
         # We only support homogenous accelerators for now.

sky/clouds/aws.py

@@ -190,7 +190,7 @@ def _make(instance_type):
             return [r]
 
         # Currently, handle a filter on accelerators only.
-        accelerators = resources.get_accelerators()
+        accelerators = resources.accelerators
         if accelerators is None:
             # No requirements to filter, so just return a default VM type.
             return _make(AWS.get_default_instance_type())

sky/clouds/azure.py

@@ -182,7 +182,7 @@ def _make(instance_type):
             return [r]
 
         # Currently, handle a filter on accelerators only.
-        accelerators = resources.get_accelerators()
+        accelerators = resources.accelerators
         if accelerators is None:
             # No requirements to filter, so just return a default VM type.
             return _make(Azure.get_default_instance_type())

sky/clouds/gcp.py

@@ -177,7 +177,7 @@ def make_deploy_resources_variables(self, resources):
             'custom_resources': None,
             'use_spot': r.use_spot,
         }
-        accelerators = r.get_accelerators()
+        accelerators = r.accelerators
         if accelerators is not None:
             assert len(accelerators) == 1, r
             acc, acc_count = list(accelerators.items())[0]

sky/resources.py

@@ -3,12 +3,23 @@
 
 from sky import clouds
 from sky import sky_logging
+from sky.clouds import service_catalog
 
 logger = sky_logging.init_logger(__name__)
 
 DEFAULT_DISK_SIZE = 256
 
 
+def _get_name_from_catalog(accelerator: str) -> str:
+    """Returns the matched accelerator name in the catalog."""
+    acc_names = list(service_catalog.list_accelerators(gpus_only=False).keys())
+    try:
+        index = [n.casefold() for n in acc_names].index(accelerator.casefold())
+    except ValueError:
+        raise ValueError(f'Invalid accelerator name: {accelerator}') from None
+    return acc_names[index]
+
+
 class Resources:
     """A cloud resource bundle.
 
@@ -69,7 +80,7 @@ def __init__(
             assert len(accelerators) == 1, accelerators
 
             acc, _ = list(accelerators.items())[0]
-            if 'tpu' in acc:
+            if 'tpu' in acc.lower():
                 if cloud is None:
                     cloud = clouds.GCP()
                 assert cloud.is_same_cloud(clouds.GCP()), 'Cloud must be GCP.'
@@ -80,7 +91,7 @@ def __init__(
                                 ' default (2.5.0)')
                     accelerator_args['tf_version'] = '2.5.0'
 
-        self.accelerators = accelerators
+        self._accelerators = self._rename_accelerators(accelerators)
         self.accelerator_args = accelerator_args
 
         self._use_spot_specified = use_spot is not None
@@ -133,20 +144,36 @@ def _try_validate_accelerators(self) -> None:
             # because e.g., the instance may have 4 GPUs, while the task
             # specifies to use 1 GPU.
 
-    def get_accelerators(self) -> Optional[Dict[str, int]]:
+    def _rename_accelerators(
+        self,
+        accelerators: Union[None, Dict[str, int]],
+    ) -> Optional[Dict[str, int]]:
+        """Renames the accelerators in a case-sensitive manner."""
+        if accelerators is not None:
+            return {_get_name_from_catalog(name): cnt \
+                    for name, cnt in accelerators.items()}
+        else:
+            return None
+
+    @property
+    def accelerators(self) -> Optional[Dict[str, int]]:
         """Returns the accelerators field directly or by inferring.
 
         For example, Resources(AWS, 'p3.2xlarge') has its accelerators field
         set to None, but this function will infer {'V100': 1} from the instance
         type.
         """
-        if self.accelerators is not None:
-            return self.accelerators
+        if self._accelerators is not None:
+            return self._accelerators
         if self.cloud is not None and self.instance_type is not None:
             return self.cloud.get_accelerators_from_instance_type(
                 self.instance_type)
         return None
 
+    @accelerators.setter
+    def accelerators(self, accelerators: Union[None, Dict[str, int]]) -> None:
+        self._accelerators = self._rename_accelerators(accelerators)
+
     def get_cost(self, seconds: float):
         """Returns cost in USD for the runtime in seconds."""
         hours = seconds / 3600
@@ -177,8 +204,8 @@ def is_same_resources(self, other: 'Resources') -> bool:
             return False
         # self.instance_type == other.instance_type
 
-        other_accelerators = other.get_accelerators()
-        accelerators = self.get_accelerators()
+        other_accelerators = other.accelerators
+        accelerators = self.accelerators
         if accelerators != other_accelerators:
             return False
         # self.accelerators == other.accelerators
@@ -204,7 +231,7 @@ def less_demanding_than(self, other: 'Resources') -> bool:
             return False
         # self.instance_type <= other.instance_type
 
-        other_accelerators = other.get_accelerators()
+        other_accelerators = other.accelerators
         if self.accelerators is not None and other_accelerators is None:
             return False