skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/jobs/server/server.py

@@ -5,11 +5,13 @@ import pathlib
 import fastapi
 
 from sky import sky_logging
+from sky.jobs import utils as managed_jobs_utils
 from sky.jobs.server import core
 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
@@ -22,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
@@ -46,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,
@@ -64,28 +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 = await api_requests.get_request_async(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,
     )
 
 
@@ -100,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
@@ -114,9 +156,9 @@ async def download_logs(
 @router.post('/pool_apply')
 async def pool_apply(request: fastapi.Request,
                      jobs_pool_apply_body: payloads.JobsPoolApplyBody) -> None:
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.pool_apply',
+        request_name=request_names.RequestName.JOBS_POOL_APPLY,
         request_body=jobs_pool_apply_body,
         func=core.pool_apply,
         schedule_type=api_requests.ScheduleType.LONG,
@@ -127,9 +169,9 @@ async def pool_apply(request: fastapi.Request,
 @router.post('/pool_down')
 async def pool_down(request: fastapi.Request,
                     jobs_pool_down_body: payloads.JobsPoolDownBody) -> None:
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.pool_down',
+        request_name=request_names.RequestName.JOBS_POOL_DOWN,
         request_body=jobs_pool_down_body,
         func=core.pool_down,
         schedule_type=api_requests.ScheduleType.SHORT,
@@ -141,9 +183,9 @@ async def pool_down(request: fastapi.Request,
 async def pool_status(
         request: fastapi.Request,
         jobs_pool_status_body: payloads.JobsPoolStatusBody) -> None:
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.pool_status',
+        request_name=request_names.RequestName.JOBS_POOL_STATUS,
         request_body=jobs_pool_status_body,
         func=core.pool_status,
         schedule_type=api_requests.ScheduleType.SHORT,
@@ -156,21 +198,25 @@ async def pool_tail_logs(
     request: fastapi.Request, log_body: payloads.JobsPoolLogsBody,
     background_tasks: fastapi.BackgroundTasks
 ) -> fastapi.responses.StreamingResponse:
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.pool_logs',
+        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 = api_requests.get_request(request.state.request_id)
+    request_task = await api_requests.get_request_async(
+        request.state.request_id, fields=['request_id'])
 
-    return stream_utils.stream_response(
+    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,
     )
 
 
@@ -188,9 +234,9 @@ async def pool_download_logs(
     # 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)
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='jobs.pool_sync_down_logs',
+        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,

sky/jobs/server/utils.py

@@ -1,17 +1,32 @@
 """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.
@@ -29,48 +44,87 @@ def check_version_mismatch_and_non_terminal_jobs() -> None:
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend)
 
-    # Get controller version and raw job table
-    code = managed_job_utils.ManagedJobCodeGen.get_version_and_job_table()
-
-    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 and jobs 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 len(output_parts) < 2 or 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]
-
-    # Rest is job table payload (preserving any newlines within it)
-    job_table_payload = output_parts[1]
+    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
-
-    # Load and filter jobs locally using existing method
-    jobs, _, _, _, _ = managed_job_utils.load_managed_job_queue(
-        job_table_payload)
+    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, show_all=False, show_user=False)
+            non_terminal_jobs,
+            pool_status=None,
+            show_all=False,
+            show_user=False)
 
         error_msg = (
             f'Controller SKYLET_VERSION ({controller_version}) does not match '