skypilot-nightly 1.0.0.dev20251009__py3-none-any.whl → 1.0.0.dev20251107__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))
@@ -116,22 +116,35 @@ def decode_jobs_queue(return_value: List[dict],) -> List[Dict[str, Any]]:
 
 
 @register_decoders('jobs.queue_v2')
-def decode_jobs_queue_v2(return_value) -> List[responses.ManagedJobRecord]:
+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'])
-    jobs = [responses.ManagedJobRecord(**job) for job in jobs]
-    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(

sky/server/requests/serializers/encoders.py

@@ -8,6 +8,8 @@ import pickle
 import typing
 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.10.7 or v0.11.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.10.7 or v0.11.0
         response_cluster['storage_mounts_metadata'] = pickle_and_encode(
-            response_cluster['storage_mounts_metadata'])
+            None)  # Always returns None.
         response.append(response_cluster)
     return response
 
@@ -121,7 +132,7 @@ 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() for job in all_jobs]
+    all_jobs = [job.model_dump(by_alias=True) for job in all_jobs]
     return encoded_all_clusters, encoded_unmanaged_clusters, all_jobs, context
 
 
@@ -148,12 +159,13 @@ def encode_jobs_queue_v2(
     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 [job.model_dump() for job in jobs]
+        return jobs_dict
     return {
-        'jobs': [job.model_dump() for job in jobs],
+        'jobs': jobs_dict,
         'total': total,
         'total_no_filter': total_no_filter,
         'status_counts': status_counts
@@ -205,10 +217,11 @@ def encode_enabled_clouds(clouds: List['clouds.Cloud']) -> List[str]:
 @register_encoder('storage_ls')
 def encode_storage_ls(
         return_value: List[responses.StorageRecord]) -> List[Dict[str, Any]]:
-    for storage_info in return_value:
+    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 [storage_info.model_dump() for storage_info in return_value]
+    return response_list
 
 
 @register_encoder('volume_list')
@@ -218,11 +231,11 @@ def encode_volume_list(
 
 
 @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')
@@ -234,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,106 @@
+"""Request execution threads management."""
+
+import concurrent.futures
+import threading
+from typing import Callable, Set
+
+from sky import exceptions
+from sky import sky_logging
+from sky.utils import atomic
+
+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, /, *args, **kwargs):
+        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

@@ -178,14 +178,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)
@@ -203,7 +205,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(
@@ -224,14 +227,67 @@ 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
+        if response.text:
+            error_msg = response.text
+
+    with ux_utils.print_exception_no_traceback():
+        raise exceptions.ServerTemporarilyUnavailableError(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()
@@ -310,11 +366,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)