From 9c8fa306605c07de346054bdc7cf4727aee42d6b Mon Sep 17 00:00:00 2001
From: Jeremy Goodsitt <jeremy.goodsitt@ipcopilot.ai>
Date: Tue, 7 May 2024 12:11:15 -0500
Subject: [PATCH] feat: refactor format of sg names

---
 sky/clouds/aws.py    | 13 +++++--------
 sky/utils/schemas.py | 33 ++++++++++++++++++++-------------
 2 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/sky/clouds/aws.py b/sky/clouds/aws.py
index 250703b2c76..dcd608b32cd 100644
--- a/sky/clouds/aws.py
+++ b/sky/clouds/aws.py
@@ -1,5 +1,6 @@
 """Amazon Web Services."""
 import enum
+import fnmatch
 import functools
 import json
 import os
@@ -400,15 +401,11 @@ def make_deploy_resources_variables(self,
 
         user_security_group = skypilot_config.get_nested(
             ('aws', 'security_group_name'), None)
-        if user_security_group is not None and not isinstance(
-                user_security_group, str):
-            for sg_name in user_security_group:
-                if cluster_name_on_cloud.startswith(
-                        sg_name) and sg_name != 'default':
-                    user_security_group = user_security_group[sg_name]
+        if user_security_group is not None and not isinstance(user_security_group, str):
+            for profile in user_security_group:
+                if fnmatch.fnmatchcase(cluster_name_on_cloud, list(profile.keys())[0]):
+                    user_security_group = list(profile.values())[0]
                     break
-                elif sg_name == 'default':
-                    user_security_group = user_security_group[sg_name]
         security_group = user_security_group
         if user_security_group is None and resources.ports is not None:
             # Already checked in Resources._try_validate_ports
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index 4cd7f1aa291..b52b6baf5f8 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -592,21 +592,28 @@ def get_config_schema():
             'additionalProperties': False,
             'properties': {
                 'security_group_name': {
-                    'oneOf': [{
-                        'type': 'string'
-                    }, {
-                        'type': 'object',
-                        'additionalProperties': False,
-                        'required': ['default'],
-                        'properties': {
-                            'sky-serve-controller': {
-                                'type': 'string',
+                    'oneOf': [
+                        {
+                            'type': 'string'
+                        },
+                        {
+                            # A list of single-element dict to pretain the order.
+                            # Example:
+                            #  security_group_name:
+                            #    - my-cluster1-*: my-security-group-1
+                            #    - my-cluster2-*: my-security-group-2
+                            #    - "*"": my-security-group-3
+                            'type': 'array',
+                            'items': {
+                                'type': 'object',
+                                'additionalProperties': {
+                                    'type': 'string'
+                                },
+                                'maxProperties': 1,
+                                'minProperties': 1,
                             },
-                            'default': {
-                                'type': 'string'
-                            }
                         }
-                    }]
+                    ]
                 },
                 **_LABELS_SCHEMA,
                 **_NETWORK_CONFIG_SCHEMA,