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

sky/server/requests/serializers/decoders.py

@@ -2,7 +2,7 @@
 import base64
 import pickle
 import typing
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from sky import jobs as managed_jobs
 from sky import models
@@ -56,10 +56,10 @@ def decode_status(
     clusters = return_value
     response = []
     for cluster in clusters:
-        cluster['handle'] = decode_and_unpickle(cluster['handle'])
+        # handle may not always be present in the response.
+        if 'handle' in cluster and cluster['handle'] is not None:
+            cluster['handle'] = decode_and_unpickle(cluster['handle'])
         cluster['status'] = status_lib.ClusterStatus(cluster['status'])
-        cluster['storage_mounts_metadata'] = decode_and_unpickle(
-            cluster['storage_mounts_metadata'])
         if 'is_managed' not in cluster:
             cluster['is_managed'] = False
         response.append(responses.StatusResponse.model_validate(cluster))
@@ -72,7 +72,7 @@ def decode_status_kubernetes(
                         List[Dict[str, Any]], Optional[str]]
 ) -> Tuple[List[kubernetes_utils.KubernetesSkyPilotClusterInfoPayload],
            List[kubernetes_utils.KubernetesSkyPilotClusterInfoPayload],
-           List[Dict[str, Any]], Optional[str]]:
+           List[responses.ManagedJobRecord], Optional[str]]:
     (encoded_all_clusters, encoded_unmanaged_clusters, all_jobs,
      context) = return_value
     all_clusters = []
