skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__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,52 @@ 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('realtime_slurm_gpu_availability')
+def encode_realtime_slurm_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/serializers/return_value_serializers.py

@@ -0,0 +1,60 @@
+"""Version-aware serializers for request return values.
+
+These serializers run at encode() time when remote_api_version is available,
+to handle backward compatibility for old clients.
+
+The existing encoders.py handles object -> dict conversion at set_return_value()
+time. This module handles dict -> JSON string serialization at encode() time,
+with version-aware field filtering for backward compatibility.
+"""
+from typing import Any, Callable, Dict
+
+import orjson
+
+from sky.server import constants as server_constants
+from sky.server import versions
+
+handlers: Dict[str, Callable[[Any], str]] = {}
+
+
+def register_serializer(*names: str):
+    """Decorator to register a version-aware serializer."""
+
+    def decorator(func):
+        for name in names:
+            if name != server_constants.DEFAULT_HANDLER_NAME:
+                name = server_constants.REQUEST_NAME_PREFIX + name
+            if name in handlers:
+                raise ValueError(f'Serializer {name} already registered: '
+                                 f'{handlers[name]}')
+            handlers[name] = func
+        return func
+
+    return decorator
+
+
+def get_serializer(name: str) -> Callable[[Any], str]:
+    """Get the serializer for a request name."""
+    return handlers.get(name, handlers[server_constants.DEFAULT_HANDLER_NAME])
+
+
+@register_serializer(server_constants.DEFAULT_HANDLER_NAME)
+def default_serializer(return_value: Any) -> str:
+    """The default serializer."""
+    return orjson.dumps(return_value).decode('utf-8')
+
+
+@register_serializer('kubernetes_node_info')
+def serialize_kubernetes_node_info(return_value: Dict[str, Any]) -> str:
+    """Serialize kubernetes node info with version compatibility.
+
+    The is_ready field was added in API version 25. Remove it for old clients
+    that don't recognize it.
+    """
+    remote_api_version = versions.get_remote_api_version()
+    if (return_value and remote_api_version is not None and
+            remote_api_version < 25):
+        # Remove is_ready field for old clients that don't recognize it
+        for node_info in return_value.get('node_info_dict', {}).values():
+            node_info.pop('is_ready', None)
+    return orjson.dumps(return_value).decode('utf-8')

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()