[Core] Added functionality in show_gpus to have an option to show all…

… regions' pricing for an accelerator (#2892) * Added functionality in show_gpus to have an option to show all regions' pricing for an accelerator * Removed all-regions from kubernetes * Removed all-regions from kubernetes * Removed all-regions from kubernetes * Fixed kubernetes unused variable issue * Fixed kubernetes unused variable issue * Fixed kubernetes unused variable issue * Fixed kubernetes unused variable issue * Update sky/cli.py Co-authored-by: Romil Bhardwaj <[email protected]> * Update sky/cli.py Co-authored-by: Romil Bhardwaj <[email protected]> * Fixed kubernetes unused variable issue --------- Co-authored-by: Romil Bhardwaj <[email protected]>
skypilot-org · Dec 30, 2023 · 6c30d79 · 6c30d79
1 parent 3eb6f2f
commit 6c30d79
Show file tree

Hide file tree

Showing 11 changed files with 84 additions and 44 deletions.
diff --git a/sky/cli.py b/sky/cli.py
@@ -3395,13 +3395,20 @@ def check(verbose: bool):
     ('The region to use. If not specified, shows accelerators from all regions.'
     ),
 )
+@click.option(
+    '--all-regions',
+    is_flag=True,
+    default=False,
+    help='Show pricing and instance details for a specified accelerator across '
+    'all regions and clouds.')
 @service_catalog.fallback_to_default_catalog
 @usage_lib.entrypoint
 def show_gpus(
         accelerator_str: Optional[str],
         all: bool,  # pylint: disable=redefined-builtin
         cloud: Optional[str],
-        region: Optional[str]):
+        region: Optional[str],
+        all_regions: Optional[bool]):
     """Show supported GPU/TPU/accelerators and their prices.
 
     The names and counts shown can be set in the ``accelerators`` field in task
@@ -3415,21 +3422,35 @@ def show_gpus(
     To show all accelerators, including less common ones and their detailed
     information, use ``sky show-gpus --all``.
 
+    To show all regions for a specified accelerator, use
+    ``sky show-gpus <accelerator> --all-regions``.
+
     Definitions of certain fields:
 
     * ``DEVICE_MEM``: Memory of a single device; does not depend on the device
       count of the instance (VM).
 
     * ``HOST_MEM``: Memory of the host instance (VM).
 
-    If ``--region`` is not specified, the price displayed for each instance
-    type is the lowest across all regions for both on-demand and spot
-    instances. There may be multiple regions with the same lowest price.
+    If ``--region`` or ``--all-regions`` is not specified, the price displayed
+    for each instance type is the lowest across all regions for both on-demand
+    and spot instances. There may be multiple regions with the same lowest
+    price.
     """
     # validation for the --region flag
     if region is not None and cloud is None:
         raise click.UsageError(
             'The --region flag is only valid when the --cloud flag is set.')
+
+    # validation for the --all-regions flag
+    if all_regions and accelerator_str is None:
+        raise click.UsageError(
+            'The --all-regions flag is only valid when an accelerator '
+            'is specified.')
+    if all_regions and region is not None:
+        raise click.UsageError(
+            '--all-regions and --region flags cannot be used simultaneously.')
+
     # This will validate 'cloud' and raise if not found.
     clouds.CLOUD_REGISTRY.from_str(cloud)
     service_catalog.validate_region_zone(region, None, clouds=cloud)
@@ -3514,7 +3535,8 @@ def _output():
                                                    quantity_filter=quantity,
                                                    region_filter=region,
                                                    clouds=cloud,
-                                                   case_sensitive=False)
+                                                   case_sensitive=False,
+                                                   all_regions=all_regions)
 
         if len(result) == 0:
             if cloud == 'kubernetes':

diff --git a/sky/clouds/service_catalog/__init__.py b/sky/clouds/service_catalog/__init__.py
@@ -27,7 +27,6 @@ def _map_clouds_catalog(clouds: CloudFilter, method_name: str, *args, **kwargs):
         # kubernetes_catalog.py
         if method_name != 'list_accelerators':
             clouds.remove('kubernetes')
-
     single = isinstance(clouds, str)
     if single:
         clouds = [clouds]  # type: ignore
