[Minor] Improve completion of accelerator name when cloud is specified (#3014)

Author: Michaelvll

sky/resources.py

@@ -509,8 +509,9 @@ def _set_accelerators(
 
             # Canonicalize the accelerator names.
             accelerators = {
-                accelerator_registry.canonicalize_accelerator_name(acc):
-                acc_count for acc, acc_count in accelerators.items()
+                accelerator_registry.canonicalize_accelerator_name(
+                    acc, self._cloud): acc_count
+                for acc, acc_count in accelerators.items()
             }
 
             acc, _ = list(accelerators.items())[0]
@@ -1311,8 +1312,9 @@ def __setstate__(self, state):
             accelerators = state.pop('_accelerators', None)
             if accelerators is not None:
                 accelerators = {
-                    accelerator_registry.canonicalize_accelerator_name(acc):
-                    acc_count for acc, acc_count in accelerators.items()
+                    accelerator_registry.canonicalize_accelerator_name(
+                        acc, cloud=None): acc_count
+                    for acc, acc_count in accelerators.items()
                 }
             state['_accelerators'] = accelerators
 

sky/utils/accelerator_registry.py

@@ -1,7 +1,13 @@
 """Accelerator registry."""
+import typing
+from typing import Optional
+
 from sky.clouds import service_catalog
 from sky.utils import ux_utils
 
+if typing.TYPE_CHECKING:
+    from sky import clouds
+
 # Canonicalized names of all accelerators (except TPUs) supported by SkyPilot.
 # NOTE: Must include accelerators supported for local clusters.
 #
@@ -67,8 +73,13 @@ def is_schedulable_non_gpu_accelerator(accelerator_name: str) -> bool:
     return False
 
 
-def canonicalize_accelerator_name(accelerator: str) -> str:
+def canonicalize_accelerator_name(accelerator: str,
+                                  cloud: Optional['clouds.Cloud']) -> str:
     """Returns the canonical accelerator name."""
+    cloud_str = None
+    if cloud is not None:
+        cloud_str = str(cloud).lower()
+
     # TPU names are always lowercase.
     if accelerator.lower().startswith('tpu-'):
         return accelerator.lower()
@@ -84,7 +95,8 @@ def canonicalize_accelerator_name(accelerator: str) -> str:
     # To cover such cases, we should search the accelerator name
     # in the service catalog.
     searched = service_catalog.list_accelerators(name_filter=accelerator,
-                                                 case_sensitive=False)
+                                                 case_sensitive=False,
+                                                 clouds=cloud_str)
     names = list(searched.keys())
 
     # Exact match.