Repository: skypilot-org/skypilot
Pull request: Added functionality to support specifying quantity of a resource in call to show-gpus
<!-- Describe the changes in this PR -->
This PR fixes issue #1915. Previously, users were not able to specify the GPU quantity that they wanted. This PR allows users to do so by running sky show-gpus V100:4 to find the cloud providers and regions that offer 4 V100 GPUs.

The output is now as follows:
```
sky show-gpus V100:8

GPU   QTY  CLOUD   INSTANCE_TYPE  DEVICE_MEM  vCPUs  HOST_MEM  HOURLY_PRICE  HOURLY_SPOT_PRICE  REGION            
V100  8    AWS     p3.16xlarge    16GB        64     488GB     $ 24.480      $ 7.344            ap-northeast-1    
V100  8    GCP     (attachable)   -           -      -         $ 19.840      $ 5.920            us-central1       
V100  8    Lambda  gpu_8x_v100    16GB        92     448GB     $ 4.400       -                  asia-northeast-1  

GPU        QTY  CLOUD  INSTANCE_TYPE       DEVICE_MEM  vCPUs  HOST_MEM  HOURLY_PRICE  HOURLY_SPOT_PRICE  REGION          
V100-32GB  8    AWS    p3dn.24xlarge       32GB        96     768GB     $ 31.212      $ 9.364            ap-northeast-1  
V100-32GB  8    Azure  Standard_ND40rs_v2  -           40     672GB     $ 22.032      $ 3.530            eastus
```

Failure output is as follows:
```
sky show-gpus V100:5

Resource 'V100' with quantity 5 not found. Try 'sky show-gpus --all' to show available accelerators and their quantities.
```

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[show-gpus] GPU count doesn't work with `sky show-gpus`
We support resource specs like `V100:4`, but `sky show-gpus V100:4` doesn't work:

```
(base) ➜  ~ sky show-gpus V100:4
Resources 'V100:4' not found. Try 'sky show-gpus --all' to show available accelerators.
```

We should allow resource count when specifying a GPU in `show-gpus` and filter the output to show resources which match. Something like:
```
(base) ➜  ~ sky show-gpus V100:4
*NOTE*: for most GCP accelerators, INSTANCE_TYPE == (attachable) means the host VM's cost is not included.

GPU   QTY  CLOUD   INSTANCE_TYPE       DEVICE_MEM  vCPUs  HOST_MEM  HOURLY_PRICE  HOURLY_SPOT_PRICE  REGION
V100  4    AWS     p3.8xlarge          16GB        32     244GB     $ 12.240      $ 3.672            ap-northeast-1
V100  4    Azure   Standard_NC24rs_v3  -           24     448GB     $ 13.460      $ 4.506            centralus
V100  4    Azure   Standard_NC24s_v3   -           24     448GB     $ 12.240      $ 4.097            centralus
V100  4    GCP     (attachable)        -           -      -         $ 9.920       $ 2.960            us-central1
```

Hints from issue discussion:
I'd like to take on this issue!

Commit message: Clean code Trial3 - fixed repetitive code

Patch to review:
diff --git a/sky/cli.py b/sky/cli.py
index c5914171f24..46fde8c2b33 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -2959,7 +2959,7 @@ def check():
 
 
 @cli.command()
