From 8486cf339af18336ab7348f8b789b8a647c733aa Mon Sep 17 00:00:00 2001
From: Justin Dong <102186064+justin-e-dong@users.noreply.github.com>
Date: Mon, 23 Sep 2024 21:45:34 -0400
Subject: [PATCH] Cleaned up details

---
 sky/clouds/service_catalog/aws_catalog.py   |  9 ++++++++-
 sky/clouds/service_catalog/azure_catalog.py | 12 +++++++++---
 sky/clouds/service_catalog/gcp_catalog.py   |  9 ++++++++-
 sky/clouds/service_catalog/ibm_catalog.py   | 11 +++++++++--
 4 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/sky/clouds/service_catalog/aws_catalog.py b/sky/clouds/service_catalog/aws_catalog.py
index 77d080c8999..240d450a9b0 100644
--- a/sky/clouds/service_catalog/aws_catalog.py
+++ b/sky/clouds/service_catalog/aws_catalog.py
@@ -238,7 +238,14 @@ def get_default_instance_type(
     instance_type_prefix = tuple(
         f'{family}.' for family in _DEFAULT_INSTANCE_FAMILY)
     df = _get_df()
-    df = df[df['InstanceType'].str.startswith(instance_type_prefix)]
+
+    default_instance_type_df = df[df['InstanceType'].str.startswith(
+        instance_type_prefix)]
+    default_instance_type = common.get_instance_type_for_cpus_mem_impl(
+        default_instance_type_df, cpus, memory_gb_or_ratio)
+    if default_instance_type is not None:
+        return default_instance_type
+
     return common.get_instance_type_for_cpus_mem_impl(df, cpus,
                                                       memory_gb_or_ratio)
 
diff --git a/sky/clouds/service_catalog/azure_catalog.py b/sky/clouds/service_catalog/azure_catalog.py
index 2d323cbac5f..bbd70d328bd 100644
--- a/sky/clouds/service_catalog/azure_catalog.py
+++ b/sky/clouds/service_catalog/azure_catalog.py
@@ -106,14 +106,20 @@ def get_default_instance_type(
         memory_gb_or_ratio = f'{_DEFAULT_MEMORY_CPU_RATIO}x'
     else:
         memory_gb_or_ratio = memory
-    df = _df[_df['InstanceType'].apply(_get_instance_family).isin(
-        _DEFAULT_INSTANCE_FAMILY)]
 
     def _filter_disk_type(instance_type: str) -> bool:
         valid, _ = Azure.check_disk_tier(instance_type, disk_tier)
         return valid
 
-    df = df.loc[df['InstanceType'].apply(_filter_disk_type)]
+    df = _df.loc[_df['InstanceType'].apply(_filter_disk_type)]
+
+    default_instance_type_df = _df[_df['InstanceType'].apply(
+        _get_instance_family).isin(_DEFAULT_INSTANCE_FAMILY)]
+    default_instance_type = common.get_instance_type_for_cpus_mem_impl(
+        default_instance_type_df, cpus, memory_gb_or_ratio)
+    if default_instance_type is not None:
+        return default_instance_type
+
     return common.get_instance_type_for_cpus_mem_impl(df, cpus,
                                                       memory_gb_or_ratio)
 
diff --git a/sky/clouds/service_catalog/gcp_catalog.py b/sky/clouds/service_catalog/gcp_catalog.py
index f861b51920e..bb7867eb4f1 100644
--- a/sky/clouds/service_catalog/gcp_catalog.py
+++ b/sky/clouds/service_catalog/gcp_catalog.py
@@ -253,13 +253,20 @@ def get_default_instance_type(
     instance_type_prefix = tuple(
         f'{family}-' for family in _DEFAULT_INSTANCE_FAMILY)
     df = _df[_df['InstanceType'].notna()]
-    df = df[df['InstanceType'].str.startswith(instance_type_prefix)]
 
     def _filter_disk_type(instance_type: str) -> bool:
         valid, _ = GCP.check_disk_tier(instance_type, disk_tier)
         return valid
 
     df = df.loc[df['InstanceType'].apply(_filter_disk_type)]
+
+    default_instance_type_df = df[df['InstanceType'].str.startswith(
+        instance_type_prefix)]
+    default_instance_type = common.get_instance_type_for_cpus_mem_impl(
+        default_instance_type_df, cpus, memory_gb_or_ratio)
+    if default_instance_type is not None:
+        return default_instance_type
+
     return common.get_instance_type_for_cpus_mem_impl(df, cpus,
                                                       memory_gb_or_ratio)
 
diff --git a/sky/clouds/service_catalog/ibm_catalog.py b/sky/clouds/service_catalog/ibm_catalog.py
index 51b4e14f569..9f0bd8297c9 100644
--- a/sky/clouds/service_catalog/ibm_catalog.py
+++ b/sky/clouds/service_catalog/ibm_catalog.py
@@ -105,8 +105,15 @@ def get_default_instance_type(
     else:
         memory_gb_or_ratio = memory
     instance_type_prefix = f'{_DEFAULT_INSTANCE_FAMILY}-'
-    df = _df[_df['InstanceType'].str.startswith(instance_type_prefix)]
-    return common.get_instance_type_for_cpus_mem_impl(df, cpus,
+
+    default_instance_type_df = _df[_df['InstanceType'].str.startswith(
+        instance_type_prefix)]
+    default_instance_type = common.get_instance_type_for_cpus_mem_impl(
+        default_instance_type_df, cpus, memory_gb_or_ratio)
+    if default_instance_type is not None:
+        return default_instance_type
+
+    return common.get_instance_type_for_cpus_mem_impl(_df, cpus,
                                                       memory_gb_or_ratio)