skypilot-nightly 1.0.0.dev20251210__py3-none-any.whl → 1.0.0.dev20260112__py3-none-any.whl

sky/serve/serve_utils.py

@@ -23,6 +23,7 @@ import filelock
 from sky import backends
 from sky import exceptions
 from sky import global_user_state
+from sky import resources as resources_lib
 from sky import sky_logging
 from sky import skypilot_config
 from sky.adaptors import common as adaptors_common
@@ -350,6 +351,13 @@ def validate_service_task(task: 'sky.Task', pool: bool) -> None:
                              f'file does not match the pool argument. '
                              f'To fix, add a valid `{field_name}` field.')
 
+    # Validate that pools do not use ordered resources
+    if pool and isinstance(task.resources, list):
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                'Ordered resources are not supported for pools. '
+                'Use `any_of` instead, or specify a single resource.')
+
     policy_description = ('on-demand'
                           if task.service.dynamic_ondemand_fallback else 'spot')
     for resource in list(task.resources):
@@ -360,22 +368,6 @@ def validate_service_task(task: 'sky.Task', pool: bool) -> None:
                                  f'{sys_name} will replenish preempted spot '
                                  f'with {policy_description} instances.')
 
-    if pool:
-        accelerators = set()
-        for resource in task.resources:
-            if resource.accelerators is not None:
-                if isinstance(resource.accelerators, str):
-                    accelerators.add(resource.accelerators)
-                elif isinstance(resource.accelerators, dict):
-                    accelerators.update(resource.accelerators.keys())
-                elif isinstance(resource.accelerators, list):
-                    accelerators.update(resource.accelerators)
-        if len(accelerators) > 1:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError('Heterogeneous clusters are not supported for '
-                                 'pools please specify one accelerator '
-                                 'for all workers.')
-
     # Try to create a spot placer from the task yaml. Check if the task yaml
     # is valid for spot placer.
     spot_placer.SpotPlacer.from_task(task.service, task)