@@ -85,6 +85,7 @@ def decode_status_kubernetes(
         cluster['status'] = status_lib.ClusterStatus(cluster['status'])
         unmanaged_clusters.append(
             kubernetes_utils.KubernetesSkyPilotClusterInfoPayload(**cluster))
+    all_jobs = [responses.ManagedJobRecord(**job) for job in all_jobs]
     return all_clusters, unmanaged_clusters, all_jobs, context
 
 
@@ -101,29 +102,49 @@ def decode_start(return_value: str) -> 'backends.CloudVmRayResourceHandle':
 
 
 @register_decoders('queue')
-def decode_queue(return_value: List[dict],) -> List[Dict[str, Any]]:
+def decode_queue(return_value: List[dict],) -> List[responses.ClusterJobRecord]:
     jobs = return_value
     for job in jobs:
         job['status'] = job_lib.JobStatus(job['status'])
-    return jobs
+    return [responses.ClusterJobRecord.model_validate(job) for job in jobs]
 
 
 @register_decoders('jobs.queue')
-def decode_jobs_queue(return_value):
+def decode_jobs_queue(return_value: List[dict],) -> List[Dict[str, Any]]:
+    # To keep backward compatibility with v0.10.2
+    return decode_jobs_queue_v2(return_value)
+
+
+@register_decoders('jobs.queue_v2')
+def decode_jobs_queue_v2(
+    return_value
+) -> Union[Tuple[List[responses.ManagedJobRecord], int, Dict[str, int], int],
+           List[responses.ManagedJobRecord]]:
     """Decode jobs queue response.
 
-    Supports legacy list, or a dict {jobs, total}.
-    - Returns list[job]
+    Supports legacy list, or a dict {jobs, total, total_no_filter,
+    status_counts}.
+
+    - Returns either list[job] or tuple(list[job], total, status_counts,
+      total_no_filter)
     """
-    # Case 1: dict shape {jobs, total}
-    if isinstance(return_value, dict) and 'jobs' in return_value:
+    # Case 1: dict shape {jobs, total, total_no_filter, status_counts}
+    if isinstance(return_value, dict):
         jobs = return_value.get('jobs', [])
+        total = return_value.get('total', len(jobs))
+        total_no_filter = return_value.get('total_no_filter', total)
+        status_counts = return_value.get('status_counts', {})
+        for job in jobs:
+            job['status'] = managed_jobs.ManagedJobStatus(job['status'])
+        jobs = [responses.ManagedJobRecord(**job) for job in jobs]
+        return jobs, total, status_counts, total_no_filter
     else:
         # Case 2: legacy list
         jobs = return_value
-    for job in jobs:
-        job['status'] = managed_jobs.ManagedJobStatus(job['status'])
-    return jobs
+        for job in jobs:
+            job['status'] = managed_jobs.ManagedJobStatus(job['status'])
+        jobs = [responses.ManagedJobRecord(**job) for job in jobs]
+        return jobs
 
 
 def _decode_serve_status(
@@ -175,14 +196,24 @@ def decode_list_accelerators(
 
 @register_decoders('storage_ls')
 def decode_storage_ls(
-        return_value: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        return_value: List[Dict[str, Any]]) -> List[responses.StorageRecord]:
     for storage_info in return_value:
         storage_info['status'] = status_lib.StorageStatus(
             storage_info['status'])
         storage_info['store'] = [
             storage.StoreType(store) for store in storage_info['store']
         ]
-    return return_value
+    return [
+        responses.StorageRecord(**storage_info) for storage_info in return_value
+    ]
+
+
+@register_decoders('volume_list')
+def decode_volume_list(
+        return_value: List[Dict[str, Any]]) -> List[responses.VolumeRecord]:
+    return [
+        responses.VolumeRecord(**volume_info) for volume_info in return_value
+    ]
 
 
 @register_decoders('job_status')

sky/server/requests/serializers/encoders.py

@@ -6,8 +6,10 @@ import base64
 import dataclasses
 import pickle
 import typing
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
+from sky import models
+from sky.catalog import common
 from sky.schemas.api import responses
 from sky.server import constants as server_constants
 from sky.utils import serialize_utils
@@ -15,7 +17,6 @@ from sky.utils import serialize_utils
 if typing.TYPE_CHECKING:
     from sky import backends
     from sky import clouds
-    from sky import models
     from sky.provision.kubernetes import utils as kubernetes_utils
 
 handlers: Dict[str, Any] = {}
@@ -60,13 +61,23 @@ def encode_status(
         clusters: List[responses.StatusResponse]) -> List[Dict[str, Any]]:
     response = []
     for cluster in clusters:
-        response_cluster = cluster.model_dump()
+        response_cluster = cluster.model_dump(exclude_none=True)
+        # These default setting is needed because last_use and status_updated_at
+        # used to be not optional.
+        # TODO(syang): remove this after v0.12.0
+        if 'last_use' not in response_cluster:
+            response_cluster['last_use'] = ''
+        if 'status_updated_at' not in response_cluster:
+            response_cluster['status_updated_at'] = 0
         response_cluster['status'] = cluster['status'].value
         handle = serialize_utils.prepare_handle_for_backwards_compatibility(
             cluster['handle'])
         response_cluster['handle'] = pickle_and_encode(handle)
+        # TODO (syang) We still need to return this field for backwards
+        # compatibility.
+        # Remove this field at or after v0.12.0
         response_cluster['storage_mounts_metadata'] = pickle_and_encode(
-            response_cluster['storage_mounts_metadata'])
+            None)  # Always returns None.
         response.append(response_cluster)
     return response
 
@@ -92,10 +103,14 @@ def encode_start(resource_handle: 'backends.CloudVmRayResourceHandle') -> str:
 
 
 @register_encoder('queue')
-def encode_queue(jobs: List[dict],) -> List[Dict[str, Any]]:
+def encode_queue(
+    jobs: List[responses.ClusterJobRecord],) -> List[Dict[str, Any]]:
+    response = []
     for job in jobs:
-        job['status'] = job['status'].value
-    return jobs
+        response_job = job.model_dump()
+        response_job['status'] = job['status'].value
+        response.append(response_job)
+    return response
 
 
 @register_encoder('status_kubernetes')
@@ -103,7 +118,7 @@ def encode_status_kubernetes(
     return_value: Tuple[
         List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
         List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
-        List[Dict[str, Any]], Optional[str]]
+        List[responses.ManagedJobRecord], Optional[str]]
 ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]],
            Optional[str]]:
     all_clusters, unmanaged_clusters, all_jobs, context = return_value
