From 77191522d61726b765c3ae77d65940612a5a6720 Mon Sep 17 00:00:00 2001
From: Jeremy Goodsitt <jeremy.goodsitt@ipcopilot.ai>
Date: Thu, 2 May 2024 00:29:36 -0400
Subject: [PATCH] refactor: use remote_identity instead of iam_instance_profile

---
 sky/backends/backend_utils.py | 11 ++++++++++-
 sky/clouds/aws.py             | 13 -------------
 sky/templates/aws-ray.yml.j2  |  4 ++--
 sky/utils/schemas.py          | 28 +++++++++-------------------
 4 files changed, 21 insertions(+), 35 deletions(-)

diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 0da2bd9ef0b2..28698cefdf4d 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -1,6 +1,7 @@
 """Util constants/functions for the backends."""
 from datetime import datetime
 import enum
+import fnmatch
 import functools
 import os
 import pathlib
@@ -798,7 +799,13 @@ def write_cluster_config(
     excluded_clouds = []
     remote_identity = skypilot_config.get_nested(
         (str(cloud).lower(), 'remote_identity'), 'LOCAL_CREDENTIALS')
-    if remote_identity == 'SERVICE_ACCOUNT':
+    if remote_identity is not None and not isinstance(
+            remote_identity, str):
+        for profile in remote_identity:
+            if fnmatch.fnmatchcase(cluster_name_on_cloud, profile):
+                remote_identity = remote_identity[profile]
+                break
+    if remote_identity != 'LOCAL_CREDENTIALS':
         if not cloud.supports_service_account_on_remote():
             raise exceptions.InvalidCloudConfigs(
                 'remote_identity: SERVICE_ACCOUNT is specified in '
@@ -881,6 +888,8 @@ def write_cluster_config(
 
                 # User-supplied instance tags.
                 'instance_tags': instance_tags,
+                # User-supplied remote_identity
+                "remote_identity": remote_identity,
                 # The reservation pools that specified by the user. This is
                 # currently only used by GCP.
                 'specific_reservations': specific_reservations,
diff --git a/sky/clouds/aws.py b/sky/clouds/aws.py
index 1962f16761fa..6008c4d080db 100644
--- a/sky/clouds/aws.py
+++ b/sky/clouds/aws.py
@@ -411,18 +411,6 @@ def make_deploy_resources_variables(self,
         else:
             security_group = DEFAULT_SECURITY_GROUP_NAME
 
-        iam_instance_profile = skypilot_config.get_nested(
-            ('aws', 'iam_instance_profile'), None)
-        if iam_instance_profile is not None and not isinstance(
-                iam_instance_profile, str):
-            for profile in iam_instance_profile:
-                if cluster_name_on_cloud.startswith(
-                        profile) and profile != 'default':
-                    iam_instance_profile = iam_instance_profile[profile]
-                    break
-                elif profile == 'default':
-                    iam_instance_profile = iam_instance_profile[profile]
-
         return {
             'instance_type': r.instance_type,
             'custom_resources': custom_resources,
@@ -430,7 +418,6 @@ def make_deploy_resources_variables(self,
             'region': region_name,
             'zones': ','.join(zone_names),
             'image_id': image_id,
-            'iam_instance_profile': iam_instance_profile,
             'security_group': security_group,
             'security_group_managed_by_skypilot':
                 str(security_group != user_security_group).lower(),
diff --git a/sky/templates/aws-ray.yml.j2 b/sky/templates/aws-ray.yml.j2
index 9fcc0e850c5e..188ccadd3f24 100644
--- a/sky/templates/aws-ray.yml.j2
+++ b/sky/templates/aws-ray.yml.j2
@@ -60,9 +60,9 @@ available_node_types:
   ray.head.default:
     resources: {}
     node_config:
-      {% if iam_instance_profile %}
+      {% if remote_identity not in ['LOCAL_CREDENTIALS', 'SERVICE_ACCOUNT'] %}
       IamInstanceProfile:
-        Name: {{iam_instance_profile}}
+        Name: {{remote_identity}}
       {% endif %}
       InstanceType: {{instance_type}}
       ImageId: {{image_id}}  # Deep Learning AMI (Ubuntu 18.04); see aws.py.
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index d0d1f73d4e41..4cd281bf3ec4 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -477,8 +477,15 @@ def get_cluster_schema():
 
 _REMOTE_IDENTITY_SCHEMA = {
     'remote_identity': {
-        'type': 'string',
-        'case_insensitive_enum': ['LOCAL_CREDENTIALS', 'SERVICE_ACCOUNT'],
+        'oneOf': [{
+            'type': 'string'
+        }, {
+            'type': 'object',
+            'required': [],
+            'additionalProperties': {
+                'type': 'string',
+            },
+        }]
     }
 }
 
@@ -518,23 +525,6 @@ def get_config_schema():
                 'security_group_name': {
                     'type': 'string'
                 },
-                'iam_instance_profile': {
-                    'oneOf': [{
-                        'type': 'string'
-                    }, {
-                        'type': 'object',
-                        'additionalProperties': False,
-                        'required': [],
-                        'properties': {
-                            'sky-serve-controller': {
-                                'type': 'string',
-                            },
-                            'default': {
-                                'type': 'string'
-                            }
-                        }
-                    }]
-                },
                 **_INSTANCE_TAGS_SCHEMA,
                 **_NETWORK_CONFIG_SCHEMA,
             }