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

sky/users/server.py

@@ -33,28 +33,38 @@ 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()
+
+    users_to_role = {}
+    for role in rbac.get_supported_roles():
+        user_ids = permission.permission_service.get_users_for_role(role)
+        for user_id in user_ids:
+            users_to_role[user_id] = role
+
     for user in user_list:
         # Filter out service accounts - they have IDs starting with "sa-"
         if user.is_service_account():
             continue
 
-        user_roles = permission.permission_service.get_user_roles(user.id)
         all_users.append({
             'id': user.id,
             'name': user.name,
             'created_at': user.created_at,
-            'role': user_roles[0] if user_roles else ''
+            'role': users_to_role.get(user.id, '')
         })
     return all_users
 
 
 @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 +80,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 +110,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 +191,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 +314,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 +378,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 +429,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 +517,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 +558,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 +594,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 +637,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/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:
@@ -331,6 +334,9 @@ class SQLiteConn(threading.local):
         """Execute the sql and commit the transaction in a sync block."""
         conn = await self._get_async_conn()
 
+        if parameters is None:
+            parameters = []
+
         def exec_and_commit(sql: str, parameters: Optional[Iterable[Any]]):
             # pylint: disable=protected-access
             conn._conn.execute(sql, parameters)
@@ -346,3 +352,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,9 +3,6 @@
 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
@@ -14,7 +11,6 @@ import filelock
 import sqlalchemy
 
 from sky import sky_logging
-from sky.skylet import constants
 
 logger = sky_logging.init_logger(__name__)
 
@@ -32,34 +28,6 @@ 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] = {}
-
-_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:
-        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)
-        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
-
 
 @contextlib.contextmanager
 def db_lock(db_name: str):

