skypilot-nightly 1.0.0.dev20250520__py3-none-any.whl → 1.0.0.dev20250522__py3-none-any.whl

sky/server/common.py

@@ -12,7 +12,7 @@ import subprocess
 import sys
 import time
 import typing
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Literal, Optional
 from urllib import parse
 import uuid
 
@@ -116,6 +116,7 @@ class ApiServerStatus(enum.Enum):
     HEALTHY = 'healthy'
     UNHEALTHY = 'unhealthy'
     VERSION_MISMATCH = 'version_mismatch'
+    NEEDS_AUTH = 'needs_auth'
 
 
 @dataclasses.dataclass
@@ -127,15 +128,21 @@ class ApiServerInfo:
     commit: Optional[str] = None
 
 
+def get_api_cookie_jar_path() -> str:
+    return os.environ.get(server_constants.API_COOKIE_FILE_ENV_VAR,
+                          server_constants.API_COOKIE_FILE_DEFAULT_LOCATION)
+
+
 def get_api_cookie_jar() -> requests.cookies.RequestsCookieJar:
     """Returns the cookie jar used by the client to access the API server."""
-    cookie_file = os.environ.get(server_constants.API_COOKIE_FILE_ENV_VAR)
     cookie_jar = requests.cookies.RequestsCookieJar()
-    if cookie_file and os.path.exists(cookie_file):
+    cookie_file = get_api_cookie_jar_path()
+    if cookie_file:
         cookie_path = pathlib.Path(cookie_file).expanduser().resolve()
-        file_cookie_jar = MozillaCookieJar(cookie_path)
-        file_cookie_jar.load()
-        cookie_jar.update(file_cookie_jar)
+        if cookie_path.exists():
+            file_cookie_jar = MozillaCookieJar(cookie_path)
+            file_cookie_jar.load()
+            cookie_jar.update(file_cookie_jar)
     return cookie_jar
 
 
@@ -196,6 +203,7 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
             response = requests.get(f'{server_url}/api/health',
                                     timeout=2.5,
                                     cookies=get_api_cookie_jar())
+            logger.debug(f'Health check status: {response.status_code}')
             if response.status_code == 200:
                 try:
                     result = response.json()
@@ -217,9 +225,24 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
                         server_info.status = ApiServerStatus.VERSION_MISMATCH
                     return server_info
                 except (json.JSONDecodeError, AttributeError) as e:
+                    # Try to check if we got redirected to a login page.
+                    for prev_response in response.history:
+                        logger.debug(f'Previous response: {prev_response.url}')
+                        # Heuristic: check if the url looks like a login page or
+                        # oauth flow.
+                        if any(key in prev_response.url
+                               for key in ['login', 'oauth2']):
+                            logger.debug(
+                                f'URL {prev_response.url} looks like '
+                                'a login page or oauth flow, so try to '
+                                'get the cookie.')
+                            return ApiServerInfo(
+                                status=ApiServerStatus.NEEDS_AUTH)
                     logger.warning('Failed to parse API server response: '
                                    f'{str(e)}')
                     return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
+            elif response.status_code == 401:
+                return ApiServerInfo(status=ApiServerStatus.NEEDS_AUTH)
             else:
                 return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
         except requests.exceptions.Timeout:
