skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/authentication.py

@@ -12,14 +12,12 @@ in ray yaml config as input,
 2. Setup the `authorized_keys` on the remote VM with the public key content,
    by cloud-init or directly using cloud provider's API.
 
-The local machine's public key should not be uploaded to the
-`~/.ssh/sky-key.pub` on the remote VM, because it will cause private/public
-key pair mismatch when the user tries to launch new VM from that remote VM
-using SkyPilot, e.g., the node is used as a jobs controller. (Lambda cloud
-is an exception, due to the limitation of the cloud provider. See the
-comments in setup_lambda_authentication)
+The local machine's public key should not be uploaded to the remote VM, because
+it will cause private/public key pair mismatch when the user tries to launch new
+VM from that remote VM using SkyPilot, e.g., the node is used as a jobs
+controller. (Lambda cloud is an exception, due to the limitation of the cloud
+provider. See the comments in setup_lambda_authentication)
 """
-import base64
 import copy
 import functools
 import os
@@ -44,10 +42,12 @@ from sky.adaptors import gcp
 from sky.adaptors import ibm
 from sky.adaptors import kubernetes
 from sky.adaptors import runpod
-from sky.clouds.utils import lambda_utils
+from sky.adaptors import vast
 from sky.provision.fluidstack import fluidstack_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.provision.lambda_cloud import lambda_utils
 from sky.utils import common_utils
+from sky.utils import config_utils
 from sky.utils import kubernetes_enums
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
@@ -61,9 +61,24 @@ logger = sky_logging.init_logger(__name__)
 
 MAX_TRIALS = 64
 # TODO(zhwu): Support user specified key pair.
-PRIVATE_SSH_KEY_PATH = '~/.ssh/sky-key'
-PUBLIC_SSH_KEY_PATH = '~/.ssh/sky-key.pub'
-_SSH_KEY_GENERATION_LOCK = '~/.sky/generated/ssh/.__internal-sky-key.lock'
+# We intentionally not have the ssh key pair to be stored in
+# ~/.sky/api_server/clients, i.e. sky.server.common.API_SERVER_CLIENT_DIR,
+# because ssh key pair need to persist across API server restarts, while
+# the former dir is empheral.
+_SSH_KEY_PATH_PREFIX = '~/.sky/clients/{user_hash}/ssh'
+
+
+def get_ssh_key_and_lock_path() -> Tuple[str, str, str]:
+    user_hash = common_utils.get_user_hash()
+    user_ssh_key_prefix = _SSH_KEY_PATH_PREFIX.format(user_hash=user_hash)
+
+    os.makedirs(os.path.expanduser(user_ssh_key_prefix),
+                exist_ok=True,
+                mode=0o700)
+    private_key_path = os.path.join(user_ssh_key_prefix, 'sky-key')
+    public_key_path = os.path.join(user_ssh_key_prefix, 'sky-key.pub')
+    lock_path = os.path.join(user_ssh_key_prefix, '.__internal-sky-key.lock')
+    return private_key_path, public_key_path, lock_path
 
 
 def _generate_rsa_key_pair() -> Tuple[str, str]:
@@ -106,16 +121,17 @@ def _save_key_pair(private_key_path: str, public_key_path: str,
 
 def get_or_generate_keys() -> Tuple[str, str]:
     """Returns the aboslute private and public key paths."""
-    private_key_path = os.path.expanduser(PRIVATE_SSH_KEY_PATH)
-    public_key_path = os.path.expanduser(PUBLIC_SSH_KEY_PATH)
+    private_key_path, public_key_path, lock_path = get_ssh_key_and_lock_path()
+    private_key_path = os.path.expanduser(private_key_path)
+    public_key_path = os.path.expanduser(public_key_path)
+    lock_path = os.path.expanduser(lock_path)
 
-    key_file_lock = os.path.expanduser(_SSH_KEY_GENERATION_LOCK)
-    lock_dir = os.path.dirname(key_file_lock)
+    lock_dir = os.path.dirname(lock_path)
     # We should have the folder ~/.sky/generated/ssh to have 0o700 permission,
     # as the ssh configs will be written to this folder as well in
     # backend_utils.SSHConfigHelper
     os.makedirs(lock_dir, exist_ok=True, mode=0o700)
-    with filelock.FileLock(key_file_lock, timeout=10):
+    with filelock.FileLock(lock_path, timeout=10):
         if not os.path.exists(private_key_path):
             public_key, private_key = _generate_rsa_key_pair()
             _save_key_pair(private_key_path, public_key_path, private_key,
@@ -270,43 +286,13 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     return configure_ssh_info(config)
 
 
-# In Azure, cloud-init script must be encoded in base64. See
-# https://learn.microsoft.com/en-us/azure/virtual-machines/custom-data
-# for more information. Here we decode it and replace the ssh user
-# and public key content, then encode it back.
-def setup_azure_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
-    _, public_key_path = get_or_generate_keys()
-    with open(public_key_path, 'r', encoding='utf-8') as f:
-        public_key = f.read().strip()
-    for node_type in config['available_node_types']:
-        node_config = config['available_node_types'][node_type]['node_config']
-        cloud_init = (
-            node_config['azure_arm_parameters']['cloudInitSetupCommands'])
-        cloud_init = base64.b64decode(cloud_init).decode('utf-8')
-        cloud_init = cloud_init.replace('skypilot:ssh_user',
-                                        config['auth']['ssh_user'])
-        cloud_init = cloud_init.replace('skypilot:ssh_public_key_content',
-                                        public_key)
-        cloud_init = base64.b64encode(
-            cloud_init.encode('utf-8')).decode('utf-8')
-        node_config['azure_arm_parameters']['cloudInitSetupCommands'] = (
-            cloud_init)
-    config_str = common_utils.dump_yaml_str(config)
-    config_str = config_str.replace('skypilot:ssh_user',
-                                    config['auth']['ssh_user'])
-    config_str = config_str.replace('skypilot:ssh_public_key_content',
-                                    public_key)
-    config = yaml.safe_load(config_str)
-    return config
-
-
 def setup_lambda_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
     get_or_generate_keys()
 
     # Ensure ssh key is registered with Lambda Cloud
     lambda_client = lambda_utils.LambdaCloudClient()
-    public_key_path = os.path.expanduser(PUBLIC_SSH_KEY_PATH)
+    _, public_key_path = get_or_generate_keys()
     with open(public_key_path, 'r', encoding='utf-8') as f:
         public_key = f.read().strip()
     prefix = f'sky-key-{common_utils.get_user_hash()}'
@@ -314,26 +300,16 @@ def setup_lambda_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     if not exists:
         lambda_client.register_ssh_key(name, public_key)
 
-    # Need to use ~ relative path because Ray uses the same
-    # path for finding the public key path on both local and head node.
-    config['auth']['ssh_public_key'] = PUBLIC_SSH_KEY_PATH
-
-    # TODO(zhwu): we need to avoid uploading the public ssh key to the
-    # nodes, as that will cause problem when the node is used as spot
-    # controller, i.e., the public and private key on the node may
-    # not match.
-    file_mounts = config['file_mounts']
-    file_mounts[PUBLIC_SSH_KEY_PATH] = PUBLIC_SSH_KEY_PATH
-    config['file_mounts'] = file_mounts
-
+    config['auth']['remote_key_name'] = name
     return config
 
 
-def setup_ibm_authentication(config):
+def setup_ibm_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     """ registers keys if they do not exist in sky folder
     and updates config file.
     keys default location: '~/.ssh/sky-key' and '~/.ssh/sky-key.pub'
     """
+    private_key_path, _ = get_or_generate_keys()
 
     def _get_unique_key_name():
         suffix_len = 10
@@ -373,17 +349,11 @@ def setup_ibm_authentication(config):
         else:
             raise Exception('Failed to register a key') from e
 
-    config['auth']['ssh_private_key'] = PRIVATE_SSH_KEY_PATH
+    config['auth']['ssh_private_key'] = private_key_path
 
     for node_type in config['available_node_types']:
         config['available_node_types'][node_type]['node_config'][
             'key_id'] = vpc_key_id
-
-    # Add public key path to file mounts
-    file_mounts = config['file_mounts']
-    file_mounts[PUBLIC_SSH_KEY_PATH] = PUBLIC_SSH_KEY_PATH
-    config['file_mounts'] = file_mounts
-
     return config
 
 
@@ -403,13 +373,19 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
         with ux_utils.print_exception_no_traceback():
             raise ValueError(str(e) + ' Please check: ~/.sky/config.yaml.') \
                 from None
-    get_or_generate_keys()
+    _, public_key_path = get_or_generate_keys()
 
     # Add the user's public key to the SkyPilot cluster.
-    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()
+    context = config['provider'].get(
+        'context', kubernetes_utils.get_current_kube_config_context_name())
+    if context == kubernetes.in_cluster_context_name():
+        # If the context is an in-cluster context name, we are running in a pod
+        # with in-cluster configuration. We need to set the context to None
+        # to use the mounted service account.
+        context = None
+    namespace = kubernetes_utils.get_namespace_from_config(config['provider'])
     k8s = kubernetes.kubernetes
     with open(public_key_path, 'r', encoding='utf-8') as f:
         public_key = f.read()
@@ -425,44 +401,71 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
         }
         custom_metadata = skypilot_config.get_nested(
             ('kubernetes', 'custom_metadata'), {})
-        kubernetes_utils.merge_dicts(custom_metadata, secret_metadata)
+        config_utils.merge_k8s_configs(secret_metadata, custom_metadata)
 
         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):
-        logger.debug(f'Key {secret_name} exists in the cluster, patching it...')
-        kubernetes.core_api().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)
+    try:
+        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(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(context).create_namespaced_secret(
+                namespace, secret)
+    except kubernetes.api_exception() as e:
+        if e.status == 409 and e.reason == 'AlreadyExists':
+            logger.debug(f'Key {secret_name} was created concurrently, '
+                         'patching it...')
+            kubernetes.core_api(context).patch_namespaced_secret(
+                secret_name, namespace, secret)
+        else:
+            raise e
 
-    ssh_jump_name = clouds.Kubernetes.SKY_SSH_JUMP_NAME
+    private_key_path, _ = get_or_generate_keys()
     if network_mode == nodeport_mode:
+        ssh_jump_name = clouds.Kubernetes.SKY_SSH_JUMP_NAME
         service_type = kubernetes_enums.KubernetesServiceType.NODEPORT
+        # 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, context,
+                                            service_type)
+        ssh_proxy_cmd = kubernetes_utils.get_ssh_proxy_command(
+            ssh_jump_name,
+            nodeport_mode,
+            private_key_path=private_key_path,
+            context=context,
+            namespace=namespace)
     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.
         kubernetes_utils.check_port_forward_mode_dependencies()
-        # Using `kubectl port-forward` creates a direct tunnel to jump pod and
-        # does not require opening any ports on Kubernetes nodes. As a result,
-        # the service can be a simple ClusterIP service which we access with
-        # `kubectl port-forward`.
-        service_type = kubernetes_enums.KubernetesServiceType.CLUSTERIP
+        # TODO(romilb): This can be further optimized. Instead of using the
+        #   head node as a jump pod for worker nodes, we can also directly
+        #   set the ssh_target to the worker node. However, that requires
+        #   changes in the downstream code to return a mapping of node IPs to
+        #   pod names (to be used as ssh_target) and updating the upstream
+        #   SSHConfigHelper to use a different ProxyCommand for each pod.
+        #   This optimization can reduce SSH time from ~0.35s to ~0.25s, tested
+        #   on GKE.
+        ssh_target = config['cluster_name'] + '-head'
+        ssh_proxy_cmd = kubernetes_utils.get_ssh_proxy_command(
+            ssh_target,
+            port_forward_mode,
+            private_key_path=private_key_path,
+            context=context,
+            namespace=namespace)
     else:
         # This should never happen because we check for this in from_str above.
         raise ValueError(f'Unsupported networking mode: {network_mode_str}')
-    # 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, service_type)
-
-    ssh_proxy_cmd = kubernetes_utils.get_ssh_proxy_command(
-        PRIVATE_SSH_KEY_PATH, ssh_jump_name, network_mode, namespace,
-        clouds.Kubernetes.PORT_FORWARD_PROXY_CMD_PATH,
-        clouds.Kubernetes.PORT_FORWARD_PROXY_CMD_TEMPLATE)
-
     config['auth']['ssh_proxy_command'] = ssh_proxy_cmd
+    config['auth']['ssh_private_key'] = private_key_path
 
     return config
 
@@ -481,15 +484,31 @@ def setup_runpod_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     return configure_ssh_info(config)
 
 
+def setup_vast_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
+    """Sets up SSH authentication for Vast.
+    - Generates a new SSH key pair if one does not exist.
+    - Adds the public SSH key to the user's Vast account.
+    """
+    _, public_key_path = get_or_generate_keys()
+    with open(public_key_path, 'r', encoding='UTF-8') as pub_key_file:
+        public_key = pub_key_file.read().strip()
+        current_key_list = vast.vast().show_ssh_keys()  # pylint: disable=assignment-from-no-return
+        # Only add an ssh key if it hasn't already been added
+        if not any(x['public_key'] == public_key for x in current_key_list):
+            vast.vast().create_ssh_key(ssh_key=public_key)
+
+    config['auth']['ssh_public_key'] = public_key_path
+    return configure_ssh_info(config)
+
+
 def setup_fluidstack_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
-    get_or_generate_keys()
+    _, public_key_path = get_or_generate_keys()
 
     client = fluidstack_utils.FluidstackClient()
-    public_key_path = os.path.expanduser(PUBLIC_SSH_KEY_PATH)
     public_key = None
     with open(public_key_path, 'r', encoding='utf-8') as f:
         public_key = f.read()
     client.get_or_add_ssh_key(public_key)
-    config['auth']['ssh_public_key'] = PUBLIC_SSH_KEY_PATH
+    config['auth']['ssh_public_key'] = public_key_path
     return configure_ssh_info(config)

sky/backends/backend.py

@@ -2,9 +2,11 @@
 import typing
 from typing import Dict, Generic, Optional
 
-import sky
 from sky.usage import usage_lib
+from sky.utils import cluster_utils
+from sky.utils import rich_utils
 from sky.utils import timeline
+from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky import resources
@@ -43,19 +45,45 @@ class Backend(Generic[_ResourceHandleType]):
     @timeline.event
     @usage_lib.messages.usage.update_runtime('provision')
     def provision(
-            self,
-            task: 'task_lib.Task',
-            to_provision: Optional['resources.Resources'],
-            dryrun: bool,
-            stream_logs: bool,
-            cluster_name: Optional[str] = None,
-            retry_until_up: bool = False) -> Optional[_ResourceHandleType]:
+        self,
+        task: 'task_lib.Task',
+        to_provision: Optional['resources.Resources'],
+        dryrun: bool,
+        stream_logs: bool,
+        cluster_name: Optional[str] = None,
+        retry_until_up: bool = False,
+        skip_unnecessary_provisioning: bool = False,
+    ) -> Optional[_ResourceHandleType]:
+        """Provisions resources for the given task.
+
+        Args:
+            task: The task to provision resources for.
+            to_provision: Resource config to provision. Should only be None if
+                cluster_name refers to an existing cluster, whose resources will
+                be used.
+            dryrun: If True, don't actually provision anything.
+            stream_logs: If True, stream additional logs to console.
+            cluster_name: Name of the cluster to provision. If None, a name will
+                be auto-generated. If the name refers to an existing cluster,
+                the existing cluster will be reused and re-provisioned.
+            retry_until_up: If True, retry provisioning until resources are
+                successfully launched.
+            skip_if_no_cluster_updates: If True, compare the cluster config to
+                the existing cluster_name's config. Skip provisioning if no
+                updates are needed for the existing cluster.
+
+        Returns:
+            A ResourceHandle object for the provisioned resources, or None if
+            dryrun is True.
+        """
         if cluster_name is None:
-            cluster_name = sky.backends.backend_utils.generate_cluster_name()
+            cluster_name = cluster_utils.generate_cluster_name()
         usage_lib.record_cluster_name_for_current_operation(cluster_name)
         usage_lib.messages.usage.update_actual_task(task)
-        return self._provision(task, to_provision, dryrun, stream_logs,
-                               cluster_name, retry_until_up)
+        with rich_utils.safe_status(ux_utils.spinner_message('Launching')):
+            return self._provision(task, to_provision, dryrun, stream_logs,
+                                   cluster_name, retry_until_up,
+                                   skip_unnecessary_provisioning)
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('sync_workdir')
@@ -76,7 +104,8 @@ class Backend(Generic[_ResourceHandleType]):
     @usage_lib.messages.usage.update_runtime('setup')
     def setup(self, handle: _ResourceHandleType, task: 'task_lib.Task',
               detach_setup: bool) -> None:
-        return self._setup(handle, task, detach_setup)
+        with rich_utils.safe_status(ux_utils.spinner_message('Running setup')):
+            return self._setup(handle, task, detach_setup)
 
     def add_storage_objects(self, task: 'task_lib.Task') -> None:
         raise NotImplementedError
@@ -96,7 +125,8 @@ class Backend(Generic[_ResourceHandleType]):
         usage_lib.record_cluster_name_for_current_operation(
             handle.get_cluster_name())
         usage_lib.messages.usage.update_actual_task(task)
-        return self._execute(handle, task, detach_run, dryrun)
+        with rich_utils.safe_status(ux_utils.spinner_message('Submitting job')):
+            return self._execute(handle, task, detach_run, dryrun)
 
     @timeline.event
     def post_execute(self, handle: _ResourceHandleType, down: bool) -> None:
@@ -121,13 +151,15 @@ class Backend(Generic[_ResourceHandleType]):
 
     # --- Implementations of the APIs ---
     def _provision(
-            self,
-            task: 'task_lib.Task',
-            to_provision: Optional['resources.Resources'],
-            dryrun: bool,
-            stream_logs: bool,
-            cluster_name: str,
-            retry_until_up: bool = False) -> Optional[_ResourceHandleType]:
+        self,
+        task: 'task_lib.Task',
+        to_provision: Optional['resources.Resources'],
+        dryrun: bool,
+        stream_logs: bool,
+        cluster_name: str,
+        retry_until_up: bool = False,
+        skip_unnecessary_provisioning: bool = False,
+    ) -> Optional[_ResourceHandleType]:
         raise NotImplementedError
 
     def _sync_workdir(self, handle: _ResourceHandleType, workdir: Path) -> None: