skypilot-nightly 1.0.0.dev20250215__py3-none-any.whl → 1.0.0.dev20250217__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 = '354bbdf3a1d031b350011bc76570cf8c009ecc4a'
+_SKYPILOT_COMMIT_SHA = '7775d44c4c91d1474982fbdb3b1e03cfbbf2385e'
 
 
 def _get_git_commit():
@@ -35,7 +35,7 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev20250215'
+__version__ = '1.0.0.dev20250217'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))
 
 
@@ -75,6 +75,7 @@ def _set_http_proxy_env_vars() -> None:
 
 
 _set_http_proxy_env_vars()
+
 # ----------------------------------------------------------------- #
 
 # Keep this order to avoid cyclic imports
@@ -85,34 +86,46 @@ from sky import clouds
 from sky.admin_policy import AdminPolicy
 from sky.admin_policy import MutatedUserRequest
 from sky.admin_policy import UserRequest
+from sky.client.sdk import api_cancel
+from sky.client.sdk import api_info
+from sky.client.sdk import api_server_logs
+from sky.client.sdk import api_start
+from sky.client.sdk import api_status
+from sky.client.sdk import api_stop
+from sky.client.sdk import autostop
+from sky.client.sdk import cancel
+from sky.client.sdk import cost_report
+from sky.client.sdk import down
+from sky.client.sdk import download_logs
+from sky.client.sdk import exec  # pylint: disable=redefined-builtin
+from sky.client.sdk import get
+from sky.client.sdk import job_status
+from sky.client.sdk import launch
+from sky.client.sdk import optimize
+from sky.client.sdk import queue
+from sky.client.sdk import start
+from sky.client.sdk import status
+from sky.client.sdk import stop
+from sky.client.sdk import storage_delete
+from sky.client.sdk import storage_ls
+from sky.client.sdk import stream_and_get
+from sky.client.sdk import tail_logs
 from sky.clouds.service_catalog import list_accelerators
-from sky.core import autostop
-from sky.core import cancel
-from sky.core import cost_report
-from sky.core import down
-from sky.core import download_logs
-from sky.core import job_status
-from sky.core import queue
-from sky.core import start
-from sky.core import status
-from sky.core import stop
-from sky.core import storage_delete
-from sky.core import storage_ls
-from sky.core import tail_logs
 from sky.dag import Dag
 from sky.data import Storage
 from sky.data import StorageMode
 from sky.data import StoreType
-from sky.execution import exec  # pylint: disable=redefined-builtin
-from sky.execution import launch
 from sky.jobs import ManagedJobStatus
 from sky.optimizer import Optimizer
-from sky.optimizer import OptimizeTarget
 from sky.resources import Resources
 from sky.skylet.job_lib import JobStatus
-from sky.skypilot_config import Config
-from sky.status_lib import ClusterStatus
 from sky.task import Task
+from sky.utils.common import OptimizeTarget
+from sky.utils.common import StatusRefreshMode
+from sky.utils.config_utils import Config
+from sky.utils.registry import CLOUD_REGISTRY
+from sky.utils.registry import JOBS_RECOVERY_STRATEGY_REGISTRY
+from sky.utils.status_lib import ClusterStatus
 
 # Aliases.
 IBM = clouds.IBM
@@ -130,7 +143,6 @@ RunPod = clouds.RunPod
 Vast = clouds.Vast
 Vsphere = clouds.Vsphere
 Fluidstack = clouds.Fluidstack
