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

sky/adaptors/kubernetes.py

@@ -1,10 +1,12 @@
 """Kubernetes adaptors"""
+import functools
 import logging
 import os
+import platform
 from typing import Any, Callable, Optional, Set
 
+from sky import sky_logging
 from sky.adaptors import common
-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
@@ -13,12 +15,23 @@ _IMPORT_ERROR_MESSAGE = ('Failed to import dependencies for Kubernetes. '
                          'Try running: pip install "skypilot[kubernetes]"')
 kubernetes = common.LazyImport('kubernetes',
                                import_error_message=_IMPORT_ERROR_MESSAGE)
+models = common.LazyImport('kubernetes.client.models',
+                           import_error_message=_IMPORT_ERROR_MESSAGE)
 urllib3 = common.LazyImport('urllib3',
                             import_error_message=_IMPORT_ERROR_MESSAGE)
+dateutil_parser = common.LazyImport('dateutil.parser',
+                                    import_error_message=_IMPORT_ERROR_MESSAGE)
 
 # Timeout to use for API calls
 API_TIMEOUT = 5
 
+# Check if KUBECONFIG is set, and use it if it is.
+DEFAULT_KUBECONFIG_PATH = '~/.kube/config'
+# From kubernetes package, keep a copy here to avoid actually importing
+# kubernetes package when parsing the KUBECONFIG env var to do credential
+# file mounts.
+ENV_KUBECONFIG_PATH_SEPARATOR = ';' if platform.system() == 'Windows' else ':'
+
 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
@@ -26,6 +39,8 @@ DEFAULT_IN_CLUSTER_REGION = 'in-cluster'
 # set to DEFAULT_IN_CLUSTER_REGION.
 IN_CLUSTER_CONTEXT_NAME_ENV_VAR = 'SKYPILOT_IN_CLUSTER_CONTEXT_NAME'
 
+logger = sky_logging.init_logger(__name__)
+
 
 def _decorate_methods(obj: Any, decorator: Callable, decoration_type: str):
     for attr_name in dir(obj):
@@ -43,7 +58,7 @@ def _decorate_methods(obj: Any, decorator: Callable, decoration_type: str):
     return obj
 
 
