skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/jobs/server/server.py

@@ -1,20 +1,20 @@
 """REST API for managed jobs."""
-import os
+
+import pathlib
 
 import fastapi
-import httpx
 
 from sky import sky_logging
+from sky.jobs import utils as managed_jobs_utils
 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 request_names
 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__)
 
@@ -24,22 +24,36 @@ router = fastapi.APIRouter()
 @router.post('/launch')
 async def launch(request: fastapi.Request,
                  jobs_launch_body: payloads.JobsLaunchBody) -> None:
-    executor.schedule_request(
+    # In consolidation mode, the jobs controller will use sky.launch on the same
+    # API server to launch the underlying job cluster. If you start run many
+    # jobs.launch requests, some may be blocked for a long time by sky.launch
+    # requests triggered by earlier jobs, which leads to confusing behavior as
+    # the jobs.launch requests trickle though. Also, since we don't have to
+    # actually launch a jobs controller sky cluster, the jobs.launch request is
+    # much quicker in consolidation mode. So we avoid the issue by just using
+    # the short executor instead - then jobs.launch will not be blocked by
+    # sky.launch.
+    consolidation_mode = managed_jobs_utils.is_consolidation_mode()
+    schedule_type = (api_requests.ScheduleType.SHORT
+                     if consolidation_mode else api_requests.ScheduleType.LONG)
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.launch',
+        request_name=request_names.RequestName.JOBS_LAUNCH,
         request_body=jobs_launch_body,
         func=core.launch,
-        schedule_type=api_requests.ScheduleType.LONG,
+        schedule_type=schedule_type,
         request_cluster_name=common.JOB_CONTROLLER_NAME,
     )
 
 
+# For backwards compatibility
+# TODO(hailong): Remove before 0.12.0.
 @router.post('/queue')
 async def queue(request: fastapi.Request,
                 jobs_queue_body: payloads.JobsQueueBody) -> None:
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.queue',
+        request_name=request_names.RequestName.JOBS_QUEUE,
         request_body=jobs_queue_body,
         func=core.queue,
         schedule_type=(api_requests.ScheduleType.LONG if jobs_queue_body.refresh
@@ -48,12 +62,27 @@ async def queue(request: fastapi.Request,
     )
 
 
+@router.post('/queue/v2')
+async def queue_v2(request: fastapi.Request,
+                   jobs_queue_body_v2: payloads.JobsQueueV2Body) -> None:
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.JOBS_QUEUE_V2,
+        request_body=jobs_queue_body_v2,
+        func=core.queue_v2_api,
+        schedule_type=(api_requests.ScheduleType.LONG
+                       if jobs_queue_body_v2.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(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.cancel',
+        request_name=request_names.RequestName.JOBS_CANCEL,
         request_body=jobs_cancel_body,
         func=core.cancel,
         schedule_type=api_requests.ScheduleType.SHORT,
@@ -66,27 +95,39 @@ async def logs(
     request: fastapi.Request, jobs_logs_body: payloads.JobsLogsBody,
     background_tasks: fastapi.BackgroundTasks
 ) -> fastapi.responses.StreamingResponse:
-    executor.schedule_request(
+    schedule_type = api_requests.ScheduleType.SHORT
+    if jobs_logs_body.refresh:
+        # When refresh is specified, the job controller might be restarted,
+        # which takes longer time to finish. We schedule it to long executor.
+        schedule_type = api_requests.ScheduleType.LONG
+    if schedule_type == api_requests.ScheduleType.SHORT:
+        executor.check_request_thread_executor_available()
+    request_task = await executor.prepare_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.logs',
+        request_name=request_names.RequestName.JOBS_LOGS,
         request_body=jobs_logs_body,
         func=core.tail_logs,
-        # TODO(aylei): We have tail logs scheduled as SHORT request, because it
-        # should be responsive. However, it can be long running if the user's
-        # job keeps running, and we should avoid it taking the SHORT worker
-        # indefinitely.
-        # When refresh is True we schedule it as LONG because a controller
-        # restart might be needed.
-        schedule_type=api_requests.ScheduleType.LONG
-        if jobs_logs_body.refresh else api_requests.ScheduleType.SHORT,
+        schedule_type=schedule_type,
         request_cluster_name=common.JOB_CONTROLLER_NAME,
     )
-    request_task = api_requests.get_request(request.state.request_id)
-
-    return stream_utils.stream_response(
+    kill_request_on_disconnect = False
+    if schedule_type == api_requests.ScheduleType.SHORT:
+        # For short request, run in the coroutine to avoid blocking
+        # short workers.
+        task = executor.execute_request_in_coroutine(request_task)
+        # Cancel the coroutine after the request is done or client disconnects
+        background_tasks.add_task(task.cancel)
+    else:
+        executor.schedule_prepared_request(request_task)
+        # When runs in long executor process, we should kill the request on
+        # disconnect to cancel the running routine.
+        kill_request_on_disconnect = True
+
+    return stream_utils.stream_response_for_long_request(
         request_id=request_task.request_id,
         logs_path=request_task.log_path,
         background_tasks=background_tasks,
+        kill_request_on_disconnect=kill_request_on_disconnect,
     )
 
 
@@ -101,9 +142,9 @@ async def download_logs(
     # 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(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.download_logs',
+        request_name=request_names.RequestName.JOBS_DOWNLOAD_LOGS,
         request_body=jobs_download_logs_body,
         func=core.download_logs,
         schedule_type=api_requests.ScheduleType.LONG
@@ -112,92 +153,92 @@ async def download_logs(
     )
 
 
-@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=5)
-                if response.is_success:
-                    break  # Connection successful, proceed with the request
-                # Raise an HTTPException here which will be caught by the
-                # following except block to retry with new connection
-                response.raise_for_status()
-            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.
-                # Assume an exception indicates that the dashboard connection
-                # is stale - remove it so that a new one is created.
-                dashboard_utils.remove_dashboard_session(user_hash)
-                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)
+@router.post('/pool_apply')
+async def pool_apply(request: fastapi.Request,
+                     jobs_pool_apply_body: payloads.JobsPoolApplyBody) -> None:
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.JOBS_POOL_APPLY,
+        request_body=jobs_pool_apply_body,
+        func=core.pool_apply,
+        schedule_type=api_requests.ScheduleType.LONG,
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )
+
+
+@router.post('/pool_down')
+async def pool_down(request: fastapi.Request,
+                    jobs_pool_down_body: payloads.JobsPoolDownBody) -> None:
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.JOBS_POOL_DOWN,
+        request_body=jobs_pool_down_body,
+        func=core.pool_down,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )
+
+
+@router.post('/pool_status')
+async def pool_status(
+        request: fastapi.Request,
+        jobs_pool_status_body: payloads.JobsPoolStatusBody) -> None:
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.JOBS_POOL_STATUS,
+        request_body=jobs_pool_status_body,
+        func=core.pool_status,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )
+
+
+@router.post('/pool_logs')
+async def pool_tail_logs(
+    request: fastapi.Request, log_body: payloads.JobsPoolLogsBody,
+    background_tasks: fastapi.BackgroundTasks
+) -> fastapi.responses.StreamingResponse:
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.JOBS_POOL_LOGS,
+        request_body=log_body,
+        func=core.pool_tail_logs,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )
+
+    request_task = await api_requests.get_request_async(
+        request.state.request_id, fields=['request_id'])
+
+    return stream_utils.stream_response_for_long_request(
+        request_id=request_task.request_id,
+        # req.log_path is derived from request_id,
+        # so it's ok to just grab the request_id in the above query.
+        logs_path=request_task.log_path,
+        background_tasks=background_tasks,
+        kill_request_on_disconnect=True,
+    )
+
+
+@router.post('/pool_sync-down-logs')
+async def pool_download_logs(
+    request: fastapi.Request,
+    download_logs_body: payloads.JobsPoolDownloadLogsBody,
+) -> None:
+    user_hash = download_logs_body.env_vars[constants.USER_ID_ENV_VAR]
+    timestamp = sky_logging.get_run_timestamp()
+    logs_dir_on_api_server = (
+        pathlib.Path(server_common.api_server_user_logs_dir_prefix(user_hash)) /
+        'pool' / f'{download_logs_body.pool_name}_{timestamp}')
+    logs_dir_on_api_server.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.
+    download_logs_body.local_dir = str(logs_dir_on_api_server)
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.JOBS_POOL_SYNC_DOWN_LOGS,
+        request_body=download_logs_body,
+        func=core.pool_sync_down_logs,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.JOB_CONTROLLER_NAME,
+    )

sky/jobs/server/utils.py