{skypilot_nightly-1.0.0.dev20250903.dist-info → skypilot_nightly-1.0.0.dev20250905.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20250903
+Version: 1.0.0.dev20250905
 Summary: SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.
 Author: SkyPilot Team
 License: Apache 2.0
@@ -145,48 +145,48 @@ Requires-Dist: grpcio>=1.63.0; extra == "server"
 Requires-Dist: protobuf<7.0.0,>=5.26.1; extra == "server"
 Requires-Dist: aiosqlite; extra == "server"
 Provides-Extra: all
-Requires-Dist: ibm-cos-sdk; extra == "all"
-Requires-Dist: ray[default]!=2.6.0,>=2.2.0; extra == "all"
-Requires-Dist: vastai-sdk>=0.1.12; extra == "all"
-Requires-Dist: aiohttp; extra == "all"
-Requires-Dist: google-cloud-storage; extra == "all"
-Requires-Dist: websockets; extra == "all"
-Requires-Dist: azure-mgmt-compute>=33.0.0; extra == "all"
 Requires-Dist: kubernetes!=32.0.0,>=20.0.0; extra == "all"
-Requires-Dist: msgraph-sdk; extra == "all"
 Requires-Dist: azure-storage-blob>=12.23.1; extra == "all"
-Requires-Dist: anyio; extra == "all"
-Requires-Dist: boto3>=1.26.1; extra == "all"
+Requires-Dist: grpcio>=1.63.0; extra == "all"
 Requires-Dist: google-api-python-client>=2.69.0; extra == "all"
+Requires-Dist: colorama<0.4.5; extra == "all"
+Requires-Dist: azure-mgmt-network>=27.0.0; extra == "all"
+Requires-Dist: sqlalchemy_adapter; extra == "all"
+Requires-Dist: aiosqlite; extra == "all"
+Requires-Dist: ibm-cos-sdk; extra == "all"
 Requires-Dist: pyopenssl<24.3.0,>=23.2.0; extra == "all"
+Requires-Dist: websockets; extra == "all"
+Requires-Dist: awscli>=1.27.10; extra == "all"
+Requires-Dist: boto3>=1.26.1; extra == "all"
+Requires-Dist: pydo>=0.3.0; extra == "all"
+Requires-Dist: google-cloud-storage; extra == "all"
 Requires-Dist: cudo-compute>=0.1.10; extra == "all"
-Requires-Dist: azure-mgmt-network>=27.0.0; extra == "all"
+Requires-Dist: azure-core>=1.24.0; extra == "all"
 Requires-Dist: pyvmomi==8.0.1.0.2; extra == "all"
+Requires-Dist: casbin; extra == "all"
+Requires-Dist: docker; extra == "all"
 Requires-Dist: passlib; extra == "all"
+Requires-Dist: ibm-vpc; extra == "all"
+Requires-Dist: oci; extra == "all"
+Requires-Dist: azure-cli>=2.65.0; extra == "all"
+Requires-Dist: nebius>=0.2.47; extra == "all"
+Requires-Dist: azure-core>=1.31.0; extra == "all"
+Requires-Dist: azure-mgmt-compute>=33.0.0; extra == "all"
 Requires-Dist: azure-identity>=1.19.0; extra == "all"
 Requires-Dist: botocore>=1.29.10; extra == "all"
-Requires-Dist: awscli>=1.27.10; extra == "all"
 Requires-Dist: python-dateutil; extra == "all"
-Requires-Dist: pyjwt; extra == "all"
-Requires-Dist: oci; extra == "all"
 Requires-Dist: ibm-cloud-sdk-core; extra == "all"
-Requires-Dist: azure-cli>=2.65.0; extra == "all"
-Requires-Dist: aiosqlite; extra == "all"
-Requires-Dist: pydo>=0.3.0; extra == "all"
-Requires-Dist: docker; extra == "all"
-Requires-Dist: azure-common; extra == "all"
-Requires-Dist: colorama<0.4.5; extra == "all"
-Requires-Dist: msrestazure; extra == "all"
-Requires-Dist: ibm-vpc; extra == "all"
-Requires-Dist: ibm-platform-services>=0.48.0; extra == "all"
-Requires-Dist: azure-core>=1.24.0; extra == "all"
-Requires-Dist: sqlalchemy_adapter; extra == "all"
-Requires-Dist: azure-core>=1.31.0; extra == "all"
-Requires-Dist: nebius>=0.2.47; extra == "all"
-Requires-Dist: casbin; extra == "all"
+Requires-Dist: ray[default]!=2.6.0,>=2.2.0; extra == "all"
 Requires-Dist: protobuf<7.0.0,>=5.26.1; extra == "all"
-Requires-Dist: grpcio>=1.63.0; extra == "all"
+Requires-Dist: ibm-platform-services>=0.48.0; extra == "all"
+Requires-Dist: msgraph-sdk; extra == "all"
+Requires-Dist: msrestazure; extra == "all"
+Requires-Dist: pyjwt; extra == "all"
+Requires-Dist: aiohttp; extra == "all"
 Requires-Dist: runpod>=1.6.1; extra == "all"
+Requires-Dist: azure-common; extra == "all"
+Requires-Dist: vastai-sdk>=0.1.12; extra == "all"
+Requires-Dist: anyio; extra == "all"
 Dynamic: author
 Dynamic: classifier
 Dynamic: description
@@ -284,21 +284,21 @@ SkyPilot supports your existing GPU, TPU, and CPU workloads, with no code change
 Install with pip:
 ```bash
 # Choose your clouds:
-pip install -U "skypilot[kubernetes,aws,gcp,azure,oci,lambda,runpod,fluidstack,paperspace,cudo,ibm,scp,nebius]"
+pip install -U "skypilot[kubernetes,aws,gcp,azure,oci,nebius,lambda,runpod,fluidstack,paperspace,cudo,ibm,scp]"
 ```
 To get the latest features and fixes, use the nightly build or [install from source](https://docs.skypilot.co/en/latest/getting-started/installation.html):
 ```bash
 # Choose your clouds:
-pip install "skypilot-nightly[kubernetes,aws,gcp,azure,oci,lambda,runpod,fluidstack,paperspace,cudo,ibm,scp,nebius]"
+pip install "skypilot-nightly[kubernetes,aws,gcp,azure,oci,nebius,lambda,runpod,fluidstack,paperspace,cudo,ibm,scp]"
 ```
 
 <p align="center">
   <img src="docs/source/_static/intro.gif" alt="SkyPilot">
 </p>
 
-Current supported infra: Kubernetes, AWS, GCP, Azure, OCI, Lambda Cloud, Fluidstack,
-RunPod, Cudo, Digital Ocean, Paperspace, Cloudflare, Samsung, IBM, Vast.ai,
-VMware vSphere, Nebius.
+Current supported infra: Kubernetes, AWS, GCP, Azure, OCI, Nebius, Lambda Cloud, RunPod, Fluidstack,
+Cudo, Digital Ocean, Paperspace, Cloudflare, Samsung, IBM, Vast.ai,
+VMware vSphere.
 <p align="center">
   <img alt="SkyPilot" src="https://raw.githubusercontent.com/skypilot-org/skypilot/master/docs/source/images/cloud-logos-light.png" width=85%>
 </p>