@@ -297,7 +320,6 @@ def _start_api_server(deploy: bool = False,
 
         log_path = os.path.expanduser(constants.API_SERVER_LOGS)
         os.makedirs(os.path.dirname(log_path), exist_ok=True)
-        cmd = f'{" ".join(args)} > {log_path} 2>&1 < /dev/null'
 
         # Start the API server process in the background and don't wait for it.
         # If this is called from a CLI invocation, we need
@@ -305,10 +327,20 @@ def _start_api_server(deploy: bool = False,
         # the API server.
         server_env = os.environ.copy()
         server_env[constants.ENV_VAR_IS_SKYPILOT_SERVER] = 'true'
-        proc = subprocess.Popen(cmd,
-                                shell=True,
-                                start_new_session=True,
-                                env=server_env)
+        with open(log_path, 'w', encoding='utf-8') as log_file:
+            # Because the log file is opened using a with statement, it may seem
+            # that the file will be closed when the with statement is exited
+            # causing the child process to be unable to write to the log file.
+            # However, Popen makes the file descriptor inheritable which means
+            # the child process will inherit its own copy of the fd,
+            # independent of the parent's fd table which enables to child
+            # process to continue writing to the log file.
+            proc = subprocess.Popen(args,
+                                    stdout=log_file,
+                                    stderr=subprocess.STDOUT,
+                                    stdin=subprocess.DEVNULL,
+                                    start_new_session=True,
+                                    env=server_env)
 
         start_time = time.time()
         while True:
@@ -360,7 +392,12 @@ def _start_api_server(deploy: bool = False,
                 f'SkyPilot API server started. {dashboard_msg}'))
 
 
-def check_server_healthy(endpoint: Optional[str] = None,) -> None:
+def check_server_healthy(
+    endpoint: Optional[str] = None
+) -> Literal[
+        # Use an incomplete list of Literals here to enforce raising for other
+        # enum values.
+        ApiServerStatus.HEALTHY, ApiServerStatus.NEEDS_AUTH]:
     """Check if the API server is healthy.
 
     Args:
@@ -370,6 +407,11 @@ def check_server_healthy(endpoint: Optional[str] = None,) -> None:
     Raises:
         RuntimeError: If the server is not healthy or the client version does
             not match the server version.
+
+    Returns:
+        ApiServerStatus: The status of the API server, unless the server is
+            unhealthy or the client version does not match the server version,
+            in which case an exception is raised.
     """
     endpoint = endpoint if endpoint is not None else get_server_url()
     api_server_info = get_api_server_status(endpoint)
@@ -432,6 +474,8 @@ def check_server_healthy(endpoint: Optional[str] = None,) -> None:
 
         hinted_for_server_install_version_mismatch = True
 
+    return api_server_status
+
 
 def _get_version_info_hint(server_info: ApiServerInfo) -> str:
     assert server_info.version is not None, 'Server version is None'
@@ -482,11 +526,13 @@ def get_skypilot_version_on_disk() -> str:
 def check_server_healthy_or_start_fn(deploy: bool = False,
                                      host: str = '127.0.0.1',
                                      foreground: bool = False):
+    api_server_status = None
     try:
-        check_server_healthy()
+        api_server_status = check_server_healthy()
     except exceptions.ApiServerConnectionError as exc:
         endpoint = get_server_url()
-        if not is_api_server_local():
+        if (not is_api_server_local() or
+                api_server_status == ApiServerStatus.NEEDS_AUTH):
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.ApiServerConnectionError(endpoint) from exc
         # Lock to prevent multiple processes from starting the server at the

sky/server/constants.py

@@ -26,6 +26,8 @@ CLUSTER_REFRESH_DAEMON_INTERVAL_SECONDS = 60
 
 # Environment variable for a file path to the API cookie file.
 API_COOKIE_FILE_ENV_VAR = f'{constants.SKYPILOT_ENV_VAR_PREFIX}API_COOKIE_FILE'
+# Default file if unset.
+API_COOKIE_FILE_DEFAULT_LOCATION = '~/.sky/cookies.txt'
 
 # The path to the dashboard build output
 DASHBOARD_DIR = os.path.join(os.path.dirname(__file__), '..', 'dashboard',

sky/server/html/token_page.html

@@ -0,0 +1,154 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>SkyPilot API Server Login</title>
+    <style>
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
+            display: flex;
+            flex-direction: column;
+            align-items: center;
+            justify-content: center;
+            min-height: 100vh;
+            margin: 0;
+            background-color: #f8f9fa;
+            color: #202124;
+            padding: 20px;
+            box-sizing: border-box;
+        }
+        .container {
+            background-color: #ffffff;
+            padding: 48px;
+            border-radius: 8px;
+            box-shadow: 0 1px 3px rgba(0,0,0,0.12), 0 1px 2px rgba(0,0,0,0.24);
+            text-align: center;
+            max-width: 600px;
+            width: 100%;
+        }
+        .logo {
+            width: 64px;
+            height: 64px;
+            margin-bottom: 20px;
+            display: inline-block;
+        }
+        .logo svg {
+            width: 100%;
+            height: 100%;
+        }
+        h1 {
+            font-size: 24px;
+            font-weight: 500;
+            margin-bottom: 20px;
+            color: #202124;
+        }
+        p {
+            font-size: 14px;
+            line-height: 1.5;
+            margin-bottom: 20px;
+            color: #5f6368;
+        }
+        .code-block {
+            background-color: #f1f3f4;
+            border: 1px solid #dadce0;
+            border-radius: 4px;
+            padding: 16px;
+            margin-top: 24px;
+            margin-bottom: 24px;
+            margin-left: auto;
+            margin-right: auto;
+            text-align: left;
+            word-break: break-all;
+            white-space: pre-wrap;
+            font-family: "SFMono-Regular", Consolas, "Liberation Mono", Menlo, Courier, monospace;
+            font-size: 13px;
+            line-height: 1.4;
+            max-width: 480px;
+        }
+        #token-box { /* Specifically for the token */
+            height: auto;
+            min-height: 6em; /* Ensure it's a reasonable size */
+            max-height: 15em; /* Prevent it from getting too large */
+            overflow-y: auto;
+        }
+        .copy-button {
+            background-color: #1a73e8;
+            color: white;
+            border: none;
+            border-radius: 4px;
+            padding: 10px 24px;
+            font-size: 14px;
+            font-weight: 500;
+            cursor: pointer;
+            transition: background-color 0.3s;
+            margin-top: 10px;
+        }
+        .copy-button:hover {
+            background-color: #287ae6;
+        }
+        .copy-button:active {
+            background-color: #1b66c9;
+        }
+        .footer-text {
+            font-size: 12px;
+            color: #5f6368;
+            margin-top: 30px;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <div class="logo">
+            <!-- SkyPilot Logo Icon -->
+            <svg viewBox="0 0 50 50" fill="none" xmlns="http://www.w3.org/2000/svg">
+                <path d="M25.1258 30.8274L19.2842 31.6783L33.8316 46.2268L31.492 37.1925L25.1258 30.8274Z" fill="#372F8A"/>
+                <path d="M46.9433 0.000976562L0.719727 13.1148L15.2661 27.6601L16.633 21.3925L10.3728 15.1323L40.183 6.74118C40.183 6.74118 46.102 0.855027 46.9444 0.00203721L46.9433 0.000976562Z" fill="#372F8A"/>
+                <path d="M40.1821 6.74021L31.4922 37.1925L33.8318 46.2257L46.9445 0C46.1022 0.85299 40.1831 6.73915 40.1831 6.73915L40.1821 6.74021Z" fill="#372F8A"/>
+                <path d="M21.3356 25.6089L19.2842 31.6783L25.1258 30.8275L30.3741 16.6011L30.3275 16.617L21.3356 25.6089Z" fill="#195D7F"/>
+                <path d="M16.632 21.3918L15.2651 27.6605L21.3357 25.6091L30.3276 16.6172L16.632 21.3918Z" fill="#39A4DD"/>
+            </svg>
+        </div>
+        <h1>Sign in to SkyPilot CLI</h1>
+        <p>You are seeing this page because a SkyPilot command requires authentication.</p>
+
+        <p>Please copy the following token and paste it into your SkyPilot CLI prompt:</p>
+        <div id="token-box" class="code-block">SKYPILOT_API_SERVER_USER_TOKEN_PLACEHOLDER</div>
+        <button id="copy-btn" class="copy-button">Copy Token</button>
+
+        <p class="footer-text">You can close this tab after copying the token.</p>
+    </div>
+
+    <script>
+        const tokenBox = document.getElementById('token-box');
+        const copyBtn  = document.getElementById('copy-btn');
+
+        function selectToken() {
+            // For <pre> or <div>, create a range to select its content
+            const range = document.createRange();
+            range.selectNodeContents(tokenBox);
+            const sel = window.getSelection();
+            sel.removeAllRanges();
+            sel.addRange(range);
+        }
+
+        // Optional: Select the token when the page loads or when token box is clicked
+        tokenBox.addEventListener('click', selectToken);
+        window.addEventListener('load', selectToken);
+
+        copyBtn.addEventListener('click', () => {
+            selectToken(); // Select the text
+            try {
+                document.execCommand('copy');
+                copyBtn.textContent = 'Copied!';
+            } catch (err) {
+                copyBtn.textContent = 'Error!';
+                console.error('Failed to copy text: ', err);
+            }
+            setTimeout(() => {
+                copyBtn.textContent = 'Copy Token';
+            }, 2000);
+        });
+    </script>
+</body>
+</html>

sky/server/requests/executor.py

@@ -20,8 +20,6 @@ See the [README.md](../README.md) for detailed architecture of the executor.
 """
 import asyncio
 import contextlib
-import contextvars
-import functools
 import multiprocessing
 import os
 import queue as queue_lib
@@ -52,6 +50,7 @@ from sky.skylet import constants
 from sky.utils import annotations
 from sky.utils import common_utils
 from sky.utils import context
+from sky.utils import context_utils
 from sky.utils import subprocess_utils
 from sky.utils import timeline
 
@@ -368,10 +367,8 @@ async def execute_request_coroutine(request: api_requests.Request):
     # 1. skypilot config is not contextual
     # 2. envs that read directly from os.environ are not contextual
     ctx.override_envs(request_body.env_vars)
-    loop = asyncio.get_running_loop()
-    pyctx = contextvars.copy_context()
-    func_call = functools.partial(pyctx.run, func, **request_body.to_kwargs())
-    fut: asyncio.Future = loop.run_in_executor(None, func_call)
+    fut: asyncio.Future = context_utils.to_thread(func,
+                                                  **request_body.to_kwargs())
 
     async def poll_task(request_id: str) -> bool:
         request = api_requests.get_request(request_id)

sky/server/requests/payloads.py

@@ -443,9 +443,9 @@ class ServeStatusBody(RequestBody):
 
 class RealtimeGpuAvailabilityRequestBody(RequestBody):
     """The request body for the realtime GPU availability endpoint."""
-    context: Optional[str]
-    name_filter: Optional[str]
-    quantity_filter: Optional[int]
+    context: Optional[str] = None
+    name_filter: Optional[str] = None
+    quantity_filter: Optional[int] = None
 
 
 class KubernetesNodeInfoRequestBody(RequestBody):

sky/server/server.py

@@ -2,9 +2,11 @@
 
 import argparse
 import asyncio
+import base64
 import contextlib
 import dataclasses
 import datetime
+import json
 import logging
 import multiprocessing
 import os
@@ -49,6 +51,7 @@ from sky.utils import admin_policy_utils
 from sky.utils import common as common_lib
 from sky.utils import common_utils
 from sky.utils import context
+from sky.utils import context_utils
 from sky.utils import dag_utils
 from sky.utils import env_options
 from sky.utils import status_lib
@@ -218,6 +221,34 @@ app.include_router(jobs_rest.router, prefix='/jobs', tags=['jobs'])
 app.include_router(serve_rest.router, prefix='/serve', tags=['serve'])
 
 
+@app.get('/token')
+async def token(request: fastapi.Request) -> fastapi.responses.HTMLResponse:
+    # Use base64 encoding to avoid having to escape anything in the HTML.
+    json_bytes = json.dumps(request.cookies).encode('utf-8')
+    base64_str = base64.b64encode(json_bytes).decode('utf-8')
+
+    html_dir = pathlib.Path(__file__).parent / 'html'
+    token_page_path = html_dir / 'token_page.html'
+    try:
+        with open(token_page_path, 'r', encoding='utf-8') as f:
+            html_content = f.read()
+    except FileNotFoundError as e:
+        raise fastapi.HTTPException(
+            status_code=500, detail='Token page template not found.') from e
+
+    html_content = html_content.replace(
+        'SKYPILOT_API_SERVER_USER_TOKEN_PLACEHOLDER', base64_str)
+
+    return fastapi.responses.HTMLResponse(
+        content=html_content,
+        headers={
+            'Cache-Control': 'no-cache, no-transform',
+            # X-Accel-Buffering: no is useful for preventing buffering issues
+            # with some reverse proxies.
+            'X-Accel-Buffering': 'no'
+        })
+
+
 @app.post('/check')
 async def check(request: fastapi.Request,
                 check_body: payloads.CheckBody) -> None:
@@ -327,25 +358,26 @@ async def validate(validate_body: payloads.ValidateBody) -> None:
     # pairs.
     logger.debug(f'Validating tasks: {validate_body.dag}')
 
+    context.initialize()
+
     def validate_dag(dag: dag_utils.dag_lib.Dag):
         # TODO: Admin policy may contain arbitrary code, which may be expensive
         # to run and may block the server thread. However, moving it into the
         # executor adds a ~150ms penalty on the local API server because of
         # added RTTs. For now, we stick to doing the validation inline in the
         # server thread.
-        dag, _ = admin_policy_utils.apply(
-            dag, request_options=validate_body.request_options)
-        # Skip validating workdir and file_mounts, as those need to be
-        # validated after the files are uploaded to the SkyPilot API server
-        # with `upload_mounts_to_api_server`.
-        dag.validate(skip_file_mounts=True, skip_workdir=True)
+        with admin_policy_utils.apply_and_use_config_in_current_request(
+                dag, request_options=validate_body.request_options) as dag:
+            # Skip validating workdir and file_mounts, as those need to be
+            # validated after the files are uploaded to the SkyPilot API server
+            # with `upload_mounts_to_api_server`.
+            dag.validate(skip_file_mounts=True, skip_workdir=True)
 
     try:
         dag = dag_utils.load_chain_dag_from_yaml_str(validate_body.dag)
-        loop = asyncio.get_running_loop()
         # Apply admin policy and validate DAG is blocking, run it in a separate
         # thread executor to avoid blocking the uvicorn event loop.
-        await loop.run_in_executor(None, validate_dag, dag)
+        await context_utils.to_thread(validate_dag, dag)
     except Exception as e:  # pylint: disable=broad-except
         raise fastapi.HTTPException(
             status_code=400, detail=exceptions.serialize_exception(e)) from e