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

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/users/server.py

@@ -33,8 +33,12 @@ USER_LOCK_TIMEOUT_SECONDS = 20
 router = fastapi.APIRouter()
 
 
+# All handlers in user handler are sync to get fastAPI run it in a
+# ThreadPoolExecutor to avoid blocking the async event loop.
+# TODO(aylei): make these async once we have the global_user_state async
+# support.
 @router.get('')
-async def users() -> List[Dict[str, Any]]:
+def users() -> List[Dict[str, Any]]:
     """Gets all users."""
     all_users = []
     user_list = global_user_state.get_all_users()
@@ -54,7 +58,7 @@ async def users() -> List[Dict[str, Any]]:
 
 
 @router.get('/role')
-async def get_current_user_role(request: fastapi.Request):
+def get_current_user_role(request: fastapi.Request):
     """Get current user's role."""
     # TODO(hailong): is there a reliable way to get the user
     # hash for the request without 'X-Auth-Request-Email' header?
@@ -70,7 +74,7 @@ async def get_current_user_role(request: fastapi.Request):
 
 
 @router.post('/create')
-async def user_create(user_create_body: payloads.UserCreateBody) -> None:
+def user_create(user_create_body: payloads.UserCreateBody) -> None:
     username = user_create_body.username
     password = user_create_body.password
     role = user_create_body.role
@@ -100,8 +104,8 @@ async def user_create(user_create_body: payloads.UserCreateBody) -> None:
 
 
 @router.post('/update')
-async def user_update(request: fastapi.Request,
-                      user_update_body: payloads.UserUpdateBody) -> None:
+def user_update(request: fastapi.Request,
+                user_update_body: payloads.UserUpdateBody) -> None:
     """Updates the user role."""
     user_id = user_update_body.user_id
     role = user_update_body.role
@@ -181,14 +185,13 @@ def _delete_user(user_id: str) -> None:
 
 
 @router.post('/delete')
-async def user_delete(user_delete_body: payloads.UserDeleteBody) -> None:
+def user_delete(user_delete_body: payloads.UserDeleteBody) -> None:
     user_id = user_delete_body.user_id
     _delete_user(user_id)
 
 
 @router.post('/import')
-async def user_import(
-        user_import_body: payloads.UserImportBody) -> Dict[str, Any]:
+def user_import(user_import_body: payloads.UserImportBody) -> Dict[str, Any]:
     """Import users from CSV content."""
     csv_content = user_import_body.csv_content
 
