skypilot-nightly 1.0.0.dev20251009__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/admin_policy.py

@@ -9,7 +9,9 @@ import pydantic
 
 import sky
 from sky import exceptions
+from sky import models
 from sky.adaptors import common as adaptors_common
+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
@@ -49,8 +51,10 @@ class _UserRequestBody(pydantic.BaseModel):
     # will be converted to JSON string, which will lose the None key.
     task: str
     skypilot_config: str
+    request_name: str
     request_options: Optional[RequestOptions] = None
     at_client_side: bool = False
+    user: str
 
 
 @dataclasses.dataclass
@@ -73,32 +77,48 @@ class UserRequest:
         skypilot_config: Global skypilot config to be used in this request.
         request_options: Request options. It is None for jobs and services.
         at_client_side: Is the request intercepted by the policy at client-side?
+        user: User who made the request.
+              Only available on the server side.
+              This value is None if at_client_side is True.
     """
     task: 'sky.Task'
     skypilot_config: 'sky.Config'
+    request_name: request_names.AdminPolicyRequestName
     request_options: Optional['RequestOptions'] = None
     at_client_side: bool = False
+    user: Optional['models.User'] = None
 
     def encode(self) -> str:
         return _UserRequestBody(
             task=yaml_utils.dump_yaml_str(self.task.to_yaml_config()),
             skypilot_config=yaml_utils.dump_yaml_str(dict(
                 self.skypilot_config)),
+            request_name=self.request_name.value,
             request_options=self.request_options,
             at_client_side=self.at_client_side,
+            user=(yaml_utils.dump_yaml_str(self.user.to_dict())
+                  if self.user is not None else ''),
         ).model_dump_json()
 
     @classmethod
     def decode(cls, body: str) -> 'UserRequest':
         user_request_body = _UserRequestBody.model_validate_json(body)
+        user_dict = yaml_utils.read_yaml_str(
+            user_request_body.user) if user_request_body.user != '' else None
+        user = models.User(
+            id=user_dict['id'],
+            name=user_dict['name']) if user_dict is not None else None
         return cls(
             task=sky.Task.from_yaml_config(
                 yaml_utils.read_yaml_all_str(user_request_body.task)[0]),
             skypilot_config=config_utils.Config.from_dict(
                 yaml_utils.read_yaml_all_str(
                     user_request_body.skypilot_config)[0]),
+            request_name=request_names.AdminPolicyRequestName(
+                user_request_body.request_name),
             request_options=user_request_body.request_options,
             at_client_side=user_request_body.at_client_side,
+            user=user,
         )
 
 

sky/authentication.py

@@ -19,31 +19,30 @@ 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
 import os
 import re
 import socket
 import subprocess
 import sys
-from typing import Any, Dict, Tuple
+from typing import Any, Dict
 import uuid
 
 import colorama
-import filelock
 
 from sky import clouds
 from sky import exceptions
-from sky import global_user_state
 from sky import sky_logging
 from sky.adaptors import gcp
 from sky.adaptors import ibm
 from sky.adaptors import runpod
 from sky.adaptors import seeweb as seeweb_adaptor
+from sky.adaptors import shadeform as shadeform_adaptor
 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.provision.primeintellect import utils as primeintellect_utils
+from sky.utils import auth_utils
 from sky.utils import common_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
@@ -56,140 +55,9 @@ logger = sky_logging.init_logger(__name__)
 # using Cloud Client Libraries for Python, where possible, for new code
 # development.
 
-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 empheral.
-_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):
-    # 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)
-    # 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:
-                raise RuntimeError(f'SSH keys not found for user {user_hash}')
-            _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.')
-
 
 def configure_ssh_info(config: Dict[str, Any]) -> Dict[str, Any]:
-    _, public_key_path = get_or_generate_keys()
+    _, public_key_path = auth_utils.get_or_generate_keys()
     with open(public_key_path, 'r', encoding='utf-8') as f:
         public_key = f.read().strip()
     config_str = yaml_utils.dump_yaml_str(config)
@@ -227,7 +95,7 @@ def parse_gcp_project_oslogin(project):
 # Retry for the GCP as sometimes there will be connection reset by peer error.
 @common_utils.retry
 def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
-    _, public_key_path = get_or_generate_keys()
+    _, public_key_path = auth_utils.get_or_generate_keys()
     config = copy.deepcopy(config)
 
     project_id = config['provider']['project_id']
@@ -352,11 +220,11 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
 def setup_lambda_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
-    get_or_generate_keys()
+    auth_utils.get_or_generate_keys()
 
     # Ensure ssh key is registered with Lambda Cloud
     lambda_client = lambda_utils.LambdaCloudClient()
-    _, public_key_path = get_or_generate_keys()
+    _, public_key_path = auth_utils.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()}'
@@ -373,7 +241,7 @@ def setup_ibm_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     and updates config file.
     keys default location: '~/.ssh/sky-key' and '~/.ssh/sky-key.pub'
     """
