skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/users/permission.py

@@ -3,7 +3,7 @@ import contextlib
 import hashlib
 import logging
 import os
-from typing import Generator, List
+from typing import Generator, List, Optional
 
 import casbin
 import filelock
@@ -14,6 +14,7 @@ from sky import models
 from sky import sky_logging
 from sky.skylet import constants
 from sky.users import rbac
+from sky.utils import annotations
 from sky.utils import common_utils
 from sky.utils.db import db_utils
 
@@ -27,14 +28,14 @@ logger = sky_logging.init_logger(__name__)
 POLICY_UPDATE_LOCK_PATH = os.path.expanduser('~/.sky/.policy_update.lock')
 POLICY_UPDATE_LOCK_TIMEOUT_SECONDS = 20
 
-_enforcer_instance = None
+_enforcer_instance: Optional['PermissionService'] = None
 
 
 class PermissionService:
     """Permission service for SkyPilot API Server."""
 
     def __init__(self):
-        self.enforcer = None
+        self.enforcer: Optional[casbin.Enforcer] = None
 
     def _lazy_initialize(self):
         if self.enforcer is not None:
@@ -42,7 +43,6 @@ class PermissionService:
         with _policy_lock():
             global _enforcer_instance
             if _enforcer_instance is None:
-                _enforcer_instance = self
                 engine = global_user_state.initialize_and_get_db()
                 db_utils.add_all_tables_to_db_sqlalchemy(
                     sqlalchemy_adapter.Base.metadata, engine)
@@ -52,11 +52,23 @@ class PermissionService:
                                           'model.conf')
                 enforcer = casbin.Enforcer(model_path, adapter)
                 self.enforcer = enforcer
+                # Only set the enforcer instance once the enforcer
+                # is successfully initialized, if we change it and then fail
+                # we will set it to None and all subsequent calls will fail.
+                _enforcer_instance = self
                 self._maybe_initialize_policies()
                 self._maybe_initialize_basic_auth_user()
             else:
+                assert _enforcer_instance is not None
                 self.enforcer = _enforcer_instance.enforcer
 
+    def _ensure_enforcer(self) -> casbin.Enforcer:
+        """Ensure enforcer is initialized and return it."""
+        self._lazy_initialize()
+        assert self.enforcer is not None, (
+            'Enforcer should be initialized after _lazy_initialize()')
+        return self.enforcer
+
     def _maybe_initialize_basic_auth_user(self) -> None:
         """Initialize basic auth user if it is enabled."""
         basic_auth = os.environ.get(constants.SKYPILOT_INITIAL_BASIC_AUTH)
@@ -72,9 +84,9 @@ class PermissionService:
                 return
             global_user_state.add_or_update_user(
                 models.User(id=user_hash, name=username, password=password))
-            self.enforcer.add_grouping_policy(user_hash,
-                                              rbac.RoleName.ADMIN.value)
-            self.enforcer.save_policy()
+            enforcer = self._ensure_enforcer()
+            enforcer.add_grouping_policy(user_hash, rbac.RoleName.ADMIN.value)
+            enforcer.save_policy()
             logger.info(f'Basic auth user {username} initialized')
 
     def _maybe_initialize_policies(self) -> None:
@@ -86,7 +98,8 @@ class PermissionService:
 
         # Check if policies are already initialized by looking for existing
         # permission policies in the enforcer
-        existing_policies = self.enforcer.get_policy()
+        enforcer = self._ensure_enforcer()
+        existing_policies = enforcer.get_policy()
 
         # If we already have policies for the expected roles, skip
         # initialization
@@ -123,7 +136,7 @@ class PermissionService:
             logger.debug('Policies not found or incomplete, initializing...')
             # Only clear p policies (permission policies),
             # keep g policies (role policies)
-            self.enforcer.remove_filtered_policy(0)
+            enforcer.remove_filtered_policy(0)
             for role, permissions in role_permissions.items():
                 if permissions['permissions'] and 'blocklist' in permissions[
                         'permissions']:
@@ -133,14 +146,14 @@ class PermissionService:
                         method = item['method']
                         logger.debug(f'Adding role policy: role={role}, '
                                      f'path={path}, method={method}')