@@ -730,7 +722,7 @@ def _get_service_status(
             for replica_info in record['replica_info']:
                 job_ids = managed_job_state.get_nonterminal_job_ids_by_pool(
                     service_name, replica_info['name'])
-                replica_info['used_by'] = job_ids[0] if job_ids else None
+                replica_info['used_by'] = job_ids
     return record
 
 
@@ -810,16 +802,112 @@ def get_ready_replicas(
     ]
 
 
-def get_next_cluster_name(service_name: str, job_id: int) -> Optional[str]:
-    """Get the next available cluster name from idle replicas.
+def _task_fits(task_resources: 'resources_lib.Resources',
+               free_resources: 'resources_lib.Resources') -> bool:
+    """Check if the task resources fit in the free resources."""
+    if not task_resources.less_demanding_than(free_resources,
+                                              check_cloud=False):
+        return False
+    if task_resources.cpus is not None:
+        if (free_resources.cpus is None or
+                task_resources.cpus > free_resources.cpus):
+            return False
+    if task_resources.memory is not None:
+        if (free_resources.memory is None or
+                task_resources.memory > free_resources.memory):
+            return False
+    return True
+
+
+def _is_empty_resource(resource: 'resources_lib.Resources') -> bool:
+    # Returns True if this resource object does not specify any resources.
+    return (resource.cpus is None and resource.memory is None and
+            resource.accelerators is None)
+
+
+def get_free_worker_resources(
+        pool: str) -> Optional[Dict[str, Optional[resources_lib.Resources]]]:
+    """Get free resources for each worker in a pool.
+
+    Args:
+        pool: Pool name (service name)
+
+    Returns:
+        Dictionary mapping cluster_name (worker) to free Resources object (or
+        None if worker is not available or has no free resources).
+    """
+
+    free_resources: Dict[str, Optional[resources_lib.Resources]] = {}
+    replicas = serve_state.get_replica_infos(pool)
+
+    for replica_info in replicas:
+        cluster_name = replica_info.cluster_name
+
+        # Get cluster handle
+        handle = replica_info.handle()
+        if handle is None or handle.launched_resources is None:
+            free_resources[cluster_name] = None
+            continue
+
+        total_resources = handle.launched_resources
+
+        # Get job IDs running on this worker
+        job_ids = managed_job_state.get_nonterminal_job_ids_by_pool(
+            pool, cluster_name)
+
+        if len(job_ids) == 0:
+            free_resources[cluster_name] = total_resources
+            continue
+
+        # Get used resources
+        # TODO(lloyd): We should batch the database calls here so that we
+        # make a single call to get all the used resources for all the jobs.
+        used_resources = managed_job_state.get_pool_worker_used_resources(
+            set(job_ids))
+        if used_resources is None:
+            # We failed to get the used resources. We should return None since
+            # we can't make any guarantees about what resources are being used.
+            logger.warning(
+                f'Failed to get used resources for cluster {cluster_name!r}')
+            return None
+
+        if _is_empty_resource(used_resources):
+            # We encountered a job that has no resources specified. We
+            # will not consider it for resource-aware scheduling so it must
+            # be scheduled on its own. To do this we will set the free
+            # worker resources to nothing by returning an empty resource
+            # object.
+            logger.debug(f'Job {job_ids} has no resources specified. '
+                         'Skipping resource-aware scheduling for cluster '
+                         f'{cluster_name!r}')
+            free_resources[cluster_name] = resources_lib.Resources()
+        else:
+            # Calculate free resources using - operator
+            free = total_resources - used_resources
+            free_resources[cluster_name] = free
+
+    return free_resources
+
+
+def get_next_cluster_name(
+    service_name: str,
+    job_id: int,
+    task_resources: Optional[typing.Union[
+        'resources_lib.Resources', typing.Set['resources_lib.Resources'],
+        typing.List['resources_lib.Resources']]] = None
+) -> Optional[str]:
+    """Get the next available cluster name from replicas with sufficient
+    resources.
 
     Args:
         service_name: The name of the service.
-        job_id: Optional job ID to associate with the acquired cluster.
-                If None, a placeholder will be used.
+        job_id: Job ID to associate with the acquired cluster.
+        task_resources: Optional task resource requirements. If provided, will
+                check if resources fit in free worker resources. Can be
+                a single Resources object or a set/list of Resources objects.
 
     Returns:
-        The cluster name if an idle replica is found, None otherwise.
+        The cluster name if a suitable replica is found, None otherwise.
     """
     # Check if service exists
     service_status = _get_service_status(service_name,
@@ -831,36 +919,126 @@ def get_next_cluster_name(service_name: str, job_id: int) -> Optional[str]:
     if not service_status['pool']:
         logger.error(f'Service {service_name!r} is not a pool.')
         return None
+
     with filelock.FileLock(get_service_filelock_path(service_name)):
+        free_resources = get_free_worker_resources(service_name)
+        logger.debug(f'Free resources: {free_resources!r}')
         logger.debug(f'Get next cluster name for pool {service_name!r}')
         ready_replicas = get_ready_replicas(service_name)
+
+        logger.debug(f'Ready replicas: {ready_replicas!r}')
+
         idle_replicas: List['replica_managers.ReplicaInfo'] = []
-        for replica_info in ready_replicas:
-            jobs_on_replica = managed_job_state.get_nonterminal_job_ids_by_pool(
-                service_name, replica_info.cluster_name)
-            # TODO(tian): Make it resources aware. Currently we allow and only
-            # allow one job per replica. In the following PR, we should:
-            #  i) When the replica is launched with `any_of` resources (
-            #     replicas can have different resources), we should check if
-            #     the resources that jobs require are available on the replica.
-            #     e.g., if a job requires A100:1 on a {L4:1, A100:1} pool, it
-            #     should only goes to replica with A100.
-            # ii) When a job only requires a subset of the resources on the
-            #     replica, each replica should be able to handle multiple jobs
-            #     at the same time. e.g., if a job requires A100:1 on a A100:8
-            #     pool, it should be able to run 4 jobs at the same time.
-            if not jobs_on_replica:
-                idle_replicas.append(replica_info)
+
+        # If task_resources is provided, use resource-aware scheduling
+        # Normalize task_resources to a list
+        if isinstance(task_resources, resources_lib.Resources):
+            task_resources_list = [task_resources]
+        elif isinstance(task_resources, (set, list)):
+            task_resources_list = list(task_resources)
+        else:
+            task_resources_list = []
+
+        # We should do resource aware scheduling if:
+        # 1. There are task resources.
+        # 2. The first task resource has some resources listed.
+        # 3. There are free resources.
+        # 4. Any free resource has some resources listed.
+        resource_aware = len(task_resources_list) > 0
+        resource_aware = (resource_aware and
+                          not _is_empty_resource(task_resources_list[0]))
+        resource_aware = resource_aware and free_resources is not None
+        if free_resources is not None:
+            for free_resource in free_resources.values():
+                if free_resource is not None and not _is_empty_resource(
+                        free_resource):
+                    resource_aware = True
+                    break
+            else:
+                resource_aware = False
+        else:
+            resource_aware = False
+
+        if resource_aware:
+            logger.debug('Doing resource aware scheduling')
+            for replica_info in ready_replicas:
+                cluster_name = replica_info.cluster_name
+                assert free_resources is not None
+                free_resources_on_worker = free_resources.get(cluster_name)
+                logger.debug(f'Free resources for cluster {cluster_name!r}: '
+                             f'{free_resources_on_worker!r}')
+
+                # Skip if worker has no free resources available
+                if free_resources_on_worker is None:
+                    logger.debug(f'Worker {cluster_name!r} has no free '
+                                 'resources')
+                    continue
+
+                # Check if any of the task resource options fit
+                fits = False
+                for task_res in task_resources_list:
+                    logger.debug(f'Task resources: {task_res!r}')
+                    if _task_fits(task_res, free_resources_on_worker):
+                        logger.debug(f'Task resources {task_res!r} fits'
+                                     ' in free resources '
+                                     f'{free_resources_on_worker!r}')
+                        fits = True
+                        break
+                    else:
+                        logger.debug(f'Task resources {task_res!r} does not fit'
+                                     ' in free resources '
+                                     f'{free_resources_on_worker!r}')
+                if fits:
+                    idle_replicas.append(replica_info)
+        # Also fall back to resource unaware scheduling if no idle replicas are
+        # found. This might be because our launched resources were improperly
+        # set. If that's the case then jobs will fail to schedule in a resource
+        # aware way because one of the resources will be `None` so we can just
+        # fallback to 1 job per replica. If we are truly resource bottlenecked
+        # then we will see that there are jobs running on the replica and will
+        # not schedule another.
+        if len(idle_replicas) == 0:
+            logger.debug('Falling back to resource unaware scheduling')
+            # Fall back to resource unaware scheduling if no task resources
+            # are provided.
+            for replica_info in ready_replicas:
+                jobs_on_replica = (
+                    managed_job_state.get_nonterminal_job_ids_by_pool(
+                        service_name, replica_info.cluster_name))
+                if not jobs_on_replica:
+                    idle_replicas.append(replica_info)
+
         if not idle_replicas:
             logger.info(f'No idle replicas found for pool {service_name!r}')
             return None
 
         # Select the first idle replica.
-        # TODO(tian): "Load balancing" policy.
         replica_info = idle_replicas[0]
         logger.info(f'Selected replica {replica_info.replica_id} with cluster '
                     f'{replica_info.cluster_name!r} for job {job_id!r} in pool '
                     f'{service_name!r}')
+
+        # If job has heterogeneous resources (any_of/ordered), update
+        # full_resources to the specific resource that was selected for this
+        # worker. This must happen before releasing the filelock to ensure
+        # atomicity with the scheduling decision.
+        if resource_aware and len(task_resources_list) > 1:
+            assert free_resources is not None
+            free_resources_on_worker = free_resources.get(
+                replica_info.cluster_name)
+            if free_resources_on_worker is not None:
+                # Find which task resource fits on this worker
+                for task_res in task_resources_list:
+                    if _task_fits(task_res, free_resources_on_worker):
+                        # Update full_resources in database to this specific
+                        # resource
+                        logger.debug(
+                            f'Updating full_resources for job {job_id!r} '
+                            f'to selected resource: {task_res!r}')
+                        managed_job_state.update_job_full_resources(
+                            job_id, task_res.to_yaml_config())
+                        break
+
         managed_job_state.set_current_cluster_name(job_id,
                                                    replica_info.cluster_name)
         return replica_info.cluster_name
@@ -1541,7 +1719,21 @@ def _format_replica_table(replica_records: List[Dict[str, Any]], show_all: bool,
         replica_status = record['status']
         status_str = replica_status.colored_str()
         used_by = record.get('used_by', None)
-        used_by_str = str(used_by) if used_by is not None else '-'
+        if used_by is None:
+            used_by_str = '-'
+        elif isinstance(used_by, str):
+            used_by_str = used_by
+        else:
+            if len(used_by) > 2:
+                used_by_str = (
+                    f'{used_by[0]}, {used_by[1]}, +{len(used_by) - 2}'
+                    ' more')
+            elif len(used_by) == 2:
+                used_by_str = f'{used_by[0]}, {used_by[1]}'
+            elif len(used_by) == 1:
+                used_by_str = str(used_by[0])
+            else:
+                used_by_str = '-'
 
         replica_handle: Optional['backends.CloudVmRayResourceHandle'] = record[
             'handle']

sky/server/common.py

@@ -17,6 +17,7 @@ import time
 import typing
 from typing import (Any, Callable, cast, Dict, Generic, Literal, Optional,
                     Tuple, TypeVar, Union)
+from urllib.request import Request
 import uuid
 
 import cachetools
@@ -147,6 +148,22 @@ def get_api_cookie_jar() -> requests.cookies.RequestsCookieJar:
     return cookie_jar
 
 
+def get_cookie_header_for_url(url: str) -> Dict[str, str]:
+    """Extract Cookie header value from a cookie jar for a specific URL"""
+    cookies = get_api_cookie_jar()
+    if not cookies:
+        return {}
+
+    # Use urllib Request to do URL-aware cookie filtering
+    request = Request(url)
+    cookies.add_cookie_header(request)
+    cookie_header = request.get_header('Cookie')
+
+    if cookie_header is None:
+        return {}
+    return {'Cookie': cookie_header}
+
+
 def set_api_cookie_jar(cookie_jar: CookieJar,
                        create_if_not_exists: bool = True) -> None:
     """Updates the file cookie jar with the given cookie jar."""

sky/server/constants.py

@@ -10,7 +10,7 @@ from sky.skylet import constants
 # based on version info is needed.
 # For more details and code guidelines, refer to:
 # https://docs.skypilot.co/en/latest/developers/CONTRIBUTING.html#backward-compatibility-guidelines
-API_VERSION = 25
+API_VERSION = 26
 
 # The minimum peer API version that the code should still work with.
 # Notes (dev):

sky/server/metrics.py

@@ -48,10 +48,12 @@ async def gpu_metrics() -> fastapi.Response:
     all_metrics: List[str] = []
     successful_contexts = 0
 
+    remote_contexts = [
+        context for context in contexts if context != 'in-cluster'
+    ]
     tasks = [
         asyncio.create_task(metrics_utils.get_metrics_for_context(context))
-        for context in contexts
-        if context != 'in-cluster'
+        for context in remote_contexts
     ]
 
     results = await asyncio.gather(*tasks, return_exceptions=True)
@@ -59,7 +61,8 @@ async def gpu_metrics() -> fastapi.Response:
     for i, result in enumerate(results):
         if isinstance(result, Exception):
             logger.error(
-                f'Failed to get metrics for context {contexts[i]}: {result}')
+                f'Failed to get metrics for context {remote_contexts[i]}: '
+                f'{result}')
         elif isinstance(result, BaseException):
             # Avoid changing behavior for non-Exception BaseExceptions
             # like KeyboardInterrupt/SystemExit: re-raise them.

sky/server/plugins.py

@@ -90,11 +90,26 @@ class RBACRule:
 class BasePlugin(abc.ABC):
     """Base class for all SkyPilot server plugins."""
 
+    @property
+    def name(self) -> Optional[str]:
+        """Plugin name for display purposes."""
+        return None
+
     @property
     def js_extension_path(self) -> Optional[str]:
         """Optional API route to the JavaScript extension to load."""
         return None
 
+    @property
+    def version(self) -> Optional[str]:
+        """Plugin version."""
+        return None
+
+    @property
+    def commit(self) -> Optional[str]:
+        """Plugin git commit hash."""
+        return None
+
     @abc.abstractmethod
     def install(self, extension_context: ExtensionContext):
         """Hook called by API server to let the plugin install itself."""
@@ -164,6 +179,7 @@ def load_plugins(extension_context: ExtensionContext):
 
     for plugin_config in config.get('plugins', []):
         class_path = plugin_config['class']
+        logger.debug(f'Loading plugins: {class_path}')
         module_path, class_name = class_path.rsplit('.', 1)
         try:
             module = importlib.import_module(module_path)

sky/server/requests/payloads.py

@@ -482,6 +482,7 @@ class VolumeApplyBody(RequestBody):
 class VolumeDeleteBody(RequestBody):
     """The request body for the volume delete endpoint."""
     names: List[str]
+    purge: bool = False
 
 
 class VolumeListBody(RequestBody):
@@ -865,3 +866,20 @@ class SlurmGpuAvailabilityRequestBody(RequestBody):
     """Request body for getting Slurm real-time GPU availability."""
     name_filter: Optional[str] = None
     quantity_filter: Optional[int] = None
+
+
+class ClusterEventsBody(RequestBody):
+    """The request body for the cluster events endpoint."""
+    cluster_name: Optional[str] = None
+    cluster_hash: Optional[str] = None
+    event_type: str  # 'STATUS_CHANGE' or 'DEBUG'
+    include_timestamps: bool = False
+    limit: Optional[
+        int] = None  # If specified, returns at most this many events
+
+
+class GetJobEventsBody(RequestBody):
+    """The request body for the get job task events endpoint."""
+    job_id: int
+    task_id: Optional[int] = None
+    limit: Optional[int] = 10  # Default to 10 most recent task events

sky/server/requests/request_names.py

@@ -31,6 +31,7 @@ class RequestName(str, enum.Enum):
     CLUSTER_JOB_LOGS = 'logs'
     CLUSTER_JOB_DOWNLOAD_LOGS = 'download_logs'
     CLUSTER_COST_REPORT = 'cost_report'
+    CLUSTER_EVENTS = 'cluster_events'
     # Storage requests
     STORAGE_LS = 'storage_ls'
     STORAGE_DELETE = 'storage_delete'
@@ -52,6 +53,7 @@ class RequestName(str, enum.Enum):
     JOBS_POOL_STATUS = 'jobs.pool_status'
     JOBS_POOL_LOGS = 'jobs.pool_logs'
     JOBS_POOL_SYNC_DOWN_LOGS = 'jobs.pool_sync_down_logs'
+    JOBS_EVENTS = 'jobs.events'
     # Serve requests
     SERVE_UP = 'serve.up'
     SERVE_UPDATE = 'serve.update'

sky/server/requests/requests.py

@@ -873,11 +873,17 @@ async def create_if_not_exists_async(request: Request) -> bool:
         f'({request_columns}) VALUES '
         f'({values_str}) ON CONFLICT(request_id) DO NOTHING RETURNING ROWID')
     request_row = request.to_row()
-    # Execute the SQL statement without getting the request lock.
-    # The request lock is used to prevent racing with cancellation codepath,
-    # but a request cannot be cancelled before it is created.
-    row = await _DB.execute_get_returning_value_async(sql_statement,
-                                                      request_row)
+    if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+        logger.debug(f'Start creating request {request.request_id}')
+    try:
+        # Execute the SQL statement without getting the request lock.
+        # The request lock is used to prevent racing with cancellation codepath,
+        # but a request cannot be cancelled before it is created.
+        row = await _DB.execute_get_returning_value_async(
+            sql_statement, request_row)
+    finally:
+        if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+            logger.debug(f'End creating request {request.request_id}')
     return True if row else False
 
 
@@ -1034,9 +1040,15 @@ _add_or_update_request_sql = (f'INSERT OR REPLACE INTO {REQUEST_TABLE} '
 def _add_or_update_request_no_lock(request: Request):
     """Add or update a REST request into the database."""
     assert _DB is not None
-    with _DB.conn:
-        cursor = _DB.conn.cursor()
-        cursor.execute(_add_or_update_request_sql, request.to_row())
+    if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+        logger.debug(f'Start adding or updating request {request.request_id}')
+    try:
+        with _DB.conn:
+            cursor = _DB.conn.cursor()
+            cursor.execute(_add_or_update_request_sql, request.to_row())
+    finally:
+        if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+            logger.debug(f'End adding or updating request {request.request_id}')
 
 
 async def _add_or_update_request_no_lock_async(request: Request):
@@ -1125,8 +1137,14 @@ async def _delete_requests(request_ids: List[str]):
     """Clean up requests by their IDs."""
     id_list_str = ','.join(repr(request_id) for request_id in request_ids)
     assert _DB is not None
-    await _DB.execute_and_commit_async(
-        f'DELETE FROM {REQUEST_TABLE} WHERE request_id IN ({id_list_str})')
+    if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+        logger.debug(f'Start deleting requests {request_ids}')
+    try:
+        await _DB.execute_and_commit_async(
+            f'DELETE FROM {REQUEST_TABLE} WHERE request_id IN ({id_list_str})')
+    finally:
+        if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+            logger.debug(f'End deleting requests {request_ids}')
 
 
 async def clean_finished_requests_with_retention(retention_seconds: int,

sky/server/requests/serializers/encoders.py

@@ -69,6 +69,11 @@ def encode_status(
             response_cluster['last_use'] = ''
         if 'status_updated_at' not in response_cluster:
             response_cluster['status_updated_at'] = 0
+        # Ensure labels is always included, defaulting to empty dict if None
+        # This is needed because exclude_none=True would exclude None labels
+        if 'labels' not in response_cluster or response_cluster.get(
+                'labels') is None:
+            response_cluster['labels'] = {}
         response_cluster['status'] = cluster['status'].value
         handle = serialize_utils.prepare_handle_for_backwards_compatibility(
             cluster['handle'])

sky/server/requests/serializers/return_value_serializers.py

@@ -50,11 +50,21 @@ def serialize_kubernetes_node_info(return_value: Dict[str, Any]) -> str:
 
     The is_ready field was added in API version 25. Remove it for old clients
     that don't recognize it.
+    The cpu_count, memory_gb, cpu_free, and memory_free_gb fields were added
+    in API version 26. Remove them for old clients that don't recognize them.
     """
     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
+    if (return_value and remote_api_version is not None):
         for node_info in return_value.get('node_info_dict', {}).values():
-            node_info.pop('is_ready', None)
+            if remote_api_version < 25:
+                # Remove is_ready field for old clients that don't recognize it
+                node_info.pop('is_ready', None)
+            if remote_api_version < 26:
+                # Remove cpu_count, memory_gb, cpu_free, and
+                # memory_free_gb fields for old clients that don't
+                # recognize them
+                node_info.pop('cpu_count', None)
+                node_info.pop('memory_gb', None)
+                node_info.pop('cpu_free', None)
+                node_info.pop('memory_free_gb', None)
     return orjson.dumps(return_value).decode('utf-8')