@@ -305,7 +308,7 @@ async def user_import(
 
 
 @router.get('/export')
-async def user_export() -> Dict[str, Any]:
+def user_export() -> Dict[str, Any]:
     """Export all users as CSV content."""
     try:
         # Get all users
@@ -369,7 +372,7 @@ def _user_lock(user_id: str) -> Generator[None, None, None]:
 
 
 @router.get('/service-account-tokens')
-async def get_service_account_tokens(
+def get_service_account_tokens(
         request: fastapi.Request) -> List[Dict[str, Any]]:
     """Get service account tokens. All users can see all tokens."""
     auth_user = request.state.auth_user
@@ -420,7 +423,7 @@ def _generate_service_account_user_id() -> str:
 
 
 @router.post('/service-account-tokens')
-async def create_service_account_token(
+def create_service_account_token(
         request: fastapi.Request,
         token_body: payloads.ServiceAccountTokenCreateBody) -> Dict[str, Any]:
     """Create a new service account token."""
@@ -508,7 +511,7 @@ async def create_service_account_token(
 
 
 @router.post('/service-account-tokens/delete')
-async def delete_service_account_token(
+def delete_service_account_token(
         request: fastapi.Request,
         token_body: payloads.ServiceAccountTokenDeleteBody) -> Dict[str, str]:
     """Delete a service account token.
@@ -549,7 +552,7 @@ async def delete_service_account_token(
 
 
 @router.post('/service-account-tokens/get-role')
-async def get_service_account_role(
+def get_service_account_role(
         request: fastapi.Request,
         role_body: payloads.ServiceAccountTokenRoleBody) -> Dict[str, Any]:
     """Get the role of a service account."""
@@ -585,7 +588,7 @@ async def get_service_account_role(
 
 
 @router.post('/service-account-tokens/update-role')
-async def update_service_account_role(
+def update_service_account_role(
         request: fastapi.Request,
         role_body: payloads.ServiceAccountTokenUpdateRoleBody
 ) -> Dict[str, str]:
@@ -628,7 +631,7 @@ async def update_service_account_role(
 
 
 @router.post('/service-account-tokens/rotate')
-async def rotate_service_account_token(
+def rotate_service_account_token(
         request: fastapi.Request,
         token_body: payloads.ServiceAccountTokenRotateBody) -> Dict[str, Any]:
     """Rotate a service account token.

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/db_utils.py

@@ -2,10 +2,12 @@
 import asyncio
 import contextlib
 import enum
+import os
+import pathlib
 import sqlite3
 import threading
 import typing
-from typing import Any, Callable, Iterable, Optional
+from typing import Any, Callable, Dict, Iterable, Optional
 
 import aiosqlite
 import aiosqlite.context
@@ -13,6 +15,7 @@ import sqlalchemy
 from sqlalchemy import exc as sqlalchemy_exc
 
 from sky import sky_logging
+from sky.skylet import constants
 
 logger = sky_logging.init_logger(__name__)
 if typing.TYPE_CHECKING:
@@ -346,3 +349,57 @@ class SQLiteConn(threading.local):
                                     ) -> Iterable[sqlite3.Row]:
         conn = await self._get_async_conn()
         return await conn.execute_fetchall(sql, parameters)
+
+    async def close(self):
+        if self._async_conn is not None:
+            await self._async_conn.close()
+        self.conn.close()
+
+
+_max_connections = 0
+_postgres_engine_cache: Dict[str, sqlalchemy.engine.Engine] = {}
+_sqlite_engine_cache: Dict[str, sqlalchemy.engine.Engine] = {}
+
+_db_creation_lock = threading.Lock()
+
+
+def set_max_connections(max_connections: int):
+    global _max_connections
+    _max_connections = max_connections
+
+
+def get_max_connections():
+    return _max_connections
+
+
+def get_engine(db_name: str):
+    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:
+        with _db_creation_lock:
+            if conn_string not in _postgres_engine_cache:
+                if _max_connections == 0:
+                    _postgres_engine_cache[conn_string] = (
+                        sqlalchemy.create_engine(
+                            conn_string, poolclass=sqlalchemy.pool.NullPool))
+                elif _max_connections == 1:
+                    _postgres_engine_cache[conn_string] = (
+                        sqlalchemy.create_engine(
+                            conn_string, poolclass=sqlalchemy.pool.StaticPool))
+                else:
+                    _postgres_engine_cache[conn_string] = (
+                        sqlalchemy.create_engine(
+                            conn_string,
+                            poolclass=sqlalchemy.pool.QueuePool,
+                            size=_max_connections,
+                            max_overflow=0))
+            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)
+        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/db/migration_utils.py

@@ -3,7 +3,6 @@
 import contextlib
 import logging
 import os
-import pathlib
 
 from alembic import command as alembic_command
 from alembic.config import Config
@@ -12,7 +11,6 @@ import filelock
 import sqlalchemy
 
 from sky import sky_logging
-from sky.skylet import constants
 
 logger = sky_logging.init_logger(__name__)
 
@@ -31,20 +29,6 @@ SERVE_VERSION = '001'
 SERVE_LOCK_PATH = '~/.sky/locks/.serve_db.lock'
 
 
-def get_engine(db_name: str):
-    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)
-    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)
-    return engine
-
-
 @contextlib.contextmanager
 def db_lock(db_name: str):
     lock_path = os.path.expanduser(f'~/.sky/locks/.{db_name}.lock')

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',