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

sky/adaptors/seeweb.py

@@ -0,0 +1,167 @@
+""" Seeweb Adaptor """
+import configparser
+import pathlib
+from typing import Optional
+
+import pydantic
+import requests  # type: ignore
+
+from sky.adaptors import common
+from sky.utils import annotations
+
+
+class SeewebError(Exception):
+    """Base exception for Seeweb adaptor errors."""
+
+
+class SeewebCredentialsFileNotFound(SeewebError):
+    """Raised when the Seeweb credentials file is missing."""
+
+
+class SeewebApiKeyMissing(SeewebError):
+    """Raised when the Seeweb API key is missing or empty."""
+
+
+class SeewebAuthenticationError(SeewebError):
+    """Raised when authenticating with Seeweb API fails."""
+
+
+_IMPORT_ERROR_MESSAGE = ('Failed to import dependencies for Seeweb.'
+                         'Try pip install "skypilot[seeweb]"')
+
+ecsapi = common.LazyImport(
+    'ecsapi',
+    import_error_message=_IMPORT_ERROR_MESSAGE,
+)
+boto3 = common.LazyImport('boto3', import_error_message=_IMPORT_ERROR_MESSAGE)
+botocore = common.LazyImport('botocore',
+                             import_error_message=_IMPORT_ERROR_MESSAGE)
+
+_LAZY_MODULES = (ecsapi, boto3, botocore)
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def check_compute_credentials() -> bool:
+    """Checks if the user has access credentials to Seeweb's compute service.
+
+    Returns True if credentials are valid; otherwise raises a SeewebError.
+    """
+    # Read API key from standard Seeweb configuration file
+    key_path = pathlib.Path('~/.seeweb_cloud/seeweb_keys').expanduser()
+    if not key_path.exists():
+        raise SeewebCredentialsFileNotFound(
+            'Missing Seeweb API key file ~/.seeweb_cloud/seeweb_keys')
+
+    parser = configparser.ConfigParser()
+    parser.read(key_path)
+    try:
+        api_key = parser['DEFAULT']['api_key'].strip()
+    except KeyError as e:
+        raise SeewebApiKeyMissing(
+            'Missing api_key in ~/.seeweb_cloud/seeweb_keys') from e
+    if not api_key:
+        raise SeewebApiKeyMissing(
+            'Empty api_key in ~/.seeweb_cloud/seeweb_keys')
+
+    # Test connection by fetching servers list to validate the key
+    try:
+        seeweb_client = ecsapi.Api(token=api_key)
+        seeweb_client.fetch_servers()
+    except Exception as e:  # pylint: disable=broad-except
+        raise SeewebAuthenticationError(
+            f'Unable to authenticate with Seeweb API: {e}') from e
+
+    return True
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+def check_storage_credentials() -> bool:
+    """Checks if the user has access credentials to Seeweb's storage service.
+
+    Mirrors compute credentials validation.
+    """
+    return check_compute_credentials()
+
+
+@common.load_lazy_modules(_LAZY_MODULES)
+@annotations.lru_cache(scope='global', maxsize=1)
+def client():
+    """Returns an authenticated ecsapi.Api object."""
+    # Create authenticated client using the same credential pattern
+    key_path = pathlib.Path('~/.seeweb_cloud/seeweb_keys').expanduser()
+    if not key_path.exists():
+        raise SeewebCredentialsFileNotFound(
+            'Missing Seeweb API key file ~/.seeweb_cloud/seeweb_keys')
+
+    parser = configparser.ConfigParser()
+    parser.read(key_path)
+    try:
+        api_key = parser['DEFAULT']['api_key'].strip()
+    except KeyError as e:
+        raise SeewebApiKeyMissing(
+            'Missing api_key in ~/.seeweb_cloud/seeweb_keys') from e
+    if not api_key:
+        raise SeewebApiKeyMissing(
+            'Empty api_key in ~/.seeweb_cloud/seeweb_keys')
+
+    api = ecsapi.Api(token=api_key)
+
+    # Monkey-patch fetch_servers to be tolerant to API schema mismatches.
+    orig_fetch_servers = api.fetch_servers
+    orig_delete_server = api.delete_server
+
+    def _tolerant_fetch_servers(
+            timeout: Optional[int] = None):  # type: ignore[override]
+        try:
+            return orig_fetch_servers(timeout=timeout)
+        except pydantic.ValidationError:
+            # Fallback path: fetch raw JSON, drop snapshot fields, then validate
+            # pylint: disable=protected-access
+            base_url = api._Api__generate_base_url()  # type: ignore
+            headers = api._Api__generate_authentication_headers(
+            )  # type: ignore
+            url = f'{base_url}/servers'
+            resp = requests.get(url, headers=headers, timeout=timeout or 15)
+            resp.raise_for_status()
+            data = resp.json()
+            try:
+                servers = data.get('server', [])
+                for s in servers:
+                    s.pop('last_restored_snapshot', None)
+            except (KeyError, TypeError, ValueError):
+                pass
+            server_list_response_cls = ecsapi._server._ServerListResponse
+            servers_response = server_list_response_cls.model_validate(data)
+            return servers_response.server
+
+    api.fetch_servers = _tolerant_fetch_servers  # type: ignore[assignment]
+
+    def _tolerant_delete_server(server_name: str,
+                                timeout: Optional[int] = None):
+        try:
+            return orig_delete_server(server_name, timeout=timeout)
+        except pydantic.ValidationError:
+            # Fallback: perform raw DELETE and interpret not_found as success
+            # pylint: disable=protected-access
+            base_url = api._Api__generate_base_url()  # type: ignore
+            headers = api._Api__generate_authentication_headers(
+            )  # type: ignore
+            url = f'{base_url}/servers/{server_name}'
+            resp = requests.delete(url, headers=headers, timeout=timeout or 15)
+            # Treat 404 as idempotent success
+            if resp.status_code == 404:
+                return None
+            # Some APIs return {status: 'not_found', message: ...}
+            try:
+                data = resp.json()
+                if isinstance(data, dict) and data.get('status') == 'not_found':
+                    return None
+            except (ValueError, TypeError):
+                pass
+            # If not clearly not_found, re-raise original behavior
+            resp.raise_for_status()
+            # Best-effort: return None to indicate deletion requested
+            return None
+
+    api.delete_server = _tolerant_delete_server  # type: ignore[assignment]
+    return api

