Repository: skypilot-org/skypilot
Pull request: Allow fall back to account ID for AWS
<!-- Describe the changes in this PR -->

Fixes #1807

We now allow the identity check fallback to AccountId if UserId does not match, because:
1.. In the case where multiple IAM users belong to a single root account, those users normally share the visibility of the VMs, so we do not need to identify them with each other. (There can be some cases, when an IAM user is given a limited permission by the admin, we may ignore that case for now, or print out a warning if the underlying userid changed for a cluster).
2. In the case where multiple users belong to an organization, those users will have different account id, so fallback works.


<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Backward compatibility
    - [x] `sky launch -c old-machine --cloud aws -t t2.micro` on master; `sky launch -c old-machine` on this PR
    - [x] `sky launch -c old-machine --cloud aws -t t2.medium` on master; `sky exec old-machine echo hi` on this PR
    - [x] `sky launch -c old-machine --cloud aws -t t2.medium echo hi` on master; `sky logs old-machine` on this PR
    - [x] `sky launch -c old-machine --cloud aws -t t2.medium echo hi` on master; `sky down old-machine` on this PR
    - [x] `sky launch -c old-machine --cloud aws -t t2.medium echo hi` on master; switch to another IAM role and `sky status -r` on this PR
- [x] All smoke tests: `pytest tests/test_smoke.py --aws` test with an AWS IAM role 


Problem statement:
[AWS] Use Account ID instead of user id for identity check
 (we can make account id as the fallback option, which means we check the user id first and then check account id if the first one does not match and print warnings), reason
In the case where multiple IAM users belong a single root account, those users normally share the visibility of the VMs, so we do not need to identity them with each other, i.e. sky status -r will be executed correctly. (There can be some cases, when an IAM user is given a limited permission for the instances it can see by the admin, we may ignore that case for now, or print out a warning if the underlying userid changed for a cluster).
In the case where the multiple users belong to an organization, those users will have different account id, so the proposed method work

Hints from issue discussion:

Commit message: Fix user identity

Patch to review:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 513615f7e45..45a7c00142a 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -1744,7 +1744,11 @@ def check_owner_identity(cluster_name: str) -> None:
     if owner_identity is None:
         global_user_state.set_owner_identity_for_cluster(
             cluster_name, current_user_identity)
