skypilot-nightly 1.0.0.dev20250902__py3-none-any.whl → 1.0.0.dev20250903__py3-none-any.whl

sky/server/server.py

@@ -24,6 +24,7 @@ import aiofiles
 import anyio
 import fastapi
 from fastapi.middleware import cors
+from sqlalchemy import pool
 import starlette.middleware.base
 import uvloop
 
@@ -1327,10 +1328,12 @@ async def provision_logs(cluster_body: payloads.ClusterNameBody,
                          tail: int = 0) -> fastapi.responses.StreamingResponse:
     """Streams the provision.log for the latest launch request of a cluster."""
     # Prefer clusters table first, then cluster_history as fallback.
-    log_path_str = global_user_state.get_cluster_provision_log_path(
+    log_path_str = await context_utils.to_thread(
+        global_user_state.get_cluster_provision_log_path,
         cluster_body.cluster_name)
     if not log_path_str:
-        log_path_str = global_user_state.get_cluster_history_provision_log_path(
+        log_path_str = await context_utils.to_thread(
+            global_user_state.get_cluster_history_provision_log_path,
             cluster_body.cluster_name)
     if not log_path_str:
         raise fastapi.HTTPException(
@@ -1429,27 +1432,29 @@ async def local_down(request: fastapi.Request) -> None:
 async def api_get(request_id: str) -> payloads.RequestPayload:
     """Gets a request with a given request ID prefix."""
     while True:
-        request_task = await requests_lib.get_request_async(request_id)
-        if request_task is None:
+        req_status = await requests_lib.get_request_status_async(request_id)
+        if req_status is None:
             print(f'No task with request ID {request_id}', flush=True)
             raise fastapi.HTTPException(
                 status_code=404, detail=f'Request {request_id!r} not found')
-        if request_task.status > requests_lib.RequestStatus.RUNNING:
-            if request_task.should_retry:
-                raise fastapi.HTTPException(
-                    status_code=503,
-                    detail=f'Request {request_id!r} should be retried')
-            request_error = request_task.get_error()
-            if request_error is not None:
-                raise fastapi.HTTPException(
-                    status_code=500, detail=request_task.encode().model_dump())
-            return request_task.encode()
-        elif (request_task.status == requests_lib.RequestStatus.RUNNING and
-              daemons.is_daemon_request_id(request_id)):
-            return request_task.encode()
+        if (req_status.status == requests_lib.RequestStatus.RUNNING and
+                daemons.is_daemon_request_id(request_id)):
+            # Daemon requests run forever, break without waiting for complete.
+            break
+        if req_status.status > requests_lib.RequestStatus.RUNNING:
+            break
         # yield control to allow other coroutines to run, sleep shortly
         # to avoid storming the DB and CPU in the meantime
         await asyncio.sleep(0.1)
+    request_task = await requests_lib.get_request_async(request_id)
+    if request_task.should_retry:
+        raise fastapi.HTTPException(
+            status_code=503, detail=f'Request {request_id!r} should be retried')
+    request_error = request_task.get_error()
+    if request_error is not None:
+        raise fastapi.HTTPException(status_code=500,
+                                    detail=request_task.encode().model_dump())
+    return request_task.encode()
 
 
 @app.get('/api/stream')
@@ -1606,10 +1611,9 @@ async def api_status(
                 requests_lib.RequestStatus.PENDING,
                 requests_lib.RequestStatus.RUNNING,
             ]
-        return [
-            request_task.readable_encode()
-            for request_task in requests_lib.get_request_tasks(status=statuses)
-        ]
+        request_tasks = await requests_lib.get_request_tasks_async(
+            req_filter=requests_lib.RequestTaskFilter(status=statuses))
+        return [r.readable_encode() for r in request_tasks]
     else:
         encoded_request_tasks = []
         for request_id in request_ids:
@@ -1808,17 +1812,20 @@ async def gpu_metrics() -> fastapi.Response:
 # === Internal APIs ===
 @app.get('/api/completion/cluster_name')
 async def complete_cluster_name(incomplete: str,) -> List[str]:
-    return global_user_state.get_cluster_names_start_with(incomplete)
+    return await context_utils.to_thread(
+        global_user_state.get_cluster_names_start_with, incomplete)
 
 
 @app.get('/api/completion/storage_name')
 async def complete_storage_name(incomplete: str,) -> List[str]:
-    return global_user_state.get_storage_names_start_with(incomplete)
+    return await context_utils.to_thread(
+        global_user_state.get_storage_names_start_with, incomplete)
 
 
 @app.get('/api/completion/volume_name')
 async def complete_volume_name(incomplete: str,) -> List[str]:
-    return global_user_state.get_volume_names_start_with(incomplete)
+    return await context_utils.to_thread(
+        global_user_state.get_volume_names_start_with, incomplete)
 
 
 @app.get('/api/completion/api_request')
@@ -1902,7 +1909,7 @@ if __name__ == '__main__':
     skyuvicorn.add_timestamp_prefix_for_server_logs()
 
     # Initialize global user state db
-    global_user_state.initialize_and_get_db()
+    global_user_state.initialize_and_get_db(pool.QueuePool)
     # Initialize request db
     requests_lib.reset_db_and_logs()
     # Restore the server user hash

sky/server/stream_utils.py

@@ -75,8 +75,10 @@ async def log_streamer(request_id: Optional[str],
         last_waiting_msg = ''
         waiting_msg = (f'Waiting for {request_task.name!r} request to be '
                        f'scheduled: {request_id}')
-        while request_task.status < requests_lib.RequestStatus.RUNNING:
-            if request_task.status_msg is not None:
+        req_status = request_task.status
+        req_msg = request_task.status_msg
+        while req_status < requests_lib.RequestStatus.RUNNING:
+            if req_msg is not None:
                 waiting_msg = request_task.status_msg
             if show_request_waiting_spinner:
                 yield status_msg.update(f'[dim]{waiting_msg}[/dim]')
@@ -91,7 +93,10 @@ async def log_streamer(request_id: Optional[str],
             # polling the DB, which can be a bottleneck for high-concurrency
             # requests.
             await asyncio.sleep(0.1)
-            request_task = await requests_lib.get_request_async(request_id)
+            status_with_msg = await requests_lib.get_request_status_async(
+                request_id, include_msg=True)
+            req_status = status_with_msg.status
+            req_msg = status_with_msg.status_msg
             if not follow:
                 break
         if show_request_waiting_spinner:
@@ -153,10 +158,13 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
         line: Optional[bytes] = await f.readline()
         if not line:
             if request_id is not None:
-                request_task = await requests_lib.get_request_async(request_id)
-                if request_task.status > requests_lib.RequestStatus.RUNNING:
-                    if (request_task.status ==
+                req_status = await requests_lib.get_request_status_async(
+                    request_id)
+                if req_status.status > requests_lib.RequestStatus.RUNNING:
+                    if (req_status.status ==
                             requests_lib.RequestStatus.CANCELLED):
+                        request_task = await requests_lib.get_request_async(
+                            request_id)
                         if request_task.should_retry:
                             buffer.append(
                                 message_utils.encode_payload(

sky/server/uvicorn.py

@@ -146,7 +146,8 @@ class Server(uvicorn.Server):
                 requests_lib.RequestStatus.PENDING,
                 requests_lib.RequestStatus.RUNNING,
             ]
-            reqs = requests_lib.get_request_tasks(status=statuses)
+            reqs = requests_lib.get_request_tasks(
+                req_filter=requests_lib.RequestTaskFilter(status=statuses))
             if not reqs:
                 break
             logger.info(f'{len(reqs)} on-going requests '

sky/templates/kubernetes-ray.yml.j2

@@ -302,7 +302,7 @@ available_node_types:
           provreq.kueue.x-k8s.io/maxRunDurationSeconds: "{{k8s_max_run_duration_seconds|string}}"
           {% endif %}
           {% endif %}
-        # https://cloud.google.com/kubernetes-engine/docs/how-to/gpu-bandwidth-gpudirect-tcpx 
+        # https://cloud.google.com/kubernetes-engine/docs/how-to/gpu-bandwidth-gpudirect-tcpx
         # Values from google cloud guide
           {% if k8s_enable_gpudirect_tcpx %}
           devices.gke.io/container.tcpx-daemon: |+
@@ -784,8 +784,8 @@ available_node_types:
                   echo "Waiting for patch package to be installed..."
                 done
                 # Apply Ray patches for progress bar fix
-                ~/.local/bin/uv pip list | grep "ray " | grep 2.9.3 2>&1 > /dev/null && { 
-                  VIRTUAL_ENV=~/skypilot-runtime python -c "from sky.skylet.ray_patches import patch; patch()" || exit 1; 
+                ~/.local/bin/uv pip list | grep "ray " | grep 2.9.3 2>&1 > /dev/null && {
+                  VIRTUAL_ENV=~/skypilot-runtime python -c "from sky.skylet.ray_patches import patch; patch()" || exit 1;
                 }
                 touch /tmp/ray_skypilot_installation_complete
                 echo "=== Ray and skypilot installation completed ==="
@@ -1202,7 +1202,7 @@ setup_commands:
     {%- endfor %}
     STEPS=("apt-ssh-setup" "runtime-setup" "env-setup")
     start_epoch=$(date +%s);
-    
+
     # Wait for SSH setup to complete before proceeding
     if [ -f /tmp/apt_ssh_setup_started ]; then
       echo "=== Logs for asynchronous SSH setup ===";
@@ -1210,7 +1210,7 @@ setup_commands:
       { tail -f -n +1 /tmp/${STEPS[0]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; until [ -f /tmp/apt_ssh_setup_complete ]; do sleep 0.5; done; kill $TAIL_PID || true; };
       [ -f /tmp/${STEPS[0]}.failed ] && { echo "Error: ${STEPS[0]} failed. Exiting."; exit 1; } || true;
     fi
-    
+
     echo "=== Logs for asynchronous ray and skypilot installation ===";
     if [ -f /tmp/skypilot_is_nimbus ]; then
       echo "=== Logs for asynchronous ray and skypilot installation ===";

sky/templates/runpod-ray.yml.j2

@@ -40,6 +40,14 @@ available_node_types:
         skypilot:ssh_public_key_content
       Preemptible: {{use_spot}}
       BidPerGPU: {{bid_per_gpu}}
+      {%- if volume_mounts and volume_mounts|length > 0 %}
+      VolumeMounts:
+      {%- for vm in volume_mounts %}
+        - VolumeNameOnCloud: {{ vm.volume_name_on_cloud }}
+          VolumeIdOnCloud: {{ vm.volume_id_on_cloud }}
+          MountPath: {{ vm.path }}
+      {%- endfor %}
+      {%- endif %}
 
 head_node_type: ray_head_default
 

sky/utils/benchmark_utils.py

@@ -0,0 +1,60 @@
+"""Utility functions for benchmarking."""
+
+import functools
+import logging
+import time
+from typing import Callable, Optional
+
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
+
+
+def log_execution_time(func: Optional[Callable] = None,
+                       *,
+                       name: Optional[str] = None,
+                       level: int = logging.DEBUG,
+                       precision: int = 4) -> Callable:
+    """Mark a function and log its execution time.
+
+    Args:
+        func: Function to decorate.
+        name: Name of the function.
+        level: Logging level.
+        precision: Number of decimal places (default: 4).
+
+    Usage:
+        from sky.utils import benchmark_utils
+
+        @benchmark_utils.log_execution_time
+        def my_function():
+            pass
+
+        @benchmark_utils.log_execution_time(name='my_module.my_function2')
+        def my_function2():
+            pass
+    """
+
+    def decorator(f: Callable) -> Callable:
+
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            nonlocal name
+            name = name or f.__name__
+            start_time = time.perf_counter()
+            try:
+                result = f(*args, **kwargs)
+                return result
+            finally:
+                end_time = time.perf_counter()
+                execution_time = end_time - start_time
+                log = (f'Method {name} executed in '
+                       f'{execution_time:.{precision}f}')
+                logger.log(level, log)
+
+        return wrapper
+
+    if func is None:
+        return decorator
+    else:
+        return decorator(func)

sky/utils/command_runner.py

@@ -41,6 +41,8 @@ RSYNC_FILTER_GITIGNORE = f'--filter=\'dir-merge,- {constants.GIT_IGNORE_FILE}\''
 # The git exclude file to support.
 GIT_EXCLUDE = '.git/info/exclude'
 RSYNC_EXCLUDE_OPTION = '--exclude-from={}'
+# Owner and group metadata is not needed for downloads.
+RSYNC_NO_OWNER_NO_GROUP_OPTION = '--no-owner --no-group'
 
 _HASH_MAX_LENGTH = 10
 _DEFAULT_CONNECT_TIMEOUT = 30
@@ -286,6 +288,8 @@ class CommandRunner:
         if prefix_command is not None:
             rsync_command.append(prefix_command)
         rsync_command += ['rsync', RSYNC_DISPLAY_OPTION]
+        if not up:
+            rsync_command.append(RSYNC_NO_OWNER_NO_GROUP_OPTION)
 
         # --filter
         # The source is a local path, so we need to resolve it.

sky/utils/db/migration_utils.py

@@ -4,6 +4,8 @@ import contextlib
 import logging
 import os
 import pathlib
+import threading
+from typing import Dict, Optional
 
 from alembic import command as alembic_command
 from alembic.config import Config
@@ -30,18 +32,32 @@ SERVE_DB_NAME = 'serve_db'
 SERVE_VERSION = '001'
 SERVE_LOCK_PATH = '~/.sky/locks/.serve_db.lock'
 
+_postgres_engine_cache: Dict[str, sqlalchemy.engine.Engine] = {}
+_sqlite_engine_cache: Dict[str, sqlalchemy.engine.Engine] = {}
 
-def get_engine(db_name: str):
+_db_creation_lock = threading.Lock()
+
+
+def get_engine(db_name: str,
+               pg_pool_class: Optional[sqlalchemy.pool.Pool] = None):
     conn_string = None
     if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
         conn_string = os.environ.get(constants.ENV_VAR_DB_CONNECTION_URI)
     if conn_string:
-        engine = sqlalchemy.create_engine(conn_string,
-                                          poolclass=sqlalchemy.NullPool)
+        if pg_pool_class is None:
+            pg_pool_class = sqlalchemy.NullPool
+        with _db_creation_lock:
+            if conn_string not in _postgres_engine_cache:
+                _postgres_engine_cache[conn_string] = sqlalchemy.create_engine(
+                    conn_string, poolclass=pg_pool_class)
+            engine = _postgres_engine_cache[conn_string]
     else:
         db_path = os.path.expanduser(f'~/.sky/{db_name}.db')
         pathlib.Path(db_path).parents[0].mkdir(parents=True, exist_ok=True)
-        engine = sqlalchemy.create_engine('sqlite:///' + db_path)
+        if db_path not in _sqlite_engine_cache:
+            _sqlite_engine_cache[db_path] = sqlalchemy.create_engine(
+                'sqlite:///' + db_path)
+        engine = _sqlite_engine_cache[db_path]
     return engine
 
 

sky/utils/resource_checker.py

@@ -269,16 +269,17 @@ def _get_active_resources(
         all_managed_jobs: List[Dict[str, Any]]
     """
 
-    def get_all_clusters():
+    def get_all_clusters() -> List[Dict[str, Any]]:
         return global_user_state.get_clusters()
 
-    def get_all_managed_jobs():
+    def get_all_managed_jobs() -> List[Dict[str, Any]]:
         # pylint: disable=import-outside-toplevel
         from sky.jobs.server import core as managed_jobs_core
         try:
-            return managed_jobs_core.queue(refresh=False,
-                                           skip_finished=True,
-                                           all_users=True)
+            filtered_jobs, _, _, _ = managed_jobs_core.queue(refresh=False,
+                                                             skip_finished=True,
+                                                             all_users=True)
+            return filtered_jobs
         except exceptions.ClusterNotUpError:
             logger.warning('All jobs should be finished.')
             return []

sky/utils/schemas.py

@@ -432,7 +432,7 @@ def get_volume_schema():
     return {
         '$schema': 'https://json-schema.org/draft/2020-12/schema',
         'type': 'object',
-        'required': ['name', 'type', 'infra'],
+        'required': ['name', 'type'],
         'additionalProperties': False,
         'properties': {
             'name': {

sky/utils/volume.py

@@ -10,6 +10,8 @@ from sky.utils import common_utils
 from sky.utils import schemas
 from sky.utils import status_lib
 
+MIN_RUNPOD_NETWORK_VOLUME_SIZE_GB = 10
+
 
 class VolumeAccessMode(enum.Enum):
     """Volume access mode."""
@@ -22,6 +24,7 @@ class VolumeAccessMode(enum.Enum):
 class VolumeType(enum.Enum):
     """Volume type."""
     PVC = 'k8s-pvc'
+    RUNPOD_NETWORK_VOLUME = 'runpod-network-volume'
 
 
 class VolumeMount:

sky/volumes/client/sdk.py

@@ -27,6 +27,34 @@ logger = sky_logging.init_logger(__name__)
 def apply(volume: volume_lib.Volume) -> server_common.RequestId[None]:
     """Creates or registers a volume.
 
+    Example:
+        .. code-block:: python
+
+            import sky.volumes
+            cfg = {
+                'name': 'pvc',
+                'type': 'k8s-pvc',
+                'size': '100GB',
+                'labels': {
+                    'key': 'value',
+                },
+            }
+            vol = sky.volumes.Volume.from_yaml_config(cfg)
+            request_id = sky.volumes.apply(vol)
+            sky.get(request_id)
+
+            or
+
+            import sky.volumes
+            vol = sky.volumes.Volume(
+                name='vol',
+                type='runpod-network-volume',
+                infra='runpod/ca/CA-MTL-1',
+                size='100GB',
+            )
+            request_id = sky.volumes.apply(vol)
+            sky.get(request_id)
+
     Args:
         volume: The volume to apply.
 

sky/volumes/server/server.py

@@ -19,10 +19,15 @@ router = fastapi.APIRouter()
 @router.get('')
 async def volume_list(request: fastapi.Request) -> None:
     """Gets the volumes."""
+    auth_user = request.state.auth_user
+    auth_user_env_vars_kwargs = {
+        'env_vars': auth_user.to_env_vars()
+    } if auth_user else {}
+    volume_list_body = payloads.VolumeListBody(**auth_user_env_vars_kwargs)
     executor.schedule_request(
         request_id=request.state.request_id,
         request_name='volume_list',
-        request_body=payloads.RequestBody(),
+        request_body=volume_list_body,
         func=core.volume_list,
         schedule_type=requests_lib.ScheduleType.SHORT,
     )
@@ -76,6 +81,11 @@ async def volume_apply(request: fastapi.Request,
         elif access_mode not in supported_access_modes:
             raise fastapi.HTTPException(
                 status_code=400, detail=f'Invalid access mode: {access_mode}')
+    elif volume_type == volume.VolumeType.RUNPOD_NETWORK_VOLUME.value:
+        if not cloud.is_same_cloud(clouds.RunPod()):
+            raise fastapi.HTTPException(
+                status_code=400,
+                detail='Runpod network volume is only supported on Runpod')
     executor.schedule_request(
         request_id=request.state.request_id,
         request_name='volume_apply',