@@ -117,13 +132,22 @@ def encode_status_kubernetes(
         encoded_cluster = dataclasses.asdict(cluster)
         encoded_cluster['status'] = encoded_cluster['status'].value
         encoded_unmanaged_clusters.append(encoded_cluster)
+    all_jobs = [job.model_dump(by_alias=True) for job in all_jobs]
     return encoded_all_clusters, encoded_unmanaged_clusters, all_jobs, context
 
 
 @register_encoder('jobs.queue')
-def encode_jobs_queue(jobs_or_tuple):
+def encode_jobs_queue(jobs: List[dict],) -> List[Dict[str, Any]]:
+    for job in jobs:
+        job['status'] = job['status'].value
+    return jobs
+
+
+@register_encoder('jobs.queue_v2')
+def encode_jobs_queue_v2(
+        jobs_or_tuple) -> Union[List[Dict[str, Any]], Dict[str, Any]]:
     # Support returning either a plain jobs list or a (jobs, total) tuple
-    status_counts = {}
+    status_counts: Dict[str, int] = {}
     if isinstance(jobs_or_tuple, tuple):
         if len(jobs_or_tuple) == 2:
             jobs, total = jobs_or_tuple
@@ -135,12 +159,13 @@ def encode_jobs_queue(jobs_or_tuple):
     else:
         jobs = jobs_or_tuple
         total = None
-    for job in jobs:
+    jobs_dict = [job.model_dump(by_alias=True) for job in jobs]
+    for job in jobs_dict:
         job['status'] = job['status'].value
     if total is None:
-        return jobs
+        return jobs_dict
     return {
-        'jobs': jobs,
+        'jobs': jobs_dict,
         'total': total,
         'total_no_filter': total_no_filter,
         'status_counts': status_counts
@@ -177,8 +202,9 @@ def encode_cost_report(
     for cluster_report in cost_report:
         if cluster_report['status'] is not None:
             cluster_report['status'] = cluster_report['status'].value
-        cluster_report['resources'] = pickle_and_encode(
-            cluster_report['resources'])
+        if 'resources' in cluster_report:
+            cluster_report['resources'] = pickle_and_encode(
+                cluster_report['resources'])
     return cost_report
 
 
@@ -190,19 +216,26 @@ def encode_enabled_clouds(clouds: List['clouds.Cloud']) -> List[str]:
 
 @register_encoder('storage_ls')
 def encode_storage_ls(
-        return_value: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-    for storage_info in return_value:
+        return_value: List[responses.StorageRecord]) -> List[Dict[str, Any]]:
+    response_list = [storage_info.model_dump() for storage_info in return_value]
+    for storage_info in response_list:
         storage_info['status'] = storage_info['status'].value
         storage_info['store'] = [store.value for store in storage_info['store']]
-    return return_value
+    return response_list
+
+
+@register_encoder('volume_list')
+def encode_volume_list(
+        return_value: List[responses.VolumeRecord]) -> List[Dict[str, Any]]:
+    return [volume_info.model_dump() for volume_info in return_value]
 
 
 @register_encoder('job_status')
-def encode_job_status(return_value: Dict[int, Any]) -> Dict[int, str]:
+def encode_job_status(return_value: Dict[int, Any]) -> Dict[str, str]:
     for job_id in return_value.keys():
         if return_value[job_id] is not None:
             return_value[job_id] = return_value[job_id].value
-    return return_value
+    return {str(k): v for k, v in return_value.items()}
 
 
 @register_encoder('kubernetes_node_info')
@@ -214,3 +247,35 @@ def encode_kubernetes_node_info(
 @register_encoder('endpoints')
 def encode_endpoints(return_value: Dict[int, str]) -> Dict[str, str]:
     return {str(k): v for k, v in return_value.items()}
+
+
+@register_encoder('realtime_kubernetes_gpu_availability')
+def encode_realtime_gpu_availability(
+    return_value: List[Tuple[str,
+                             List[Any]]]) -> List[Tuple[str, List[List[Any]]]]:
+    # Convert RealtimeGpuAvailability namedtuples to lists
+    # for JSON serialization.
+    encoded = []
+    for context, gpu_list in return_value:
+        converted_gpu_list = []
+        for gpu in gpu_list:
+            assert isinstance(gpu, models.RealtimeGpuAvailability), (
+                f'Expected RealtimeGpuAvailability, got {type(gpu)}')
+            converted_gpu_list.append(list(gpu))
+        encoded.append((context, converted_gpu_list))
+    return encoded
+
+
+@register_encoder('list_accelerators')
+def encode_list_accelerators(
+        return_value: Dict[str, List[Any]]) -> Dict[str, Any]:
+    encoded: Dict[str, Any] = {}
+    for accelerator_name, instances in return_value.items():
+        # Convert InstanceTypeInfo namedtuples to lists for JSON serialization.
+        converted_instances: List[Any] = []
+        for instance in instances:
+            assert isinstance(instance, common.InstanceTypeInfo), (
+                f'Expected InstanceTypeInfo, got {type(instance)}')
+            converted_instances.append(list(instance))
+        encoded[accelerator_name] = converted_instances
+    return encoded

sky/server/requests/threads.py

@@ -0,0 +1,117 @@
+"""Request execution threads management."""
+
+import concurrent.futures
+import sys
+import threading
+from typing import Callable, Set, TypeVar
+
+from sky import exceptions
+from sky import sky_logging
+from sky.utils import atomic
+
+# pylint: disable=ungrouped-imports
+if sys.version_info >= (3, 10):
+    from typing import ParamSpec
+else:
+    from typing_extensions import ParamSpec
+
+_P = ParamSpec('_P')
+_T = TypeVar('_T')
+
+logger = sky_logging.init_logger(__name__)
+
+
+class OnDemandThreadExecutor(concurrent.futures.Executor):
+    """An executor that creates a new thread for each task and destroys it
+    after the task is completed.
+
+    Note(dev):
+    We raise an error instead of queuing the request if the limit is reached, so
+    that:
+    1. the request might be handled by other processes that have idle workers
+       upon retry;
+    2. if not, then users can be clearly hinted that they need to scale the API
+       server to support higher concurrency.
+    So this executor is only suitable for carefully selected cases where the
+    error can be properly handled by caller. To make this executor general, we
+    need to support configuring the queuing behavior (exception or queueing).
+    """
+
+    def __init__(self, name: str, max_workers: int):
+        self.name: str = name
+        self.max_workers: int = max_workers
+        self.running: atomic.AtomicInt = atomic.AtomicInt(0)
+        self._shutdown: bool = False
+        self._shutdown_lock: threading.Lock = threading.Lock()
+        self._threads: Set[threading.Thread] = set()
+        self._threads_lock: threading.Lock = threading.Lock()
+
+    def _cleanup_thread(self, thread: threading.Thread):
+        with self._threads_lock:
+            self._threads.discard(thread)
+
+    def _task_wrapper(self, fn: Callable, fut: concurrent.futures.Future, /,
+                      *args, **kwargs):
+        try:
+            result = fn(*args, **kwargs)
+            fut.set_result(result)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.debug(f'Executor [{self.name}] error executing {fn}: {e}')
+            fut.set_exception(e)
+        finally:
+            self.running.decrement()
+            self._cleanup_thread(threading.current_thread())
+
+    def check_available(self, borrow: bool = False) -> int:
+        """Check if there are available workers.
+
+        Args:
+            borrow: If True, the caller borrow a worker from the executor.
+                The caller is responsible for returning the worker to the
+                executor after the task is completed.
+        """
+        count = self.running.increment()
+        if count > self.max_workers:
+            self.running.decrement()
+            raise exceptions.ConcurrentWorkerExhaustedError(
+                f'Maximum concurrent workers {self.max_workers} of threads '
+                f'executor [{self.name}] reached')
+        if not borrow:
+            self.running.decrement()
+        return count
+
+    def submit(self, fn: Callable[_P, _T], *args: _P.args,
+               **kwargs: _P.kwargs) -> 'concurrent.futures.Future[_T]':
+        with self._shutdown_lock:
+            if self._shutdown:
+                raise RuntimeError(
+                    'Cannot submit task after executor is shutdown')
+            count = self.check_available(borrow=True)
+            fut: concurrent.futures.Future = concurrent.futures.Future()
+            # Name is assigned for debugging purpose, duplication is fine
+            thread = threading.Thread(target=self._task_wrapper,
+                                      name=f'{self.name}-{count}',
+                                      args=(fn, fut, *args),
+                                      kwargs=kwargs,
+                                      daemon=True)
+            with self._threads_lock:
+                self._threads.add(thread)
+            try:
+                thread.start()
+            except Exception as e:
+                self.running.decrement()
+                self._cleanup_thread(thread)
+                fut.set_exception(e)
+                raise
+            assert thread.ident is not None, 'Thread should be started'
+            return fut
+
+    def shutdown(self, wait=True):
+        with self._shutdown_lock:
+            self._shutdown = True
+        if not wait:
+            return
+        with self._threads_lock:
+            threads = list(self._threads)
+        for t in threads:
+            t.join()

sky/server/rest.py

@@ -4,11 +4,14 @@ import asyncio
 import contextlib
 import contextvars
 import functools
+import html
+import re
 import time
 import typing
 from typing import Any, Callable, cast, Optional, TypeVar
 
 import colorama
+import urllib3.exceptions
 
 from sky import exceptions
 from sky import sky_logging
@@ -31,7 +34,15 @@ else:
 
 F = TypeVar('F', bound=Callable[..., Any])
 
-_RETRY_CONTEXT = contextvars.ContextVar('retry_context', default=None)
+
+class RetryContext:
+
+    def __init__(self):
+        self.line_processed = 0
+
+
+_RETRY_CONTEXT: contextvars.ContextVar[Optional[RetryContext]] = (
+    contextvars.ContextVar('retry_context', default=None))
 
 _session = requests.Session()
 # Tune connection pool size, otherwise the default max is just 10.
@@ -53,13 +64,11 @@ _session.headers[constants.VERSION_HEADER] = (
 _transient_errors = [
     requests.exceptions.RequestException,
     ConnectionError,
+    urllib3.exceptions.HTTPError,
 ]
 
-
-class RetryContext:
-
-    def __init__(self):
-        self.line_processed = 0
+_HTML_TITLE_RE = re.compile(r'<title[^>]*>(.*?)</title>',
+                            re.IGNORECASE | re.DOTALL)
 
 
 @contextlib.contextmanager
@@ -176,14 +185,16 @@ def _retry_on_server_unavailable(max_wait_seconds: int = 600,
     Notes(dev):
     """
 
+    def _readable_error_msg(message: str) -> str:
+        return (f'{colorama.Fore.YELLOW}API server is temporarily '
+                f'unavailable: {message}.\nRetrying...'
+                f'{colorama.Style.RESET_ALL}')
+
     def decorator(func: F) -> F:
 
         @functools.wraps(func)
         def wrapper(*args, **kwargs) -> Any:
-            msg = (
-                f'{colorama.Fore.YELLOW}API server is temporarily unavailable: '
-                'upgrade in progress. Waiting to resume...'
-                f'{colorama.Style.RESET_ALL}')
+
             backoff = common_utils.Backoff(
                 initial_backoff=initial_backoff,
                 max_backoff_factor=max_backoff_factor)
@@ -201,7 +212,8 @@ def _retry_on_server_unavailable(max_wait_seconds: int = 600,
                         # stop the status spinner before retrying func() to
                         # avoid the status spinner get stuck if the func() runs
                         # for a long time without update status, e.g. sky logs.
-                        with rich_utils.client_status(msg):
+                        with rich_utils.client_status(
+                                _readable_error_msg(e.message)):
                             if time.time() - start_time > max_wait_seconds:
                                 # pylint: disable=line-too-long
                                 raise exceptions.ServerTemporarilyUnavailableError(
@@ -222,14 +234,98 @@ def _retry_on_server_unavailable(max_wait_seconds: int = 600,
 
 
 def handle_server_unavailable(response: 'requests.Response') -> None:
-    if response.status_code == 503:
-        # TODO(aylei): Hacky, depends on how nginx controller handles backends
-        # with no ready endpoints. Should use self-defined status code or header
-        # to distinguish retryable server error from general 503 errors.
-        with ux_utils.print_exception_no_traceback():
-            raise exceptions.ServerTemporarilyUnavailableError(
-                'SkyPilot API server is temporarily unavailable. '
-                'Please try again later.')
+    """Handle 503 (Service Unavailable) error
+
+    The client get 503 error in the following cases:
+    1. The reverse proxy cannot find any ready backend endpoints to serve the
+       request, e.g. when there is and rolling-update.
+    2. The skypilot API server has temporary resource issue, e.g. when the
+       cucurrency of the handling process is exhausted.
+
+    We expect the caller (CLI or SDK) retry on these cases and show clear wait
+    message to the user to let user decide whether keep waiting or abort the
+    request.
+    """
+    if response.status_code != 503:
+        return
+
+    # error_msg = 'SkyPilot API server is temporarily unavailable. '
+    error_msg = ''
+    try:
+        response_data = response.json()
+        if 'detail' in response_data:
+            error_msg = response_data['detail']
+    except Exception:  # pylint: disable=broad-except
+        error_msg = handle_response_text(response)
+
+    with ux_utils.print_exception_no_traceback():
+        raise exceptions.ServerTemporarilyUnavailableError(error_msg)
+
+
+def handle_response_text(response: 'requests.Response') -> str:
+    """Handle the plaintext response to get the error message
+
+    There is a special handling for html content which might be returned
+    by the reverse proxy to make the error message more user-friendly.
+    """
+    error_msg = ''
+    if isinstance(response, str):
+        text, headers = response, {}
+    else:
+        text = getattr(response, 'text', '')
+        headers = getattr(response, 'headers', {}) or {}
+    if not isinstance(text, str):
+        text = str(text) if text is not None else ''
+    if not text:
+        return ''
+    content_type = headers.get('Content-Type', '')
+    is_html = isinstance(content_type, str) and 'html' in (content_type.lower())
+    if not is_html:
+        stripped = text.lstrip()
+        is_html = stripped.startswith('<') and '<title' in stripped.lower()
+    if is_html:
+        match = _HTML_TITLE_RE.search(text)
+        if match:
+            title = html.unescape(match.group(1)).strip()
+            if title:
+                error_msg = title
+    if not error_msg:
+        error_msg = text
+    return error_msg
+
+
+async def handle_server_unavailable_async(
+        response: 'aiohttp.ClientResponse') -> None:
+    """Async version: Handle 503 (Service Unavailable) error
+
+    The client get 503 error in the following cases:
+    1. The reverse proxy cannot find any ready backend endpoints to serve the
+       request, e.g. when there is and rolling-update.
+    2. The skypilot API server has temporary resource issue, e.g. when the
+       cucurrency of the handling process is exhausted.
+
+    We expect the caller (CLI or SDK) retry on these cases and show clear wait
+    message to the user to let user decide whether keep waiting or abort the
+    request.
+    """
+    if response.status != 503:
+        return
+
+    error_msg = ''
+    try:
+        response_data = await response.json()
+        if 'detail' in response_data:
+            error_msg = response_data['detail']
+    except Exception:  # pylint: disable=broad-except
+        try:
+            text = await response.text()
+            if text:
+                error_msg = text
+        except Exception:  # pylint: disable=broad-except
+            pass
+
+    with ux_utils.print_exception_no_traceback():
+        raise exceptions.ServerTemporarilyUnavailableError(error_msg)
 
 
 @_retry_on_server_unavailable()
@@ -308,11 +404,7 @@ async def request_without_retry_async(session: 'aiohttp.ClientSession',
         response = await session.request(method, url, **kwargs)
 
         # Handle server unavailability (503 status) - same as sync version
-        if response.status == 503:
-            with ux_utils.print_exception_no_traceback():
-                raise exceptions.ServerTemporarilyUnavailableError(
-                    'SkyPilot API server is temporarily unavailable. '
-                    'Please try again later.')
+        await handle_server_unavailable_async(response)
 
         # Set remote API version and version from headers - same as sync version
         remote_api_version = response.headers.get(constants.API_VERSION_HEADER)