[Tests] Fix smoke test for GCP disk tier best (#3992)

* [Tests] Fix smoke test for GCP disk tier `best` * Update tests/test_smoke.py Co-authored-by: Zhanghao Wu <[email protected]> * test multiple possibilities for gcp disk tier * Update tests/test_smoke.py Co-authored-by: Zhanghao Wu <[email protected]> --------- Co-authored-by: Zhanghao Wu <[email protected]>
skypilot-org · Sep 27, 2024 · 5cef688 · 5cef688
1 parent 4e46cf4
commit 5cef688
Showing 1 changed file with 27 additions and 16 deletions.
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
@@ -3427,26 +3427,37 @@ def _get_aws_query_command(region, instance_id, field, expected):
 @pytest.mark.gcp
 def test_gcp_disk_tier():
     for disk_tier in list(resources_utils.DiskTier):
-        type = GCP._get_disk_type(disk_tier)
+        disk_types = [GCP._get_disk_type(disk_tier)]
         name = _get_cluster_name() + '-' + disk_tier.value
         name_on_cloud = common_utils.make_cluster_name_on_cloud(
             name, sky.GCP.max_cluster_name_length())
         region = 'us-west2'
-        test = Test(
-            'gcp-disk-tier-' + disk_tier.value,
-            [
-                f'sky launch -y -c {name} --cloud gcp --region {region} '
-                f'--disk-tier {disk_tier.value} echo "hello sky"',
-                f'name=`gcloud compute instances list --filter='
-                f'"labels.ray-cluster-name:{name_on_cloud}" '
-                '--format="value(name)"`; '
-                f'gcloud compute disks list --filter="name=$name" '
-                f'--format="value(type)" | grep {type} '
-            ],
-            f'sky down -y {name}',
-            timeout=6 * 60,  # 6 mins  (it takes around ~3 mins)
-        )
-        run_one_test(test)
+        instance_type_options = ['']
+        if disk_tier == resources_utils.DiskTier.BEST:
+            # Ultra disk tier requires n2 instance types to have more than 64 CPUs.
+            # If using default instance type, it will only enable the high disk tier.
+            disk_types = [
+                GCP._get_disk_type(resources_utils.DiskTier.HIGH),
+                GCP._get_disk_type(resources_utils.DiskTier.ULTRA),
+            ]
+            instance_type_options = ['', '--instance-type n2-standard-64']
+        for disk_type, instance_type_option in zip(disk_types,
+                                                   instance_type_options):
+            test = Test(
+                'gcp-disk-tier-' + disk_tier.value,
+                [
+                    f'sky launch -y -c {name} --cloud gcp --region {region} '
+                    f'--disk-tier {disk_tier.value} {instance_type_option} ',
+                    f'name=`gcloud compute instances list --filter='
+                    f'"labels.ray-cluster-name:{name_on_cloud}" '
+                    '--format="value(name)"`; '
+                    f'gcloud compute disks list --filter="name=$name" '
+                    f'--format="value(type)" | grep {disk_type} '
+                ],
+                f'sky down -y {name}',
+                timeout=6 * 60,  # 6 mins  (it takes around ~3 mins)
+            )
+            run_one_test(test)
 
 
 @pytest.mark.azure