sky/adaptors/shadeform.py

@@ -0,0 +1,89 @@
+"""Shadeform cloud adaptor."""
+
+import functools
+import socket
+from typing import Any, Dict, List, Optional
+
+import requests
+
+from sky import sky_logging
+from sky.provision.shadeform import shadeform_utils
+from sky.utils import common_utils
+
+logger = sky_logging.init_logger(__name__)
+
+_shadeform_sdk = None
+
+
+def import_package(func):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        global _shadeform_sdk
+        if _shadeform_sdk is None:
+            try:
+                import shadeform as _shadeform  # pylint: disable=import-outside-toplevel
+                _shadeform_sdk = _shadeform
+            except ImportError:
+                raise ImportError(
+                    'Failed to import dependencies for Shadeform. '
+                    'Try pip install "skypilot[shadeform]"') from None
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+@import_package
+def shadeform():
+    """Return the shadeform package."""
+    return _shadeform_sdk
+
+
+def list_ssh_keys() -> List[Dict[str, Any]]:
+    """List all SSH keys in Shadeform account."""
+    try:
+        response = shadeform_utils.get_ssh_keys()
+        return response.get('ssh_keys', [])
+    except (ValueError, KeyError, requests.exceptions.RequestException) as e:
+        logger.warning(f'Failed to list SSH keys from Shadeform: {e}')
+        return []
+
+
+def add_ssh_key_to_shadeform(public_key: str) -> Optional[str]:
+    """Add SSH key to Shadeform if it doesn't already exist.
+
+    Args:
+        public_key: The SSH public key string.
+
+    Returns:
+        The name of the key if added successfully, None otherwise.
+    """
+    try:
+        # Check if key already exists
+        existing_keys = list_ssh_keys()
+        key_exists = False
+        key_id = None
+        for key in existing_keys:
+            if key.get('public_key') == public_key:
+                key_exists = True
+                key_id = key.get('id')
+                break
+
+        if key_exists:
+            logger.info('SSH key already exists in Shadeform account')
+            return key_id
+
+        # Generate a unique key name
+        hostname = socket.gethostname()
+        key_name = f'skypilot-{hostname}-{common_utils.get_user_hash()[:8]}'
+
+        # Add the key
+        response = shadeform_utils.add_ssh_key(name=key_name,
+                                               public_key=public_key)
+        key_id = response['id']
+        logger.info(f'Added SSH key to Shadeform: {key_name, key_id}')
+        return key_id
+
+    except (ValueError, KeyError, requests.exceptions.RequestException) as e:
+        logger.warning(f'Failed to add SSH key to Shadeform: {e}')
+        return None

sky/admin_policy.py

@@ -4,12 +4,26 @@ import dataclasses
 import typing
 from typing import Optional
 
+import colorama
+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
+from sky.utils import yaml_utils
+
 if typing.TYPE_CHECKING:
-    import sky
+    import requests
+else:
+    requests = adaptors_common.LazyImport('requests')
 
 
-@dataclasses.dataclass
-class RequestOptions:
+class RequestOptions(pydantic.BaseModel):
     """Request options for admin policy.
 
     Args:
@@ -21,11 +35,28 @@ class RequestOptions:
         dryrun: Is the request a dryrun?
     """
     cluster_name: Optional[str]
+    # Keep these two fields for backward compatibility. The values are copied
+    # from task.resources.autostop_config, so that legacy admin policy plugins
+    # can still read the correct autostop config from request options before
+    # we drop the compatibility.
+    # TODO(aylei): remove these fields after 0.12.0
     idle_minutes_to_autostop: Optional[int]
     down: bool
     dryrun: bool
 
 
