skypilot-nightly 1.0.0.dev20250624__py3-none-any.whl → 1.0.0.dev20250626__py3-none-any.whl

sky/resources.py

@@ -30,6 +30,9 @@ from sky.utils import resources_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
 
+if typing.TYPE_CHECKING:
+    from sky.volumes import volume as volume_lib
+
 logger = sky_logging.init_logger(__name__)
 
 _DEFAULT_DISK_SIZE_GB = 256
@@ -289,7 +292,8 @@ class Resources:
                 self._job_recovery = job_recovery
 
         if disk_size is not None:
-            self._disk_size = int(parse_memory_resource(disk_size, 'disk_size'))
+            self._disk_size = int(
+                resources_utils.parse_memory_resource(disk_size, 'disk_size'))
         else:
             self._disk_size = _DEFAULT_DISK_SIZE_GB
 
@@ -707,11 +711,11 @@ class Resources:
             self._memory = None
             return
 
-        memory = parse_memory_resource(str(memory),
-                                       'memory',
-                                       ret_type=float,
-                                       allow_plus=True,
-                                       allow_x=True)
+        memory = resources_utils.parse_memory_resource(str(memory),
+                                                       'memory',
+                                                       ret_type=float,
+                                                       allow_plus=True,
+                                                       allow_x=True)
         self._memory = memory
         if memory.endswith(('+', 'x')):
             # 'x' is used internally for make sure our resources used by
@@ -1465,11 +1469,15 @@ class Resources:
     def get_spot_str(self) -> str:
         return '[Spot]' if self.use_spot else ''
 
