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

sky/adaptors/kubernetes.py

@@ -1,11 +1,12 @@
 """Kubernetes adaptors"""
-
-# pylint: disable=import-outside-toplevel
-
+import logging
 import os
+from typing import Any, Callable, Optional, Set
 
 from sky.adaptors import common
-from sky.utils import env_options
+from sky.sky_logging import set_logging_level
+from sky.utils import annotations
+from sky.utils import common_utils
 from sky.utils import ux_utils
 
 _IMPORT_ERROR_MESSAGE = ('Failed to import dependencies for Kubernetes. '
@@ -15,117 +16,140 @@ kubernetes = common.LazyImport('kubernetes',
 urllib3 = common.LazyImport('urllib3',
                             import_error_message=_IMPORT_ERROR_MESSAGE)
 
-_configured = False
-_core_api = None
-_auth_api = None
-_networking_api = None
-_custom_objects_api = None
-_node_api = None
-_apps_api = None
-_api_client = None
-
 # Timeout to use for API calls
 API_TIMEOUT = 5
 
+DEFAULT_IN_CLUSTER_REGION = 'in-cluster'
+# The name for the environment variable that stores the in-cluster context name
+# for Kubernetes clusters. This is used to associate a name with the current
+# context when running with in-cluster auth. If not set, the context name is
+# set to DEFAULT_IN_CLUSTER_REGION.
+IN_CLUSTER_CONTEXT_NAME_ENV_VAR = 'SKYPILOT_IN_CLUSTER_CONTEXT_NAME'
+
+
+def _decorate_methods(obj: Any, decorator: Callable, decoration_type: str):
+    for attr_name in dir(obj):
+        attr = getattr(obj, attr_name)
+        # Skip methods starting with '__' since they are invoked through one
+        # of the main methods, which are already decorated.
+        if callable(attr) and not attr_name.startswith('__'):
+            decorated_types: Set[str] = getattr(attr, '_sky_decorator_types',
+                                                set())
+            if decoration_type not in decorated_types:
+                decorated_attr = decorator(attr)
+                decorated_attr._sky_decorator_types = (  # pylint: disable=protected-access
+                    decorated_types | {decoration_type})
+                setattr(obj, attr_name, decorated_attr)
+    return obj
+
+
+def _api_logging_decorator(logger: str, level: int):
+    """Decorator to set logging level for API calls.
+
+    This is used to suppress the verbose logging from urllib3 when calls to the
+    Kubernetes API timeout.
+    """
+
+    def decorated_api(api):
+
+        def wrapped(*args, **kwargs):
+            obj = api(*args, **kwargs)
+            _decorate_methods(obj, set_logging_level(logger, level), 'api_log')
+            return obj
+
+        return wrapped
+
+    return decorated_api
 
-def _load_config():
-    global _configured
-    if _configured:
-        return
-    try:
-        # Load in-cluster config if running in a pod
-        # Kubernetes set environment variables for service discovery do not
-        # show up in SkyPilot tasks. For now, we work around by using
-        # DNS name instead of environment variables.
-        # See issue: https://github.com/skypilot-org/skypilot/issues/2287
-        os.environ['KUBERNETES_SERVICE_HOST'] = 'kubernetes.default.svc'
-        os.environ['KUBERNETES_SERVICE_PORT'] = '443'
-        kubernetes.config.load_incluster_config()
-    except kubernetes.config.config_exception.ConfigException:
+
+def _load_config(context: Optional[str] = None):
+    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+    def _load_config_from_kubeconfig(context: Optional[str] = None):
         try:
-            kubernetes.config.load_kube_config()
+            kubernetes.config.load_kube_config(context=context)
         except kubernetes.config.config_exception.ConfigException as e:
-            suffix = ''
-            if env_options.Options.SHOW_DEBUG_INFO.get():
-                suffix += f' Error: {str(e)}'
+            suffix = common_utils.format_exception(e, use_bracket=True)
             # Check if exception was due to no current-context
             if 'Expected key current-context' in str(e):
-                err_str = ('Failed to load Kubernetes configuration. '
-                           'Kubeconfig does not contain any valid context(s).'
-                           f'{suffix}\n'
-                           '    If you were running a local Kubernetes '
-                           'cluster, run `sky local up` to start the cluster.')
+                err_str = (
+                    f'Failed to load Kubernetes configuration for {context!r}. '
+                    'Kubeconfig does not contain any valid context(s).'
+                    f'\n{suffix}\n'
+                    '    If you were running a local Kubernetes '
+                    'cluster, run `sky local up` to start the cluster.')
             else:
-                err_str = ('Failed to load Kubernetes configuration. '
-                           'Please check if your kubeconfig file exists at '
-                           f'~/.kube/config and is valid.{suffix}')
+                err_str = (
+                    f'Failed to load Kubernetes configuration for {context!r}. '
+                    'Please check if your kubeconfig file exists at '
+                    f'~/.kube/config and is valid.\n{suffix}')
             err_str += '\nTo disable Kubernetes for SkyPilot: run `sky check`.'
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(err_str) from None
-    _configured = True
-
-
-def core_api():
-    global _core_api
-    if _core_api is None:
-        _load_config()
-        _core_api = kubernetes.client.CoreV1Api()
-
-    return _core_api
-
-
-def auth_api():
-    global _auth_api
-    if _auth_api is None:
-        _load_config()
-        _auth_api = kubernetes.client.RbacAuthorizationV1Api()
-
-    return _auth_api
 
+    if context == in_cluster_context_name() or context is None:
+        try:
+            # Load in-cluster config if running in a pod and context is None.
+            # Kubernetes set environment variables for service discovery do not
+            # show up in SkyPilot tasks. For now, we work around by using
+            # DNS name instead of environment variables.
+            # See issue: https://github.com/skypilot-org/skypilot/issues/2287
+            os.environ['KUBERNETES_SERVICE_HOST'] = 'kubernetes.default.svc'
+            os.environ['KUBERNETES_SERVICE_PORT'] = '443'
+            kubernetes.config.load_incluster_config()
+        except kubernetes.config.config_exception.ConfigException:
+            _load_config_from_kubeconfig()
+    else:
+        _load_config_from_kubeconfig(context)
 
-def networking_api():
-    global _networking_api
-    if _networking_api is None:
-        _load_config()
-        _networking_api = kubernetes.client.NetworkingV1Api()
 
-    return _networking_api
+@_api_logging_decorator('urllib3', logging.ERROR)
+@annotations.lru_cache(scope='request')
+def core_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.CoreV1Api()
 
 
-def custom_objects_api():
-    global _custom_objects_api
-    if _custom_objects_api is None:
-        _load_config()
-        _custom_objects_api = kubernetes.client.CustomObjectsApi()
+@_api_logging_decorator('urllib3', logging.ERROR)
+@annotations.lru_cache(scope='request')
+def auth_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.RbacAuthorizationV1Api()
 
-    return _custom_objects_api
 
+@_api_logging_decorator('urllib3', logging.ERROR)
+@annotations.lru_cache(scope='request')
+def networking_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.NetworkingV1Api()
 
-def node_api():
-    global _node_api
-    if _node_api is None:
-        _load_config()
-        _node_api = kubernetes.client.NodeV1Api()
 
-    return _node_api
+@_api_logging_decorator('urllib3', logging.ERROR)
+@annotations.lru_cache(scope='request')
+def custom_objects_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.CustomObjectsApi()
 
 
-def apps_api():
-    global _apps_api
-    if _apps_api is None:
-        _load_config()
-        _apps_api = kubernetes.client.AppsV1Api()
+@_api_logging_decorator('urllib3', logging.ERROR)
+@annotations.lru_cache(scope='global')
+def node_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.NodeV1Api()
 
-    return _apps_api
 
+@_api_logging_decorator('urllib3', logging.ERROR)
+@annotations.lru_cache(scope='request')
+def apps_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.AppsV1Api()
 
-def api_client():
-    global _api_client
-    if _api_client is None:
-        _load_config()
-        _api_client = kubernetes.client.ApiClient()
 
-    return _api_client
+@_api_logging_decorator('urllib3', logging.ERROR)
+@annotations.lru_cache(scope='request')
+def api_client(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.ApiClient()
 
 
 def api_exception():
@@ -142,3 +166,13 @@ def max_retry_error():
 
 def stream():
     return kubernetes.stream.stream
+
+
+def in_cluster_context_name() -> Optional[str]:
+    """Returns the name of the in-cluster context from the environment.
+
+    If the environment variable is not set, returns the default in-cluster
+    context name.
+    """
+    return (os.environ.get(IN_CLUSTER_CONTEXT_NAME_ENV_VAR) or
+            DEFAULT_IN_CLUSTER_REGION)

sky/adaptors/nebius.py

@@ -0,0 +1,100 @@
+"""Nebius cloud adaptor."""
+import os
+
+from sky.adaptors import common
+
+NEBIUS_TENANT_ID_FILENAME = 'NEBIUS_TENANT_ID.txt'
+NEBIUS_IAM_TOKEN_FILENAME = 'NEBIUS_IAM_TOKEN.txt'
+NEBIUS_PROJECT_ID_FILENAME = 'NEBIUS_PROJECT_ID.txt'
+NEBIUS_TENANT_ID_PATH = '~/.nebius/' + NEBIUS_TENANT_ID_FILENAME
+NEBIUS_IAM_TOKEN_PATH = '~/.nebius/' + NEBIUS_IAM_TOKEN_FILENAME
+NEBIUS_PROJECT_ID_PATH = '~/.nebius/' + NEBIUS_PROJECT_ID_FILENAME
+
+MAX_RETRIES_TO_DISK_CREATE = 120
+MAX_RETRIES_TO_INSTANCE_STOP = 120
+MAX_RETRIES_TO_INSTANCE_START = 120
+MAX_RETRIES_TO_INSTANCE_READY = 240
+
+MAX_RETRIES_TO_DISK_DELETE = 120
+MAX_RETRIES_TO_INSTANCE_WAIT = 120  # Maximum number of retries
+
+POLL_INTERVAL = 5
+
+_iam_token = None
+_tenant_id = None
+_project_id = None
+
+nebius = common.LazyImport(
+    'nebius',
+    import_error_message='Failed to import dependencies for Nebius AI Cloud. '
+    'Try running: pip install "skypilot[nebius]"',
+    # https://github.com/grpc/grpc/issues/37642 to avoid spam in console
+    set_loggers=lambda: os.environ.update({'GRPC_VERBOSITY': 'NONE'}))
+
+
+def request_error():
+    return nebius.aio.service_error.RequestError
+
+
+def compute():
+    # pylint: disable=import-outside-toplevel
+    from nebius.api.nebius.compute import v1 as compute_v1
+    return compute_v1
+
+
+def iam():
+    # pylint: disable=import-outside-toplevel
+    from nebius.api.nebius.iam import v1 as iam_v1
+    return iam_v1
+
+
+def nebius_common():
+    # pylint: disable=import-outside-toplevel
+    from nebius.api.nebius.common import v1 as common_v1
+    return common_v1
+
+
+def vpc():
+    # pylint: disable=import-outside-toplevel
+    from nebius.api.nebius.vpc import v1 as vpc_v1
+    return vpc_v1
+
+
+def get_iam_token():
+    global _iam_token
+    if _iam_token is None:
+        try:
+            with open(os.path.expanduser(NEBIUS_IAM_TOKEN_PATH),
+                      encoding='utf-8') as file:
+                _iam_token = file.read().strip()
+        except FileNotFoundError:
+            return None
+    return _iam_token
+
+
+def get_project_id():
+    global _project_id
+    if _project_id is None:
+        try:
+            with open(os.path.expanduser(NEBIUS_PROJECT_ID_PATH),
+                      encoding='utf-8') as file:
+                _project_id = file.read().strip()
+        except FileNotFoundError:
+            return None
+    return _project_id
+
+
+def get_tenant_id():
+    global _tenant_id
+    if _tenant_id is None:
+        try:
+            with open(os.path.expanduser(NEBIUS_TENANT_ID_PATH),
+                      encoding='utf-8') as file:
+                _tenant_id = file.read().strip()
+        except FileNotFoundError:
+            return None
+    return _tenant_id
+
+
+def sdk():
+    return nebius.sdk.SDK(credentials=get_iam_token())

sky/adaptors/oci.py

@@ -1,8 +1,17 @@
 """Oracle OCI cloud adaptor"""
 
+import functools
+import logging
 import os
 
 from sky.adaptors import common
+from sky.clouds.utils import oci_utils
+
+# Suppress OCI circuit breaker logging before lazy import, because
+# oci modules prints additional message during imports, i.e., the
+# set_logger in the LazyImport called after imports will not take
+# effect.
+logging.getLogger('oci.circuit_breaker').setLevel(logging.WARNING)
 
 CONFIG_PATH = '~/.oci/config'
 ENV_VAR_OCI_CONFIG = 'OCI_CONFIG'
@@ -23,10 +32,16 @@ def get_config_file() -> str:
 
 def get_oci_config(region=None, profile='DEFAULT'):
     conf_file_path = get_config_file()
+    if not profile or profile == 'DEFAULT':
+        config_profile = oci_utils.oci_config.get_profile()
+    else:
+        config_profile = profile
+
     oci_config = oci.config.from_file(file_location=conf_file_path,
-                                      profile_name=profile)
+                                      profile_name=config_profile)
     if region is not None:
         oci_config['region'] = region
+
     return oci_config
 
 
@@ -47,6 +62,29 @@ def get_identity_client(region=None, profile='DEFAULT'):
     return oci.identity.IdentityClient(get_oci_config(region, profile))
 
 
+def get_object_storage_client(region=None, profile='DEFAULT'):
+    return oci.object_storage.ObjectStorageClient(
+        get_oci_config(region, profile))
+
+
 def service_exception():
     """OCI service exception."""
     return oci.exceptions.ServiceError
+
+
+def with_oci_env(f):
+
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        # pylint: disable=line-too-long
+        enter_env_cmds = [
+            'conda info --envs | grep "sky-oci-cli-env" || conda create -n sky-oci-cli-env python=3.10 -y',
+            '. $(conda info --base 2> /dev/null)/etc/profile.d/conda.sh > /dev/null 2>&1 || true',
+            'conda activate sky-oci-cli-env', 'pip install oci-cli',
+            'export OCI_CLI_SUPPRESS_FILE_PERMISSIONS_WARNING=True'
+        ]
+        operation_cmd = [f(*args, **kwargs)]
+        leave_env_cmds = ['conda deactivate']
+        return ' && '.join(enter_env_cmds + operation_cmd + leave_env_cmds)
+
+    return wrapper

sky/adaptors/vast.py

@@ -0,0 +1,29 @@
+"""Vast cloud adaptor."""
+
+import functools
+
+_vast_sdk = None
+
+
+def import_package(func):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        global _vast_sdk
+
+        if _vast_sdk is None:
+            try:
+                import vastai_sdk as _vast  # pylint: disable=import-outside-toplevel
+                _vast_sdk = _vast.VastAI()
+            except ImportError as e:
+                raise ImportError(f'Fail to import dependencies for vast: {e}\n'
+                                  'Try pip install "skypilot[vast]"') from None
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+@import_package
+def vast():
+    """Return the vast package."""
+    return _vast_sdk

sky/admin_policy.py

@@ -0,0 +1,101 @@
+"""Interface for admin-defined policy for user requests."""
+import abc
+import dataclasses
+import typing
+from typing import Optional
+
+if typing.TYPE_CHECKING:
+    import sky
+
+
+@dataclasses.dataclass
+class RequestOptions:
+    """Request options for admin policy.
+
+    Args:
+        cluster_name: Name of the cluster to create/reuse. It is None if not
+            specified by the user.
+        idle_minutes_to_autostop: Autostop setting requested by a user. The
+            cluster will be set to autostop after this many minutes of idleness.
+        down: If true, use autodown rather than autostop.
+        dryrun: Is the request a dryrun?
+    """
+    cluster_name: Optional[str]
+    idle_minutes_to_autostop: Optional[int]
+    down: bool
+    dryrun: bool
+
+
+@dataclasses.dataclass
+class UserRequest:
+    """A user request.
+
+    A "user request" is defined as a `sky launch / exec` command or its API
+    equivalent.
+
+    `sky jobs launch / serve up` involves multiple launch requests, including
+    the launch of controller and clusters for a job (which can have multiple
+    tasks if it is a pipeline) or service replicas. Each launch is a separate
+    request.
+
+    This class wraps the underlying task, the global skypilot config used to run
+    a task, and the request options.
+
+    Args:
+        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.
+    """
+    task: 'sky.Task'
+    skypilot_config: 'sky.Config'
+    request_options: Optional['RequestOptions'] = None
+
+
+@dataclasses.dataclass
+class MutatedUserRequest:
+    task: 'sky.Task'
+    skypilot_config: 'sky.Config'
+
+
+# pylint: disable=line-too-long
+class AdminPolicy:
+    """Abstract interface of an admin-defined policy for all user requests.
+
+    Admins can implement a subclass of AdminPolicy with the following signature:
+
+        import sky
+
+        class SkyPilotPolicyV1(sky.AdminPolicy):
+            def validate_and_mutate(user_request: UserRequest) -> MutatedUserRequest:
+                ...
+                return MutatedUserRequest(task=..., skypilot_config=...)
+
+    The policy can mutate both task and skypilot_config. Admins then distribute
+    a simple module that contains this implementation, installable in a way
+    that it can be imported by users from the same Python environment where
+    SkyPilot is running.
+
+    Users can register a subclass of AdminPolicy in the SkyPilot config file
+    under the key 'admin_policy', e.g.
+
+        admin_policy: my_package.SkyPilotPolicyV1
+    """
+
+    @classmethod
+    @abc.abstractmethod
+    def validate_and_mutate(cls,
+                            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')