+class _UserRequestBody(pydantic.BaseModel):
+    """Auxiliary model to validate and serialize a user request."""
+    # We have to use serialized YAML string, instead of a dict, because dict
+    # 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
 class UserRequest:
     """A user request.
@@ -45,20 +76,101 @@ class UserRequest:
         task: User specified task.
         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,
+        )
+
+
+class _MutatedUserRequestBody(pydantic.BaseModel):
+    """Auxiliary model to validate and serialize a user request."""
+    task: str
+    skypilot_config: str
 
 
 @dataclasses.dataclass
 class MutatedUserRequest:
+    """Mutated user request."""
+
     task: 'sky.Task'
     skypilot_config: 'sky.Config'
 
+    def encode(self) -> str:
+        return _MutatedUserRequestBody(
+            task=yaml_utils.dump_yaml_str(self.task.to_yaml_config()),
+            skypilot_config=yaml_utils.dump_yaml_str(dict(
+                self.skypilot_config),)).model_dump_json()
+
+    @classmethod
+    def decode(cls, mutated_user_request_body: str,
+               original_request: UserRequest) -> 'MutatedUserRequest':
+        mutated_user_request_body = _MutatedUserRequestBody.model_validate_json(
+            mutated_user_request_body)
+        task = sky.Task.from_yaml_config(
+            yaml_utils.read_yaml_all_str(mutated_user_request_body.task)[0])
+        # Some internal Task fields are not serialized. We need to manually
+        # restore them from the original request.
+        task.managed_job_dag = original_request.task.managed_job_dag
+        task.service_name = original_request.task.service_name
+        return cls(task=task,
+                   skypilot_config=config_utils.Config.from_dict(
+                       yaml_utils.read_yaml_all_str(
+                           mutated_user_request_body.skypilot_config)[0],))
+
+
+class PolicyInterface:
+    """Interface for admin-defined policy for user requests."""
+
+    @abc.abstractmethod
+    def apply(self, user_request: UserRequest) -> MutatedUserRequest:
+        """Apply the admin policy to the user request."""
+
+    def __str__(self):
+        return f'{self.__class__.__name__}'
+
 
 # pylint: disable=line-too-long
-class AdminPolicy:
+class AdminPolicy(PolicyInterface):
     """Abstract interface of an admin-defined policy for all user requests.
 
     Admins can implement a subclass of AdminPolicy with the following signature:
@@ -99,3 +211,74 @@ class AdminPolicy:
         """
         raise NotImplementedError(
             'Your policy must implement validate_and_mutate')
+
+    def apply(self, user_request: UserRequest) -> MutatedUserRequest:
+        return self.validate_and_mutate(user_request)
+
+
+class PolicyTemplate(PolicyInterface):
+    """Admin policy template that can be instantiated to create a policy."""
+
+    @abc.abstractmethod
+    def validate_and_mutate(self,
+                            user_request: UserRequest) -> MutatedUserRequest:
+        """Validates and mutates the user request and returns mutated request.
+
+        Args:
+            user_request: The user request to validate and mutate.
+                UserRequest contains (sky.Task, sky.Config)
+
+        Returns:
+            MutatedUserRequest: The mutated user request.
+
+        Raises:
+            Exception to throw if the user request failed the validation.
+        """
+        raise NotImplementedError(
+            'Your policy must implement validate_and_mutate')
+
+    def apply(self, user_request: UserRequest) -> MutatedUserRequest:
+        return self.validate_and_mutate(user_request)
+
+
+class RestfulAdminPolicy(PolicyTemplate):
+    """Admin policy that calls a RESTful API for validation."""
+
+    def __init__(self, policy_url: str):
+        super().__init__()
+        self.policy_url = policy_url
+
+    def validate_and_mutate(self,
+                            user_request: UserRequest) -> MutatedUserRequest:
+        try:
+            response = requests.post(
+                self.policy_url,
+                json=user_request.encode(),
+                headers={'Content-Type': 'application/json'},
+                # TODO(aylei): make this configurable
+                timeout=30)
+            if response.status_code == 400:
+                raise exceptions.UserRequestRejectedByPolicy(
+                    f'{colorama.Fore.RED}User request is rejected by admin '
+                    f'policy {self.policy_url}{colorama.Fore.RESET}: '
+                    f'{response.text}')
+            response.raise_for_status()
+        except requests.exceptions.RequestException as e:
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.RestfulPolicyError(
+                    f'Failed to call admin policy URL '
+                    f'{self.policy_url}: {e}') from None
+
+        try:
+            mutated_user_request = MutatedUserRequest.decode(
+                response.json(), user_request)
+        except Exception as e:  # pylint: disable=broad-except
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.RestfulPolicyError(
+                    f'Failed to decode response from admin policy URL '
+                    f'{self.policy_url}: {common_utils.format_exception(e, use_bracket=True)}'
+                ) from None
+        return mutated_user_request
+
+    def __repr__(self):
+        return f'RestfulAdminPolicy(policy_url={self.policy_url})'