-                        self.enforcer.add_policy(role, path, method)
+                        enforcer.add_policy(role, path, method)
                         policy_updated = True
 
             for workspace_name, users in workspace_policy_permissions.items():
                 for user in users:
                     logger.debug(f'Initializing workspace policy: user={user}, '
                                  f'workspace={workspace_name}')
-                    self.enforcer.add_policy(user, workspace_name, '*')
+                    enforcer.add_policy(user, workspace_name, '*')
                     policy_updated = True
             logger.debug('Policies initialized successfully')
         else:
@@ -153,7 +166,7 @@ class PermissionService:
             policy_updated = policy_updated or user_added
 
         if policy_updated:
-            self.enforcer.save_policy()
+            enforcer.save_policy()
 
     def add_user_if_not_exists(self, user_id: str) -> None:
         """Add user role relationship."""
@@ -167,34 +180,35 @@ class PermissionService:
         Returns:
             True if the user was added, False otherwise.
         """
-        user_roles = self.enforcer.get_roles_for_user(user_id)
+        enforcer = self._ensure_enforcer()
+        user_roles = enforcer.get_roles_for_user(user_id)
         if not user_roles:
-            self.enforcer.add_grouping_policy(user_id, rbac.get_default_role())
+            enforcer.add_grouping_policy(user_id, rbac.get_default_role())
             return True
         return False
 
     def delete_user(self, user_id: str) -> None:
         """Delete user role relationship."""
-        self._lazy_initialize()
         with _policy_lock():
             # Get current roles
             self._load_policy_no_lock()
             # Avoid calling get_user_roles, as it will require the lock.
-            current_roles = self.enforcer.get_roles_for_user(user_id)
+            enforcer = self._ensure_enforcer()
+            current_roles = enforcer.get_roles_for_user(user_id)
             if not current_roles:
                 logger.debug(f'User {user_id} has no roles')
                 return
-            self.enforcer.remove_grouping_policy(user_id, current_roles[0])
-            self.enforcer.save_policy()
+            enforcer.remove_grouping_policy(user_id, current_roles[0])
+            enforcer.save_policy()
 
     def update_role(self, user_id: str, new_role: str) -> None:
         """Update user role relationship."""
-        self._lazy_initialize()
         with _policy_lock():
             # Get current roles
             self._load_policy_no_lock()
             # Avoid calling get_user_roles, as it will require the lock.
-            current_roles = self.enforcer.get_roles_for_user(user_id)
+            enforcer = self._ensure_enforcer()
+            current_roles = enforcer.get_roles_for_user(user_id)
             if not current_roles:
                 logger.debug(f'User {user_id} has no roles')
             else:
@@ -203,11 +217,11 @@ class PermissionService:
                 if current_role == new_role:
                     logger.debug(f'User {user_id} already has role {new_role}')
                     return
-                self.enforcer.remove_grouping_policy(user_id, current_role)
+                enforcer.remove_grouping_policy(user_id, current_role)
 
             # Update user role
-            self.enforcer.add_grouping_policy(user_id, new_role)
-            self.enforcer.save_policy()
+            enforcer.add_grouping_policy(user_id, new_role)
+            enforcer.save_policy()
 
     def get_user_roles(self, user_id: str) -> List[str]:
         """Get all roles for a user.
