Repository: skypilot-org/skypilot
Pull request: Fix assertion on Azure's A100/V100 and fuzzy resource search
Fixes https://github.com/sky-proj/sky/issues/367. @gmittal 

Problem statement:
Fuzzy resource search
I type something like this: `sky gpunode --cloud azure --gpus A100`, but sky throws an exception (as expected, since this not a valid configuration). It would be really nice to show some kind of "nearest neighbor" resource that is valid and still accomplishes my goal of launching an A100 instance on Azure.

Example:
```
$ sky gpunode --cloud azure --gpus A100
Invalid resource configuration. Did you mean --cloud azure --gpus A100:8?
```

Unable to provision A100 cluster
@infwinston mentioned that he has a 1 line fix to solve this. This is hi-pri for our users so we should get this fixed as soon as we can.

```
sky ❯ gnode -c a100 --cloud azure --gpus A100:8
Traceback (most recent call last):
  File "/Users/gautam/miniconda3/envs/sky/bin/sky", line 33, in <module>
    sys.exit(load_entry_point('sky', 'console_scripts', 'sky')())
  File "/Users/gautam/miniconda3/envs/sky/lib/python3.8/site-packages/click/core.py", line 1128, in __call__
    return self.main(*args, **kwargs)
  File "/Users/gautam/miniconda3/envs/sky/lib/python3.8/site-packages/click/core.py", line 1053, in main
    rv = self.invoke(ctx)
  File "/Users/gautam/miniconda3/envs/sky/lib/python3.8/site-packages/click/core.py", line 1659, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/Users/gautam/miniconda3/envs/sky/lib/python3.8/site-packages/click/core.py", line 1395, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/Users/gautam/miniconda3/envs/sky/lib/python3.8/site-packages/click/core.py", line 754, in invoke
    return __callback(*args, **kwargs)
  File "/Users/gautam/Desktop/workbench/sky/sky/cli.py", line 1167, in gpunode
    _create_and_ssh_into_node(
  File "/Users/gautam/Desktop/workbench/sky/sky/cli.py", line 311, in _create_and_ssh_into_node
    dag = sky.optimize(dag)
  File "/Users/gautam/Desktop/workbench/sky/sky/optimizer.py", line 82, in optimize
    optimized_dag, unused_best_plan = Optimizer._optimize_cost(
  File "/Users/gautam/Desktop/workbench/sky/sky/optimizer.py", line 195, in _optimize_cost
    _fill_in_launchable_resources(
  File "/Users/gautam/Desktop/workbench/sky/sky/optimizer.py", line 404, in _fill_in_launchable_resources
    resources] = resources.cloud.get_feasible_launchable_resources(
  File "/Users/gautam/Desktop/workbench/sky/sky/clouds/azure.py", line 185, in get_feasible_launchable_resources
    instance_type = service_catalog.get_instance_type_for_accelerator(
  File "/Users/gautam/Desktop/workbench/sky/sky/clouds/service_catalog/__init__.py", line 119, in get_instance_type_for_accelerator
    return _map_clouds_catalog(clouds, 'get_instance_type_for_accelerator',
  File "/Users/gautam/Desktop/workbench/sky/sky/clouds/service_catalog/__init__.py", line 36, in _map_clouds_catalog
    results.append(method(*args, **kwargs))
  File "/Users/gautam/Desktop/workbench/sky/sky/clouds/service_catalog/azure_catalog.py", line 36, in get_instance_type_for_accelerator
    return common.get_instance_type_for_accelerator_impl(
  File "/Users/gautam/Desktop/workbench/sky/sky/clouds/service_catalog/common.py", line 101, in get_instance_type_for_accelerator_impl
    assert its == ['g3.4xlarge', 'g3s.xlarge'
AssertionError: ['Standard_ND96amsr_A100_v4', 'Standard_ND96asr_v4']
```

Hints from issue discussion:
Fixing it now.

Commit message: Fix assertion on Azure's A100/V100

Patch to review:
diff --git a/sky/clouds/service_catalog/common.py b/sky/clouds/service_catalog/common.py
index 4ab38819633..e4684c7d96e 100644
--- a/sky/clouds/service_catalog/common.py
+++ b/sky/clouds/service_catalog/common.py
@@ -95,11 +95,17 @@ def get_instance_type_for_accelerator_impl(
         # and count. Throw so we can manually investigate. The current
         # whitelist consists of:
         if len(instance_types) == 2:
+            # FIXME: refactor below implementation
             # - M60, offered by AWS g3s.xl and g3.4xl
             # - "Promo" instance types offered by Azure
+            # - A100 instances (80GB or 40GB A100)
+            # - V100 instances (w or w/o RDMA)
             its = sorted(instance_types)
-            assert its == ['g3.4xlarge', 'g3s.xlarge'
-                          ] or its[0] + '_Promo' == its[1], its
+            assert its == [
+                'g3.4xlarge', 'g3s.xlarge'
+            ] or its[0] + '_Promo' == its[1] or its == [
+                'Standard_ND96amsr_A100_v4', 'Standard_ND96asr_v4'
+            ] or its == ['Standard_NC24rs_v3', 'Standard_NC24s_v3'], its
         elif len(instance_types) == 4:
             its = sorted(instance_types)
             assert its == [


Code review comments:
diff_hunk: @@ -95,11 +95,17 @@ def get_instance_type_for_accelerator_impl(
         # and count. Throw so we can manually investigate. The current
         # whitelist consists of:
         if len(instance_types) == 2:
+            # FIXME: refactor below implementation
             # - M60, offered by AWS g3s.xl and g3.4xl
             # - "Promo" instance types offered by Azure
+            # - A100 instances (80GB or 40GB A100)
+            # - V100 instances (w or w/o RDMA)
             its = sorted(instance_types)
-            assert its == ['g3.4xlarge', 'g3s.xlarge'
-                          ] or its[0] + '_Promo' == its[1], its
+            assert its == [
+                'g3.4xlarge', 'g3s.xlarge'
+            ] or its[0] + '_Promo' == its[1] or its == [
+                'Standard_ND96amsr_A100_v4', 'Standard_ND96asr_v4'
path: sky/clouds/service_catalog/common.py
text: @user1:
So we will choose the 40GB A100 by default since it is cheaper?

@author:
Yes that's the logic for now if user only specifies `A100:8`. But I guess we should probably hint users that there are multiple choices of A100.

@user1:
Maybe we can add a TODO here? We have V100 vs V100-32GB by changing the acc_name in the catalog fetching code. Maybe we can do similar thing for A100.
https://github.com/sky-proj/sky/blob/f212f1e4598f4f187ab7bf711d0585a8b5af358e/sky/clouds/service_catalog/data_fetchers/fetch_aws.py#L114-L115

@author:
Ah you're right. I should probably follow this. But still users may not easily find out there are `V100-32GB` or `A100-80GB` options. I think a slightly annoying but unambiguous way is to prompt two options `V100-16GB` or `V100-32GB` when user type `V100` and let them choose one?

@user1:
The user can find it by `sky show-gpus`. I think it would be still fine to use `V100` for `V100-16GB` by default, otherwise, it may slightly increase the learning curve for new users?

@author:
Yeah it should be fine for now.

@author:
just realized a problem: Azure doesn't offer instances with 8x V100-16GB. so if user types `V100:8 --cloud azure` sky will say `No launchable resource`.