skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/jobs/server/server.py

@@ -0,0 +1,190 @@
+"""REST API for managed jobs."""
+import os
+
+import fastapi
+import httpx
+
+from sky import sky_logging
+from sky.jobs.server import core
+from sky.jobs.server import dashboard_utils
+from sky.server import common as server_common
+from sky.server import stream_utils
+from sky.server.requests import executor
+from sky.server.requests import payloads
+from sky.server.requests import requests as api_requests
+from sky.skylet import constants
+from sky.utils import common
+from sky.utils import common_utils
+
+logger = sky_logging.init_logger(__name__)
+
+router = fastapi.APIRouter()
+
+
+@router.post('/launch')
+async def launch(request: fastapi.Request,
+                 jobs_launch_body: payloads.JobsLaunchBody) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.launch',
+        request_body=jobs_launch_body,
+        func=core.launch,
+        schedule_type=api_requests.ScheduleType.LONG,
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )
+
+
+@router.post('/queue')
+async def queue(request: fastapi.Request,
+                jobs_queue_body: payloads.JobsQueueBody) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.queue',
+        request_body=jobs_queue_body,
+        func=core.queue,
+        schedule_type=(api_requests.ScheduleType.LONG if jobs_queue_body.refresh
+                       else api_requests.ScheduleType.SHORT),
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )
+
+
+@router.post('/cancel')
+async def cancel(request: fastapi.Request,
+                 jobs_cancel_body: payloads.JobsCancelBody) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.cancel',
+        request_body=jobs_cancel_body,
+        func=core.cancel,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )
+
+
+@router.post('/logs')
+async def logs(
+    request: fastapi.Request, jobs_logs_body: payloads.JobsLogsBody,
+    background_tasks: fastapi.BackgroundTasks
+) -> fastapi.responses.StreamingResponse:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.logs',
+        request_body=jobs_logs_body,
+        func=core.tail_logs,
+        schedule_type=api_requests.ScheduleType.SHORT
+        if jobs_logs_body.refresh else api_requests.ScheduleType.LONG,
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )
+    request_task = api_requests.get_request(request.state.request_id)
+
+    return stream_utils.stream_response(
+        request_id=request_task.request_id,
+        logs_path=request_task.log_path,
+        background_tasks=background_tasks,
+    )
+
+
+@router.post('/download_logs')
+async def download_logs(
+        request: fastapi.Request,
+        jobs_download_logs_body: payloads.JobsDownloadLogsBody) -> None:
+    user_hash = jobs_download_logs_body.env_vars[constants.USER_ID_ENV_VAR]
+    logs_dir_on_api_server = server_common.api_server_user_logs_dir_prefix(
+        user_hash)
+    logs_dir_on_api_server.expanduser().mkdir(parents=True, exist_ok=True)
+    # We should reuse the original request body, so that the env vars, such as
+    # user hash, are kept the same.
+    jobs_download_logs_body.local_dir = str(logs_dir_on_api_server)
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='jobs.download_logs',
+        request_body=jobs_download_logs_body,
+        func=core.download_logs,
+        schedule_type=api_requests.ScheduleType.LONG
+        if jobs_download_logs_body.refresh else api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )
+
+
+@router.get('/dashboard')
+async def dashboard(request: fastapi.Request,
+                    user_hash: str) -> fastapi.Response:
+    # TODO(cooperc): Support showing only jobs for a specific user.
+
+    # FIX(zhwu/cooperc/eric): Fix log downloading (assumes global
+    # /download_log/xx route)
+
+    # Note: before #4717, each user had their own controller, and thus their own
+    # dashboard. Now, all users share the same controller, so this isn't really
+    # necessary. TODO(cooperc): clean up.
+
+    # TODO: Put this in an executor to avoid blocking the main server thread.
+    # It can take a long time if it needs to check the controller status.
+
+    # Find the port for the dashboard of the user
+    os.environ[constants.USER_ID_ENV_VAR] = user_hash
+    server_common.reload_for_new_request(client_entrypoint=None,
+                                         client_command=None,
+                                         using_remote_api_server=False)
+    logger.info(f'Starting dashboard for user hash: {user_hash}')
+
+    with dashboard_utils.get_dashboard_lock_for_user(user_hash):
+        max_retries = 3
+        for attempt in range(max_retries):
+            port, pid = dashboard_utils.get_dashboard_session(user_hash)
+            if port == 0 or attempt > 0:
+                # Let the client know that we are waiting for starting the
+                # dashboard.
+                try:
+                    port, pid = core.start_dashboard_forwarding()
+                except Exception as e:  # pylint: disable=broad-except
+                    # We catch all exceptions to gracefully handle unknown
+                    # errors and raise an HTTPException to the client.
+                    msg = (
+                        'Dashboard failed to start: '
+                        f'{common_utils.format_exception(e, use_bracket=True)}')
+                    logger.error(msg)
+                    raise fastapi.HTTPException(status_code=503, detail=msg)
+                dashboard_utils.add_dashboard_session(user_hash, port, pid)
+
+            # Assuming the dashboard is forwarded to localhost on the API server
+            dashboard_url = f'http://localhost:{port}'
+            try:
+                # Ping the dashboard to check if it's still running
+                async with httpx.AsyncClient() as client:
+                    response = await client.request('GET',
+                                                    dashboard_url,
+                                                    timeout=1)
+                break  # Connection successful, proceed with the request
+            except Exception as e:  # pylint: disable=broad-except
+                # We catch all exceptions to gracefully handle unknown
+                # errors and retry or raise an HTTPException to the client.
+                msg = (
+                    f'Dashboard connection attempt {attempt + 1} failed with '
+                    f'{common_utils.format_exception(e, use_bracket=True)}')
+                logger.info(msg)
+                if attempt == max_retries - 1:
+                    raise fastapi.HTTPException(status_code=503, detail=msg)
+
+    # Create a client session to forward the request
+    try:
+        async with httpx.AsyncClient() as client:
+            # Make the request and get the response
+            response = await client.request(
+                method='GET',
+                url=f'{dashboard_url}',
+                headers=request.headers.raw,
+            )
+
+            # Create a new response with the content already read
+            content = await response.aread()
+            return fastapi.Response(
+                content=content,
+                status_code=response.status_code,
+                headers=dict(response.headers),
+                media_type=response.headers.get('content-type'))
+    except Exception as e:
+        msg = (f'Failed to forward request to dashboard: '
+               f'{common_utils.format_exception(e, use_bracket=True)}')
+        logger.error(msg)
+        raise fastapi.HTTPException(status_code=502, detail=msg)