@@ -61,6 +60,7 @@ def list_accelerators(
     quantity_filter: Optional[int] = None,
     clouds: CloudFilter = None,
     case_sensitive: bool = True,
+    all_regions: bool = False,
 ) -> 'Dict[str, List[common.InstanceTypeInfo]]':
     """List the names of all accelerators offered by Sky.
 
@@ -72,7 +72,7 @@ def list_accelerators(
     """
     results = _map_clouds_catalog(clouds, 'list_accelerators', gpus_only,
                                   name_filter, region_filter, quantity_filter,
-                                  case_sensitive)
+                                  case_sensitive, all_regions)
     if not isinstance(results, list):
         results = [results]
     ret: Dict[str,

diff --git a/sky/clouds/service_catalog/aws_catalog.py b/sky/clouds/service_catalog/aws_catalog.py
@@ -291,12 +291,13 @@ def list_accelerators(
         name_filter: Optional[str],
         region_filter: Optional[str],
         quantity_filter: Optional[int],
-        case_sensitive: bool = True
-) -> Dict[str, List[common.InstanceTypeInfo]]:
+        case_sensitive: bool = True,
+        all_regions: bool = False) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in AWS offering accelerators."""
     return common.list_accelerators_impl('AWS', _get_df(), gpus_only,
                                          name_filter, region_filter,
-                                         quantity_filter, case_sensitive)
+                                         quantity_filter, case_sensitive,
+                                         all_regions)
 
 
 def get_image_id_from_tag(tag: str, region: Optional[str]) -> Optional[str]:

diff --git a/sky/clouds/service_catalog/azure_catalog.py b/sky/clouds/service_catalog/azure_catalog.py
@@ -169,9 +169,9 @@ def list_accelerators(
         name_filter: Optional[str],
         region_filter: Optional[str],
         quantity_filter: Optional[int],
-        case_sensitive: bool = True
-) -> Dict[str, List[common.InstanceTypeInfo]]:
+        case_sensitive: bool = True,
+        all_regions: bool = False) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in Azure offering GPUs."""
     return common.list_accelerators_impl('Azure', _df, gpus_only, name_filter,
                                          region_filter, quantity_filter,
-                                         case_sensitive)
+                                         case_sensitive, all_regions)
diff --git a/sky/clouds/service_catalog/common.py b/sky/clouds/service_catalog/common.py
@@ -452,6 +452,7 @@ def list_accelerators_impl(
     region_filter: Optional[str],
     quantity_filter: Optional[int],
     case_sensitive: bool = True,
+    all_regions: bool = False,
 ) -> Dict[str, List[InstanceTypeInfo]]:
     """Lists accelerators offered in a cloud service catalog.
 
@@ -500,12 +501,20 @@ def list_accelerators_impl(
     grouped = df.groupby('AcceleratorName')
 
     def make_list_from_df(rows):
-        # Only keep the lowest prices across regions.
-        rows = rows.groupby([
-            'InstanceType', 'AcceleratorName', 'AcceleratorCount', 'vCPUs',
-            'MemoryGiB'
-        ],
-                            dropna=False).aggregate('min').reset_index()
+        if all_regions:
+            # Keep all regions.
+            rows = rows.groupby([
+                'InstanceType', 'AcceleratorName', 'AcceleratorCount', 'vCPUs',
+                'MemoryGiB', 'Region'
+            ],
+                                dropna=False).aggregate('min').reset_index()
+        else:
+            # Only keep the lowest prices across regions.
+            rows = rows.groupby([
+                'InstanceType', 'AcceleratorName', 'AcceleratorCount', 'vCPUs',
+                'MemoryGiB'
+            ],
+                                dropna=False).aggregate('min').reset_index()
         ret = rows.apply(
             lambda row: InstanceTypeInfo(
                 cloud,

diff --git a/sky/clouds/service_catalog/gcp_catalog.py b/sky/clouds/service_catalog/gcp_catalog.py
@@ -369,11 +369,12 @@ def list_accelerators(
     region_filter: Optional[str] = None,
     quantity_filter: Optional[int] = None,
     case_sensitive: bool = True,
+    all_regions: bool = False,
 ) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in GCP offering GPUs."""
     results = common.list_accelerators_impl('GCP', _df, gpus_only, name_filter,
                                             region_filter, quantity_filter,
-                                            case_sensitive)
+                                            case_sensitive, all_regions)
 
     # Remove GPUs that are unsupported by SkyPilot.
     new_results = {}