@@ -222,15 +236,15 @@ class PermissionService:
         Returns:
             A list of role names that the user has.
         """
-        self._lazy_initialize()
         self._load_policy_no_lock()
-        return self.enforcer.get_roles_for_user(user_id)
+        enforcer = self._ensure_enforcer()
+        return enforcer.get_roles_for_user(user_id)
 
     def get_users_for_role(self, role: str) -> List[str]:
         """Get all users for a role."""
-        self._lazy_initialize()
         self._load_policy_no_lock()
-        return self.enforcer.get_users_for_role(role)
+        enforcer = self._ensure_enforcer()
+        return enforcer.get_users_for_role(role)
 
     def check_endpoint_permission(self, user_id: str, path: str,
                                   method: str) -> bool:
@@ -241,19 +255,22 @@ class PermissionService:
         # it is a hot path in every request. It is ok to have a stale policy,
         # as long as it is eventually consistent.
         # self._load_policy_no_lock()
-        self._lazy_initialize()
-        return self.enforcer.enforce(user_id, path, method)
+        enforcer = self._ensure_enforcer()
+        return enforcer.enforce(user_id, path, method)
 
     def _load_policy_no_lock(self):
         """Load policy from storage."""
-        self.enforcer.load_policy()
+        enforcer = self._ensure_enforcer()
+        enforcer.load_policy()
 
     def load_policy(self):
         """Load policy from storage with lock."""
-        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
 
+    # Right now, not a lot of users are using multiple workspaces,
+    # so 5 should be more than enough.
+    @annotations.lru_cache(scope='request', maxsize=5)
     def check_workspace_permission(self, user_id: str,
                                    workspace_name: str) -> bool:
         """Check workspace permission.
@@ -266,7 +283,6 @@ class PermissionService:
         For public workspaces, the permission is granted via a wildcard policy
         ('*').
         """
-        self._lazy_initialize()
         if os.getenv(constants.ENV_VAR_IS_SKYPILOT_SERVER) is None:
             # When it is not on API server, we allow all users to access all
             # workspaces, as the workspace check has been done on API server.
@@ -279,7 +295,8 @@ class PermissionService:
         # r.act == p.act
         # This means if there's a policy ('*', workspace_name, '*'), it will
         # match any user
-        result = self.enforcer.enforce(user_id, workspace_name, '*')
+        enforcer = self._ensure_enforcer()
+        result = enforcer.enforce(user_id, workspace_name, '*')
         logger.debug(f'Workspace permission check: user={user_id}, '
                      f'workspace={workspace_name}, result={result}')
         return result
@@ -323,13 +340,13 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
-        self._lazy_initialize()
         with _policy_lock():
+            enforcer = self._ensure_enforcer()
             for user in users:
                 logger.debug(f'Adding workspace policy: user={user}, '
                              f'workspace={workspace_name}')
-                self.enforcer.add_policy(user, workspace_name, '*')
-            self.enforcer.save_policy()
+                enforcer.add_policy(user, workspace_name, '*')
+            enforcer.save_policy()
 
     def update_workspace_policy(self, workspace_name: str,
                                 users: List[str]) -> None:
@@ -341,24 +358,24 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
-        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
+            enforcer = self._ensure_enforcer()
             # Remove all existing policies for this workspace
-            self.enforcer.remove_filtered_policy(1, workspace_name)
+            enforcer.remove_filtered_policy(1, workspace_name)
             # Add new policies
             for user in users:
                 logger.debug(f'Updating workspace policy: user={user}, '
                              f'workspace={workspace_name}')
-                self.enforcer.add_policy(user, workspace_name, '*')
-            self.enforcer.save_policy()
+                enforcer.add_policy(user, workspace_name, '*')
+            enforcer.save_policy()
 
     def remove_workspace_policy(self, workspace_name: str) -> None:
         """Remove workspace policy."""
-        self._lazy_initialize()
         with _policy_lock():
-            self.enforcer.remove_filtered_policy(1, workspace_name)
-            self.enforcer.save_policy()
+            enforcer = self._ensure_enforcer()
+            enforcer.remove_filtered_policy(1, workspace_name)
+            enforcer.save_policy()
 
 
 @contextlib.contextmanager

sky/utils/accelerator_registry.py

@@ -3,6 +3,7 @@ import typing
 from typing import List, Optional
 
 from sky import catalog
+from sky.catalog import common as catalog_common
 from sky.utils import rich_utils
 from sky.utils import ux_utils
 
@@ -34,8 +35,8 @@ if typing.TYPE_CHECKING:
 
 # Use a cached version of accelerators to cloud mapping, so that we don't have
 # to download and read the catalog file for every cloud locally.
-_accelerator_df = catalog.common.read_catalog('common/accelerators.csv')
-_memory_df = catalog.common.read_catalog('common/metadata.csv')
+_accelerator_df = catalog_common.read_catalog('common/accelerators.csv')
+_memory_df = catalog_common.read_catalog('common/metadata.csv')
 
 # List of non-GPU accelerators that are supported by our backend for job queue
 # scheduling.
@@ -107,10 +108,12 @@ def canonicalize_accelerator_name(accelerator: str,
     if not names and cloud_str in ['Kubernetes', None]:
         with rich_utils.safe_status(
                 ux_utils.spinner_message('Listing accelerators on Kubernetes')):
+            # Only search for Kubernetes to reduce the lookup cost.
+            # For other clouds, the catalog has been searched in previous steps.
             searched = catalog.list_accelerators(
                 name_filter=accelerator,
                 case_sensitive=False,
-                clouds=cloud_str,
+                clouds='Kubernetes',
             )
         names = list(searched.keys())
         if accelerator in names:

sky/utils/admin_policy_utils.py

@@ -2,8 +2,9 @@
 import contextlib
 import copy
 import importlib
+import typing
 from typing import Iterator, Optional, Tuple, Union
-import urllib.parse
+from urllib import parse as urlparse
 
 import colorama
 
@@ -13,17 +14,21 @@ from sky import exceptions
 from sky import sky_logging
 from sky import skypilot_config
 from sky import task as task_lib
+from sky.server.requests import request_names
 from sky.utils import common_utils
 from sky.utils import config_utils
 from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
+if typing.TYPE_CHECKING:
+    from sky import models
+
 
 def _is_url(policy_string: str) -> bool:
     """Check if the policy string is a URL."""
     try:
-        parsed = urllib.parse.urlparse(policy_string)
+        parsed = urlparse.urlparse(policy_string)
         return parsed.scheme in ('http', 'https')
     except Exception:  # pylint: disable=broad-except
         return False
@@ -73,6 +78,7 @@ def _get_policy_impl(
 @contextlib.contextmanager
 def apply_and_use_config_in_current_request(
     entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
+    request_name: request_names.AdminPolicyRequestName,
     request_options: Optional[admin_policy.RequestOptions] = None,
     at_client_side: bool = False,
 ) -> Iterator['dag_lib.Dag']:
@@ -86,7 +92,8 @@ def apply_and_use_config_in_current_request(
     Refer to `apply()` for more details.
     """
     original_config = skypilot_config.to_dict()