@@ -0,0 +1,136 @@
+"""Utility functions for managed jobs."""
+import typing
+
+from sky import backends
+from sky import exceptions
+from sky import sky_logging
+from sky.adaptors import common as adaptors_common
+from sky.backends import backend_utils
+from sky.backends import cloud_vm_ray_backend
+from sky.jobs import utils as managed_job_utils
+from sky.skylet import constants as skylet_constants
+from sky.utils import controller_utils
+
+logger = sky_logging.init_logger(__name__)
+
+if typing.TYPE_CHECKING:
+    from sky.schemas.generated import managed_jobsv1_pb2
+else:
+    managed_jobsv1_pb2 = adaptors_common.LazyImport(
+        'sky.schemas.generated.managed_jobsv1_pb2')
+
+_MANAGED_JOB_FIELDS_TO_GET = [
+    'job_id', 'task_id', 'workspace', 'job_name', 'task_name', 'resources',
+    'submitted_at', 'end_at', 'job_duration', 'recovery_count', 'status', 'pool'
+]
+
+
+def check_version_mismatch_and_non_terminal_jobs() -> None:
+    """Check if controller has version mismatch and non-terminal jobs exist.
+    Raises:
+        ValueError: If there's a version mismatch and non-terminal jobs exist.
+        sky.exceptions.ClusterNotUpError: If the controller is not accessible.
+    """
+    # Get the current local SKYLET_VERSION
+    local_version = skylet_constants.SKYLET_VERSION
+
+    # Get controller handle (works the same in both normal and
+    # consolidation mode)
+    jobs_controller_type = controller_utils.Controllers.JOBS_CONTROLLER
+    handle = backend_utils.is_controller_accessible(
+        controller=jobs_controller_type,
+        stopped_message='Jobs controller is not running.')
+
+    backend = backend_utils.get_backend_from_handle(handle)
+    assert isinstance(backend, backends.CloudVmRayBackend)
+
+    use_legacy = not handle.is_grpc_enabled_with_flag
+
+    if not use_legacy:
+        try:
+            version_request = managed_jobsv1_pb2.GetVersionRequest()
+            version_response = backend_utils.invoke_skylet_with_retries(
+                lambda: cloud_vm_ray_backend.SkyletClient(
+                    handle.get_grpc_channel(
+                    )).get_managed_job_controller_version(version_request))
+            controller_version = version_response.controller_version
+
+            job_table_request = managed_jobsv1_pb2.GetJobTableRequest(
+                skip_finished=True,
+                fields=managed_jobsv1_pb2.Fields(
+                    fields=_MANAGED_JOB_FIELDS_TO_GET),
+            )
+            job_table_response = backend_utils.invoke_skylet_with_retries(
+                lambda: cloud_vm_ray_backend.SkyletClient(
+                    handle.get_grpc_channel()).get_managed_job_table(
+                        job_table_request))
+            jobs = managed_job_utils.decode_managed_job_protos(
+                job_table_response.jobs)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        # Get controller version and raw job table
+        code = managed_job_utils.ManagedJobCodeGen.get_version()
+
+        returncode, output, stderr = backend.run_on_head(handle,
+                                                         code,
+                                                         require_outputs=True,
+                                                         stream_logs=False,
+                                                         separate_stderr=True)
+
+        if returncode != 0:
+            logger.error(output + stderr)
+            raise ValueError('Failed to check controller version with '
+                             f'returncode: {returncode}.\n{output + stderr}')
+
+        # Parse the output to extract controller version (split only on first
+        # newline)
+        output_parts = output.strip().split('\n', 1)
+
+        # Extract controller version from first line
+        if not output_parts[0].startswith('controller_version:'):
+            raise ValueError(
+                f'Expected controller version in first line, got: {output}')
+
+        controller_version = output_parts[0].split(':', 1)[1]
+
+        code = managed_job_utils.ManagedJobCodeGen.get_job_table(
+            skip_finished=True, fields=_MANAGED_JOB_FIELDS_TO_GET)
+        returncode, job_table_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+
+        if returncode != 0:
+            logger.error(job_table_payload + stderr)
+            raise ValueError('Failed to fetch managed jobs with returncode: '
+                             f'{returncode}.\n{job_table_payload + stderr}')
+
+        jobs, _, _, _, _ = (
+            managed_job_utils.load_managed_job_queue(job_table_payload))
+
+    # Process locally: check version match and filter non-terminal jobs
+    version_matches = (controller_version == local_version or
+                       int(controller_version) > 17)
+    non_terminal_jobs = [job for job in jobs if not job['status'].is_terminal()]
+    has_non_terminal_jobs = len(non_terminal_jobs) > 0
+
+    if not version_matches and has_non_terminal_jobs:
+        # Format job table locally using the same method as queue()
+        formatted_job_table = managed_job_utils.format_job_table(
+            non_terminal_jobs,
+            pool_status=None,
+            show_all=False,
+            show_user=False)
+
+        error_msg = (
+            f'Controller SKYLET_VERSION ({controller_version}) does not match '
+            f'current version ({local_version}), and there are non-terminal '
+            'jobs on the controller. Please wait for all jobs to complete or '
+            'cancel them before launching new jobs with the updated version.'
+            f'\n\nCurrent non-terminal jobs:\n{formatted_job_table}')
+
+        raise ValueError(error_msg)