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

sky/server/requests/serializers/decoders.py

@@ -2,17 +2,17 @@
 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
-from sky.clouds.service_catalog import common
+from sky.catalog import common
 from sky.data import storage
 from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.schemas.api import responses
 from sky.serve import serve_state
 from sky.server import constants as server_constants
 from sky.skylet import job_lib
-from sky.utils import registry
 from sky.utils import status_lib
 
 if typing.TYPE_CHECKING:
@@ -51,13 +51,19 @@ def default_decode_handler(return_value: Any) -> Any:
 
 
 @register_decoders('status')
-def decode_status(return_value: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+def decode_status(
+        return_value: List[Dict[str, Any]]) -> List[responses.StatusResponse]:
     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'])
-
-    return clusters
+        if 'is_managed' not in cluster:
+            cluster['is_managed'] = False
+        response.append(responses.StatusResponse.model_validate(cluster))
+    return response
 
 
 @register_decoders('status_kubernetes')
@@ -66,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 = []
@@ -79,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
 
 
@@ -95,24 +102,53 @@ 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: List[dict],) -> List[Dict[str, Any]]:
-    jobs = return_value
-    for job in jobs:
-        job['status'] = managed_jobs.ManagedJobStatus(job['status'])
-    return jobs
-
-
-@register_decoders('serve.status')
-def decode_serve_status(return_value: List[dict]) -> List[Dict[str, Any]]:
-    service_statuses = return_value
+    # 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, 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, 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
+
+
+def _decode_serve_status(
+        service_statuses: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     for service_status in service_statuses:
         service_status['status'] = serve_state.ServiceStatus(
             service_status['status'])
@@ -123,6 +159,16 @@ def decode_serve_status(return_value: List[dict]) -> List[Dict[str, Any]]:
     return service_statuses
 
 
+@register_decoders('serve.status')
+def decode_serve_status(return_value: List[dict]) -> List[Dict[str, Any]]:
+    return _decode_serve_status(return_value)
+
+
+@register_decoders('jobs.pool_status')
+def decode_jobs_pool_status(return_value: List[dict]) -> List[Dict[str, Any]]:
+    return _decode_serve_status(return_value)
+
+
 @register_decoders('cost_report')
 def decode_cost_report(
         return_value: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
@@ -135,16 +181,6 @@ def decode_cost_report(
     return return_value
 
 
-@register_decoders('enabled_clouds')
-def decode_enabled_clouds(return_value: List[str]) -> List['clouds.Cloud']:
-    clouds = []
-    for cloud_name in return_value:
-        cloud = registry.CLOUD_REGISTRY.from_str(cloud_name)
-        assert cloud is not None, return_value
-        clouds.append(cloud)
-    return clouds
-
-
 @register_decoders('list_accelerators')
 def decode_list_accelerators(
     return_value: Dict[str, List[List[Any]]]
@@ -160,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')
@@ -190,3 +236,8 @@ def decode_job_status(
 def decode_kubernetes_node_info(
         return_value: Dict[str, Any]) -> models.KubernetesNodesInfo:
     return models.KubernetesNodesInfo.from_dict(return_value)
+
+
+@register_decoders('endpoints')
+def decode_endpoints(return_value: Dict[int, str]) -> Dict[int, str]:
+    return {int(k): v for k, v in return_value.items()}

sky/server/requests/serializers/encoders.py

@@ -6,14 +6,17 @@ 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
 
 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] = {}
@@ -21,6 +24,9 @@ handlers: Dict[str, Any] = {}
 
 def pickle_and_encode(obj: Any) -> str:
     try:
+        # Apply backwards compatibility processing at the lowest level
+        # to catch any handles that might have bypassed the encoders
+        obj = serialize_utils.prepare_handle_for_backwards_compatibility(obj)
         return base64.b64encode(pickle.dumps(obj)).decode('utf-8')
     except TypeError as e:
         raise ValueError(f'Failed to pickle object: {obj}') from e
@@ -51,13 +57,29 @@ def default_encoder(return_value: Any) -> Any:
 
 
 @register_encoder('status')
-def encode_status(clusters: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+def encode_status(
+        clusters: List[responses.StatusResponse]) -> List[Dict[str, Any]]:
+    response = []
     for cluster in clusters:
-        cluster['status'] = cluster['status'].value
-        cluster['handle'] = pickle_and_encode(cluster['handle'])
-        cluster['storage_mounts_metadata'] = pickle_and_encode(
-            cluster['storage_mounts_metadata'])
-    return clusters
+        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(
+            None)  # Always returns None.
+        response.append(response_cluster)
+    return response
 
 
 @register_encoder('launch', 'exec', 'jobs.launch')
@@ -65,6 +87,7 @@ def encode_launch(
     job_id_handle: Tuple[Optional[int], Optional['backends.ResourceHandle']]
 ) -> Dict[str, Any]:
     job_id, handle = job_id_handle
+    handle = serialize_utils.prepare_handle_for_backwards_compatibility(handle)
     return {
         'job_id': job_id,
         'handle': pickle_and_encode(handle),
@@ -73,14 +96,21 @@ def encode_launch(
 
 @register_encoder('start')
 def encode_start(resource_handle: 'backends.CloudVmRayResourceHandle') -> str:
+    resource_handle = (
+        serialize_utils.prepare_handle_for_backwards_compatibility(
+            resource_handle))
     return pickle_and_encode(resource_handle)
 
 
 @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')
@@ -88,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
@@ -102,6 +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(by_alias=True) for job in all_jobs]
     return encoded_all_clusters, encoded_unmanaged_clusters, all_jobs, context
 
 
@@ -112,25 +143,68 @@ def encode_jobs_queue(jobs: List[dict],) -> List[Dict[str, Any]]:
     return jobs
 
 
-@register_encoder('serve.status')
-def encode_serve_status(
+@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: Dict[str, int] = {}
+    if isinstance(jobs_or_tuple, tuple):
+        if len(jobs_or_tuple) == 2:
+            jobs, total = jobs_or_tuple
+            total_no_filter = total
+        elif len(jobs_or_tuple) == 4:
+            jobs, total, status_counts, total_no_filter = jobs_or_tuple
+        else:
+            raise ValueError(f'Invalid jobs tuple: {jobs_or_tuple}')
+    else:
+        jobs = jobs_or_tuple
+        total = None
+    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_dict
+    return {
+        'jobs': jobs_dict,
+        'total': total,
+        'total_no_filter': total_no_filter,
+        'status_counts': status_counts
+    }
+
+
+def _encode_serve_status(
         service_statuses: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     for service_status in service_statuses:
         service_status['status'] = service_status['status'].value
         for replica_info in service_status.get('replica_info', []):
             replica_info['status'] = replica_info['status'].value
-            replica_info['handle'] = pickle_and_encode(replica_info['handle'])
+            handle = serialize_utils.prepare_handle_for_backwards_compatibility(
+                replica_info['handle'])
+            replica_info['handle'] = pickle_and_encode(handle)
     return service_statuses
 
 
+@register_encoder('serve.status')
+def encode_serve_status(
+        service_statuses: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    return _encode_serve_status(service_statuses)
+
+
+@register_encoder('jobs.pool_status')
+def encode_jobs_pool_status(
+        pool_statuses: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    return _encode_serve_status(pool_statuses)
+
+
 @register_encoder('cost_report')
 def encode_cost_report(
         cost_report: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     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
 
 
@@ -142,22 +216,66 @@ 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')
 def encode_kubernetes_node_info(
         return_value: 'models.KubernetesNodesInfo') -> Dict[str, Any]:
     return return_value.to_dict()
+
+
+@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()