-    elif owner_identity != current_user_identity:
+    else:
+        assert isinstance(owner_identity, list)
+        for owner, current in zip(owner_identity, current_user_identity):
+            if owner == current:
+                return  # The user identity matches.
         with ux_utils.print_exception_no_traceback():
             raise exceptions.ClusterOwnerIdentityMismatchError(
                 f'{cluster_name!r} ({cloud}) is owned by account '
diff --git a/sky/clouds/aws.py b/sky/clouds/aws.py
index b29d1380ee5..35634458e29 100644
--- a/sky/clouds/aws.py
+++ b/sky/clouds/aws.py
@@ -444,8 +444,14 @@ def _is_current_identity_sso(cls) -> bool:
         return 'sso' in proc.stdout.decode().split()
 
     @classmethod
-    def get_current_user_identity(cls) -> Optional[str]:
-        """Returns the identity of the user on this cloud.
+    def get_current_user_identity(cls) -> Optional[List[str]]:
+        """Returns a list of identities of the user on this cloud.
+
+        Returns:
+            A list of strings that uniquely identifies the user on this cloud.
+            For identity check, we will fallback through the list of strings
+            until we find a match, and print a warning if we fail for the
+            first string.
 
         Raises:
             exceptions.CloudUserIdentityError: if the user identity cannot be
@@ -453,16 +459,26 @@ def get_current_user_identity(cls) -> Optional[str]:
         """
         try:
             sts = aws.client('sts')
-            # The caller identity contains 3 fields: UserId, AccountId, Arn.
+            # The caller identity contains 3 fields: UserId, Account, Arn.
             # 'UserId' is unique across all AWS entity, which looks like
             # "AROADBQP57FF2AEXAMPLE:role-session-name"
-            # 'AccountId' can be shared by multiple users under the same
+            # 'Account' can be shared by multiple users under the same
             # organization
             # 'Arn' is the full path to the user, which can be reused when
             # the user is deleted and recreated.
             # Refer to https://docs.aws.amazon.com/cli/latest/reference/sts/get-caller-identity.html # pylint: disable=line-too-long
             # and https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_variables.html#principaltable # pylint: disable=line-too-long
-            user_id = sts.get_caller_identity()['UserId']
+            user_info = sts.get_caller_identity()
+            # Allow fallback to AccountId if UserId does not match, because:
+            # 1. In the case where multiple IAM users belong a single root account,
+            # those users normally share the visibility of the VMs, so we do not
+            # need to identity them with each other. (There can be some cases,
+            # when an IAM user is given a limited permission by the admin, we may
+            # ignore that case for now, or print out a warning if the underlying
+            # userid changed for a cluster).
+            # 2. In the case where the multiple users belong to an organization,
+            # those users will have different account id, so fallback works.
+            user_ids = [user_info['UserId'], user_info['Account']]
         except aws.botocore_exceptions().NoCredentialsError:
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.CloudUserIdentityError(
@@ -507,7 +523,7 @@ def get_current_user_identity(cls) -> Optional[str]:
                     f'Failed to get AWS user.\n'
                     f'  Reason: {common_utils.format_exception(e, use_bracket=True)}.'
                 ) from None
-        return user_id
+        return user_ids
 
     def get_credential_file_mounts(self) -> Dict[str, str]:
         # TODO(skypilot): ~/.aws/credentials is required for users using multiple clouds.
diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py
index e2e69768ee3..827711af6a2 100644
--- a/sky/clouds/azure.py
+++ b/sky/clouds/azure.py
@@ -348,7 +348,7 @@ def accelerator_in_region_or_zone(self,
 
     @classmethod
     @functools.lru_cache(maxsize=1)  # Cache since getting identity is slow.
-    def get_current_user_identity(cls) -> Optional[str]:
+    def get_current_user_identity(cls) -> Optional[List[str]]:
         """Returns the cloud user identity."""
         # This returns the user's email address + [subscription_id].
         retry_cnt = 0
@@ -390,7 +390,7 @@ def get_current_user_identity(cls) -> Optional[str]:
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.CloudUserIdentityError(
                     'Failed to get Azure project ID.') from e
-        return f'{account_email} [subscription_id={project_id}]'
+        return [f'{account_email} [subscription_id={project_id}]']
 
     @classmethod
     def get_project_id(cls, dryrun: bool = False) -> str:
diff --git a/sky/clouds/cloud.py b/sky/clouds/cloud.py
index 54ae403dbed..c7abe5b27d8 100644
--- a/sky/clouds/cloud.py
+++ b/sky/clouds/cloud.py
@@ -296,7 +296,7 @@ def check_credentials(cls) -> Tuple[bool, Optional[str]]:
         raise NotImplementedError
 
     @classmethod
-    def get_current_user_identity(cls) -> Optional[str]:
+    def get_current_user_identity(cls) -> Optional[List[str]]:
         """(Advanced) Returns currently active user identity of this cloud.
 
         The user "identity" is associated with each SkyPilot cluster they
@@ -319,6 +319,15 @@ def get_current_user_identity(cls) -> Optional[str]:
         resources are used when the user invoked each cloud's default
         CLI/API.
 
+        The returned identity is a list of strings. The list is in the
+        order of strictness, i.e., the first element is the most strict
+        identity, and the last element is the least strict identity.
+        This is useful for the IAM roles under the same root account
+        share the same set of permission. In that case, we assume
+        although the user can switch between different IAM roles, the
+        resources they can access are the same. Therefore, we can ignore
+        the difference between the IAM roles and only use the root account.
+
         Example identities (see cloud implementations):
             - AWS: unique aws:user_id
             - GCP: email address + project ID
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index 87bea4b3dc4..a12347662d8 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -571,9 +571,10 @@ def check_credentials(cls) -> Tuple[bool, Optional[str]]:
         if len(ret_permissions) < len(GCP_PREMISSION_CHECK_LIST):
             diffs = set(GCP_PREMISSION_CHECK_LIST).difference(
                 set(ret_permissions))
+            identity_str = identity[0] if identity else None
             return False, (
                 'The following permissions are not enabled for the current '
-                f'GCP identity ({identity}):\n    '
+                f'GCP identity ({identity_str}):\n    '
                 f'{diffs}\n    '
                 'For more details, visit: https://skypilot.readthedocs.io/en/latest/reference/faq.html#what-are-the-required-iam-permissons-on-gcp-for-skypilot')  # pylint: disable=line-too-long
         return True, None
@@ -606,7 +607,7 @@ def get_credential_file_mounts(self) -> Dict[str, str]:
 
     @classmethod
     @functools.lru_cache(maxsize=1)  # Cache since getting identity is slow.
-    def get_current_user_identity(cls) -> Optional[str]:
+    def get_current_user_identity(cls) -> Optional[List[str]]:
         """Returns the email address + project id of the active user."""
         try:
             account = _run_output('gcloud auth list --filter=status:ACTIVE '
@@ -627,7 +628,7 @@ def get_current_user_identity(cls) -> Optional[str]:
                     '--format="value(account)"` and ensure it correctly '
                     'returns the current user.')
         try:
-            return f'{account} [project_id={cls.get_project_id()}]'
+            project_id = cls.get_project_id()
         except Exception as e:  # pylint: disable=broad-except
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.CloudUserIdentityError(
@@ -636,6 +637,7 @@ def get_current_user_identity(cls) -> Optional[str]:
                     '  Reason: '
                     f'{common_utils.format_exception(e, use_bracket=True)}'
                 ) from e
+        return [f'{account} [project_id={project_id}]']
 
     def instance_type_exists(self, instance_type):
         return service_catalog.instance_type_exists(instance_type, 'gcp')
diff --git a/sky/clouds/lambda_cloud.py b/sky/clouds/lambda_cloud.py
index 14531c50cc5..ea8a8ab442c 100644
--- a/sky/clouds/lambda_cloud.py
+++ b/sky/clouds/lambda_cloud.py
@@ -229,7 +229,7 @@ def get_credential_file_mounts(self) -> Dict[str, str]:
         }
 
     @classmethod