-    def make_deploy_variables(self, cluster_name: resources_utils.ClusterName,
-                              region: clouds.Region,
-                              zones: Optional[List[clouds.Zone]],
-                              num_nodes: int,
-                              dryrun: bool) -> Dict[str, Optional[str]]:
+    def make_deploy_variables(
+        self,
+        cluster_name: resources_utils.ClusterName,
+        region: clouds.Region,
+        zones: Optional[List[clouds.Zone]],
+        num_nodes: int,
+        dryrun: bool,
+        volume_mounts: Optional[List['volume_lib.VolumeMount']] = None,
+    ) -> Dict[str, Optional[str]]:
         """Converts planned sky.Resources to resource variables.
 
         These variables are divided into two categories: cloud-specific and
@@ -1491,7 +1499,7 @@ class Resources:
         # Cloud specific variables
         assert self.cloud is not None, 'Cloud must be specified'
         cloud_specific_variables = self.cloud.make_deploy_resources_variables(
-            self, cluster_name, region, zones, num_nodes, dryrun)
+            self, cluster_name, region, zones, num_nodes, dryrun, volume_mounts)
 
         # TODO(andyl): Should we print some warnings if users' envs share
         # same names with the cloud specific variables, but not enabled
@@ -2291,67 +2299,3 @@ def parse_time_minutes(time: str) -> int:
                 continue
 
     raise ValueError(f'Invalid time format: {time}')
-
-
-def parse_memory_resource(resource_qty_str: Union[str, int, float],
-                          field_name: str,
-                          ret_type: type = int,
-                          unit: str = 'gb',
-                          allow_plus: bool = False,
-                          allow_x: bool = False,
-                          allow_rounding: bool = False) -> str:
-    """Returns memory size in chosen units given a resource quantity string.
-
-    Args:
-        resource_qty_str: Resource quantity string
-        unit: Unit to convert to
-        allow_plus: Whether to allow '+' prefix
-        allow_x: Whether to allow 'x' suffix
-    """
-    assert unit in constants.MEMORY_SIZE_UNITS, f'Invalid unit: {unit}'
-
-    error_msg = f'"{field_name}" field should be a <int><b|k|m|g|t|p><+?>,'\
-                f' got {resource_qty_str}'
-
-    resource_str = str(resource_qty_str)
-
-    # Handle plus and x suffixes, x is only used internally for jobs controller
-    plus = ''
-    if resource_str.endswith('+'):
-        if allow_plus:
-            resource_str = resource_str[:-1]
-            plus = '+'
-        else:
-            raise ValueError(error_msg)
-
-    x = ''
-    if resource_str.endswith('x'):
-        if allow_x:
-            resource_str = resource_str[:-1]
-            x = 'x'
-        else:
-            raise ValueError(error_msg)
-
-    try:
-        # We assume it is already in the wanted units to maintain backwards
-        # compatibility
-        ret_type(resource_str)
-        return f'{resource_str}{plus}{x}'
-    except ValueError:
-        pass
-
-    resource_str = resource_str.lower()
-    for mem_unit, multiplier in constants.MEMORY_SIZE_UNITS.items():
-        if resource_str.endswith(mem_unit):
-            try:
-                value = ret_type(resource_str[:-len(mem_unit)])
-                converted = (value * multiplier /
-                             constants.MEMORY_SIZE_UNITS[unit])
-                if not allow_rounding and ret_type(converted) != converted:
-                    raise ValueError(error_msg)
-                converted = ret_type(converted)
-                return f'{converted}{plus}{x}'
-            except ValueError:
-                continue
-
-    raise ValueError(error_msg)

sky/serve/client/sdk.py

@@ -5,9 +5,9 @@ from typing import List, Optional, Union
 
 import click
 
-from sky.adaptors import common as adaptors_common
 from sky.client import common as client_common
 from sky.server import common as server_common
+from sky.server import rest
 from sky.server.requests import payloads
 from sky.usage import usage_lib
 from sky.utils import admin_policy_utils
@@ -17,12 +17,8 @@ from sky.utils import dag_utils
 if typing.TYPE_CHECKING:
     import io
 
-    import requests
-
     import sky
     from sky.serve import serve_utils
-else:
-    requests = adaptors_common.LazyImport('requests')
 
 
 @context.contextual
@@ -78,7 +74,7 @@ def up(
             task=dag_str,
             service_name=service_name,
         )
-        response = requests.post(
+        response = rest.post(
             f'{server_common.get_server_url()}/serve/up',
             json=json.loads(body.model_dump_json()),
             timeout=(5, None),
@@ -140,7 +136,7 @@ def update(
             mode=mode,
         )
 
-        response = requests.post(
+        response = rest.post(
             f'{server_common.get_server_url()}/serve/update',
             json=json.loads(body.model_dump_json()),
             timeout=(5, None),
@@ -182,7 +178,7 @@ def down(
         all=all,
         purge=purge,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/down',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -217,7 +213,7 @@ def terminate_replica(service_name: str, replica_id: int,
         replica_id=replica_id,
         purge=purge,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/terminate-replica',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -290,7 +286,7 @@ def status(
         exceptions.ClusterNotUpError: if the sky serve controller is not up.
     """
     body = payloads.ServeStatusBody(service_names=service_names,)
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/status',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -301,6 +297,7 @@ def status(
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
+@rest.retry_on_server_unavailable()
 def tail_logs(service_name: str,
               target: Union[str, 'serve_utils.ServiceComponent'],
               replica_id: Optional[int] = None,
@@ -376,7 +373,7 @@ def tail_logs(service_name: str,
         replica_id=replica_id,
         follow=follow,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/logs',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -384,7 +381,10 @@ def tail_logs(service_name: str,
         cookies=server_common.get_api_cookie_jar(),
     )
     request_id = server_common.get_request_id(response)
-    sdk.stream_response(request_id, response, output_stream)
+    return sdk.stream_response(request_id=request_id,
+                               response=response,
+                               output_stream=output_stream,
+                               resumable=True)
 
 
 @usage_lib.entrypoint
@@ -436,7 +436,7 @@ def sync_down_logs(service_name: str,
         targets=targets,
         replica_ids=replica_ids,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/sync-down-logs',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),

sky/serve/server/core.py

@@ -28,6 +28,7 @@ from sky.utils import command_runner
 from sky.utils import common
 from sky.utils import common_utils
 from sky.utils import controller_utils
+from sky.utils import dag_utils
 from sky.utils import rich_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
@@ -139,10 +140,13 @@ def up(
                              f'{constants.CLUSTER_NAME_VALID_REGEX}')
 
     serve_utils.validate_service_task(task)
+    dag = dag_utils.convert_entrypoint_to_dag(task)
+    dag.resolve_and_validate_volumes()
     # Always apply the policy again here, even though it might have been applied
     # in the CLI. This is to ensure that we apply the policy to the final DAG
     # and get the mutated config.
-    dag, mutated_user_config = admin_policy_utils.apply(task)
+    dag, mutated_user_config = admin_policy_utils.apply(dag)
+    dag.pre_mount_volumes()
     task = dag.tasks[0]
 
     with rich_utils.safe_status(

sky/server/common.py

@@ -9,11 +9,13 @@ import json
 import os
 import pathlib
 import re
+import shutil
 import subprocess
 import sys
+import tempfile
 import time
 import typing
-from typing import Any, Dict, Literal, Optional, Tuple
+from typing import Any, Dict, Literal, Optional, Tuple, Union
 from urllib import parse
 import uuid
 
@@ -27,6 +29,7 @@ from sky import skypilot_config
 from sky.adaptors import common as adaptors_common
 from sky.data import data_utils
 from sky.server import constants as server_constants
+from sky.server import rest
 from sky.skylet import constants
 from sky.usage import usage_lib
 from sky.utils import annotations
@@ -240,9 +243,9 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
     server_url = endpoint if endpoint is not None else get_server_url()
     while time_out_try_count <= RETRY_COUNT_ON_TIMEOUT:
         try:
-            response = requests.get(f'{server_url}/api/health',
-                                    timeout=2.5,
-                                    cookies=get_api_cookie_jar())
+            response = rest.get(f'{server_url}/api/health',
+                                timeout=2.5,
+                                cookies=get_api_cookie_jar())
         except requests.exceptions.Timeout:
             if time_out_try_count == RETRY_COUNT_ON_TIMEOUT:
                 return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
@@ -327,6 +330,8 @@ def get_request_id(response: 'requests.Response') -> RequestId:
 def _start_api_server(deploy: bool = False,
                       host: str = '127.0.0.1',
                       foreground: bool = False,
+                      metrics: bool = False,
+                      metrics_port: Optional[int] = None,
                       enable_basic_auth: bool = False):
     """Starts a SkyPilot API server locally."""
     server_url = get_server_url(host)
@@ -357,10 +362,13 @@ def _start_api_server(deploy: bool = False,
             args += ['--deploy']
         if host is not None:
             args += [f'--host={host}']
+        if metrics_port is not None:
+            args += [f'--metrics-port={metrics_port}']
 
         if foreground:
             # Replaces the current process with the API server
             os.environ[constants.ENV_VAR_IS_SKYPILOT_SERVER] = 'true'
+            _set_metrics_env_var(os.environ, metrics, deploy)
             if enable_basic_auth:
                 os.environ[constants.ENV_VAR_ENABLE_BASIC_AUTH] = 'true'
             os.execvp(args[0], args)
@@ -368,6 +376,10 @@ def _start_api_server(deploy: bool = False,
         log_path = os.path.expanduser(constants.API_SERVER_LOGS)
         os.makedirs(os.path.dirname(log_path), exist_ok=True)
 
+        # For spawn mode, copy the environ to avoid polluting the SDK process.
+        server_env = os.environ.copy()
+        server_env[constants.ENV_VAR_IS_SKYPILOT_SERVER] = 'true'
+        _set_metrics_env_var(server_env, metrics, deploy)
         # Start the API server process in the background and don't wait for it.
         # If this is called from a CLI invocation, we need
         # start_new_session=True so that SIGINT on the CLI will not also kill
@@ -437,6 +449,26 @@ def _start_api_server(deploy: bool = False,
                 f'SkyPilot API server started. {dashboard_msg}'))
 
 
+def _set_metrics_env_var(env: Union[Dict[str, str], os._Environ], metrics: bool,
+                         deploy: bool):
+    """Sets the metrics environment variables.
+
+    Args:
+        env: The environment variables to set.
+        metrics: Whether to enable metrics.
+        deploy: Whether the server is running in deploy mode, which means
+            multiple processes might be running.
+    """
+    if metrics:
+        env[constants.ENV_VAR_SERVER_METRICS_ENABLED] = 'true'
+        if deploy:
+            metrics_dir = os.path.join(tempfile.gettempdir(), 'metrics')
+            shutil.rmtree(metrics_dir, ignore_errors=True)
+            os.makedirs(metrics_dir, exist_ok=True)
+            # Refer to https://prometheus.github.io/client_python/multiprocess/
+            env['PROMETHEUS_MULTIPROC_DIR'] = metrics_dir
+
+
 def check_server_healthy(
     endpoint: Optional[str] = None
 ) -> Tuple[Literal[
@@ -571,6 +603,8 @@ def get_skypilot_version_on_disk() -> str:
 def check_server_healthy_or_start_fn(deploy: bool = False,
                                      host: str = '127.0.0.1',
                                      foreground: bool = False,
+                                     metrics: bool = False,
+                                     metrics_port: Optional[int] = None,
                                      enable_basic_auth: bool = False):
     api_server_status = None
     try:
@@ -592,7 +626,8 @@ def check_server_healthy_or_start_fn(deploy: bool = False,
             # have started the server while we were waiting for the lock.
             api_server_info = get_api_server_status(endpoint)
             if api_server_info.status == ApiServerStatus.UNHEALTHY:
-                _start_api_server(deploy, host, foreground, enable_basic_auth)
+                _start_api_server(deploy, host, foreground, metrics,
+                                  metrics_port, enable_basic_auth)
 
 
 def check_server_healthy_or_start(func):

sky/server/constants.py

@@ -7,7 +7,7 @@ from sky.skylet import constants
 # API server version, whenever there is a change in API server that requires a
 # restart of the local API server or error out when the client does not match
 # the server version.
-API_VERSION = '9'
+API_VERSION = '10'
 
 # Prefix for API request names.
 REQUEST_NAME_PREFIX = 'sky.'
@@ -22,6 +22,10 @@ API_SERVER_REQUEST_DB_PATH = '~/.sky/api_server/requests.db'
 # background.
 CLUSTER_REFRESH_DAEMON_INTERVAL_SECONDS = 60
 
+# The interval (seconds) for the volume status to be refreshed in the
+# background.
+VOLUME_REFRESH_DAEMON_INTERVAL_SECONDS = 60
+
 # Environment variable for a file path to the API cookie file.
 # Keep in sync with websocket_proxy.py
 API_COOKIE_FILE_ENV_VAR = f'{constants.SKYPILOT_ENV_VAR_PREFIX}API_COOKIE_FILE'

sky/server/metrics.py

@@ -0,0 +1,105 @@
+"""Instrumentation for the API server."""
+
+import asyncio
+import os
+import time
+
+import fastapi
+from prometheus_client import generate_latest
+from prometheus_client import multiprocess
+import prometheus_client as prom
+import starlette.middleware.base
+import uvicorn
+
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
+
+# Total number of API server requests, grouped by path, method, and status.
+sky_apiserver_requests_total = prom.Counter(
+    'sky_apiserver_requests_total',
+    'Total number of API server requests',
+    ['path', 'method', 'status'],
+)
+
+# Time spent processing API server requests, grouped by path, method, and
+# status.
+sky_apiserver_request_duration_seconds = prom.Histogram(
+    'sky_apiserver_request_duration_seconds',
+    'Time spent processing API server requests',
+    ['path', 'method', 'status'],
+    buckets=(0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
+             float('inf')),
+)
+
+metrics_app = fastapi.FastAPI()
+
+
+@metrics_app.get('/metrics')
+async def metrics() -> fastapi.Response:
+    """Expose aggregated Prometheus metrics from all worker processes."""
+    if os.environ.get('PROMETHEUS_MULTIPROC_DIR'):
+        # In multiprocess mode, we need to collect metrics from all processes.
+        registry = prom.CollectorRegistry()
+        multiprocess.MultiProcessCollector(registry)
+        data = generate_latest(registry)
+    else:
+        data = generate_latest()
+    return fastapi.Response(content=data,
+                            media_type=prom.CONTENT_TYPE_LATEST,
+                            headers={'Cache-Control': 'no-cache'})
+
+
+def run_metrics_server(host: str, port: int):
+    metrics_config = uvicorn.Config(
+        'sky.server.metrics:metrics_app',
+        host=host,
+        port=port,
+        workers=1,
+    )
+    metrics_server_instance = uvicorn.Server(metrics_config)
+    asyncio.run(metrics_server_instance.serve())
+
+
+def _get_status_code_group(status_code: int) -> str:
+    """Group status codes into classes (2xx, 5xx) to reduce cardinality."""
+    return f'{status_code // 100}xx'
+
+
+def _is_streaming_api(path: str) -> bool:
+    """Check if the path is a streaming API."""
+    path = path.rstrip('/')
+    return path.endswith('/logs') or path.endswith('/api/stream')
+
+
+class PrometheusMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
+    """Middleware to collect Prometheus metrics for HTTP requests."""
+
+    async def dispatch(self, request: fastapi.Request, call_next):
+        path = request.url.path
+        logger.info(f'PROM Middleware Request: {request}, {request.url.path}')
+        streaming = _is_streaming_api(path)
+        if not streaming:
+            # Exclude streaming APIs, the duration is not meaningful.
+            # TODO(aylei): measure the duration of async execution instead.
+            start_time = time.time()
+        method = request.method
+        status_code_group = ''
+
+        try:
+            response = await call_next(request)
+            status_code_group = _get_status_code_group(response.status_code)
+        except Exception:  # pylint: disable=broad-except
+            status_code_group = '5xx'
+            raise
+        finally:
+            sky_apiserver_requests_total.labels(path=path,
+                                                method=method,
+                                                status=status_code_group).inc()
+            if not streaming:
+                duration = time.time() - start_time
+                sky_apiserver_request_duration_seconds.labels(
+                    path=path, method=method,
+                    status=status_code_group).observe(duration)
+
+        return response

sky/server/requests/executor.py

@@ -149,10 +149,25 @@ class RequestWorker:
         self.schedule_type = schedule_type
         self.garanteed_parallelism = config.garanteed_parallelism
         self.burstable_parallelism = config.burstable_parallelism
+        self._thread: Optional[threading.Thread] = None
+        self._cancel_event = threading.Event()
 
     def __str__(self) -> str:
         return f'Worker(schedule_type={self.schedule_type.value})'
 
+    def run_in_background(self) -> None:
+        # Thread dispatcher is sufficient for current scale, refer to
+        # tests/load_tests/test_queue_dispatcher.py for more details.
+        # Use daemon thread for automatic cleanup.
+        thread = threading.Thread(target=self.run, daemon=True)
+        thread.start()
+        self._thread = thread
+
+    def cancel(self) -> None:
+        if self._thread is not None:
+            self._cancel_event.set()
+            self._thread.join()
+
     def process_request(self, executor: process.BurstableExecutor,
                         queue: RequestQueue) -> None:
         try:
@@ -219,7 +234,7 @@ class RequestWorker:
                 burst_workers=self.burstable_parallelism,
                 initializer=executor_initializer,
                 initargs=(proc_group,))
-            while True:
+            while not self._cancel_event.is_set():
                 self.process_request(executor, queue)
         # TODO(aylei): better to distinct between KeyboardInterrupt and SIGTERM.
         except KeyboardInterrupt:
@@ -539,15 +554,21 @@ def schedule_request(request_id: str,
         enqueue()
 
 
-def start(config: server_config.ServerConfig) -> List[multiprocessing.Process]:
+def start(
+    config: server_config.ServerConfig
+) -> Tuple[Optional[multiprocessing.Process], List[RequestWorker]]:
     """Start the request workers.
 
     Request workers run in background, schedule the requests and delegate the
     request execution to executor processes.
+
+    Returns:
+        A tuple of the queue server process and the list of request worker
+        threads.
     """
     global queue_backend
     queue_backend = config.queue_backend
-    sub_procs = []
+    queue_server = None
     # Setup the queues.
     if queue_backend == server_config.QueueBackend.MULTIPROCESSING:
         logger.info('Creating shared request queues')
@@ -564,7 +585,6 @@ def start(config: server_config.ServerConfig) -> List[multiprocessing.Process]:
         queue_server = multiprocessing.Process(
             target=mp_queue.start_queue_manager, args=(queue_names, port))
         queue_server.start()
-        sub_procs.append(queue_server)
         mp_queue.wait_for_queues_to_be_ready(queue_names,
                                              queue_server,
                                              port=port)
@@ -577,20 +597,16 @@ def start(config: server_config.ServerConfig) -> List[multiprocessing.Process]:
 
     logger.info('Request queues created')
 
-    def run_worker_in_background(worker: RequestWorker):
-        # Thread dispatcher is sufficient for current scale, refer to
-        # tests/load_tests/test_queue_dispatcher.py for more details.
-        # Use daemon thread for automatic cleanup.
-        thread = threading.Thread(target=worker.run, daemon=True)
-        thread.start()
-
+    workers = []
     # Start a worker for long requests.
     long_worker = RequestWorker(schedule_type=api_requests.ScheduleType.LONG,
                                 config=config.long_worker_config)
-    run_worker_in_background(long_worker)
+    long_worker.run_in_background()
+    workers.append(long_worker)
 
     # Start a worker for short requests.
     short_worker = RequestWorker(schedule_type=api_requests.ScheduleType.SHORT,
                                  config=config.short_worker_config)
-    run_worker_in_background(short_worker)
-    return sub_procs
+    short_worker.run_in_background()
+    workers.append(short_worker)
+    return queue_server, workers

sky/server/requests/payloads.py

@@ -368,6 +368,22 @@ class StorageBody(RequestBody):
     name: str
 
 
+class VolumeApplyBody(RequestBody):
+    """The request body for the volume apply endpoint."""
+    name: str
+    volume_type: str
+    cloud: str
+    region: Optional[str] = None
+    zone: Optional[str] = None
+    size: Optional[str] = None
+    config: Optional[Dict[str, Any]] = None
+
+
+class VolumeDeleteBody(RequestBody):
+    """The request body for the volume delete endpoint."""
+    names: List[str]
+
+
 class EndpointsBody(RequestBody):
     """The request body for the endpoint."""
     cluster: str