-@click.argument('gpu_name', required=False)
+@click.argument('accelerator_str', required=False)
 @click.option('--all',
               '-a',
               is_flag=True,
@@ -2980,11 +2980,11 @@ def check():
 @service_catalog.use_default_catalog
 @usage_lib.entrypoint
 def show_gpus(
-        gpu_name: Optional[str],
+        accelerator_str: Optional[str],
         all: bool,  # pylint: disable=redefined-builtin
         cloud: Optional[str],
         region: Optional[str]):
-    """Show supported GPU/TPU/accelerators and their prices.
+    """Show supported GPU/TPU/accelerators and their= prices.
 
     The names and counts shown can be set in the ``accelerators`` field in task
     YAMLs, or in the ``--gpus`` flag in CLI commands. For example, if this
@@ -3016,7 +3016,7 @@ def show_gpus(
     clouds.CLOUD_REGISTRY.from_str(cloud)
     service_catalog.validate_region_zone(region, None, clouds=cloud)
     show_all = all
-    if show_all and gpu_name is not None:
+    if show_all and accelerator_str is not None:
         raise click.UsageError('--all is only allowed without a GPU name.')
 
     def _list_to_str(lst):
@@ -3030,7 +3030,7 @@ def _output():
         other_table = log_utils.create_table(
             ['OTHER_GPU', 'AVAILABLE_QUANTITIES'])
 
-        if gpu_name is None:
+        if accelerator_str is None:
             result = service_catalog.list_accelerator_counts(
                 gpus_only=True,
                 clouds=cloud,
@@ -3063,16 +3063,35 @@ def _output():
                 return
 
         # Show detailed accelerator information
+
+        name, quantity = accelerator_str.split(':')[0], None
+        has_quantity = accelerator_str.count(':') > 0
+        if has_quantity:
+            quantity = int(accelerator_str.split(':')[1])
         result = service_catalog.list_accelerators(gpus_only=True,
-                                                   name_filter=gpu_name,
+                                                   name_filter=name,
+                                                   quantity=quantity,
                                                    region_filter=region,
                                                    clouds=cloud)
         if len(result) == 0:
-            yield f'Resources \'{gpu_name}\' not found. '
+            yield f'Resources \'{name}\' not found. '
             yield 'Try \'sky show-gpus --all\' '
             yield 'to show available accelerators.'
             return
 
+        num_accelerators_found = 0
+        num_unique_accelerators = len([1 for i in result if len(result[i]) > 0])
+        for i in result:
+            num_accelerators_found += len(result[i])
+        if num_accelerators_found == 0:
+            yield f'Resource \'{name}\', with '
+            yield f'requested quantity {quantity}, '
+            yield 'not found. '
+            yield 'Try \'sky show-gpus --all\' '
+            yield 'to show available accelerators '
+            yield 'and their quantities.'
+            return
+
         yield '*NOTE*: for most GCP accelerators, '
         yield 'INSTANCE_TYPE == (attachable) means '
         yield 'the host VM\'s cost is not included.\n\n'
@@ -3113,6 +3132,7 @@ def _output():
                 spot_price_str = f'$ {item.spot_price:.3f}' if not pd.isna(
                     item.spot_price) else '-'
                 region_str = item.region if not pd.isna(item.region) else '-'
+
                 accelerator_table_vals = [
                     item.accelerator_name,
                     item.accelerator_count,
@@ -3127,10 +3147,10 @@ def _output():
                 if not show_all:
                     accelerator_table_vals.append(region_str)
                 accelerator_table.add_row(accelerator_table_vals)
-
-            if i != 0:
+            if i != 0 and num_unique_accelerators > 0:
                 yield '\n\n'
             yield from accelerator_table.get_string()
+            num_unique_accelerators -= 1
 
     if show_all:
         click.echo_via_pager(_output())
diff --git a/sky/clouds/service_catalog/__init__.py b/sky/clouds/service_catalog/__init__.py
index 81884ddd0e6..3a387cae91f 100644
--- a/sky/clouds/service_catalog/__init__.py
+++ b/sky/clouds/service_catalog/__init__.py
@@ -51,6 +51,7 @@ def _map_clouds_catalog(clouds: CloudFilter, method_name: str, *args, **kwargs):
 def list_accelerators(
     gpus_only: bool = True,
     name_filter: Optional[str] = None,
+    quantity: Optional[int] = None,
     region_filter: Optional[str] = None,
     clouds: CloudFilter = None,
     case_sensitive: bool = True,
@@ -71,7 +72,14 @@ def list_accelerators(
               List['common.InstanceTypeInfo']] = collections.defaultdict(list)
     for result in results:
         for gpu, items in result.items():
-            ret[gpu] += items
+            if quantity:
+                new_items = []
+                for item in items:
+                    if item.accelerator_count == quantity:
+                        new_items.append(item)
+                ret[gpu] += new_items
+            else:
+                ret[gpu] += items
     return dict(ret)
 
 


Code review comments:
diff_hunk: @@ -3063,16 +3063,35 @@ def _output():
                 return
 
         # Show detailed accelerator information
+
+        name, quantity = accelerator_str.split(':')[0], None
+        has_quantity = accelerator_str.count(':') > 0
+        if has_quantity:
+            quantity = int(accelerator_str.split(':')[1])
path: sky/cli.py
text: @user1:
We can improve readability and add validation for input here:
```suggestion
        accelerator_split = accelerator_str.split(':')
        if len(accelerator_split) > 2:
            raise click.UsageError(
                f'Invalid accelerator string {accelerator_str}. '
                'Expected format: <accelerator_name>[:<quantity>].')
        if len(accelerator_split) == 2:
            name = accelerator_split[0]
            # Check if quantity is valid
            try:
                quantity = int(accelerator_split[1])
            except ValueError:
                raise click.UsageError(
                    f'Invalid accelerator quantity {accelerator_split[1]}. '
                    'Expected an integer.')
        else:
            name, quantity = accelerator_str, None
```

@author:
Added validation as suggested!

diff_hunk: @@ -71,7 +72,14 @@ def list_accelerators(
               List['common.InstanceTypeInfo']] = collections.defaultdict(list)
     for result in results:
         for gpu, items in result.items():
-            ret[gpu] += items
+            if quantity:
+                new_items = []
+                for item in items:
+                    if item.accelerator_count == quantity:
+                        new_items.append(item)
+                ret[gpu] += new_items
+            else:
+                ret[gpu] += items
path: sky/clouds/service_catalog/__init__.py
text: @user1:
My bad, I should have clarified earlier - we should implement this filtering further downstream in [list_accelerators_impl](https://github.com/skypilot-org/skypilot/blob/9dca02b2c6b0290d46d1ab03829879cee9949d8c/sky/clouds/service_catalog/common.py#L423). That was `results` returned in L67 should already have GPUs filtered by quantity. You may need to add a `quantity_filter` arg to list_accelerators_impl. 

@author:
Pushed all the filtering down into list_accelerators_impl and added the quantity_filter to all the list_accelerator functions in all the cloud catalog files

diff_hunk: @@ -3127,10 +3147,10 @@ def _output():
                 if not show_all:
                     accelerator_table_vals.append(region_str)
                 accelerator_table.add_row(accelerator_table_vals)
-
-            if i != 0:
+            if i != 0 and num_unique_accelerators > 0:
                 yield '\n\n'
             yield from accelerator_table.get_string()
+            num_unique_accelerators -= 1
path: sky/cli.py
text: @user1:
Should not be required once we update `list_accelerators` to return empty dict if resources do not match pass the quantity filter.

diff_hunk: @@ -2980,11 +2980,11 @@ def check():
 @service_catalog.use_default_catalog
 @usage_lib.entrypoint
 def show_gpus(
-        gpu_name: Optional[str],
+        accelerator_str: Optional[str],
         all: bool,  # pylint: disable=redefined-builtin
         cloud: Optional[str],
         region: Optional[str]):
-    """Show supported GPU/TPU/accelerators and their prices.
+    """Show supported GPU/TPU/accelerators and their= prices.
path: sky/cli.py
text: @user1:
`=` unintended?

diff_hunk: @@ -3113,6 +3132,7 @@ def _output():
                 spot_price_str = f'$ {item.spot_price:.3f}' if not pd.isna(
                     item.spot_price) else '-'
                 region_str = item.region if not pd.isna(item.region) else '-'
+
path: sky/cli.py
text: @user1:
Unintended?

diff_hunk: @@ -51,6 +51,7 @@ def _map_clouds_catalog(clouds: CloudFilter, method_name: str, *args, **kwargs):
 def list_accelerators(
     gpus_only: bool = True,
     name_filter: Optional[str] = None,
+    quantity: Optional[int] = None,
path: sky/clouds/service_catalog/__init__.py
text: @user1:
Can we call this `quantity_filter`?

diff_hunk: @@ -3063,16 +3063,35 @@ def _output():
                 return
 
         # Show detailed accelerator information
+
+        name, quantity = accelerator_str.split(':')[0], None
+        has_quantity = accelerator_str.count(':') > 0
+        if has_quantity:
+            quantity = int(accelerator_str.split(':')[1])
         result = service_catalog.list_accelerators(gpus_only=True,
-                                                   name_filter=gpu_name,
+                                                   name_filter=name,
+                                                   quantity=quantity,
                                                    region_filter=region,
                                                    clouds=cloud)
         if len(result) == 0:
-            yield f'Resources \'{gpu_name}\' not found. '
+            yield f'Resources \'{name}\' not found. '
             yield 'Try \'sky show-gpus --all\' '
             yield 'to show available accelerators.'
             return
 
+        num_accelerators_found = 0
+        num_unique_accelerators = len([1 for i in result if len(result[i]) > 0])
+        for i in result:
+            num_accelerators_found += len(result[i])
+        if num_accelerators_found == 0:
path: sky/cli.py
text: @user1:
We shouldn't need these checks - result should be empty if no accelerators match the quantity_filter. We should modify `list_accelerators` have this behavior.  

@author:
Removed all this unnecessary filtering