-    def get_current_user_identity(cls) -> Optional[str]:
+    def get_current_user_identity(cls) -> Optional[List[str]]:
         # TODO(ewzeng): Implement get_current_user_identity for Lambda
         return None
 
diff --git a/sky/clouds/service_catalog/aws_catalog.py b/sky/clouds/service_catalog/aws_catalog.py
index 0639a584a7f..2d6082b600f 100644
--- a/sky/clouds/service_catalog/aws_catalog.py
+++ b/sky/clouds/service_catalog/aws_catalog.py
@@ -87,8 +87,9 @@ def _fetch_and_apply_az_mapping(df: pd.DataFrame) -> pd.DataFrame:
         with the zone name (e.g. us-east-1a).
     """
     try:
-        user_identity = aws.AWS.get_current_user_identity()
-        assert user_identity is not None, 'user_identity is None'
+        user_identity_list = aws.AWS.get_current_user_identity()
+        assert user_identity_list, user_identity_list
+        user_identity = user_identity_list[0]
         aws_user_hash = hashlib.md5(user_identity.encode()).hexdigest()[:8]
     except exceptions.CloudUserIdentityError:
         glob_name = common.get_catalog_path('aws/az_mappings-*.csv')
diff --git a/sky/global_user_state.py b/sky/global_user_state.py
index 1f848b11010..7ebb40e02a8 100644
--- a/sky/global_user_state.py
+++ b/sky/global_user_state.py
@@ -468,11 +468,12 @@ def _set_cluster_usage_intervals(cluster_hash: str,
 
 
 def set_owner_identity_for_cluster(cluster_name: str,
-                                   owner_identity: Optional[str]) -> None:
+                                   owner_identity: Optional[List[str]]) -> None:
     if owner_identity is None:
         return
+    owner_identity_str = json.dumps(owner_identity)
     _DB.cursor.execute('UPDATE clusters SET owner=(?) WHERE name=(?)',
-                       (owner_identity, cluster_name))
+                       (owner_identity_str, cluster_name))
 
     count = _DB.cursor.rowcount
     _DB.conn.commit()
@@ -503,6 +504,24 @@ def get_launched_resources_from_cluster_hash(
         return num_nodes, launched_resources
 
 
+def _load_owner(record_owner: Optional[str]) -> Optional[List[str]]:
+    if record_owner is None:
+        return None
+    try:
+        result = json.loads(record_owner)
+        if result is not None and not isinstance(result, list):
+            # Backwards compatibility for old records, which were stored as
+            # a string instead of a list. It is possible that json.loads
+            # will parse the string with all numbers as an int, so we need
+            # to convert it back to a list of strings.
+            return [str(result)]
+        return result
+    except json.JSONDecodeError:
+        # Backwards compatibility for old records, which were stored as
+        # a string instead of a list.
+        return [record_owner]
+
+
 def get_cluster_from_name(
         cluster_name: Optional[str]) -> Optional[Dict[str, Any]]:
     rows = _DB.cursor.execute('SELECT * FROM clusters WHERE name=(?)',
@@ -522,7 +541,7 @@ def get_cluster_from_name(
             'status': ClusterStatus[status],
             'autostop': autostop,
             'to_down': bool(to_down),
-            'owner': owner,
+            'owner': _load_owner(owner),
             'metadata': json.loads(metadata),
             'cluster_hash': cluster_hash,
         }
@@ -546,7 +565,7 @@ def get_clusters() -> List[Dict[str, Any]]:
             'status': ClusterStatus[status],
             'autostop': autostop,
             'to_down': bool(to_down),
-            'owner': owner,
+            'owner': _load_owner(owner),
             'metadata': json.loads(metadata),
             'cluster_hash': cluster_hash,
         }


Code review comments:
diff_hunk: @@ -1744,7 +1744,11 @@ def check_owner_identity(cluster_name: str) -> None:
     if owner_identity is None:
         global_user_state.set_owner_identity_for_cluster(
             cluster_name, current_user_identity)
-    elif owner_identity != current_user_identity:
+    else:
+        assert isinstance(owner_identity, list)
+        for owner, current in zip(owner_identity, current_user_identity):
+            if owner == current:
+                return  # The user identity matches.
path: sky/backends/backend_utils.py
text: @user1:
Do we want to assert current_user_identity is a list, and/or has the same length?

What happens if lengths are different?

@author:
The `current_user_identity` should be list guaranteed by the type annotation.

Added a comment for the case where the lengths are different. Wdyt?
```
# It is OK if the owner identity is shorter, which will happen when
# the cluster is launched before #1808. In that case, we only check
# the same length (zip will stop at the shorter one).
```

diff_hunk: @@ -319,6 +319,15 @@ def get_current_user_identity(cls) -> Optional[str]:
         resources are used when the user invoked each cloud's default
         CLI/API.
 
+        The returned identity is a list of strings. The list is in the
+        order of strictness, i.e., the first element is the most strict
+        identity, and the last element is the least strict identity.
+        This is useful for the IAM roles under the same root account
+        share the same set of permission. In that case, we assume
+        although the user can switch between different IAM roles, the
+        resources they can access are the same. Therefore, we can ignore
+        the difference between the IAM roles and only use the root account.
+
         Example identities (see cloud implementations):
             - AWS: unique aws:user_id
path: sky/clouds/cloud.py
text: @user1:
update these 3 lines

@author:
Done. Thanks!