-optimize = Optimizer.optimize
 
 __all__ = [
     '__version__',
@@ -161,11 +173,13 @@ __all__ = [
     'ClusterStatus',
     'JobStatus',
     'ManagedJobStatus',
+    'StatusRefreshMode',
     # APIs
     'Dag',
     'Task',
     'Resources',
-    # execution APIs
+    # core APIs
+    'optimize',
     'launch',
     'exec',
     # core APIs
@@ -184,9 +198,21 @@ __all__ = [
     # core APIs Storage Management
     'storage_ls',
     'storage_delete',
+    # API server APIs
+    'get',
+    'stream_and_get',
+    'api_status',
+    'api_cancel',
+    'api_info',
+    'api_start',
+    'api_stop',
+    'api_server_logs',
     # Admin Policy
     'UserRequest',
     'MutatedUserRequest',
     'AdminPolicy',
     'Config',
+    # Registry
+    'CLOUD_REGISTRY',
+    'JOBS_RECOVERY_STRATEGY_REGISTRY',
 ]

sky/adaptors/aws.py

@@ -35,6 +35,7 @@ import time
 from typing import Any, Callable
 
 from sky.adaptors import common
+from sky.utils import annotations
 from sky.utils import common_utils
 
 _IMPORT_ERROR_MESSAGE = ('Failed to import dependencies for AWS. '
@@ -59,7 +60,7 @@ class _ThreadLocalLRUCache(threading.local):
 
     def __init__(self, maxsize=32):
         super().__init__()
-        self.cache = functools.lru_cache(maxsize=maxsize)
+        self.cache = annotations.lru_cache(scope='global', maxsize=maxsize)
 
 
 def _thread_local_lru_cache(maxsize=32):

sky/adaptors/azure.py

@@ -3,7 +3,6 @@
 # pylint: disable=import-outside-toplevel
 import asyncio
 import datetime
-import functools
 import logging
 import threading
 import time
@@ -14,6 +13,7 @@ from sky import exceptions as sky_exceptions
 from sky import sky_logging
 from sky.adaptors import common
 from sky.skylet import constants
+from sky.utils import annotations
 from sky.utils import common_utils
 from sky.utils import ux_utils
 
@@ -33,7 +33,7 @@ _MAX_RETRY_FOR_GET_SUBSCRIPTION_ID = 5
 
 
 @common.load_lazy_modules(modules=_LAZY_MODULES)
-@functools.lru_cache()
+@annotations.lru_cache(scope='global', maxsize=1)
 def get_subscription_id() -> str:
     """Get the default subscription id."""
     from azure.common import credentials
@@ -69,7 +69,7 @@ def exceptions():
     return azure_exceptions
 
 
-@functools.lru_cache()
+@annotations.lru_cache(scope='global')
 @common.load_lazy_modules(modules=_LAZY_MODULES)
 def azure_mgmt_models(name: str):
     if name == 'compute':
@@ -83,7 +83,7 @@ def azure_mgmt_models(name: str):
 # We should keep the order of the decorators having 'lru_cache' followed
 # by 'load_lazy_modules' as we need to make sure a caller can call
 # 'get_client.cache_clear', which is a function provided by 'lru_cache'
-@functools.lru_cache()
+@annotations.lru_cache(scope='global')
 @common.load_lazy_modules(modules=_LAZY_MODULES)
 def get_client(name: str,
                subscription_id: Optional[str] = None,

sky/adaptors/cloudflare.py

@@ -2,12 +2,12 @@
 # pylint: disable=import-outside-toplevel
 
 import contextlib
-import functools
 import os
 import threading
 from typing import Dict, Optional, Tuple
 
 from sky.adaptors import common
+from sky.utils import annotations
 from sky.utils import ux_utils
 
 _IMPORT_ERROR_MESSAGE = ('Failed to import dependencies for Cloudflare.'
@@ -62,7 +62,7 @@ def get_r2_credentials(boto3_session):
 # lru_cache() is thread-safe and it will return the same session object
 # for different threads.
 # Reference: https://docs.python.org/3/library/functools.html#functools.lru_cache # pylint: disable=line-too-long
-@functools.lru_cache()
+@annotations.lru_cache(scope='global')
 def session():
     """Create an AWS session."""
     # Creating the session object is not thread-safe for boto3,
@@ -76,7 +76,7 @@ def session():
         return session_
 
 
-@functools.lru_cache()
+@annotations.lru_cache(scope='global')
 def resource(resource_name: str, **kwargs):
     """Create a Cloudflare resource.
 
@@ -102,7 +102,7 @@ def resource(resource_name: str, **kwargs):
         **kwargs)
 
 
-@functools.lru_cache()
+@annotations.lru_cache(scope='global')
 def client(service_name: str, region):
     """Create an CLOUDFLARE client of a certain service.
 

sky/adaptors/kubernetes.py

@@ -1,11 +1,11 @@
 """Kubernetes adaptors"""
-import functools
 import logging
 import os
 from typing import Any, Callable, Optional, Set
 
 from sky.adaptors import common
 from sky.sky_logging import set_logging_level
+from sky.utils import annotations
 from sky.utils import env_options
 from sky.utils import ux_utils
 
@@ -106,49 +106,49 @@ def _load_config(context: Optional[str] = None):
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-@functools.lru_cache()
+@annotations.lru_cache(scope='request')
 def core_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.CoreV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-@functools.lru_cache()
+@annotations.lru_cache(scope='request')
 def auth_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.RbacAuthorizationV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-@functools.lru_cache()
+@annotations.lru_cache(scope='request')
 def networking_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.NetworkingV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-@functools.lru_cache()
+@annotations.lru_cache(scope='request')
 def custom_objects_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.CustomObjectsApi()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-@functools.lru_cache()
+@annotations.lru_cache(scope='global')
 def node_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.NodeV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-@functools.lru_cache()
+@annotations.lru_cache(scope='request')
 def apps_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.AppsV1Api()
 
 
 @_api_logging_decorator('urllib3', logging.ERROR)
-@functools.lru_cache()
+@annotations.lru_cache(scope='request')
 def api_client(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.ApiClient()

sky/authentication.py

@@ -12,12 +12,11 @@ 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 copy
 import functools
@@ -48,6 +47,7 @@ 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,
@@ -276,7 +292,7 @@ def setup_lambda_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
     # 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()}'
@@ -284,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 jobs
-    # 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
@@ -343,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
 
 
@@ -373,10 +373,9 @@ 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
     context = config['provider'].get(
@@ -386,9 +385,7 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
         # with in-cluster configuration. We need to set the context to None
         # to use the mounted service account.
         context = None
-    namespace = config['provider'].get(
-        'namespace',
-        kubernetes_utils.get_kube_config_context_namespace(context))
+    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()
@@ -404,7 +401,7 @@ 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),
@@ -468,6 +465,7 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
         # This should never happen because we check for this in from_str above.
         raise ValueError(f'Unsupported networking mode: {network_mode_str}')
     config['auth']['ssh_proxy_command'] = ssh_proxy_cmd
+    config['auth']['ssh_private_key'] = private_key_path
 
     return config
 
@@ -499,19 +497,18 @@ def setup_vast_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
         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_SSH_KEY_PATH
+    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,8 +2,8 @@
 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
@@ -77,7 +77,7 @@ class Backend(Generic[_ResourceHandleType]):
             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)
         with rich_utils.safe_status(ux_utils.spinner_message('Launching')):