skypilot-nightly 1.0.0.dev20250526__py3-none-any.whl → 1.0.0.dev20250528__py3-none-any.whl

sky/server/requests/serializers/decoders.py

@@ -12,7 +12,6 @@ from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.serve import serve_state
 from sky.server import constants as server_constants
 from sky.skylet import job_lib
-from sky.utils import registry
 from sky.utils import status_lib
 
 if typing.TYPE_CHECKING:
@@ -135,16 +134,6 @@ def decode_cost_report(
     return return_value
 
 
-@register_decoders('enabled_clouds')
-def decode_enabled_clouds(return_value: List[str]) -> List['clouds.Cloud']:
-    clouds = []
-    for cloud_name in return_value:
-        cloud = registry.CLOUD_REGISTRY.from_str(cloud_name)
-        assert cloud is not None, return_value
-        clouds.append(cloud)
-    return clouds
-
-
 @register_decoders('list_accelerators')
 def decode_list_accelerators(
     return_value: Dict[str, List[List[Any]]]

sky/server/server.py

@@ -127,6 +127,11 @@ class AuthProxyMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
 
     async def dispatch(self, request: fastapi.Request, call_next):
         auth_user = _get_auth_user_header(request)
+
+        # Add user to database if auth_user is present
+        if auth_user is not None:
+            global_user_state.add_or_update_user(auth_user)
+
         body = await request.body()
         if auth_user and body:
             try:
@@ -137,10 +142,16 @@ class AuthProxyMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
                 logger.debug(f'Overriding user for {request.state.request_id}: '
                              f'{auth_user.name}, {auth_user.id}')
                 if 'env_vars' in original_json:
-                    original_json['env_vars'][
-                        constants.USER_ID_ENV_VAR] = auth_user.id
-                    original_json['env_vars'][
-                        constants.USER_ENV_VAR] = auth_user.name
+                    if isinstance(original_json.get('env_vars'), dict):
+                        original_json['env_vars'][
+                            constants.USER_ID_ENV_VAR] = auth_user.id
+                        original_json['env_vars'][
+                            constants.USER_ENV_VAR] = auth_user.name
+                    else:
+                        logger.warning(
+                            f'"env_vars" in request body is not a dictionary '
+                            f'for request {request.state.request_id}. '
+                            'Skipping user info injection into body.')
                 request._body = json.dumps(original_json).encode('utf-8')  # pylint: disable=protected-access
         return await call_next(request)
 
@@ -262,10 +273,7 @@ app.include_router(workspaces_rest.router,
 
 @app.get('/token')
 async def token(request: fastapi.Request) -> fastapi.responses.HTMLResponse:
-    # If we have auth info, save this user to the database.
     user = _get_auth_user_header(request)
-    if user is not None:
-        global_user_state.add_or_update_user(user)
 
     token_data = {
         'v': 1,  # Token version number, bump for backwards incompatible.
@@ -315,12 +323,14 @@ async def check(request: fastapi.Request,
 
 @app.get('/enabled_clouds')
 async def enabled_clouds(request: fastapi.Request,
-                         workspace: Optional[str] = None) -> None:
+                         workspace: Optional[str] = None,
+                         expand: bool = False) -> None:
     """Gets enabled clouds on the server."""
     executor.schedule_request(
         request_id=request.state.request_id,
         request_name='enabled_clouds',
-        request_body=payloads.EnabledCloudsBody(workspace=workspace),
+        request_body=payloads.EnabledCloudsBody(workspace=workspace,
+                                                expand=expand),
         func=core.enabled_clouds,
         schedule_type=requests_lib.ScheduleType.SHORT,
     )
@@ -411,6 +421,10 @@ async def validate(validate_body: payloads.ValidateBody) -> None:
     logger.debug(f'Validating tasks: {validate_body.dag}')
 
     context.initialize()
+    ctx = context.get()
+    assert ctx is not None
+    # TODO(aylei): generalize this to all requests without a db record.
+    ctx.override_envs(validate_body.env_vars)
 
     def validate_dag(dag: dag_utils.dag_lib.Dag):
         # TODO: Admin policy may contain arbitrary code, which may be expensive
@@ -1189,13 +1203,10 @@ async def health(request: fastapi.Request) -> Dict[str, Any]:
 
 
 @app.websocket('/kubernetes-pod-ssh-proxy')
-async def kubernetes_pod_ssh_proxy(
-    websocket: fastapi.WebSocket,
-    cluster_name_body: payloads.ClusterNameBody = fastapi.Depends()
-) -> None:
+async def kubernetes_pod_ssh_proxy(websocket: fastapi.WebSocket,
+                                   cluster_name: str) -> None:
     """Proxies SSH to the Kubernetes pod with websocket."""
     await websocket.accept()
-    cluster_name = cluster_name_body.cluster_name
     logger.info(f'WebSocket connection accepted for cluster: {cluster_name}')
 
     cluster_records = core.status(cluster_name, all_users=True)

sky/setup_files/dependencies.py

@@ -57,6 +57,7 @@ install_requires = [
     'httpx',
     'setproctitle',
     'sqlalchemy',
+    'psycopg2-binary',
 ]
 
 local_ray = [

sky/skylet/constants.py

@@ -408,3 +408,7 @@ SKY_USER_FILE_PATH = '~/.sky/generated'
 ENV_VAR_IS_SKYPILOT_SERVER = 'IS_SKYPILOT_SERVER'
 
 SKYPILOT_DEFAULT_WORKSPACE = 'default'
+
+# Experimental - may be deprecated in the future without notice.
+SKYPILOT_API_SERVER_DB_URL_ENV_VAR: str = (
+    f'{SKYPILOT_ENV_VAR_PREFIX}API_SERVER_DB_URL')

sky/skypilot_config.py

@@ -299,6 +299,10 @@ def get_nested(keys: Tuple[str, ...],
 def get_workspace_cloud(cloud: str,
                         workspace: Optional[str] = None) -> config_utils.Config:
     """Returns the workspace config."""
+    # TODO(zhwu): Instead of just returning the workspace specific config, we
+    # should return the config that already merges the global config, so that
+    # the caller does not need to manually merge the global config with
+    # the workspace specific config.
     if workspace is None:
         workspace = get_active_workspace()
     clouds = get_nested(keys=(

sky/utils/db_utils.py

@@ -88,60 +88,48 @@ def add_column_to_table_sqlalchemy(
     session: 'Session',
     table_name: str,
     column_name: str,
-    column_type: str,
+    column_type: sqlalchemy.types.TypeEngine,
+    default_statement: Optional[str] = None,
     copy_from: Optional[str] = None,
     value_to_replace_existing_entries: Optional[Any] = None,
 ):
     """Add a column to a table."""
-    dialect = session.bind.dialect
-    if dialect.name == SQLAlchemyDialect.SQLITE.value:
-        try:
+    # column type may be different for different dialects.
+    # for example, sqlite uses BLOB for LargeBinary
+    # while postgres uses BYTEA.
+    column_type_str = column_type.compile(dialect=session.bind.dialect)
+    default_statement_str = (f' {default_statement}'
+                             if default_statement is not None else '')
+    try:
+        session.execute(
+            sqlalchemy.text(f'ALTER TABLE {table_name} '
+                            f'ADD COLUMN {column_name} {column_type_str}'
+                            f'{default_statement_str}'))
+        if copy_from is not None:
             session.execute(
-                sqlalchemy.text(f'ALTER TABLE {table_name} '
-                                f'ADD COLUMN {column_name} {column_type}'))
-            if copy_from is not None:
-                session.execute(
-                    sqlalchemy.text(f'UPDATE {table_name} '
-                                    f'SET {column_name} = {copy_from}'))
-            if value_to_replace_existing_entries is not None:
-                session.execute(
-                    sqlalchemy.text(f'UPDATE {table_name} '
-                                    f'SET {column_name} = :replacement_value '
-                                    f'WHERE {column_name} IS NULL'),
-                    {'replacement_value': value_to_replace_existing_entries})
-        except sqlalchemy_exc.OperationalError as e:
-            if 'duplicate column name' in str(e):
-                pass
-            else:
-                raise
-    elif dialect.name == SQLAlchemyDialect.POSTGRESQL.value:
-        # TODO(syang) support postgres dialect
-        session.rollback()
-        raise ValueError('Unsupported database dialect')
-    else:
-        session.rollback()
-        raise ValueError('Unsupported database dialect')
+                sqlalchemy.text(f'UPDATE {table_name} '
+                                f'SET {column_name} = {copy_from}'))
+        if value_to_replace_existing_entries is not None:
+            session.execute(
+                sqlalchemy.text(f'UPDATE {table_name} '
+                                f'SET {column_name} = :replacement_value '
+                                f'WHERE {column_name} IS NULL'),
+                {'replacement_value': value_to_replace_existing_entries})
+    #sqlite
+    except sqlalchemy_exc.OperationalError as e:
+        if 'duplicate column name' in str(e):
+            pass
+        else:
+            raise
+    #postgressql
+    except sqlalchemy_exc.ProgrammingError as e:
+        if 'already exists' in str(e):
+            pass
+        else:
+            raise
     session.commit()
 
 
-def rename_column(
-    cursor: 'sqlite3.Cursor',
-    conn: 'sqlite3.Connection',
-    table_name: str,
-    old_name: str,
-    new_name: str,
-):
-    """Rename a column in a table."""
-    # NOTE: This only works for sqlite3 >= 3.25.0. Be careful to use this.
-
-    for row in cursor.execute(f'PRAGMA table_info({table_name})'):
-        if row[1] == old_name:
-            cursor.execute(f'ALTER TABLE {table_name} '
-                           f'RENAME COLUMN {old_name} to {new_name}')
-            break
-    conn.commit()
-
-
 class SQLiteConn(threading.local):
     """Thread-local connection to the sqlite3 database."""
 

sky/utils/kubernetes/exec_kubeconfig_converter.py

@@ -5,6 +5,9 @@ the 'command' field in the exec configuration, leaving only the executable name.
 This is useful when moving between different environments where auth plugin
 executables might be installed in different locations.
 
+For Nebius kubeconfigs, it also changes the --profile argument to 'sky' to
+ensure compatibility with SkyPilot's expected profile configuration.
+
 It assumes the target environment has the auth executable available in PATH.
 If not, you'll need to update your environment container to include the auth
 executable in PATH.
@@ -21,6 +24,8 @@ import yaml
 def strip_auth_plugin_paths(kubeconfig_path: str, output_path: str):
     """Strip path information from exec plugin commands in a kubeconfig file.
 
+    For Nebius kubeconfigs, also changes the --profile argument to 'sky'.
+
     Args:
         kubeconfig_path (str): Path to the input kubeconfig file
         output_path (str): Path where the modified kubeconfig will be saved
@@ -40,6 +45,20 @@ def strip_auth_plugin_paths(kubeconfig_path: str, output_path: str):
                 exec_info['command'] = executable
                 updated = True
 
+            # Handle Nebius kubeconfigs: change --profile to 'sky'
+            if executable == 'nebius' or current_command == 'nebius':
+                args = exec_info.get('args', [])
+                if args and '--profile' in args:
+                    try:
+                        profile_index = args.index('--profile')
+                        if profile_index + 1 < len(args):
+                            old_profile = args[profile_index + 1]
+                            if old_profile != 'sky':
+                                args[profile_index + 1] = 'sky'
+                                updated = True
+                    except ValueError:
+                        pass  # --profile not found in args
+
     if updated:
         with open(output_path, 'w', encoding='utf-8') as file:
             yaml.safe_dump(config, file)

sky/utils/schemas.py

@@ -1044,6 +1044,25 @@ def get_config_schema():
                 },
             }
         },
+        'ssh': {
+            'type': 'object',
+            'required': [],
+            'additionalProperties': False,
+            'properties': {
+                'allowed_node_pools': {
+                    'type': 'array',
+                    'items': {
+                        'type': 'string',
+                    },
+                },
+                'pod_config': {
+                    'type': 'object',
+                    'required': [],
+                    # Allow arbitrary keys since validating pod spec is hard
+                    'additionalProperties': True,
+                },
+            }
+        },
         'oci': {
             'type': 'object',
             'required': [],
@@ -1177,12 +1196,13 @@ def get_config_schema():
 
     allowed_workspace_cloud_names = list(
         service_catalog.ALL_CLOUDS) + ['cloudflare']
-    # Create pattern for non-GCP clouds (all clouds except gcp)
-    non_gcp_clouds = [
+    # Create pattern for not supported clouds, i.e.
+    # all clouds except gcp, kubernetes, ssh
+    not_supported_clouds = [
         cloud for cloud in allowed_workspace_cloud_names
-        if cloud.lower() != 'gcp'
+        if cloud.lower() not in ['gcp', 'kubernetes', 'ssh']
     ]
-    non_gcp_cloud_regex = '|'.join(non_gcp_clouds)
+    not_supported_cloud_regex = '|'.join(not_supported_clouds)
     workspaces_schema = {
         'type': 'object',
         'required': [],
@@ -1192,7 +1212,7 @@ def get_config_schema():
             'additionalProperties': False,
             'patternProperties': {
                 # Pattern for non-GCP clouds - only allows 'disabled' property
-                f'^({non_gcp_cloud_regex})$': {
+                f'^({not_supported_cloud_regex})$': {
                     'type': 'object',
                     'additionalProperties': False,
                     'properties': {
@@ -1217,6 +1237,38 @@ def get_config_schema():
                     },
                     'additionalProperties': False,
                 },
+                'ssh': {
+                    'type': 'object',
+                    'required': [],
+                    'properties': {
+                        'allowed_node_pools': {
+                            'type': 'array',
+                            'items': {
+                                'type': 'string',
+                            },
+                        },
+                        'disabled': {
+                            'type': 'boolean'
+                        },
+                    },
+                    'additionalProperties': False,
+                },
+                'kubernetes': {
+                    'type': 'object',
+                    'required': [],
+                    'properties': {
+                        'allowed_contexts': {
+                            'type': 'array',
+                            'items': {
+                                'type': 'string',
+                            },
+                        },
+                        'disabled': {
+                            'type': 'boolean'
+                        },
+                    },
+                    'additionalProperties': False,
+                },
             },
         },
     }

sky/utils/subprocess_utils.py

@@ -246,11 +246,10 @@ def kill_process_with_grace_period(proc: GenericProcess,
         # The child process may have already been terminated.
         return
     except psutil.TimeoutExpired:
-        # Pass to finally to force kill the process.
-        pass
-    finally:
         logger.debug(f'Process {proc.pid} did not terminate after '
                      f'{grace_period} seconds')
+        # Continue to finally to force kill the process.
+    finally:
         # Attempt to force kill if the normal termination fails
         if not force:
             logger.debug(f'Force killing process {proc.pid}')

sky/workspaces/core.py

@@ -13,6 +13,7 @@ from sky import skypilot_config
 from sky.skylet import constants
 from sky.usage import usage_lib
 from sky.utils import common_utils
+from sky.utils import config_utils
 from sky.utils import schemas
 
 logger = sky_logging.init_logger(__name__)
@@ -88,70 +89,103 @@ def _check_workspace_has_no_active_resources(workspace_name: str,
     Raises:
         ValueError: If the workspace has active clusters or managed jobs.
     """
+    _check_workspaces_have_no_active_resources([(workspace_name, operation)])
 
-    def check_clusters():
-        # Check for active clusters
-        all_clusters = global_user_state.get_clusters()
-        workspace_clusters = [
-            cluster for cluster in all_clusters
-            if (cluster.get('workspace', constants.SKYPILOT_DEFAULT_WORKSPACE)
-                == workspace_name)
-        ]
-        return workspace_clusters
 
-    def check_managed_jobs():
-        # Check for active managed jobs using the jobs controller
-        # pylint: disable=import-outside-toplevel
-        from sky.jobs.server import core as managed_jobs_core
+def _check_workspaces_have_no_active_resources(
+        workspace_operations: list) -> None:
+    """Check if workspaces have active clusters or managed jobs.
 
-        try:
-            # Get active managed jobs from the jobs controller
-            # (skip_finished=True)
-            managed_jobs = managed_jobs_core.queue(refresh=False,
-                                                   skip_finished=True,
-                                                   all_users=True)
-
-            workspace_active_jobs = [
-                job for job in managed_jobs
-                if job.get('workspace', constants.SKYPILOT_DEFAULT_WORKSPACE) ==
-                workspace_name
-            ]
+    Args:
+        workspace_operations: List of tuples (workspace_name, operation) where
+            operation is 'update' or 'delete'.
 
-            return workspace_active_jobs
+    Raises:
+        ValueError: If any workspace has active clusters or managed jobs.
+            The error message will include all workspaces with issues.
+    """
+    if not workspace_operations:
+        return
+
+    def get_all_clusters():
+        return global_user_state.get_clusters()
 
+    def get_all_managed_jobs():
+        # 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)
         except exceptions.ClusterNotUpError:
-            # If we can't check managed jobs (e.g., controller not running),
-            # log a warning but don't fail the operation
             logger.warning('All jobs should be finished in workspace.')
             return []
 
-    # Run both checks in parallel
+    # Fetch both clusters and jobs in parallel
     with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
-        cluster_future = executor.submit(check_clusters)
-        jobs_future = executor.submit(check_managed_jobs)
+        clusters_future = executor.submit(get_all_clusters)
+        jobs_future = executor.submit(get_all_managed_jobs)
+
+        all_clusters = clusters_future.result()
+        all_managed_jobs = jobs_future.result()
 
-        # Wait for both to complete
-        workspace_clusters = cluster_future.result()
-        workspace_active_jobs = jobs_future.result()
+    # Collect all error messages instead of raising immediately
+    error_messages = []
+
+    # Check each workspace against the fetched data
+    for workspace_name, operation in workspace_operations:
+        # Filter clusters for this workspace
+        workspace_clusters = [
+            cluster for cluster in all_clusters
+            if (cluster.get('workspace', constants.SKYPILOT_DEFAULT_WORKSPACE)
+                == workspace_name)
+        ]
 
-    # Check results
-    if workspace_clusters:
-        active_cluster_names = [
-            cluster['name'] for cluster in workspace_clusters
+        # Filter managed jobs for this workspace
+        workspace_active_jobs = [
+            job for job in all_managed_jobs
+            if job.get('workspace', constants.SKYPILOT_DEFAULT_WORKSPACE) ==
+            workspace_name
         ]
-        cluster_list = ', '.join(active_cluster_names)
-        raise ValueError(
-            f'Cannot {operation} workspace {workspace_name!r} because it has '
-            f'{len(workspace_clusters)} active cluster(s): {cluster_list}. '
-            f'Please terminate these clusters first.')
 
-    if workspace_active_jobs:
-        job_names = [job['job_id'] for job in workspace_active_jobs]
-        job_list = ', '.join(job_names)
-        raise ValueError(
-            f'Cannot {operation} workspace {workspace_name!r} because it has '
-            f'{len(workspace_active_jobs)} active managed job(s): '
-            f'{job_list}. Please cancel these jobs first.')
+        # Collect error messages for this workspace
+        workspace_errors = []
+
+        if workspace_clusters:
+            active_cluster_names = [
+                cluster['name'] for cluster in workspace_clusters
+            ]
+            cluster_list = ', '.join(active_cluster_names)
+            workspace_errors.append(
+                f'{len(workspace_clusters)} active cluster(s): {cluster_list}')
+
+        if workspace_active_jobs:
+            job_names = [job['job_id'] for job in workspace_active_jobs]
+            job_list = ', '.join(job_names)
+            workspace_errors.append(
+                f'{len(workspace_active_jobs)} active managed job(s): '
+                f'{job_list}')
+
+        # If this workspace has issues, add to overall error messages
+        if workspace_errors:
+            workspace_error_summary = ' and '.join(workspace_errors)
+            error_messages.append(
+                f'Cannot {operation} workspace {workspace_name!r} because it '
+                f'has {workspace_error_summary}.')
+
+    # If we collected any errors, raise them all together
+    if error_messages:
+        if len(error_messages) == 1:
+            # Single workspace error
+            full_message = error_messages[
+                0] + ' Please terminate these resources first.'
+        else:
+            # Multiple workspace errors
+            full_message = (f'Cannot proceed due to active resources in '
+                            f'{len(error_messages)} workspace(s):\n' +
+                            '\n'.join(f'• {msg}' for msg in error_messages) +
+                            '\nPlease terminate these resources first.')
+        raise ValueError(full_message)
 
 
 def _validate_workspace_config(workspace_name: str,
@@ -293,3 +327,105 @@ def delete_workspace(workspace_name: str) -> Dict[str, Any]:
 
     # Use the internal helper function to save
     return _update_workspaces_config(delete_workspace_fn)
+
+
+# =========================
+# = Config Management =
+# =========================
+
+
+@usage_lib.entrypoint
+def get_config() -> Dict[str, Any]:
+    """Returns the entire SkyPilot configuration.
+
+    Returns:
+        The complete SkyPilot configuration as a dictionary.
+    """
+    return skypilot_config.to_dict()
+
+
+@usage_lib.entrypoint
+def update_config(config: Dict[str, Any]) -> Dict[str, Any]:
+    """Updates the entire SkyPilot configuration.
+
+    Args:
+        config: The new configuration to save.
+
+    Returns:
+        The updated configuration.
+
+    Raises:
+        ValueError: If the configuration is invalid, or if there are
+            active clusters or managed jobs in workspaces being modified.
+        FileNotFoundError: If the config file cannot be found.
+        PermissionError: If the config file cannot be written.
+    """
+    # Validate the configuration using the schema
+    try:
+        common_utils.validate_schema(config, schemas.get_config_schema(),
+                                     'Invalid SkyPilot configuration: ')
+    except exceptions.InvalidSkyPilotConfigError as e:
+        raise ValueError(str(e)) from e
+
+    # Check for API server changes and validate them
+    current_config = skypilot_config.to_dict()
+
+    current_endpoint = current_config.get('api_server', {}).get('endpoint')
+    new_endpoint = config.get('api_server', {}).get('endpoint')
+    if current_endpoint != new_endpoint:
+        raise ValueError('API server endpoint should not be changed to avoid '
+                         'unexpected behavior.')
+
+    # Check for workspace changes and validate them
+    current_workspaces = current_config.get('workspaces', {})
+    new_workspaces = config.get('workspaces', {})
+
+    # Collect all workspaces that need to be checked for active resources
+    workspaces_to_check = []
+
+    # Check each workspace that is being modified
+    for workspace_name, new_workspace_config in new_workspaces.items():
+        current_workspace_config = current_workspaces.get(workspace_name, {})
+
+        # If workspace configuration is changing, validate and mark for checking
+        if current_workspace_config != new_workspace_config:
+            _validate_workspace_config(workspace_name, new_workspace_config)
+            workspaces_to_check.append((workspace_name, 'update'))
+
+    # Check for workspace deletions
+    for workspace_name in current_workspaces:
+        if workspace_name not in new_workspaces:
+            # Workspace is being deleted
+            if workspace_name == constants.SKYPILOT_DEFAULT_WORKSPACE:
+                raise ValueError(f'Cannot delete the default workspace '
+                                 f'{constants.SKYPILOT_DEFAULT_WORKSPACE!r}.')
+            workspaces_to_check.append((workspace_name, 'delete'))
+
+    # Check all workspaces for active resources in one efficient call
+    _check_workspaces_have_no_active_resources(workspaces_to_check)
+
+    # Use file locking to prevent race conditions
+    lock_path = skypilot_config.get_skypilot_config_lock_path()
+    try:
+        with filelock.FileLock(lock_path,
+                               _WORKSPACE_CONFIG_LOCK_TIMEOUT_SECONDS):
+            # Convert to config_utils.Config and save
+            config_obj = config_utils.Config.from_dict(config)
+            skypilot_config.update_config_no_lock(config_obj)
+    except filelock.Timeout as e:
+        raise RuntimeError(
+            f'Failed to update configuration due to a timeout '
+            f'when trying to acquire the lock at {lock_path}. This may '
+            'indicate another SkyPilot process is currently updating the '
+            'configuration. Please try again or manually remove the lock '
+            f'file if you believe it is stale.') from e
+
+    # Validate the configuration by running sky check
+    try:
+        sky_check.check(quiet=True)
+    except Exception as e:  # pylint: disable=broad-except
+        logger.warning(f'Configuration saved but '
+                       f'validation check failed: {e}')
+        # Don't fail the update if the check fails, just warn
+
+    return config