skypilot-nightly 1.0.0.dev20251029__py3-none-any.whl → 1.0.0.dev20251101__py3-none-any.whl

sky/server/server.py

@@ -6,6 +6,7 @@ import base64
 from concurrent.futures import ThreadPoolExecutor
 import contextlib
 import datetime
+from enum import IntEnum
 import hashlib
 import json
 import multiprocessing
@@ -15,6 +16,7 @@ import posixpath
 import re
 import resource
 import shutil
+import struct
 import sys
 import threading
 import traceback
@@ -62,6 +64,7 @@ from sky.server.auth import oauth2_proxy
 from sky.server.requests import executor
 from sky.server.requests import payloads
 from sky.server.requests import preconditions
+from sky.server.requests import request_names
 from sky.server.requests import requests as requests_lib
 from sky.skylet import constants
 from sky.ssh_node_pools import server as ssh_node_pools_rest
@@ -460,7 +463,7 @@ async def schedule_on_boot_check_async():
     try:
         await executor.schedule_request_async(
             request_id='skypilot-server-on-boot-check',
-            request_name='check',
+            request_name=request_names.RequestName.CHECK,
             request_body=payloads.CheckBody(),
             func=sky_check.check,
             schedule_type=requests_lib.ScheduleType.SHORT,
@@ -732,7 +735,7 @@ async def check(request: fastapi.Request,
     """Checks enabled clouds."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='check',
+        request_name=request_names.RequestName.CHECK,
         request_body=check_body,
         func=sky_check.check,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -746,7 +749,7 @@ async def enabled_clouds(request: fastapi.Request,
     """Gets enabled clouds on the server."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='enabled_clouds',
+        request_name=request_names.RequestName.ENABLED_CLOUDS,
         request_body=payloads.EnabledCloudsBody(workspace=workspace,
                                                 expand=expand),
         func=core.enabled_clouds,
@@ -762,7 +765,8 @@ async def realtime_kubernetes_gpu_availability(
     """Gets real-time Kubernetes GPU availability."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='realtime_kubernetes_gpu_availability',
+        request_name=request_names.RequestName.
+        REALTIME_KUBERNETES_GPU_AVAILABILITY,
         request_body=realtime_gpu_availability_body,
         func=core.realtime_kubernetes_gpu_availability,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -777,7 +781,7 @@ async def kubernetes_node_info(
     """Gets Kubernetes nodes information and hints."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='kubernetes_node_info',
+        request_name=request_names.RequestName.KUBERNETES_NODE_INFO,
         request_body=kubernetes_node_info_body,
         func=kubernetes_utils.get_kubernetes_node_info,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -789,7 +793,7 @@ async def status_kubernetes(request: fastapi.Request) -> None:
     """Gets Kubernetes status."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='status_kubernetes',
+        request_name=request_names.RequestName.STATUS_KUBERNETES,
         request_body=payloads.RequestBody(),
         func=core.status_kubernetes,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -803,7 +807,7 @@ async def list_accelerators(
     """Gets list of accelerators from cloud catalog."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='list_accelerators',
+        request_name=request_names.RequestName.LIST_ACCELERATORS,
         request_body=list_accelerator_counts_body,
         func=catalog.list_accelerators,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -818,7 +822,7 @@ async def list_accelerator_counts(
     """Gets list of accelerator counts from cloud catalog."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='list_accelerator_counts',
+        request_name=request_names.RequestName.LIST_ACCELERATOR_COUNTS,
         request_body=list_accelerator_counts_body,
         func=catalog.list_accelerator_counts,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -875,7 +879,7 @@ async def optimize(optimize_body: payloads.OptimizeBody,
     """Optimizes the user's DAG."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='optimize',
+        request_name=request_names.RequestName.OPTIMIZE,
         request_body=optimize_body,
         ignore_return_value=True,
         func=core.optimize,
@@ -1085,7 +1089,7 @@ async def launch(launch_body: payloads.LaunchBody,
     logger.info(f'Launching request: {request_id}')
     await executor.schedule_request_async(
         request_id,
-        request_name='launch',
+        request_name=request_names.RequestName.CLUSTER_LAUNCH,
         request_body=launch_body,
         func=execution.launch,
         schedule_type=requests_lib.ScheduleType.LONG,
@@ -1101,7 +1105,7 @@ async def exec(request: fastapi.Request, exec_body: payloads.ExecBody) -> None:
     cluster_name = exec_body.cluster_name
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='exec',
+        request_name=request_names.RequestName.CLUSTER_EXEC,
         request_body=exec_body,
         func=execution.exec,
         precondition=preconditions.ClusterStartCompletePrecondition(
@@ -1119,7 +1123,7 @@ async def stop(request: fastapi.Request,
     """Stops a cluster."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='stop',
+        request_name=request_names.RequestName.CLUSTER_STOP,
         request_body=stop_body,
         func=core.stop,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1139,7 +1143,7 @@ async def status(
             detail='Server is shutting down, please try again later.')
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='status',
+        request_name=request_names.RequestName.CLUSTER_STATUS,
         request_body=status_body,
         func=core.status,
         schedule_type=(requests_lib.ScheduleType.LONG if
@@ -1154,7 +1158,7 @@ async def endpoints(request: fastapi.Request,
     """Gets the endpoint for a given cluster and port number (endpoint)."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='endpoints',
+        request_name=request_names.RequestName.CLUSTER_ENDPOINTS,
         request_body=endpoint_body,
         func=core.endpoints,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1168,7 +1172,7 @@ async def down(request: fastapi.Request,
     """Tears down a cluster."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='down',
+        request_name=request_names.RequestName.CLUSTER_DOWN,
         request_body=down_body,
         func=core.down,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1182,7 +1186,7 @@ async def start(request: fastapi.Request,
     """Restarts a cluster."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='start',
+        request_name=request_names.RequestName.CLUSTER_START,
         request_body=start_body,
         func=core.start,
         schedule_type=requests_lib.ScheduleType.LONG,
@@ -1196,7 +1200,7 @@ async def autostop(request: fastapi.Request,
     """Schedules an autostop/autodown for a cluster."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='autostop',
+        request_name=request_names.RequestName.CLUSTER_AUTOSTOP,
         request_body=autostop_body,
         func=core.autostop,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1210,7 +1214,7 @@ async def queue(request: fastapi.Request,
     """Gets the job queue of a cluster."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='queue',
+        request_name=request_names.RequestName.CLUSTER_QUEUE,
         request_body=queue_body,
         func=core.queue,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1224,7 +1228,7 @@ async def job_status(request: fastapi.Request,
     """Gets the status of a job."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='job_status',
+        request_name=request_names.RequestName.CLUSTER_JOB_STATUS,
         request_body=job_status_body,
         func=core.job_status,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1238,7 +1242,7 @@ async def cancel(request: fastapi.Request,
     """Cancels jobs on a cluster."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='cancel',
+        request_name=request_names.RequestName.CLUSTER_JOB_CANCEL,
         request_body=cancel_body,
         func=core.cancel,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1258,7 +1262,7 @@ async def logs(
     executor.check_request_thread_executor_available()
     request_task = await executor.prepare_request_async(
         request_id=request.state.request_id,
-        request_name='logs',
+        request_name=request_names.RequestName.CLUSTER_JOB_LOGS,
         request_body=cluster_job_body,
         func=core.tail_logs,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1289,7 +1293,7 @@ async def download_logs(
     cluster_jobs_body.local_dir = str(logs_dir_on_api_server)
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='download_logs',
+        request_name=request_names.RequestName.CLUSTER_JOB_DOWNLOAD_LOGS,
         request_body=cluster_jobs_body,
         func=core.download_logs,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1440,7 +1444,7 @@ async def cost_report(request: fastapi.Request,
     """Gets the cost report of a cluster."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='cost_report',
+        request_name=request_names.RequestName.CLUSTER_COST_REPORT,
         request_body=cost_report_body,
         func=core.cost_report,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1452,7 +1456,7 @@ async def storage_ls(request: fastapi.Request) -> None:
     """Gets the storages."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='storage_ls',
+        request_name=request_names.RequestName.STORAGE_LS,
         request_body=payloads.RequestBody(),
         func=core.storage_ls,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1465,7 +1469,7 @@ async def storage_delete(request: fastapi.Request,
     """Deletes a storage."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='storage_delete',
+        request_name=request_names.RequestName.STORAGE_DELETE,
         request_body=storage_body,
         func=core.storage_delete,
         schedule_type=requests_lib.ScheduleType.LONG,
@@ -1478,7 +1482,7 @@ async def local_up(request: fastapi.Request,
     """Launches a Kubernetes cluster on API server."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='local_up',
+        request_name=request_names.RequestName.LOCAL_UP,
         request_body=local_up_body,
         func=core.local_up,
         schedule_type=requests_lib.ScheduleType.LONG,
@@ -1491,7 +1495,7 @@ async def local_down(request: fastapi.Request,
     """Tears down the Kubernetes cluster started by local_up."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='local_down',
+        request_name=request_names.RequestName.LOCAL_DOWN,
         request_body=local_down_body,
         func=core.local_down,
         schedule_type=requests_lib.ScheduleType.LONG,
@@ -1699,7 +1703,7 @@ async def api_cancel(request: fastapi.Request,
     """Cancels requests."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='api_cancel',
+        request_name=request_names.RequestName.API_CANCEL,
         request_body=request_cancel_body,
         func=requests_lib.kill_requests_with_prefix,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -1804,16 +1808,31 @@ async def health(request: fastapi.Request) -> responses.APIHealthResponse:
         basic_auth_enabled=os.environ.get(constants.ENV_VAR_ENABLE_BASIC_AUTH,
                                           'false').lower() == 'true',
         user=user if user is not None else None,
+        service_account_token_enabled=(os.environ.get(
+            constants.ENV_VAR_ENABLE_SERVICE_ACCOUNTS,
+            'false').lower() == 'true'),
     )
 
 
+class KubernetesSSHMessageType(IntEnum):
+    REGULAR_DATA = 0
+    PINGPONG = 1
+    LATENCY_MEASUREMENT = 2
+
+
 @app.websocket('/kubernetes-pod-ssh-proxy')
-async def kubernetes_pod_ssh_proxy(websocket: fastapi.WebSocket,
-                                   cluster_name: str) -> None:
+async def kubernetes_pod_ssh_proxy(
+        websocket: fastapi.WebSocket,
+        cluster_name: str,
+        client_version: Optional[int] = None) -> None:
     """Proxies SSH to the Kubernetes pod with websocket."""
     await websocket.accept()
     logger.info(f'WebSocket connection accepted for cluster: {cluster_name}')
 
+    timestamps_supported = client_version is not None and client_version > 21
+    logger.info(f'Websocket timestamps supported: {timestamps_supported}, \
+        client_version = {client_version}')
+
     # Run core.status in another thread to avoid blocking the event loop.
     with ThreadPoolExecutor(max_workers=1) as thread_pool_executor:
         cluster_records = await context_utils.to_thread_with_executor(
@@ -1868,6 +1887,42 @@ async def kubernetes_pod_ssh_proxy(websocket: fastapi.WebSocket,
         async def websocket_to_ssh():
             try:
                 async for message in websocket.iter_bytes():
+                    if timestamps_supported:
+                        type_size = struct.calcsize('!B')
+                        message_type = struct.unpack('!B',
+                                                     message[:type_size])[0]
+                        if (message_type ==
+                                KubernetesSSHMessageType.REGULAR_DATA):
+                            # Regular data - strip type byte and forward to SSH
+                            message = message[type_size:]
+                        elif message_type == KubernetesSSHMessageType.PINGPONG:
+                            # PING message - respond with PONG (type 1)
+                            ping_id_size = struct.calcsize('!I')
+                            if len(message) != type_size + ping_id_size:
+                                raise ValueError('Invalid PING message '
+                                                 f'length: {len(message)}')
+                            # Return the same PING message, so that the client
+                            # can measure the latency.
+                            await websocket.send_bytes(message)
+                            continue
+                        elif (message_type ==
+                              KubernetesSSHMessageType.LATENCY_MEASUREMENT):
+                            # Latency measurement from client
+                            latency_size = struct.calcsize('!Q')
+                            if len(message) != type_size + latency_size:
+                                raise ValueError(
+                                    'Invalid latency measurement '
+                                    f'message length: {len(message)}')
+                            avg_latency_ms = struct.unpack(
+                                '!Q',
+                                message[type_size:type_size + latency_size])[0]
+                            latency_seconds = avg_latency_ms / 1000
+                            metrics_utils.SKY_APISERVER_WEBSOCKET_SSH_LATENCY_SECONDS.labels(pid=os.getpid()).observe(latency_seconds)  # pylint: disable=line-too-long
+                            continue
+                        else:
+                            # Unknown message type.
+                            raise ValueError(
+                                f'Unknown message type: {message_type}')
                     writer.write(message)
                     try:
                         await writer.drain()
@@ -1898,6 +1953,11 @@ async def kubernetes_pod_ssh_proxy(websocket: fastapi.WebSocket,
                             nonlocal ssh_failed
                             ssh_failed = True
                         break
+                    if timestamps_supported:
+                        # Prepend message type byte (0 = regular data)
+                        message_type_bytes = struct.pack(
+                            '!B', KubernetesSSHMessageType.REGULAR_DATA.value)
+                        data = message_type_bytes + data
                     await websocket.send_bytes(data)
             except Exception:  # pylint: disable=broad-except
                 pass
@@ -1937,7 +1997,7 @@ async def all_contexts(request: fastapi.Request) -> None:
 
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='all_contexts',
+        request_name=request_names.RequestName.ALL_CONTEXTS,
         request_body=payloads.RequestBody(),
         func=core.get_all_contexts,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -2051,7 +2111,6 @@ if __name__ == '__main__':
     # Serve metrics on a separate port to isolate it from the application APIs:
     # metrics port will not be exposed to the public network typically.
     parser.add_argument('--metrics-port', default=9090, type=int)
-    parser.add_argument('--start-with-python', action='store_true')
     cmd_args = parser.parse_args()
     if cmd_args.port == cmd_args.metrics_port:
         logger.error('port and metrics-port cannot be the same, exiting.')
@@ -2066,9 +2125,18 @@ if __name__ == '__main__':
         logger.error(f'Port {cmd_args.port} is not available, exiting.')
         raise RuntimeError(f'Port {cmd_args.port} is not available')
 
-    if not cmd_args.start_with_python:
-        # Maybe touch the signal file on API server startup.
-        managed_job_utils.is_consolidation_mode(on_api_restart=True)
+    # Maybe touch the signal file on API server startup. Do it again here even
+    # if we already touched it in the sky/server/common.py::_start_api_server.
+    # This is because the sky/server/common.py::_start_api_server function call
+    # is running outside the skypilot API server process tree. The process tree
+    # starts within that function (see the `subprocess.Popen` call in
+    # sky/server/common.py::_start_api_server). When pg is used, the
+    # _start_api_server function will not load the config file from db, which
+    # will ignore the consolidation mode config. Here, inside the process tree,
+    # we already reload the config as a server (with env var _start_api_server),
+    # so we will respect the consolidation mode config.
+    # Refers to #7717 for more details.
+    managed_job_utils.is_consolidation_mode(on_api_restart=True)
 
     # Show the privacy policy if it is not already shown. We place it here so
     # that it is shown only when the API server is started.

sky/skylet/constants.py

@@ -100,7 +100,7 @@ TASK_ID_LIST_ENV_VAR = f'{SKYPILOT_ENV_VAR_PREFIX}TASK_IDS'
 # cluster yaml is updated.
 #
 # TODO(zongheng,zhanghao): make the upgrading of skylet automatic?
-SKYLET_VERSION = '23'
+SKYLET_VERSION = '25'
 # The version of the lib files that skylet/jobs use. Whenever there is an API
 # change for the job_lib or log_lib, we need to bump this version, so that the
 # user can be notified to update their SkyPilot version on the remote cluster.
@@ -422,6 +422,8 @@ SKIPPED_CLIENT_OVERRIDE_KEYS: List[Tuple[str, ...]] = [
     #   but the configs won't be applied)
     ('jobs', 'controller', 'consolidation_mode'),
     ('serve', 'controller', 'consolidation_mode'),
+    ('jobs', 'controller', 'controller_logs_gc_retention_hours'),
+    ('jobs', 'controller', 'task_logs_gc_retention_hours'),
 ]
 
 # Constants for Azure blob storage
@@ -548,3 +550,6 @@ ENV_VAR_LOOP_LAG_THRESHOLD_MS = (SKYPILOT_ENV_VAR_PREFIX +
 
 ARM64_ARCH = 'arm64'
 X86_64_ARCH = 'x86_64'
+
+SSH_DISABLE_LATENCY_MEASUREMENT_ENV_VAR = (
+    f'{SKYPILOT_ENV_VAR_PREFIX}SSH_DISABLE_LATENCY_MEASUREMENT')

sky/skylet/events.py

@@ -326,8 +326,15 @@ class AutostopEvent(SkyletEvent):
         cluster_name_on_cloud = cluster_config['cluster_name']
         is_cluster_multinode = cluster_config['max_workers'] > 0
 
+        # Clear AWS credentials from environment to force boto3 to use IAM
+        # role attached to the instance (lowest priority in credential chain).
+        # This allows the cluster to stop/terminate itself using its IAM role.
         os.environ.pop('AWS_ACCESS_KEY_ID', None)
         os.environ.pop('AWS_SECRET_ACCESS_KEY', None)
+        os.environ.pop('AWS_SESSION_TOKEN', None)
+        # Point boto3 to /dev/null to skip reading credentials from files.
+        os.environ['AWS_SHARED_CREDENTIALS_FILE'] = '/dev/null'
+        os.environ['AWS_CONFIG_FILE'] = '/dev/null'
 
         # Stop the ray autoscaler to avoid scaling up, during
         # stopping/terminating of the cluster.

sky/skylet/services.py

@@ -407,7 +407,9 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
         context: grpc.ServicerContext
     ) -> managed_jobsv1_pb2.GetJobTableResponse:
         try:
-            accessible_workspaces = list(request.accessible_workspaces)
+            accessible_workspaces = (
+                list(request.accessible_workspaces.workspaces)
+                if request.HasField('accessible_workspaces') else None)
             job_ids = (list(request.job_ids.ids)
                        if request.HasField('job_ids') else None)
             user_hashes: Optional[List[Optional[str]]] = None
@@ -419,6 +421,8 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
                     user_hashes.append(None)
             statuses = (list(request.statuses.statuses)
                         if request.HasField('statuses') else None)
+            fields = (list(request.fields.fields)
+                      if request.HasField('fields') else None)
             job_queue = managed_job_utils.get_managed_job_queue(
                 skip_finished=request.skip_finished,
                 accessible_workspaces=accessible_workspaces,
@@ -432,7 +436,9 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
                 page=request.page if request.HasField('page') else None,
                 limit=request.limit if request.HasField('limit') else None,
                 user_hashes=user_hashes,
-                statuses=statuses)
+                statuses=statuses,
+                fields=fields,
+            )
             jobs = job_queue['jobs']
             total = job_queue['total']
             total_no_filter = job_queue['total_no_filter']
@@ -440,7 +446,16 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
 
             jobs_info = []
             for job in jobs:
+                converted_metadata = None
+                metadata = job.get('metadata')
+                if metadata:
+                    converted_metadata = {
+                        k: v for k, v in metadata.items() if v is not None
+                    }
                 job_info = managed_jobsv1_pb2.ManagedJobInfo(
+                    # The `spot.job_id`, which can be used to identify
+                    # different tasks for the same job
+                    _job_id=job.get('_job_id'),
                     job_id=job.get('job_id'),
                     task_id=job.get('task_id'),
                     job_name=job.get('job_name'),
@@ -468,11 +483,7 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
                     end_at=job.get('end_at'),
                     user_yaml=job.get('user_yaml'),
                     entrypoint=job.get('entrypoint'),
-                    metadata={
-                        k: v
-                        for k, v in job.get('metadata', {}).items()
-                        if v is not None
-                    },
+                    metadata=converted_metadata,
                     pool=job.get('pool'),
                     pool_hash=job.get('pool_hash'))
                 jobs_info.append(job_info)

sky/ssh_node_pools/server.py

@@ -7,6 +7,7 @@ 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 request_names
 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
@@ -101,7 +102,7 @@ async def deploy_ssh_node_pool(request: fastapi.Request,
         ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=False)
         await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_up',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_UP,
             request_body=ssh_up_body,
             func=sky_core.ssh_up,
             schedule_type=requests_lib.ScheduleType.LONG,
@@ -126,7 +127,7 @@ async def deploy_ssh_node_pool_general(
     try:
         await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_up',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_UP,
             request_body=ssh_up_body,
             func=sky_core.ssh_up,
             schedule_type=requests_lib.ScheduleType.LONG,
@@ -152,7 +153,7 @@ async def down_ssh_node_pool(request: fastapi.Request,
         ssh_up_body = payloads.SSHUpBody(infra=pool_name, cleanup=True)
         await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_down',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_DOWN,
             request_body=ssh_up_body,
             func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
             schedule_type=requests_lib.ScheduleType.LONG,
@@ -180,7 +181,7 @@ async def down_ssh_node_pool_general(
         ssh_up_body.cleanup = True
         await executor.schedule_request_async(
             request_id=request.state.request_id,
-            request_name='ssh_down',
+            request_name=request_names.RequestName.SSH_NODE_POOLS_DOWN,
             request_body=ssh_up_body,
             func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
             schedule_type=requests_lib.ScheduleType.LONG,

sky/task.py

@@ -1,6 +1,5 @@
 """Task: a coarse-grained stage in an application."""
 import collections
-import inspect
 import json
 import os
 import re
@@ -29,10 +28,6 @@ from sky.utils import yaml_utils
 
 logger = sky_logging.init_logger(__name__)
 
-# A lambda generating commands (node rank_i, node addrs -> cmd_i).
-CommandGen = Callable[[int, List[str]], Optional[str]]
-CommandOrCommandGen = Union[str, CommandGen]
-
 _VALID_NAME_REGEX = '[a-zA-Z0-9]+(?:[._-]{1,2}[a-zA-Z0-9]+)*'
 _VALID_NAME_DESCR = ('ASCII characters and may contain lowercase and'
                      ' uppercase letters, digits, underscores, periods,'
@@ -236,7 +231,7 @@ class Task:
         name: Optional[str] = None,
         *,
         setup: Optional[Union[str, List[str]]] = None,
-        run: Optional[Union[CommandOrCommandGen, List[str]]] = None,
+        run: Optional[Union[str, List[str]]] = None,
         envs: Optional[Dict[str, str]] = None,
         secrets: Optional[Dict[str, str]] = None,
         workdir: Optional[Union[str, Dict[str, Any]]] = None,
@@ -349,7 +344,7 @@ class Task:
         self._volumes = volumes or {}
 
         # concatenate commands if given as list
-        def _concat(commands):
+        def _concat(commands: Optional[Union[str, List[str]]]) -> Optional[str]:
             if isinstance(commands, list):
                 return '\n'.join(commands)
             return commands
@@ -447,42 +442,9 @@ class Task:
 
     def validate_run(self):
         """Validates if the run command is valid."""
-        if callable(self.run):
-            run_sig = inspect.signature(self.run)
-            # Check that run is a function with 2 arguments.
-            if len(run_sig.parameters) != 2:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(_RUN_FN_CHECK_FAIL_MSG.format(run_sig))
-
-            type_list = [int, List[str]]
-            # Check annotations, if exists
-            for i, param in enumerate(run_sig.parameters.values()):
-                if param.annotation != inspect.Parameter.empty:
-                    if param.annotation != type_list[i]:
-                        with ux_utils.print_exception_no_traceback():
-                            raise ValueError(
-                                _RUN_FN_CHECK_FAIL_MSG.format(run_sig))
-
-            # Check self containedness.
-            run_closure = inspect.getclosurevars(self.run)
-            if run_closure.nonlocals:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'run command generator must be self contained. '
-                        f'Found nonlocals: {run_closure.nonlocals}')
-            if run_closure.globals:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'run command generator must be self contained. '
-                        f'Found globals: {run_closure.globals}')
-            if run_closure.unbound:
-                # Do not raise an error here. Import statements, which are
-                # allowed, will be considered as unbounded.
-                pass
-        elif self.run is not None and not isinstance(self.run, str):
+        if self.run is not None and not isinstance(self.run, str):
             with ux_utils.print_exception_no_traceback():
-                raise ValueError('run must be either a shell script (str) or '
-                                 f'a command generator ({CommandGen}). '
+                raise ValueError('run must be a shell script (str). '
                                  f'Got {type(self.run)}')
 
     def expand_and_validate_file_mounts(self):

sky/templates/kubernetes-ray.yml.j2

@@ -1059,7 +1059,7 @@ available_node_types:
                 # Also, skip the jobs that are waiting to be scheduled as those does not have a controller process running.
                 # For SkyServe, this will be None and every service will be recovered. This is because SkyServe
                 # will delete the service from the database after it is terminated so everything in the database is running.
-                ALL_IN_PROGRESS_JOBS=$({{sky_python_cmd}} -c "from sky.jobs import state; jobs = state.get_managed_jobs(); print(' '.join({str(job['job_id']) for job in jobs if job['schedule_state'] not in [state.ManagedJobScheduleState.DONE, state.ManagedJobScheduleState.WAITING]}) if jobs else None)")
+                ALL_IN_PROGRESS_JOBS=$({{sky_python_cmd}} -c "from sky.jobs import state; jobs, _ = state.get_managed_jobs_with_filters(fields=['job_id', 'schedule_state']); print(' '.join({str(job['job_id']) for job in jobs if job['schedule_state'] not in [state.ManagedJobScheduleState.DONE, state.ManagedJobScheduleState.WAITING]}) if jobs else None)")
                 if [ "$ALL_IN_PROGRESS_JOBS" != "None" ]; then
                   read -ra ALL_IN_PROGRESS_JOBS_SEQ <<< "$ALL_IN_PROGRESS_JOBS"
                 fi