-    private_key_path, _ = get_or_generate_keys()
+    private_key_path, _ = auth_utils.get_or_generate_keys()
 
     def _get_unique_key_name():
         suffix_len = 10
@@ -382,7 +250,7 @@ def setup_ibm_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     client = ibm.client(region=config['provider']['region'])
     resource_group_id = config['provider']['resource_group_id']
 
-    _, public_key_path = get_or_generate_keys()
+    _, public_key_path = auth_utils.get_or_generate_keys()
     with open(os.path.abspath(os.path.expanduser(public_key_path)),
               'r',
               encoding='utf-8') as file:
@@ -424,7 +292,7 @@ def setup_ibm_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     context = kubernetes_utils.get_context_from_config(config['provider'])
     namespace = kubernetes_utils.get_namespace_from_config(config['provider'])
-    private_key_path, _ = get_or_generate_keys()
+    private_key_path, _ = auth_utils.get_or_generate_keys()
     # 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()
@@ -455,7 +323,7 @@ def setup_runpod_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     - Generates a new SSH key pair if one does not exist.
     - Adds the public SSH key to the user's RunPod account.
     """
-    _, public_key_path = get_or_generate_keys()
+    _, public_key_path = auth_utils.get_or_generate_keys()
     with open(public_key_path, 'r', encoding='UTF-8') as pub_key_file:
         public_key = pub_key_file.read().strip()
         runpod.runpod.cli.groups.ssh.functions.add_ssh_key(public_key)
@@ -468,7 +336,7 @@ def setup_vast_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     - 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()
+    _, public_key_path = auth_utils.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
@@ -482,7 +350,7 @@ def setup_vast_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
 def setup_fluidstack_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
-    _, public_key_path = get_or_generate_keys()
+    _, public_key_path = auth_utils.get_or_generate_keys()
 
     client = fluidstack_utils.FluidstackClient()
     public_key = None
@@ -495,7 +363,7 @@ def setup_fluidstack_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
 def setup_hyperbolic_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     """Sets up SSH authentication for Hyperbolic."""
-    _, public_key_path = get_or_generate_keys()
+    _, public_key_path = auth_utils.get_or_generate_keys()
     with open(public_key_path, 'r', encoding='utf-8') as f:
         public_key = f.read().strip()
 
@@ -511,6 +379,48 @@ def setup_hyperbolic_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     return configure_ssh_info(config)
 
 
+def setup_shadeform_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
+    """Sets up SSH authentication for Shadeform.
+    - Generates a new SSH key pair if one does not exist.
+    - Adds the public SSH key to the user's Shadeform account.
+
+    Note: This assumes there is a Shadeform Python SDK available.
+    If no official SDK exists, this function would need to use direct API calls.
+    """
+
+    _, public_key_path = auth_utils.get_or_generate_keys()
+    ssh_key_id = None
+
+    with open(public_key_path, 'r', encoding='utf-8') as f:
+        public_key = f.read().strip()
+
+    try:
+        # Add SSH key to Shadeform using our utility functions
+        ssh_key_id = shadeform_adaptor.add_ssh_key_to_shadeform(public_key)
+
+    except ImportError as e:
+        # If required dependencies are missing
+        logger.warning(
+            f'Failed to add Shadeform SSH key due to missing dependencies: '
+            f'{e}. Manually configure SSH keys in your Shadeform account.')
+
+    except Exception as e:
+        logger.warning(f'Failed to set up Shadeform authentication: {e}')
+        raise exceptions.CloudUserIdentityError(
+            'Failed to set up SSH authentication for Shadeform. '
+            f'Please ensure your Shadeform credentials are configured: {e}'
+        ) from e
+
+    if ssh_key_id is None:
+        raise Exception('Failed to add SSH key to Shadeform')
+
+    # Configure SSH info in the config
+    config['auth']['ssh_public_key'] = public_key_path
+    config['auth']['ssh_key_id'] = ssh_key_id
+
+    return configure_ssh_info(config)
+
+
 def setup_primeintellect_authentication(
         config: Dict[str, Any]) -> Dict[str, Any]:
     """Sets up SSH authentication for Prime Intellect.
@@ -518,7 +428,7 @@ def setup_primeintellect_authentication(
     - Adds the public SSH key to the user's Prime Intellect account.
     """
     # Ensure local SSH keypair exists and fetch public key content
-    _, public_key_path = get_or_generate_keys()
+    _, public_key_path = auth_utils.get_or_generate_keys()
     with open(public_key_path, 'r', encoding='utf-8') as f:
         public_key = f.read().strip()
 
@@ -538,10 +448,10 @@ def setup_primeintellect_authentication(
 def setup_seeweb_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     """Registers the public key with Seeweb and notes the remote name."""
     # 1. local key pair
-    get_or_generate_keys()
+    auth_utils.get_or_generate_keys()
 
     # 2. public key
-    _, public_key_path = get_or_generate_keys()
+    _, public_key_path = auth_utils.get_or_generate_keys()
     with open(public_key_path, 'r', encoding='utf-8') as f:
         public_key = f.read().strip()