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

sky/server/uvicorn.py

@@ -3,17 +3,165 @@
 This module is a wrapper around uvicorn to customize the behavior of the
 server.
 """
-import functools
+import asyncio
 import os
+import signal
 import threading
-from typing import Optional
+import time
+from types import FrameType
+from typing import Optional, Union
 
+import filelock
 import uvicorn
 from uvicorn.supervisors import multiprocess
 
+from sky import sky_logging
+from sky.server import state
+from sky.server.requests import requests as requests_lib
+from sky.skylet import constants
 from sky.utils import context_utils
 from sky.utils import subprocess_utils
 
+logger = sky_logging.init_logger(__name__)
+
+# File lock path for coordinating graceful shutdown across processes
+_GRACEFUL_SHUTDOWN_LOCK_PATH = '/tmp/skypilot_graceful_shutdown.lock'
+
+# Interval to check for on-going requests.
+_WAIT_REQUESTS_INTERVAL_SECONDS = 5
+
+# Timeout for waiting for on-going requests to finish.
+try:
+    _WAIT_REQUESTS_TIMEOUT_SECONDS = int(
+        os.environ.get(constants.GRACE_PERIOD_SECONDS_ENV_VAR, '60'))
+except ValueError:
+    _WAIT_REQUESTS_TIMEOUT_SECONDS = 60
+
+# TODO(aylei): use decorator to register requests that need to be proactively
+# cancelled instead of hardcoding here.
+_RETRIABLE_REQUEST_NAMES = [
+    'sky.logs',
+    'sky.jobs.logs',
+    'sky.serve.logs',
+]
+
+
+class Server(uvicorn.Server):
+    """Server wrapper for uvicorn.
+
+    Extended functionalities:
+    - Handle exit signal and perform custom graceful shutdown.
+    - Run the server process with contextually aware.
+    """
+
+    def __init__(self, config: uvicorn.Config):
+        super().__init__(config=config)
+        self.exiting: bool = False
+
+    def handle_exit(self, sig: int, frame: Union[FrameType, None]) -> None:
+        """Handle exit signal.
+
+        When a server process receives a SIGTERM or SIGINT signal, a graceful
+        shutdown will be initiated. If a SIGINT signal is received again, the
+        server will be forcefully shutdown.
+        """
+        if self.exiting and sig == signal.SIGINT:
+            # The server has been siganled to exit and recieved a SIGINT again,
+            # do force shutdown.
+            logger.info('Force shutdown.')
+            self.should_exit = True
+            super().handle_exit(sig, frame)
+            return
+        if not self.exiting:
+            self.exiting = True
+            # Perform graceful shutdown in a separate thread to avoid blocking
+            # the main thread.
+            threading.Thread(target=self._graceful_shutdown,
+                             args=(sig, frame),
+                             daemon=True).start()
+
+    def _graceful_shutdown(self, sig: int, frame: Union[FrameType,
+                                                        None]) -> None:
+        """Perform graceful shutdown."""
+        # Block new requests so that we can wait until all on-going requests
+        # are finished. Note that /api/$verb operations are still allowed in
+        # this stage to ensure the client can still operate the on-going
+        # requests, e.g. /api/logs, /api/cancel, etc.
+        logger.info('Block new requests being submitted in worker '
+                    f'{os.getpid()}.')
+        state.set_block_requests(True)
+        # Ensure the shutting_down are set on all workers before next step.
+        # TODO(aylei): hacky, need a reliable solution.
+        time.sleep(1)
+
+        lock = filelock.FileLock(_GRACEFUL_SHUTDOWN_LOCK_PATH)
+        # Elect a coordinator process to handle on-going requests check
+        with lock.acquire():
+            logger.info(f'Worker {os.getpid()} elected as shutdown coordinator')
+            self._wait_requests()
+
+        logger.info('Shutting down server...')
+        self.should_exit = True
+        super().handle_exit(sig, frame)
+
+    def _wait_requests(self) -> None:
+        """Wait until all on-going requests are finished or cancelled."""
+        start_time = time.time()
+        while True:
+            statuses = [
+                requests_lib.RequestStatus.PENDING,
+                requests_lib.RequestStatus.RUNNING,
+            ]
+            reqs = requests_lib.get_request_tasks(status=statuses)
+            if not reqs:
+                break
+            logger.info(f'{len(reqs)} on-going requests '
+                        'found, waiting for them to finish...')
+            # Proactively cancel internal requests and logs requests since
+            # they can run for infinite time.
+            internal_request_ids = [
+                d.id for d in requests_lib.INTERNAL_REQUEST_DAEMONS
+            ]
+            if time.time() - start_time > _WAIT_REQUESTS_TIMEOUT_SECONDS:
+                logger.warning('Timeout waiting for on-going requests to '
+                               'finish, cancelling all on-going requests.')
+                for req in reqs:
+                    self.interrupt_request_for_retry(req.request_id)
+                break
+            interrupted = 0
+            for req in reqs:
+                if req.request_id in internal_request_ids:
+                    self.interrupt_request_for_retry(req.request_id)
+                    interrupted += 1
+                elif req.name in _RETRIABLE_REQUEST_NAMES:
+                    self.interrupt_request_for_retry(req.request_id)
+                    interrupted += 1
+                # TODO(aylei): interrupt pending requests to accelerate the
+                # shutdown.
+            # If some requests are not interrupted, wait for them to finish,
+            # otherwise we just check again immediately to accelerate the
+            # shutdown process.
+            if interrupted < len(reqs):
+                time.sleep(_WAIT_REQUESTS_INTERVAL_SECONDS)
+
+    def interrupt_request_for_retry(self, request_id: str) -> None:
+        """Interrupt a request for retry."""
+        with requests_lib.update_request(request_id) as req:
+            if req is None:
+                return
+            if req.pid is not None:
+                os.kill(req.pid, signal.SIGTERM)
+            req.status = requests_lib.RequestStatus.CANCELLED
+            req.should_retry = True
+        logger.info(
+            f'Request {request_id} interrupted and will be retried by client.')
+
+    def run(self, *args, **kwargs):
+        """Run the server process."""
+        context_utils.hijack_sys_attrs()
+        with self.capture_signals():
+            asyncio.run(self.serve(*args, **kwargs))
+
 
 def run(config: uvicorn.Config):
     """Run unvicorn server."""
@@ -22,28 +170,20 @@ def run(config: uvicorn.Config):
         # in uvicorn. Since we do not use reload now, simply
         # guard by an exception.
         raise ValueError('Reload is not supported yet.')
-    server = uvicorn.Server(config=config)
-    run_server_process = functools.partial(_run_server_process, server)
+    server = Server(config=config)
     try:
         if config.workers is not None and config.workers > 1:
             sock = config.bind_socket()
-            SlowStartMultiprocess(config,
-                                  target=run_server_process,
+            SlowStartMultiprocess(config, target=server.run,
                                   sockets=[sock]).run()
         else:
-            run_server_process()
+            server.run()
     finally:
         # Copied from unvicorn.run()
         if config.uds and os.path.exists(config.uds):
             os.remove(config.uds)
 
 
-def _run_server_process(server: uvicorn.Server, *args, **kwargs):
-    """Run the server process with contextually aware."""
-    context_utils.hijack_sys_attrs()
-    server.run(*args, **kwargs)
-
-
 class SlowStartMultiprocess(multiprocess.Multiprocess):
     """Uvicorn Multiprocess wrapper with slow start.
 

sky/setup_files/dependencies.py

@@ -62,6 +62,8 @@ install_requires = [
     # the client-side actually not importing them.
     'casbin',
     'sqlalchemy_adapter',
+    # Required for API server metrics
+    'prometheus_client>=0.8.0',
     'passlib',
 ]
 

sky/skylet/constants.py

@@ -413,6 +413,13 @@ SKY_USER_FILE_PATH = '~/.sky/generated'
 # Environment variable that is set to 'true' if this is a skypilot server.
 ENV_VAR_IS_SKYPILOT_SERVER = 'IS_SKYPILOT_SERVER'
 
+# Environment variable that is set to 'true' if metrics are enabled.
+ENV_VAR_SERVER_METRICS_ENABLED = 'SKY_API_SERVER_METRICS_ENABLED'
+
+# Environment variable that is used as the DB connection string for the
+# skypilot server.
+ENV_VAR_DB_CONNECTION_URI = (f'{SKYPILOT_ENV_VAR_PREFIX}DB_CONNECTION_URI')
+
 # Environment variable that is set to 'true' if basic
 # authentication is enabled in the API server.
 ENV_VAR_ENABLE_BASIC_AUTH = 'ENABLE_BASIC_AUTH'
@@ -449,20 +456,29 @@ TIME_PATTERN: str = (
 
 MEMORY_SIZE_UNITS = {
     'kb': 2**10,
+    'ki': 2**10,
     'mb': 2**20,
+    'mi': 2**20,
     'gb': 2**30,
+    'gi': 2**30,
     'tb': 2**40,
+    'ti': 2**40,
     'pb': 2**50,
+    'pi': 2**50,
 }
 
 MEMORY_SIZE_PATTERN = (
     '^[0-9]+('
-    f'{"|".join([unit.lower() for unit in MEMORY_SIZE_UNITS])}'
-    ')?$/i')
-MEMORY_SIZE_PLUS_PATTERN = f'{MEMORY_SIZE_PATTERN[:-3]}+?$/i'
+    f'{"|".join([unit.lower() for unit in MEMORY_SIZE_UNITS])}|'
+    f'{"|".join([unit.upper() for unit in MEMORY_SIZE_UNITS])}|'
+    f'{"|".join([unit[0].upper() + unit[1:] for unit in MEMORY_SIZE_UNITS if len(unit) > 1])}'  # pylint: disable=line-too-long
+    ')?$')
+
+LAST_USE_TRUNC_LENGTH = 25
 
 MIN_PRIORITY = -1000
 MAX_PRIORITY = 1000
 DEFAULT_PRIORITY = 0
 