-    dag, mutated_config = apply(entrypoint, request_options, at_client_side)
+    dag, mutated_config = apply(entrypoint, request_name, request_options,
+                                at_client_side)
     if mutated_config != original_config:
         with skypilot_config.replace_skypilot_config(mutated_config):
             yield dag
@@ -96,6 +103,7 @@ def apply_and_use_config_in_current_request(
 
 def apply(
     entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
+    request_name: request_names.AdminPolicyRequestName,
     request_options: Optional[admin_policy.RequestOptions] = None,
     at_client_side: bool = False,
 ) -> Tuple['dag_lib.Dag', config_utils.Config]:
@@ -126,9 +134,13 @@ def apply(
     if policy is None:
         return dag, skypilot_config.to_dict()
 
+    user = None
     if at_client_side:
         logger.info(f'Applying client admin policy: {policy}')
     else:
+        # When being called by the server, the middleware has set the
+        # current user and this information is available at this point.
+        user = common_utils.get_current_user()
         logger.info(f'Applying server admin policy: {policy}')
     config = copy.deepcopy(skypilot_config.to_dict())
     mutated_dag = dag_lib.Dag()
@@ -136,8 +148,9 @@ def apply(
 
     mutated_config = None
     for task in dag.tasks:
-        user_request = admin_policy.UserRequest(task, config, request_options,
-                                                at_client_side)
+        user_request = admin_policy.UserRequest(task, config, request_name,
+                                                request_options, at_client_side,
+                                                user)
         try:
             mutated_user_request = policy.apply(user_request)
         # Avoid duplicate exception wrapping.

sky/utils/annotations.py

@@ -3,6 +3,7 @@
 import functools
 from typing import Callable, Literal, TypeVar
 
+import cachetools
 from typing_extensions import ParamSpec
 
 # Whether the current process is a SkyPilot API server process.
@@ -56,6 +57,27 @@ def lru_cache(scope: Literal['global', 'request'], *lru_cache_args,
     return decorator
 
 
+def ttl_cache(scope: Literal['global', 'request'], *ttl_cache_args,
+              **ttl_cache_kwargs) -> Callable:
+    """TTLCache decorator for functions.
+
+    This decorator allows us to track which functions need to be reloaded for a
+    new request using the scope argument.
+    """
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+        if scope == 'global':
+            return cachetools.cached(
+                cachetools.TTLCache(*ttl_cache_args, **ttl_cache_kwargs))(func)
+        else:
+            cached_func = cachetools.cached(
+                cachetools.TTLCache(*ttl_cache_args, **ttl_cache_kwargs))(func)
+            _FUNCTIONS_NEED_RELOAD_CACHE.append(cached_func)
+            return cached_func
+
+    return decorator
+
+
 def clear_request_level_cache():
     """Clear the request-level cache."""
     for func in _FUNCTIONS_NEED_RELOAD_CACHE:

sky/utils/asyncio_utils.py

@@ -0,0 +1,78 @@
+"""Asyncio utilities."""
+
+import asyncio
+import functools
+from typing import Set
+
+_background_tasks: Set[asyncio.Task] = set()
+
+
+def shield(func):
+    """Shield the decorated async function from cancellation.
+
+    If the outer coroutine is cancelled, the inner decorated function
+    will be protected from cancellation by asyncio.shield(). And we will
+    maintain a reference to the the inner task to avoid it get GCed before
+    it is done.
+
+    For example, filelock.AsyncFileLock is not cancellation safe. The
+    following code:
+
+        async def fn_with_lock():
+            async with filelock.AsyncFileLock('lock'):
+                await asyncio.sleep(1)
+
+    is equivalent to:
+
+        # The lock may leak if the cancellation happens in
+        # lock.acquire() or lock.release()
+        async def fn_with_lock():
+            lock = filelock.AsyncFileLock('lock')
+            await lock.acquire()
+            try:
+                await asyncio.sleep(1)
+            finally:
+                await lock.release()
+
+    Shilding the function ensures there is no cancellation will happen in the
+    function, thus the lock will be released properly:
+
+        @shield
+        async def fn_with_lock()
+
+    Note that the resource acquisition and release should usually be protected
+    in one @shield block but not separately, e.g.:
+
+        lock = filelock.AsyncFileLock('lock')
+
+        @shield
+        async def acquire():
+            await lock.acquire()
+
+        @shield
+        async def release():
+            await lock.release()
+
+        async def fn_with_lock():
+            await acquire()
+            try:
+                do_something()
+            finally:
+                await release()
+
+    The above code is not safe because if `fn_with_lock` is cancelled,
+    `acquire()` and `release()` will be executed in the background
+    concurrently and causes race conditions.
+    """
+
+    @functools.wraps(func)
+    async def async_wrapper(*args, **kwargs):
+        task = asyncio.create_task(func(*args, **kwargs))
+        try:
+            return await asyncio.shield(task)
+        except asyncio.CancelledError:
+            _background_tasks.add(task)
+            task.add_done_callback(lambda _: _background_tasks.discard(task))
+            raise
+
+    return async_wrapper

sky/utils/atomic.py

@@ -1,4 +1,4 @@
-"""Atomic structures and utilties."""
+"""Atomic structures and utilities."""
 
 import threading
 

sky/utils/auth_utils.py

@@ -0,0 +1,153 @@
+"""Utils for managing SkyPilot SSH key pairs."""
+
+import functools
+import os
+from typing import Tuple
+
+import filelock
+
+from sky import global_user_state
+from sky import sky_logging
+from sky.utils import common_utils
+
+logger = sky_logging.init_logger(__name__)
+
+MAX_TRIALS = 64
+# TODO(zhwu): Support user specified key pair.
+# 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 ephemeral.
+_SSH_KEY_PATH_PREFIX = '~/.sky/clients/{user_hash}/ssh'
+
+
+def get_ssh_key_and_lock_path(user_hash: str) -> Tuple[str, str, str]:
+    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]:
+    # Keep the import of the cryptography local to avoid expensive
+    # third-party imports when not needed.
+    # pylint: disable=import-outside-toplevel
+    from cryptography.hazmat.backends import default_backend
+    from cryptography.hazmat.primitives import serialization
+    from cryptography.hazmat.primitives.asymmetric import rsa
+
+    key = rsa.generate_private_key(backend=default_backend(),
+                                   public_exponent=65537,
+                                   key_size=2048)
+
+    private_key = key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.TraditionalOpenSSL,
+        encryption_algorithm=serialization.NoEncryption()).decode(
+            'utf-8').strip()
+
+    public_key = key.public_key().public_bytes(
+        serialization.Encoding.OpenSSH,
+        serialization.PublicFormat.OpenSSH).decode('utf-8').strip()
+
+    return public_key, private_key
+
+
+def _save_key_pair(private_key_path: str, public_key_path: str,
+                   private_key: str, public_key: str) -> None:
+    key_dir = os.path.dirname(private_key_path)
+    os.makedirs(key_dir, exist_ok=True, mode=0o700)
+
+    with open(
+            private_key_path,
+            'w',
+            encoding='utf-8',
+            opener=functools.partial(os.open, mode=0o600),
+    ) as f:
+        f.write(private_key)
+
+    with open(public_key_path,
+              'w',
+              encoding='utf-8',
+              opener=functools.partial(os.open, mode=0o644)) as f:
+        f.write(public_key)
+
+
+def get_or_generate_keys() -> Tuple[str, str]:
+    """Returns the absolute private and public key paths."""
+    user_hash = common_utils.get_user_hash()
+    private_key_path, public_key_path, lock_path = get_ssh_key_and_lock_path(
+        user_hash)
+    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)
+
+    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(lock_path, timeout=10):
+        if not os.path.exists(private_key_path):
+            ssh_public_key, ssh_private_key, exists = (
+                global_user_state.get_ssh_keys(user_hash))
+            if not exists:
+                ssh_public_key, ssh_private_key = _generate_rsa_key_pair()
+                global_user_state.set_ssh_keys(user_hash, ssh_public_key,
+                                               ssh_private_key)
+            _save_key_pair(private_key_path, public_key_path, ssh_private_key,
+                           ssh_public_key)
+    assert os.path.exists(public_key_path), (
+        'Private key found, but associated public key '
+        f'{public_key_path} does not exist.')
+    return private_key_path, public_key_path
+
+
+def create_ssh_key_files_from_db(private_key_path: str) -> bool:
+    """Creates the ssh key files from the database.
+
+    Returns:
+        True if the ssh key files are created successfully, False otherwise.
+    """
+    # Assume private key path is in the format of
+    # ~/.sky/clients/<user_hash>/ssh/sky-key
+    separated_path = os.path.normpath(private_key_path).split(os.path.sep)
+    assert separated_path[-1] == 'sky-key'
+    assert separated_path[-2] == 'ssh'
+    user_hash = separated_path[-3]
+
+    private_key_path_generated, public_key_path, lock_path = (
+        get_ssh_key_and_lock_path(user_hash))
+    assert private_key_path == os.path.expanduser(private_key_path_generated), (
+        f'Private key path {private_key_path} does not '
+        'match the generated path '
+        f'{os.path.expanduser(private_key_path_generated)}')
+    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)
+    lock_dir = os.path.dirname(lock_path)
+
+    if os.path.exists(private_key_path) and os.path.exists(public_key_path):
+        return True
+    # 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(lock_path, timeout=10):
+        if not os.path.exists(private_key_path):
+            ssh_public_key, ssh_private_key, exists = (
+                global_user_state.get_ssh_keys(user_hash))
+            if not exists:
+                logger.debug(f'SSH keys not found for user {user_hash}')
+                return False
+            _save_key_pair(private_key_path, public_key_path, ssh_private_key,
+                           ssh_public_key)
+    assert os.path.exists(public_key_path), (
+        'Private key found, but associated public key '
+        f'{public_key_path} does not exist.')
+    return True