diff --git a/sky/clouds/service_catalog/ibm_catalog.py b/sky/clouds/service_catalog/ibm_catalog.py
@@ -85,16 +85,17 @@ def get_region_zones_for_instance_type(instance_type: str,
 
 
 def list_accelerators(
-        gpus_only: bool,
-        name_filter: Optional[str],
-        region_filter: Optional[str],
-        quantity_filter: Optional[int],
-        case_sensitive: bool = True
+    gpus_only: bool,
+    name_filter: Optional[str],
+    region_filter: Optional[str],
+    quantity_filter: Optional[int],
+    case_sensitive: bool = True,
+    all_regions: bool = False,
 ) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in IBM offering accelerators."""
     return common.list_accelerators_impl('IBM', _df, gpus_only, name_filter,
                                          region_filter, quantity_filter,
-                                         case_sensitive)
+                                         case_sensitive, all_regions)
 
 
 def get_default_instance_type(cpus: Optional[str] = None,

diff --git a/sky/clouds/service_catalog/kubernetes_catalog.py b/sky/clouds/service_catalog/kubernetes_catalog.py
@@ -34,13 +34,15 @@ def is_image_tag_valid(tag: str, region: Optional[str]) -> bool:
 
 
 def list_accelerators(
-        gpus_only: bool,
-        name_filter: Optional[str],
-        region_filter: Optional[str],
-        quantity_filter: Optional[int],
-        case_sensitive: bool = True
+    gpus_only: bool,
+    name_filter: Optional[str],
+    region_filter: Optional[str],
+    quantity_filter: Optional[int],
+    case_sensitive: bool = True,
+    all_regions: bool = False,
 ) -> Dict[str, List[common.InstanceTypeInfo]]:
     k8s_cloud = Kubernetes()
+    del all_regions  # unused
     if not any(
             map(k8s_cloud.is_same_cloud, global_user_state.get_enabled_clouds())
     ) or not kubernetes_utils.check_credentials()[0]:

diff --git a/sky/clouds/service_catalog/lambda_catalog.py b/sky/clouds/service_catalog/lambda_catalog.py
@@ -126,13 +126,14 @@ def get_region_zones_for_instance_type(instance_type: str,
 
 
 def list_accelerators(
-        gpus_only: bool,
-        name_filter: Optional[str],
-        region_filter: Optional[str],
-        quantity_filter: Optional[int],
-        case_sensitive: bool = True
+    gpus_only: bool,
+    name_filter: Optional[str],
+    region_filter: Optional[str],
+    quantity_filter: Optional[int],
+    case_sensitive: bool = True,
+    all_regions: bool = False,
 ) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in Lambda offering GPUs."""
     return common.list_accelerators_impl('Lambda', _df, gpus_only, name_filter,
                                          region_filter, quantity_filter,
-                                         case_sensitive)
+                                         case_sensitive, all_regions)
diff --git a/sky/clouds/service_catalog/oci_catalog.py b/sky/clouds/service_catalog/oci_catalog.py
@@ -167,16 +167,18 @@ def get_region_zones_for_instance_type(instance_type: str,
 
 
 def list_accelerators(
-        gpus_only: bool,
-        name_filter: Optional[str],
-        region_filter: Optional[str],
-        quantity_filter: Optional[int],
-        case_sensitive: bool = True
+    gpus_only: bool,
+    name_filter: Optional[str],
+    region_filter: Optional[str],
+    quantity_filter: Optional[int],
+    case_sensitive: bool = True,
+    all_regions: bool = False,
 ) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in OCI offering GPUs."""
     return common.list_accelerators_impl('OCI', _get_df(), gpus_only,
                                          name_filter, region_filter,
-                                         quantity_filter, case_sensitive)
+                                         quantity_filter, case_sensitive,
+                                         all_regions)
 
 
 def get_vcpus_mem_from_instance_type(

diff --git a/sky/clouds/service_catalog/scp_catalog.py b/sky/clouds/service_catalog/scp_catalog.py
@@ -127,11 +127,12 @@ def list_accelerators(
     region_filter: Optional[str],
     quantity_filter: Optional[int],
     case_sensitive: bool = True,
+    all_regions: bool = False,
 ) -> Dict[str, List[common.InstanceTypeInfo]]:
     """Returns all instance types in SCP offering GPUs."""
     return common.list_accelerators_impl('scp', _df, gpus_only, name_filter,
                                          region_filter, quantity_filter,
-                                         case_sensitive)
+                                         case_sensitive, all_regions)
 
 
 def get_image_id_from_tag(tag: str, region: Optional[str]) -> Optional[str]: