skypilot-nightly 1.0.0.dev20240910__py3-none-any.whl → 1.0.0.dev20240911__py3-none-any.whl

sky/__init__.py

@@ -5,7 +5,7 @@ from typing import Optional
 import urllib.request
 
 # Replaced with the current commit when building the wheels.
-_SKYPILOT_COMMIT_SHA = 'db9aeb029d07cf7452751b4658af46e187985a98'
+_SKYPILOT_COMMIT_SHA = 'bad7dabf372f3512679098f63e4cfca8bc9b0870'
 
 
 def _get_git_commit():
@@ -35,7 +35,7 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev20240910'
+__version__ = '1.0.0.dev20240911'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))
 
 

sky/adaptors/kubernetes.py

@@ -1,10 +1,8 @@
 """Kubernetes adaptors"""
-
-# pylint: disable=import-outside-toplevel
-
+import functools
 import logging
 import os
-from typing import Any, Callable, Set
+from typing import Any, Callable, Optional, Set
 
 from sky.adaptors import common
 from sky.sky_logging import set_logging_level
@@ -18,15 +16,6 @@ kubernetes = common.LazyImport('kubernetes',
 urllib3 = common.LazyImport('urllib3',
                             import_error_message=_IMPORT_ERROR_MESSAGE)
 
-_configured = False
-_core_api = None
-_auth_api = None
-_networking_api = None
-_custom_objects_api = None
-_node_api = None
-_apps_api = None
-_api_client = None
-
 # Timeout to use for API calls
 API_TIMEOUT = 5
 
@@ -66,10 +55,7 @@ def _api_logging_decorator(logger: str, level: int):
     return decorated_api
 
 
-def _load_config():
-    global _configured
-    if _configured:
-        return
+def _load_config(context: Optional[str] = None):
     urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
     try:
         # Load in-cluster config if running in a pod
@@ -82,7 +68,7 @@ def _load_config():
         kubernetes.config.load_incluster_config()
     except kubernetes.config.config_exception.ConfigException:
         try:
-            kubernetes.config.load_kube_config()
+            kubernetes.config.load_kube_config(context=context)
         except kubernetes.config.config_exception.ConfigException as e:
             suffix = ''
             if env_options.Options.SHOW_DEBUG_INFO.get():
@@ -101,76 +87,55 @@ def _load_config():
             err_str += '\nTo disable Kubernetes for SkyPilot: run `sky check`.'
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(err_str) from None
-    _configured = True
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-def core_api():
-    global _core_api
-    if _core_api is None:
-        _load_config()
-        _core_api = kubernetes.client.CoreV1Api()
-    return _core_api
+@functools.lru_cache()
+def core_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.CoreV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-def auth_api():
-    global _auth_api
-    if _auth_api is None:
-        _load_config()
-        _auth_api = kubernetes.client.RbacAuthorizationV1Api()
-
-    return _auth_api
+@functools.lru_cache()
+def auth_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.RbacAuthorizationV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-def networking_api():
-    global _networking_api
-    if _networking_api is None:
-        _load_config()
-        _networking_api = kubernetes.client.NetworkingV1Api()
-
-    return _networking_api
+@functools.lru_cache()
+def networking_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.NetworkingV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-def custom_objects_api():
-    global _custom_objects_api
-    if _custom_objects_api is None:
-        _load_config()
-        _custom_objects_api = kubernetes.client.CustomObjectsApi()
-
-    return _custom_objects_api
+@functools.lru_cache()
+def custom_objects_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.CustomObjectsApi()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-def node_api():
-    global _node_api
-    if _node_api is None:
-        _load_config()
-        _node_api = kubernetes.client.NodeV1Api()
-
-    return _node_api
+@functools.lru_cache()
+def node_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.NodeV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-def apps_api():
-    global _apps_api
-    if _apps_api is None:
-        _load_config()
-        _apps_api = kubernetes.client.AppsV1Api()
-
-    return _apps_api
+@functools.lru_cache()
+def apps_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.AppsV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-def api_client():
-    global _api_client
-    if _api_client is None:
-        _load_config()
-        _api_client = kubernetes.client.ApiClient()
-
-    return _api_client
+@functools.lru_cache()
+def api_client(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.ApiClient()
 
 
 def api_exception():

sky/authentication.py

@@ -378,7 +378,11 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     public_key_path = os.path.expanduser(PUBLIC_SSH_KEY_PATH)
     secret_name = clouds.Kubernetes.SKY_SSH_KEY_SECRET_NAME
     secret_field_name = clouds.Kubernetes().ssh_key_secret_field_name
-    namespace = kubernetes_utils.get_current_kube_config_context_namespace()
+    namespace = config['provider'].get(
+        'namespace',
+        kubernetes_utils.get_current_kube_config_context_namespace())
+    context = config['provider'].get(
+        'context', kubernetes_utils.get_current_kube_config_context_name())
     k8s = kubernetes.kubernetes
     with open(public_key_path, 'r', encoding='utf-8') as f:
         public_key = f.read()
@@ -399,14 +403,14 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
         secret = k8s.client.V1Secret(
             metadata=k8s.client.V1ObjectMeta(**secret_metadata),
             string_data={secret_field_name: public_key})
-    if kubernetes_utils.check_secret_exists(secret_name, namespace):
+    if kubernetes_utils.check_secret_exists(secret_name, namespace, context):
         logger.debug(f'Key {secret_name} exists in the cluster, patching it...')
-        kubernetes.core_api().patch_namespaced_secret(secret_name, namespace,
-                                                      secret)
+        kubernetes.core_api(context).patch_namespaced_secret(
+            secret_name, namespace, secret)
     else:
         logger.debug(
             f'Key {secret_name} does not exist in the cluster, creating it...')
-        kubernetes.core_api().create_namespaced_secret(namespace, secret)
+        kubernetes.core_api(context).create_namespaced_secret(namespace, secret)
 
     private_key_path, _ = get_or_generate_keys()
     if network_mode == nodeport_mode:
@@ -415,13 +419,14 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
         # Setup service for SSH jump pod. We create the SSH jump service here
         # because we need to know the service IP address and port to set the
         # ssh_proxy_command in the autoscaler config.
-        kubernetes_utils.setup_ssh_jump_svc(ssh_jump_name, namespace,
+        kubernetes_utils.setup_ssh_jump_svc(ssh_jump_name, namespace, context,
                                             service_type)
         ssh_proxy_cmd = kubernetes_utils.get_ssh_proxy_command(
             ssh_jump_name,
             nodeport_mode,
             private_key_path=private_key_path,
-            namespace=namespace)
+            namespace=namespace,
+            context=context)
     elif network_mode == port_forward_mode:
         # Using `kubectl port-forward` creates a direct tunnel to the pod and
         # does not require a ssh jump pod.

sky/backends/backend_utils.py

@@ -1558,58 +1558,65 @@ def check_owner_identity(cluster_name: str) -> None:
         return
 
     cloud = handle.launched_resources.cloud
-    current_user_identity = cloud.get_current_user_identity()
+    user_identities = cloud.get_user_identities()
     owner_identity = record['owner']
-    if current_user_identity is None:
+    if user_identities is None:
         # Skip the check if the cloud does not support user identity.
         return
     # The user identity can be None, if the cluster is created by an older
     # version of SkyPilot. In that case, we set the user identity to the
-    # current one.
+    # current active one.
     # NOTE: a user who upgrades SkyPilot and switches to a new cloud identity
     # immediately without `sky status --refresh` first, will cause a leakage
     # of the existing cluster. We deem this an acceptable tradeoff mainly
     # because multi-identity is not common (at least at the moment).
     if owner_identity is None:
         global_user_state.set_owner_identity_for_cluster(
-            cluster_name, current_user_identity)
+            cluster_name, user_identities[0])
     else:
         assert isinstance(owner_identity, list)
         # 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).
-        for i, (owner,
-                current) in enumerate(zip(owner_identity,
-                                          current_user_identity)):
-            # Clean up the owner identity for the backslash and newlines, caused
-            # by the cloud CLI output, e.g. gcloud.
-            owner = owner.replace('\n', '').replace('\\', '')
-            if owner == current:
-                if i != 0:
-                    logger.warning(
-                        f'The cluster was owned by {owner_identity}, but '
-                        f'a new identity {current_user_identity} is activated. We still '
-                        'allow the operation as the two identities are likely to have '
-                        'the same access to the cluster. Please be aware that this can '
-                        'cause unexpected cluster leakage if the two identities are not '
-                        'actually equivalent (e.g., belong to the same person).'
-                    )
-                if i != 0 or len(owner_identity) != len(current_user_identity):
-                    # We update the owner of a cluster, when:
-                    # 1. The strictest identty (i.e. the first one) does not
-                    # match, but the latter ones match.
-                    # 2. The length of the two identities are different, which
-                    # will only happen when the cluster is launched before #1808.
-                    # Update the user identity to avoid showing the warning above
-                    # again.
-                    global_user_state.set_owner_identity_for_cluster(
-                        cluster_name, current_user_identity)
-                return  # The user identity matches.
+        for identity in user_identities:
+            for i, (owner, current) in enumerate(zip(owner_identity, identity)):
+                # Clean up the owner identity for the backslash and newlines, caused
+                # by the cloud CLI output, e.g. gcloud.
+                owner = owner.replace('\n', '').replace('\\', '')
+                if owner == current:
+                    if i != 0:
+                        logger.warning(
+                            f'The cluster was owned by {owner_identity}, but '
+                            f'a new identity {identity} is activated. We still '
+                            'allow the operation as the two identities are '
+                            'likely to have the same access to the cluster. '
+                            'Please be aware that this can cause unexpected '
+                            'cluster leakage if the two identities are not '
+                            'actually equivalent (e.g., belong to the same '
+                            'person).')
+                    if i != 0 or len(owner_identity) != len(identity):
+                        # We update the owner of a cluster, when:
+                        # 1. The strictest identty (i.e. the first one) does not
+                        # match, but the latter ones match.
+                        # 2. The length of the two identities are different,
+                        # which will only happen when the cluster is launched
+                        # before #1808. Update the user identity to avoid
+                        # showing the warning above again.
+                        global_user_state.set_owner_identity_for_cluster(
+                            cluster_name, identity)
+                    return  # The user identity matches.
+        # Generate error message if no match found
+        if len(user_identities) == 1:
+            err_msg = f'the activated identity is {user_identities[0]!r}.'
+        else:
+            err_msg = (f'available identities are {user_identities!r}.')
+        if cloud.is_same_cloud(clouds.Kubernetes()):
+            err_msg += (' Check your kubeconfig file and make sure the '
+                        'correct context is available.')
         with ux_utils.print_exception_no_traceback():
             raise exceptions.ClusterOwnerIdentityMismatchError(
                 f'{cluster_name!r} ({cloud}) is owned by account '
-                f'{owner_identity!r}, but the activated account '
-                f'is {current_user_identity!r}.')
+                f'{owner_identity!r}, but ' + err_msg)
 
 
 def tag_filter_for_cluster(cluster_name: str) -> Dict[str, str]:

sky/backends/cloud_vm_ray_backend.py

@@ -1945,7 +1945,7 @@ class RetryingVmProvisioner(object):
                 if dryrun:
                     cloud_user = None
                 else:
-                    cloud_user = to_provision.cloud.get_current_user_identity()
+                    cloud_user = to_provision.cloud.get_active_user_identity()
 
                 requested_features = self._requested_features.copy()
                 # Skip stop feature for Kubernetes and RunPod controllers.

sky/check.py

@@ -44,7 +44,7 @@ def check(
         if ok:
             enabled_clouds.append(cloud_repr)
             if verbose and cloud is not cloudflare:
-                activated_account = cloud.get_current_user_identity_str()
+                activated_account = cloud.get_active_user_identity_str()
                 if activated_account is not None:
                     echo(f'    Activated account: {activated_account}')
             if reason is not None:

sky/clouds/aws.py

@@ -547,7 +547,7 @@ class AWS(clouds.Cloud):
         # Checks if AWS credentials 1) exist and 2) are valid.
         # https://stackoverflow.com/questions/53548737/verify-aws-credentials-with-boto3
         try:
-            identity_str = cls.get_current_user_identity_str()
+            identity_str = cls.get_active_user_identity_str()
         except exceptions.CloudUserIdentityError as e:
             return False, str(e)
 
@@ -584,7 +584,7 @@ class AWS(clouds.Cloud):
         else:
             # This file is required because it is required by the VMs launched on
             # other clouds to access private s3 buckets and resources like EC2.
-            # `get_current_user_identity` does not guarantee this file exists.
+            # `get_active_user_identity` does not guarantee this file exists.
             if not static_credential_exists:
                 return (False, '~/.aws/credentials does not exist. ' +
                         cls._STATIC_CREDENTIAL_HELP_STR)
@@ -648,7 +648,7 @@ class AWS(clouds.Cloud):
             return AWSIdentityType.SHARED_CREDENTIALS_FILE
 
     @classmethod
-    def get_current_user_identity(cls) -> Optional[List[str]]:
+    def get_user_identities(cls) -> Optional[List[List[str]]]:
         """Returns a [UserId, Account] list that uniquely identifies the user.
 
         These fields come from `aws sts get-caller-identity`. We permit the same
@@ -752,11 +752,13 @@ class AWS(clouds.Cloud):
                     f'Failed to get AWS user.\n'
                     f'  Reason: {common_utils.format_exception(e, use_bracket=True)}.'
                 ) from None
-        return user_ids
+        # TODO: Return a list of identities in the profile when we support
+        #   automatic switching for AWS. Currently we only support one identity.
+        return [user_ids]
 
     @classmethod
-    def get_current_user_identity_str(cls) -> Optional[str]:
-        user_identity = cls.get_current_user_identity()
+    def get_active_user_identity_str(cls) -> Optional[str]:
+        user_identity = cls.get_active_user_identity()
         if user_identity is None:
             return None
         identity_str = f'{user_identity[0]} [account={user_identity[1]}]'

sky/clouds/azure.py

@@ -483,7 +483,7 @@ class Azure(clouds.Cloud):
         # If Azure is properly logged in, this will return the account email
         # address + subscription ID.
         try:
-            cls.get_current_user_identity()
+            cls.get_active_user_identity()
         except exceptions.CloudUserIdentityError as e:
             return False, (f'Getting user\'s Azure identity failed.{help_str}\n'
                            f'{cls._INDENT_PREFIX}Details: '
@@ -516,7 +516,7 @@ class Azure(clouds.Cloud):
 
     @classmethod
     @functools.lru_cache(maxsize=1)  # Cache since getting identity is slow.
-    def get_current_user_identity(cls) -> Optional[List[str]]:
+    def get_user_identities(cls) -> Optional[List[List[str]]]:
         """Returns the cloud user identity."""
         # This returns the user's email address + [subscription_id].
         retry_cnt = 0
@@ -558,11 +558,13 @@ class Azure(clouds.Cloud):
             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}]']
+        # TODO: Return a list of identities in the profile when we support
+        #   automatic switching for Az. Currently we only support one identity.
+        return [[f'{account_email} [subscription_id={project_id}]']]
 
     @classmethod
-    def get_current_user_identity_str(cls) -> Optional[str]:
-        user_identity = cls.get_current_user_identity()
+    def get_active_user_identity_str(cls) -> Optional[str]:
+        user_identity = cls.get_active_user_identity()
         if user_identity is None:
             return None
         return user_identity[0]

sky/clouds/cloud.py

@@ -441,11 +441,11 @@ class Cloud:
 
     # TODO(zhwu): Make the return type immutable.
     @classmethod
-    def get_current_user_identity(cls) -> Optional[List[str]]:
-        """(Advanced) Returns currently active user identity of this cloud.
+    def get_user_identities(cls) -> Optional[List[List[str]]]:
+        """(Advanced) Returns all available user identities of this cloud.
 
         The user "identity" is associated with each SkyPilot cluster they
-        creates. This is used in protecting cluster operations, such as
+        create. This is used in protecting cluster operations, such as
         provision, teardown and status refreshing, in a multi-identity
         scenario, where the same user/device can switch between different
         cloud identities. We check that the user identity matches before:
@@ -453,10 +453,16 @@ class Cloud:
             - Stopping/tearing down a cluster
             - Refreshing the status of a cluster
 
-        Design choice: we allow the operations that can correctly work with
-        a different user identity, as a user should have full control over
-        all their clusters (no matter which identity it belongs to), e.g.,
-        submitting jobs, viewing logs, auto-stopping, etc.
+        Design choices:
+          1. We allow the operations that can correctly work with a different
+             user identity, as a user should have full control over all their
+             clusters (no matter which identity it belongs to), e.g.,
+             submitting jobs, viewing logs, auto-stopping, etc.
+          2. A cloud implementation can optionally switch between different
+             identities if required for cluster operations. In this case,
+             the cloud implementation should return multiple identities
+             as a list. E.g., our Kubernetes implementation can use multiple
+             kubeconfig contexts to switch between different identities.
 
         The choice of what constitutes an identity is up to each cloud's
         implementation. In general, to suffice for the above purposes,
@@ -464,24 +470,34 @@ class Cloud:
         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
+        An 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.
         When performing an identity check between the current active identity
         and the owner identity associated with a cluster, we compare the two
         lists in order: if a position does not match, we go to the next. To
-        see an example, see the docstring of the AWS.get_current_user_identity.
-
+        see an example, see the docstring of the AWS.get_user_identities.
 
         Example identities (see cloud implementations):
             - AWS: [UserId, AccountId]
             - GCP: [email address + project ID]
             - Azure: [email address + subscription ID]
+            - Kubernetes: [context name]
+
+        Example return values:
+            - AWS: [[UserId, AccountId]]
+            - GCP: [[email address + project ID]]
+            - Azure: [[email address + subscription ID]]
+            - Kubernetes: [[current active context], [context 2], ...]
 
         Returns:
             None if the cloud does not have a concept of user identity
             (access protection will be disabled for these clusters);
-            otherwise the currently active user identity.
+            otherwise a list of available identities with the current active
+            identity being the first element. Most clouds have only one identity
+            available, so the returned list will only have one element: the
+            current active identity.
+
         Raises:
             exceptions.CloudUserIdentityError: If the user identity cannot be
                 retrieved.
@@ -489,13 +505,26 @@ class Cloud:
         return None
 
     @classmethod
-    def get_current_user_identity_str(cls) -> Optional[str]:
-        """Returns a user friendly representation of the current identity."""
-        user_identity = cls.get_current_user_identity()
+    def get_active_user_identity_str(cls) -> Optional[str]:
+        """Returns a user friendly representation of the active identity."""
+        user_identity = cls.get_active_user_identity()
         if user_identity is None:
             return None
         return ', '.join(user_identity)
 
+    @classmethod
+    def get_active_user_identity(cls) -> Optional[List[str]]:
+        """Returns currently active user identity of this cloud
+
+        See get_user_identities for definition of user identity.
+
+        Returns:
+            None if the cloud does not have a concept of user identity;
+            otherwise the current active identity.
+        """
+        identities = cls.get_user_identities()
+        return identities[0] if identities is not None else None
+
     def get_credential_file_mounts(self) -> Dict[str, str]:
         """Returns the files necessary to access this cloud.
 

sky/clouds/cudo.py

@@ -328,7 +328,7 @@ class Cudo(clouds.Cloud):
         }
 
     @classmethod
-    def get_current_user_identity(cls) -> Optional[List[str]]:
+    def get_user_identities(cls) -> Optional[List[List[str]]]:
         # NOTE: used for very advanced SkyPilot functionality
         # Can implement later if desired
         return None

sky/clouds/fluidstack.py

@@ -291,8 +291,8 @@ class Fluidstack(clouds.Cloud):
         return {filename: filename for filename in _CREDENTIAL_FILES}
 
     @classmethod
-    def get_current_user_identity(cls) -> Optional[List[str]]:
-        # TODO(mjibril): Implement get_current_user_identity for Fluidstack
+    def get_user_identities(cls) -> Optional[List[List[str]]]:
+        # TODO(mjibril): Implement get_active_user_identity for Fluidstack
         return None
 
     def instance_type_exists(self, instance_type: str) -> bool:

sky/clouds/gcp.py

@@ -715,7 +715,7 @@ class GCP(clouds.Cloud):
             project_id = cls.get_project_id()
 
             # Check if the user is activated.
-            identity = cls.get_current_user_identity()
+            identity = cls.get_active_user_identity()
         except (auth.exceptions.DefaultCredentialsError,
                 exceptions.CloudUserIdentityError) as e:
             # See also: https://stackoverflow.com/a/53307505/1165051
@@ -826,16 +826,19 @@ class GCP(clouds.Cloud):
     @classmethod
     def _get_identity_type(cls) -> Optional[GCPIdentityType]:
         try:
-            account = cls.get_current_user_identity()[0]
+            account = cls.get_active_user_identity()
         except exceptions.CloudUserIdentityError:
             return None
-        if GCPIdentityType.SERVICE_ACCOUNT.value in account:
+        if account is None:
+            return None
+        assert account is not None
+        if GCPIdentityType.SERVICE_ACCOUNT.value in account[0]:
             return GCPIdentityType.SERVICE_ACCOUNT
         return GCPIdentityType.SHARED_CREDENTIALS_FILE
 
     @classmethod
     @functools.lru_cache(maxsize=1)  # Cache since getting identity is slow.
-    def get_current_user_identity(cls) -> List[str]:
+    def get_user_identities(cls) -> List[List[str]]:
         """Returns the email address + project id of the active user."""
         try:
             account = _run_output('gcloud auth list --filter=status:ACTIVE '
@@ -866,11 +869,13 @@ class GCP(clouds.Cloud):
                     '  Reason: '
                     f'{common_utils.format_exception(e, use_bracket=True)}'
                 ) from e
-        return [f'{account} [project_id={project_id}]']
+        # TODO: Return a list of identities in the profile when we support
+        #   automatic switching for GCP. Currently we only support one identity.
+        return [[f'{account} [project_id={project_id}]']]
 
     @classmethod
-    def get_current_user_identity_str(cls) -> Optional[str]:
-        user_identity = cls.get_current_user_identity()
+    def get_active_user_identity_str(cls) -> Optional[str]:
+        user_identity = cls.get_active_user_identity()
         if user_identity is None:
             return None
         return user_identity[0].replace('\n', '')