-def _api_logging_decorator(logger: str, level: int):
+def _api_logging_decorator(logger_src: 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
@@ -54,7 +69,9 @@ def _api_logging_decorator(logger: str, level: int):
 
         def wrapped(*args, **kwargs):
             obj = api(*args, **kwargs)
-            _decorate_methods(obj, set_logging_level(logger, level), 'api_log')
+            _decorate_methods(obj,
+                              sky_logging.set_logging_level(logger_src, level),
+                              'api_log')
             return obj
 
         return wrapped
@@ -62,31 +79,61 @@ def _api_logging_decorator(logger: str, level: int):
     return decorated_api
 
 
+def _get_config_file() -> str:
+    # Kubernetes load the kubeconfig from the KUBECONFIG env var on
+    # package initialization. So we have to reload the KUBECOFNIG env var
+    # everytime in case the KUBECONFIG env var is changed.
+    return os.environ.get('KUBECONFIG', '~/.kube/config')
+
+
 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(context=context)
+            kubernetes.config.load_kube_config(config_file=_get_config_file(),
+                                               context=context)
         except kubernetes.config.config_exception.ConfigException as e:
             suffix = common_utils.format_exception(e, use_bracket=True)
             context_name = '(current-context)' if context is None else context
+            is_ssh_node_pool = False
+            if context_name.startswith('ssh-'):
+                context_name = common_utils.removeprefix(context_name, 'ssh-')
+                is_ssh_node_pool = 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 for '
-                           f'{context_name!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.')
+                if is_ssh_node_pool:
+                    context_name = common_utils.removeprefix(
+                        context_name, 'ssh-')
+                    err_str = ('Failed to load SSH Node Pool configuration for '
+                               f'{context_name!r}.\n'
+                               '    Run `sky ssh up --infra {context_name}` to '
+                               'set up or repair the cluster.')
+                else:
+                    err_str = (
+                        'Failed to load Kubernetes configuration for '
+                        f'{context_name!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:
                 kubeconfig_path = os.environ.get('KUBECONFIG', '~/.kube/config')
-                err_str = (
-                    f'Failed to load Kubernetes configuration for '
-                    f'{context_name!r}. Please check if your kubeconfig file '
-                    f'exists at {kubeconfig_path} and is valid.\n{suffix}')
-            err_str += '\nTo disable Kubernetes for SkyPilot: run `sky check`.'
-            if context is None:  # kubernetes defaults to current-context.
+                if is_ssh_node_pool:
+                    err_str = (
+                        f'Failed to load SSH Node Pool configuration for '
+                        f'{context_name!r}. Run `sky ssh up --infra '
+                        f'{context_name}` to set up or repair the cluster.')
+                else:
+                    err_str = (
+                        'Failed to load Kubernetes configuration for '
+                        f'{context_name!r}. Please check if your kubeconfig '
+                        f'file exists at {kubeconfig_path} and is valid.'
+                        f'\n{suffix}\n')
+            if is_ssh_node_pool:
+                err_str += (f'\nTo disable SSH Node Pool {context_name!r}: '
+                            'run `sky check`.')
+            else:
                 err_str += (
                     '\nHint: Kubernetes attempted to query the current-context '
                     'set in kubeconfig. Check if the current-context is valid.')
@@ -100,8 +147,11 @@ def _load_config(context: Optional[str] = None):
             # 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'
+            # Only set if not already present (preserving existing values)
+            if 'KUBERNETES_SERVICE_HOST' not in os.environ:
+                os.environ['KUBERNETES_SERVICE_HOST'] = 'kubernetes.default.svc'
+            if 'KUBERNETES_SERVICE_PORT' not in os.environ:
+                os.environ['KUBERNETES_SERVICE_PORT'] = '443'
             kubernetes.config.load_incluster_config()
         except kubernetes.config.config_exception.ConfigException:
             _load_config_from_kubeconfig()
@@ -109,8 +159,65 @@ def _load_config(context: Optional[str] = None):
         _load_config_from_kubeconfig(context)
 
 
+def list_kube_config_contexts():
+    return kubernetes.config.list_kube_config_contexts(_get_config_file())
+
+
+class ClientWrapper:
+    """Wrapper around the kubernetes API clients.
+
+    This is needed because we cache kubernetes.client.ApiClient and other typed
+    clients (e.g. kubernetes.client.CoreV1Api) and lru_cache.cache_clear() does
+    not call close() on the client to cleanup external resources like
+    semaphores. This decorator wraps the client with __del__ to ensure the
+    external state of kubernetes clients are properly cleaned up on GC.
+    """
+
+    def __init__(self, client):
+        self._client = client
+
+    def __getattr__(self, name):
+        """Delegate to the underlying client"""
+        return getattr(self._client, name)
+
+    def __del__(self):
+        """Clean up the underlying client"""
+        try:
+            real_client = None
+            if isinstance(self._client, kubernetes.client.ApiClient):
+                real_client = self._client
+            elif isinstance(self._client, kubernetes.watch.Watch):
+                real_client = getattr(self._client, '_api_client', None)
+            else:
+                # Otherwise, the client is a typed client, the typed client
+                # is generated by codegen and all of them should have an
+                # 'api_client' attribute referring to the real client.
+                real_client = getattr(self._client, 'api_client', None)
+            if real_client is not None:
+                real_client.close()
+            else:
+                # logger may already be cleaned up during __del__ at shutdown
+                if logger is not None:
+                    logger.debug(f'No client found for {self._client}')
+        except Exception as e:  # pylint: disable=broad-except
+            if logger is not None:
+                logger.debug(f'Error closing Kubernetes client: {e}')
+
+
+def wrap_kubernetes_client(func):
+    """Wraps kubernetes API clients for proper cleanup."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        obj = func(*args, **kwargs)
+        return ClientWrapper(obj)
+
+    return wrapper
+
+
 @_api_logging_decorator('urllib3', logging.ERROR)
 @annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
 def core_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.CoreV1Api()
@@ -118,6 +225,15 @@ def core_api(context: Optional[str] = None):
 
 @_api_logging_decorator('urllib3', logging.ERROR)
 @annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
+def storage_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.StorageV1Api()
+
+
+@_api_logging_decorator('urllib3', logging.ERROR)
+@annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
 def auth_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.RbacAuthorizationV1Api()
@@ -125,6 +241,7 @@ def auth_api(context: Optional[str] = None):
 
 @_api_logging_decorator('urllib3', logging.ERROR)
 @annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
 def networking_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.NetworkingV1Api()
@@ -132,6 +249,7 @@ def networking_api(context: Optional[str] = None):
 
 @_api_logging_decorator('urllib3', logging.ERROR)
 @annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
 def custom_objects_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.CustomObjectsApi()
@@ -139,6 +257,7 @@ def custom_objects_api(context: Optional[str] = None):
 
 @_api_logging_decorator('urllib3', logging.ERROR)
 @annotations.lru_cache(scope='global')
+@wrap_kubernetes_client
 def node_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.NodeV1Api()
@@ -146,6 +265,7 @@ def node_api(context: Optional[str] = None):
 
 @_api_logging_decorator('urllib3', logging.ERROR)
 @annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
 def apps_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.AppsV1Api()
@@ -153,6 +273,7 @@ def apps_api(context: Optional[str] = None):
 
 @_api_logging_decorator('urllib3', logging.ERROR)
 @annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
 def batch_api(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.BatchV1Api()
@@ -160,6 +281,7 @@ def batch_api(context: Optional[str] = None):
 
 @_api_logging_decorator('urllib3', logging.ERROR)
 @annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
 def api_client(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.client.ApiClient()
@@ -167,6 +289,15 @@ def api_client(context: Optional[str] = None):
 
 @_api_logging_decorator('urllib3', logging.ERROR)
 @annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
+def custom_resources_api(context: Optional[str] = None):
+    _load_config(context)
+    return kubernetes.client.CustomObjectsApi()
+
+
+@_api_logging_decorator('urllib3', logging.ERROR)
+@annotations.lru_cache(scope='request')
+@wrap_kubernetes_client
 def watch(context: Optional[str] = None):
     _load_config(context)
     return kubernetes.watch.Watch()

sky/adaptors/nebius.py

@@ -1,17 +1,106 @@
 """Nebius cloud adaptor."""
+import asyncio
 import os
 import threading
+from typing import Any, Awaitable, List, Optional
 
+from sky import sky_logging
+from sky import skypilot_config
 from sky.adaptors import common
 from sky.utils import annotations
 from sky.utils import ux_utils
 
-NEBIUS_TENANT_ID_FILENAME = 'NEBIUS_TENANT_ID.txt'
-NEBIUS_IAM_TOKEN_FILENAME = 'NEBIUS_IAM_TOKEN.txt'
-NEBIUS_CREDENTIALS_FILENAME = 'credentials.json'
-NEBIUS_TENANT_ID_PATH = '~/.nebius/' + NEBIUS_TENANT_ID_FILENAME
-NEBIUS_IAM_TOKEN_PATH = '~/.nebius/' + NEBIUS_IAM_TOKEN_FILENAME
-NEBIUS_CREDENTIALS_PATH = '~/.nebius/' + NEBIUS_CREDENTIALS_FILENAME
+# Default read timeout for nebius SDK
+READ_TIMEOUT = 10
+
+logger = sky_logging.init_logger(__name__)
+
+_loop_lock = threading.Lock()
+_loop = None
+
+
+def _get_event_loop() -> asyncio.AbstractEventLoop:
+    """Get event loop for nebius sdk."""
+    global _loop
+
+    if _loop is not None:
+        return _loop
+
+    with _loop_lock:
+        if _loop is None:
+            # Create a new event loop in a dedicated thread
+            _loop = asyncio.new_event_loop()
+            threading.Thread(target=_loop.run_forever, daemon=True).start()
+
+        return _loop
+
+
+def sync_call(awaitable: Awaitable[Any]) -> Any:
+    """Synchronously run an awaitable in coroutine.
+
+    This wrapper is used to workaround:
+    https://github.com/nebius/pysdk/issues/76
+
+    Uses a dedicated background event loop to avoid conflicts
+    with existing asyncio contexts and prevent BlockingIOError.
+    """
+    loop = _get_event_loop()
+    future = asyncio.run_coroutine_threadsafe(_coro(awaitable), loop)
+    return future.result()
+
+
+async def _coro(awaitable: Awaitable[Any]) -> Any:
+    """Wrapper coroutine for awaitable."""
+    return await awaitable
+
+
+def tenant_id_path() -> str:
+    return '~/.nebius/NEBIUS_TENANT_ID.txt'
+
+
+def iam_token_path() -> str:
+    return '~/.nebius/NEBIUS_IAM_TOKEN.txt'
+
+
+def domain_path() -> str:
+    return '~/.nebius/NEBIUS_DOMAIN.txt'
+
+
+def credentials_path() -> str:
+    workspace_path = skypilot_config.get_workspace_cloud('nebius').get(
+        'credentials_file_path', None)
+    if workspace_path is not None:
+        return workspace_path
+    return _get_default_credentials_path()
+
+
+def _get_workspace_credentials_path() -> Optional[str]:
+    """Get credentials path if explicitly set in workspace config."""
+    workspace_cred_path = skypilot_config.get_workspace_cloud('nebius').get(
+        'credentials_file_path', None)
+    return workspace_cred_path
+
+
+def _get_default_credentials_path() -> str:
+    """Get the default credentials path."""
+    return '~/.nebius/credentials.json'
+
+
+def api_domain() -> Optional[str]:
+    domain_in_ws_config = skypilot_config.get_workspace_cloud('nebius').get(
+        'domain', None)
+    if domain_in_ws_config is not None:
+        return domain_in_ws_config
+    domain_in_config = skypilot_config.get_effective_region_config(
+        cloud='nebius', region=None, keys=('domain',), default_value=None)
+    if domain_in_config is not None:
+        return domain_in_config
+    try:
+        with open(os.path.expanduser(domain_path()), encoding='utf-8') as file:
+            return file.read().strip()
+    except FileNotFoundError:
+        return None
+
 
 DEFAULT_REGION = 'eu-north1'
 
@@ -62,6 +151,12 @@ def iam():
     return iam_v1
 
 
+def billing():
+    # pylint: disable=import-outside-toplevel
+    from nebius.api.nebius.billing import v1alpha1 as billing_v1alpha1
+    return billing_v1alpha1
+
+
 def nebius_common():
     # pylint: disable=import-outside-toplevel
     from nebius.api.nebius.common import v1 as common_v1
@@ -74,39 +169,79 @@ def vpc():
     return vpc_v1
 
 
-@annotations.lru_cache(scope='request')
 def get_iam_token():
     try:
-        with open(os.path.expanduser(NEBIUS_IAM_TOKEN_PATH),
+        with open(os.path.expanduser(iam_token_path()),
                   encoding='utf-8') as file:
             return file.read().strip()
     except FileNotFoundError:
         return None
 
 
-@annotations.lru_cache(scope='request')
 def is_token_or_cred_file_exist():
-    return (os.path.exists(os.path.expanduser(NEBIUS_IAM_TOKEN_PATH)) or
-            os.path.exists(os.path.expanduser(NEBIUS_CREDENTIALS_PATH)))
+    return (os.path.exists(os.path.expanduser(iam_token_path())) or
+            os.path.exists(os.path.expanduser(credentials_path())))
 
 
-@annotations.lru_cache(scope='request')
 def get_tenant_id():
+    tenant_id_in_ws_config = skypilot_config.get_workspace_cloud('nebius').get(
+        'tenant_id', None)
+    if tenant_id_in_ws_config is not None:
+        return tenant_id_in_ws_config
+    tenant_id_in_config = skypilot_config.get_effective_region_config(
+        cloud='nebius', region=None, keys=('tenant_id',), default_value=None)
+    if tenant_id_in_config is not None:
+        return tenant_id_in_config
     try:
-        with open(os.path.expanduser(NEBIUS_TENANT_ID_PATH),
+        with open(os.path.expanduser(tenant_id_path()),
                   encoding='utf-8') as file:
             return file.read().strip()
     except FileNotFoundError:
         return None
 
 
-@annotations.lru_cache(scope='request')
 def sdk():
+    """Create the Nebius SDK with the correct credentials.
+
+    The order of priority is:
+    1. Credentials file specified in workspace config, if set
+    2. IAM token file, if set
+    3. Default credentials path
+    """
+    # 1. Check if credentials path is set in workspace config (highest priority)
+    workspace_cred_path = _get_workspace_credentials_path()
+    if workspace_cred_path is not None:
+        # Check if token is also available and warn
+        token = get_iam_token()
+        if token is not None:
+            logger.warning(
+                f'Both workspace credentials file ({workspace_cred_path}) and '
+                f'IAM token file ({iam_token_path()}) are available. Using '
+                'workspace credentials file.')
+        return _sdk(None, workspace_cred_path)
+
+    # 2. Check for IAM token file (second priority)
     token = get_iam_token()
     if token is not None:
-        return nebius.sdk.SDK(credentials=token)
-    return nebius.sdk.SDK(
-        credentials_file_name=os.path.expanduser(NEBIUS_CREDENTIALS_PATH))
+        return _sdk(token, None)
+
+    # 3. Fall back to default credentials path (lowest priority)
+    default_cred_path = _get_default_credentials_path()
+    return _sdk(None, default_cred_path)
+
+
+@annotations.lru_cache(scope='request')
+def _sdk(token: Optional[str], cred_path: Optional[str]):
+    # Exactly one of token or cred_path must be provided
+    assert (token is None) != (cred_path is None), (token, cred_path)
+    if token is not None:
+        return nebius.sdk.SDK(credentials=token, domain=api_domain())
+    if cred_path is not None:
+        return nebius.sdk.SDK(
+            credentials_file_name=os.path.expanduser(cred_path),
+            domain=api_domain(),
+        )
+    raise ValueError('Either token or credentials file path must be provided')
 
 
 def get_nebius_credentials(boto3_session):
@@ -184,3 +319,21 @@ def botocore_exceptions():
     # pylint: disable=import-outside-toplevel
     from botocore import exceptions
     return exceptions
+
+
+def get_credential_file_paths() -> List[str]:
+    """Get the list of credential file paths based on current configuration."""
+    paths = {
+        # Always include tenant ID and IAM token paths
+        tenant_id_path(),
+        iam_token_path(),
+    }
+
+    # Add workspace-specific credentials path if set
+    workspace_cred_path = _get_workspace_credentials_path()
+    if workspace_cred_path is not None:
+        paths.add(workspace_cred_path)
+    # Always add default path in case it's needed for fallback
+    paths.add(_get_default_credentials_path())
+
+    return list(paths)

sky/adaptors/primeintellect.py

@@ -0,0 +1 @@
+"""Prime Intellect cloud adaptor."""

sky/adaptors/runpod.py

@@ -1,8 +1,76 @@
 """RunPod cloud adaptor."""
 
+import os
+import time
+from typing import Any, Dict, Optional
+
 from sky.adaptors import common
 
 runpod = common.LazyImport(
     'runpod',
     import_error_message='Failed to import dependencies for RunPod. '
     'Try running: pip install "skypilot[runpod]"')
+
+# Lazy imports
+requests = common.LazyImport('requests')
+
+_REST_BASE = 'https://rest.runpod.io/v1'
+_MAX_RETRIES = 3
+_TIMEOUT = 10
+
+
+def _get_api_key() -> str:
+    api_key = getattr(runpod, 'api_key', None)
+    if not api_key:
+        # Fallback to env if SDK global not set
+        api_key = os.environ.get('RUNPOD_API_KEY')
+    if not api_key:
+        raise RuntimeError(
+            'RunPod API key is not set. Please set runpod.api_key '
+            'or RUNPOD_API_KEY.')
+    return str(api_key)
+
+
+def rest_request(method: str,
+                 path: str,
+                 json: Optional[Dict[str, Any]] = None) -> Any:
+    url = f'{_REST_BASE}{path}'
+    headers = {
+        'Authorization': f'Bearer {_get_api_key()}',
+        'Content-Type': 'application/json',
+    }
+    attempt = 0
+    while True:
+        attempt += 1
+        try:
+            resp = requests.request(method,
+                                    url,
+                                    headers=headers,
+                                    json=json,
+                                    timeout=_TIMEOUT)
+        except Exception as e:  # pylint: disable=broad-except
+            # Retry on transient network errors
+            if attempt >= _MAX_RETRIES:
+                raise RuntimeError(f'RunPod REST network error: {e}') from e
+            time.sleep(1)
+            continue
+
+        # Retry on 5xx and 429
+        if resp.status_code >= 500 or resp.status_code == 429:
+            if attempt >= _MAX_RETRIES:
+                raise RuntimeError(
+                    f'RunPod REST error {resp.status_code}: {resp.text}')
+            time.sleep(1)
+            continue
+
+        if resp.status_code >= 400:
+            # Non-retryable client error
+            raise RuntimeError(
+                f'RunPod REST error {resp.status_code}: {resp.text}')
+
+        if resp.text:
+            try:
+                return resp.json()
+            except Exception:  # pylint: disable=broad-except
+                return resp.text
+        return None