Add user identity to cluster status to avoid leakage when switching account

.github/workflows/pytest.yml

@@ -53,4 +53,4 @@ jobs:
           pip install pytest pytest-xdist pytest-env>=0.6
 
       - name: Run tests with pytest
-        run: SKYPILOT_DISABLE_USAGE_COLLECTION=1 pytest ${{ matrix.test-path }}
+        run: SKYPILOT_DISABLE_USAGE_COLLECTION=1 SKYPILOT_SKIP_CLOUD_IDENTITY_CHECK=1 pytest ${{ matrix.test-path }}

sky/adaptors/aws.py

@@ -2,15 +2,15 @@
 
 # pylint: disable=import-outside-toplevel
 
-from functools import wraps
+import functools
 
 boto3 = None
 botocore = None
 
 
 def import_package(func):
 
-    @wraps(func)
+    @functools.wraps(func)
     def wrapper(*args, **kwargs):
         global boto3, botocore
         if boto3 is None or botocore is None:
@@ -27,17 +27,16 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+@functools.lru_cache()
 @import_package
-def client(service_name: str, **kwargs):
-    """Create an AWS client of a certain service.
-
-    Args:
-        service_name: AWS service name (e.g., 's3', 'ec2').
-        kwargs: Other options.
-    """
-    return boto3.client(service_name, **kwargs)
+def session():
+    """Create an AWS session."""
+    # functools.lru_cache() is used to cache the session object
+    # for each thread.
+    return boto3.session.Session()
 
 
+@functools.lru_cache()
 @import_package
 def resource(resource_name: str, **kwargs):
     """Create an AWS resource.
@@ -46,13 +45,26 @@ def resource(resource_name: str, **kwargs):
         resource_name: AWS resource name (e.g., 's3').
         kwargs: Other options.
     """
-    return boto3.resource(resource_name, **kwargs)
+    # Need to use the resource retrieved from the per-thread session
+    # to avoid thread-safety issues (Directly creating the client
+    # with boto3.resource() is not thread-safe).
+    # Reference: https://stackoverflow.com/a/59635814
+    return session().resource(resource_name, **kwargs)
 
 
-@import_package
-def session():
-    """Create an AWS session."""
-    return boto3.Session()
+@functools.lru_cache()
+def client(service_name: str, **kwargs):
+    """Create an AWS client of a certain service.
+
+    Args:
+        service_name: AWS service name (e.g., 's3', 'ec2').
+        kwargs: Other options.
+    """
+    # Need to use the client retrieved from the per-thread session
+    # to avoid thread-safety issues (Directly creating the client
+    # with boto3.client() is not thread-safe).
+    # Reference: https://stackoverflow.com/a/59635814
+    return session().client(service_name, **kwargs)
 
 
 @import_package

sky/adaptors/azure.py

@@ -28,3 +28,10 @@ def get_subscription_id() -> str:
     """Get the default subscription id."""
     from azure.common import credentials
     return credentials.get_cli_profile().get_subscription_id()
+
+
+@import_package
+def get_current_account_user() -> str:
+    """Get the default account user."""
+    from azure.common import credentials
+    return credentials.get_cli_profile().get_current_account_user()