+GRACE_PERIOD_SECONDS_ENV_VAR = SKYPILOT_ENV_VAR_PREFIX + 'GRACE_PERIOD_SECONDS'
 COST_REPORT_DEFAULT_DAYS = 30

sky/skypilot_config.py

@@ -564,7 +564,10 @@ def _reload_config_as_server() -> None:
     _set_loaded_config_path(None)
 
     server_config_path = _resolve_server_config_path()
+    db_url_from_env = os.environ.get(constants.ENV_VAR_DB_CONNECTION_URI)
     server_config = _get_config_from_path(server_config_path)
+    if db_url_from_env:
+        server_config.set_nested(('db',), db_url_from_env)
 
     if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
         logger.debug(f'server config: \n'

sky/ssh_node_pools/__init__.py

@@ -0,0 +1 @@
+"""SSH Node Pool management package."""

sky/ssh_node_pools/core.py

@@ -0,0 +1,133 @@
+"""SSH Node Pool management core functionality."""
+import os
+from pathlib import Path
+from typing import Any, Dict, List
+
+import yaml
+
+
+class SSHNodePoolManager:
+    """Manager for SSH Node Pool configurations."""
+
+    def __init__(self):
+        self.config_path = Path.home() / '.sky' / 'ssh_node_pools.yaml'
+        self.keys_dir = Path.home() / '.sky' / 'ssh_keys'
+        self.keys_dir.mkdir(parents=True, exist_ok=True)
+
+    def get_all_pools(self) -> Dict[str, Any]:
+        """Read all SSH Node Pool configurations from YAML file."""
+        if not self.config_path.exists():
+            return {}
+
+        try:
+            with open(self.config_path, 'r', encoding='utf-8') as f:
+                return yaml.safe_load(f) or {}
+        except Exception as e:
+            raise RuntimeError(
+                f'Failed to read SSH Node Pool config: {e}') from e
+
+    def save_all_pools(self, pools_config: Dict[str, Any]) -> None:
+        """Write SSH Node Pool configurations to YAML file."""
+        try:
+            self.config_path.parent.mkdir(parents=True, exist_ok=True)
+            with open(self.config_path, 'w', encoding='utf-8') as f:
+                yaml.dump(pools_config, f, default_flow_style=False)
+        except Exception as e:
+            raise RuntimeError(
+                f'Failed to save SSH Node Pool config: {e}') from e
+
+    def update_pools(self, pools_config: Dict[str, Any]) -> None:
+        """Update SSH Node Pool configurations."""
+        all_pools = self.get_all_pools()
+        all_pools.update(pools_config)
+        self.save_all_pools(all_pools)
+
+    def add_or_update_pool(self, pool_name: str,
+                           pool_config: Dict[str, Any]) -> None:
+        """Add or update a single SSH Node Pool configuration."""
+        # Validate pool configuration
+        self._validate_pool_config(pool_config)
+
+        all_pools = self.get_all_pools()
+        all_pools[pool_name] = pool_config
+        self.save_all_pools(all_pools)
+
+    def delete_pool(self, pool_name: str) -> bool:
+        """Delete a SSH Node Pool configuration."""
+        all_pools = self.get_all_pools()
+        if pool_name in all_pools:
+            del all_pools[pool_name]
+            self.save_all_pools(all_pools)
+            return True
+        return False
+
+    def save_ssh_key(self, key_name: str, key_content: str) -> str:
+        """Save SSH private key to ~/.sky/ssh_keys/ directory."""
+        # Validate key name
+        if not key_name or '/' in key_name or key_name.startswith('.'):
+            raise ValueError('Invalid key name')
+
+        key_path = self.keys_dir / key_name
+        try:
+            with open(key_path, 'w', encoding='utf-8') as f:
+                f.write(key_content)
+            os.chmod(key_path, 0o600)  # Set secure permissions
+            return str(key_path)
+        except Exception as e:
+            raise RuntimeError(f'Failed to save SSH key: {e}') from e
+
+    def list_ssh_keys(self) -> List[str]:
+        """List available SSH key files."""
+        if not self.keys_dir.exists():
+            return []
+        try:
+            return [f.name for f in self.keys_dir.iterdir() if f.is_file()]
+        except Exception:  # pylint: disable=broad-except
+            return []
+
+    def _validate_pool_config(self, config: Dict[str, Any]) -> None:
+        """Validate SSH Node Pool configuration."""
+        if 'hosts' not in config:
+            raise ValueError('Pool configuration must include `hosts`')
+
+        if not isinstance(config['hosts'], list) or not config['hosts']:
+            raise ValueError('`hosts` must be a non-empty list')
+
+        # Validate user field
+        if not config.get('user', '').strip():
+            raise ValueError('Pool configuration must include `user`')
+
+        # Validate authentication - must have either identity_file or password
+        if not config.get('identity_file') and not config.get('password'):
+            raise ValueError('Pool configuration must include '
+                             'either `identity_file` or `password`')
+
+
+def get_all_pools() -> Dict[str, Any]:
+    """Get all SSH Node Pool configurations."""
+    manager = SSHNodePoolManager()
+    return manager.get_all_pools()
+
+
+def update_pools(pools_config: Dict[str, Any]) -> None:
+    """Update SSH Node Pool configurations."""
+    manager = SSHNodePoolManager()
+    manager.update_pools(pools_config)
+
+
+def delete_pool(pool_name: str) -> bool:
+    """Delete a SSH Node Pool configuration."""
+    manager = SSHNodePoolManager()
+    return manager.delete_pool(pool_name)
+
+
+def upload_ssh_key(key_name: str, key_content: str) -> str:
+    """Upload SSH private key."""
+    manager = SSHNodePoolManager()
+    return manager.save_ssh_key(key_name, key_content)
+
+
+def list_ssh_keys() -> List[str]:
+    """List available SSH keys."""
+    manager = SSHNodePoolManager()
+    return manager.list_ssh_keys()

sky/ssh_node_pools/server.py

@@ -0,0 +1,232 @@
+"""SSH Node Pool management API endpoints."""
+import re
+from typing import Any, Dict, List
+
+import fastapi
+
+from sky import core as sky_core
+from sky.server.requests import executor
+from sky.server.requests import payloads
+from sky.server.requests import requests as requests_lib
+from sky.ssh_node_pools import core as ssh_node_pools_core
+from sky.utils import common_utils
+
+router = fastapi.APIRouter()
+
+
+@router.get('')
+async def get_ssh_node_pools() -> Dict[str, Any]:
+    """Get all SSH Node Pool configurations."""
+    try:
+        return ssh_node_pools_core.get_all_pools()
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=
+            f'Failed to get SSH Node Pools: {common_utils.format_exception(e)}')
+
+
+@router.post('')
+async def update_ssh_node_pools(pools_config: Dict[str, Any]) -> Dict[str, str]:
+    """Update SSH Node Pool configurations."""
+    try:
+        ssh_node_pools_core.update_pools(pools_config)
+        return {'status': 'success'}
+    except Exception as e:
+        raise fastapi.HTTPException(status_code=400,
+                                    detail=f'Failed to update SSH Node Pools:'
+                                    f' {common_utils.format_exception(e)}')
+
+
+@router.delete('/{pool_name}')
+async def delete_ssh_node_pool(pool_name: str) -> Dict[str, str]:
+    """Delete a SSH Node Pool configuration."""
+    try:
+        if ssh_node_pools_core.delete_pool(pool_name):
+            return {'status': 'success'}
+        else:
+            raise fastapi.HTTPException(
+                status_code=404,
+                detail=f'SSH Node Pool `{pool_name}` not found')
+    except fastapi.HTTPException:
+        raise
+    except Exception as e:
+        raise fastapi.HTTPException(status_code=500,
+                                    detail='Failed to delete SSH Node Pool: '
+                                    f'{common_utils.format_exception(e)}')
+
+
+@router.post('/keys')
+async def upload_ssh_key(request: fastapi.Request) -> Dict[str, str]:
+    """Upload SSH private key."""
+    try:
+        form = await request.form()
+        key_name = form.get('key_name')
+        key_file = form.get('key_file')
+
+        if not key_name or not key_file:
+            raise fastapi.HTTPException(status_code=400,
+                                        detail='Missing key_name or key_file')
+
+        key_content = await key_file.read()
+        key_path = ssh_node_pools_core.upload_ssh_key(key_name,
+                                                      key_content.decode())
+
+        return {'status': 'success', 'key_path': key_path}
+    except fastapi.HTTPException:
+        raise
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=
+            f'Failed to upload SSH key: {common_utils.format_exception(e)}')
+
+
+@router.get('/keys')
+async def list_ssh_keys() -> List[str]:
+    """List available SSH keys."""
+    try:
+        return ssh_node_pools_core.list_ssh_keys()
+    except Exception as e:
+        exception_msg = common_utils.format_exception(e)
+        raise fastapi.HTTPException(
+            status_code=500, detail=f'Failed to list SSH keys: {exception_msg}')
+
+
+@router.post('/{pool_name}/deploy')
+async def deploy_ssh_node_pool(request: fastapi.Request,
+                               pool_name: str) -> Dict[str, str]:
+    """Deploy SSH Node Pool using existing ssh_up functionality."""
+    try:
+        ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=False)
+        executor.schedule_request(
+            request_id=request.state.request_id,
+            request_name='ssh_up',
+            request_body=ssh_up_body,
+            func=sky_core.ssh_up,
+            schedule_type=requests_lib.ScheduleType.LONG,
+        )
+
+        return {
+            'status': 'success',
+            'request_id': request.state.request_id,
+            'message': f'SSH Node Pool `{pool_name}` deployment started'
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(status_code=500,
+                                    detail=f'Failed to deploy SSH Node Pool: '
+                                    f'{common_utils.format_exception(e)}')
+
+
+@router.post('/deploy')
+async def deploy_ssh_node_pool_general(
+        request: fastapi.Request,
+        ssh_up_body: payloads.SSHUpBody) -> Dict[str, str]:
+    """Deploys all SSH Node Pools."""
+    try:
+        executor.schedule_request(
+            request_id=request.state.request_id,
+            request_name='ssh_up',
+            request_body=ssh_up_body,
+            func=sky_core.ssh_up,
+            schedule_type=requests_lib.ScheduleType.LONG,
+        )
+
+        pool_name = ssh_up_body.infra or 'default'
+        return {
+            'status': 'success',
+            'request_id': request.state.request_id,
+            'message': f'SSH Node Pool `{pool_name}` deployment started'
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(status_code=500,
+                                    detail=f'Failed to deploy SSH Node Pool: '
+                                    f'{common_utils.format_exception(e)}')
+
+
+@router.post('/{pool_name}/down')
+async def down_ssh_node_pool(request: fastapi.Request,
+                             pool_name: str) -> Dict[str, str]:
+    """Cleans up a SSH Node Pools."""
+    try:
+        ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=True)
+        executor.schedule_request(
+            request_id=request.state.request_id,
+            request_name='ssh_down',
+            request_body=ssh_up_body,
+            func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
+            schedule_type=requests_lib.ScheduleType.LONG,
+        )
+
+        return {
+            'status': 'success',
+            'request_id': request.state.request_id,
+            'message': f'SSH Node Pool `{pool_name}` teardown started'
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=f'Failed to tear down SSH Node Pool: '
+            f'{common_utils.format_exception(e)}')
+
+
+@router.post('/down')
+async def down_ssh_node_pool_general(
+        request: fastapi.Request,
+        ssh_up_body: payloads.SSHUpBody) -> Dict[str, str]:
+    """Cleans up all SSH Node Pools."""
+    try:
+        # Set cleanup=True for down operation
+        ssh_up_body.cleanup = True
+        executor.schedule_request(
+            request_id=request.state.request_id,
+            request_name='ssh_down',
+            request_body=ssh_up_body,
+            func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
+            schedule_type=requests_lib.ScheduleType.LONG,
+        )
+
+        pool_name = ssh_up_body.infra or 'default'
+        return {
+            'status': 'success',
+            'request_id': request.state.request_id,
+            'message': f'SSH Node Pool `{pool_name}` teardown started'
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=f'Failed to tear down SSH Node Pool: '
+            f'{common_utils.format_exception(e)}')
+
+
+@router.get('/{pool_name}/status')
+async def get_ssh_node_pool_status(pool_name: str) -> Dict[str, str]:
+    """Get the status of a specific SSH Node Pool."""
+    try:
+        # Call ssh_status to check the context
+        context_name = f'ssh-{pool_name}'
+        is_ready, reason = sky_core.ssh_status(context_name)
+
+        # Strip ANSI escape codes from the reason text
+        def strip_ansi_codes(text):
+            if not text:
+                return text
+            # Remove ANSI escape sequences (color codes, formatting, etc.)
+            text = re.sub(r'\x1b\[[0-9;]*m', '', text)
+            # Remove 'disabled. Reason: ' prefix if present
+            text = text.replace('disabled. Reason: ', '')
+            return text
+
+        cleaned_reason = strip_ansi_codes(reason) if reason else reason
+
+        return {
+            'pool_name': pool_name,
+            'context_name': context_name,
+            'status': 'Ready' if is_ready else 'Not Ready',
+            'reason': cleaned_reason
+        }
+    except Exception as e:
+        raise fastapi.HTTPException(
+            status_code=500,
+            detail=f'Failed to get SSH Node Pool status: '
+            f'{common_utils.format_exception(e)}')