skypilot-nightly 1.0.0.dev20250919__py3-none-any.whl → 1.0.0.dev20250925__py3-none-any.whl

sky/execution.py

@@ -265,8 +265,10 @@ def _execute_dag(
 
     cluster_exists = False
     if cluster_name is not None:
-        cluster_record = global_user_state.get_cluster_from_name(cluster_name)
-        cluster_exists = cluster_record is not None
+        # We use launched_at to check if the cluster exists, because this
+        # db query is faster than get_cluster_from_name.
+        cluster_exists = global_user_state.cluster_with_name_exists(
+            cluster_name)
         # TODO(woosuk): If the cluster exists, print a warning that
         # `cpus` and `memory` are not used as a job scheduling constraint,
         # unlike `gpus`.

sky/global_user_state.py

@@ -24,6 +24,7 @@ from sqlalchemy import exc as sqlalchemy_exc
 from sqlalchemy import orm
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.dialects import sqlite
+from sqlalchemy.ext import asyncio as sql_async
 from sqlalchemy.ext import declarative
 
 from sky import models
@@ -51,6 +52,7 @@ _ENABLED_CLOUDS_KEY_PREFIX = 'enabled_clouds_'
 _ALLOWED_CLOUDS_KEY_PREFIX = 'allowed_clouds_'
 
 _SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
+_SQLALCHEMY_ENGINE_ASYNC: Optional[sql_async.AsyncEngine] = None
 _SQLALCHEMY_ENGINE_LOCK = threading.Lock()
 
 DEFAULT_CLUSTER_EVENT_RETENTION_HOURS = 24.0
@@ -296,6 +298,20 @@ def create_table(engine: sqlalchemy.engine.Engine):
         migration_utils.GLOBAL_USER_STATE_VERSION)
 
 
+def initialize_and_get_db_async() -> sql_async.AsyncEngine:
+    global _SQLALCHEMY_ENGINE_ASYNC
+    if _SQLALCHEMY_ENGINE_ASYNC is not None:
+        return _SQLALCHEMY_ENGINE_ASYNC
+    with _SQLALCHEMY_ENGINE_LOCK:
+        if _SQLALCHEMY_ENGINE_ASYNC is not None:
+            return _SQLALCHEMY_ENGINE_ASYNC
+
+        _SQLALCHEMY_ENGINE_ASYNC = db_utils.get_engine('state',
+                                                       async_engine=True)
+    initialize_and_get_db()
+    return _SQLALCHEMY_ENGINE_ASYNC
+
+
 # We wrap the sqlalchemy engine initialization in a thread
 # lock to ensure that multiple threads do not initialize the
 # engine which could result in a rare race condition where
@@ -321,6 +337,22 @@ def initialize_and_get_db() -> sqlalchemy.engine.Engine:
         return _SQLALCHEMY_ENGINE
 
 
+def _init_db_async(func):
+    """Initialize the async database."""
+
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        if _SQLALCHEMY_ENGINE_ASYNC is None:
+            # this may happen multiple times since there is no locking
+            # here but thats fine, this is just a short circuit for the
+            # common case.
+            await context_utils.to_thread(initialize_and_get_db_async)
+
+        return await func(*args, **kwargs)
+
+    return wrapper
+
+
 def _init_db(func):
     """Initialize the database."""
 
@@ -1010,29 +1042,68 @@ def get_handle_from_cluster_name(
     assert _SQLALCHEMY_ENGINE is not None
     assert cluster_name is not None, 'cluster_name cannot be None'
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(cluster_table).filter_by(name=cluster_name).first()
+        row = (session.query(
+            cluster_table.c.handle).filter_by(name=cluster_name).first())
     if row is None:
         return None
     return pickle.loads(row.handle)
 
 
+@_init_db_async
+@metrics_lib.time_me
+async def get_status_from_cluster_name_async(
+        cluster_name: str) -> Optional[status_lib.ClusterStatus]:
+    """Get the status of a cluster."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    assert cluster_name is not None, 'cluster_name cannot be None'
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.select(cluster_table.c.status).where(
+                cluster_table.c.name == cluster_name))
+        row = result.first()
+
+        if row is None:
+            return None
+        return status_lib.ClusterStatus(row[0])
+
+
 @_init_db
 @metrics_lib.time_me
-def get_glob_cluster_names(cluster_name: str) -> List[str]:
+def get_status_from_cluster_name(
+        cluster_name: str) -> Optional[status_lib.ClusterStatus]:
+    assert _SQLALCHEMY_ENGINE is not None
+    assert cluster_name is not None, 'cluster_name cannot be None'
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        row = session.query(
+            cluster_table.c.status).filter_by(name=cluster_name).first()
+    if row is None:
+        return None
+    return status_lib.ClusterStatus[row.status]
+
+
+@_init_db
+@metrics_lib.time_me
+def get_glob_cluster_names(
+        cluster_name: str,
+        workspaces_filter: Optional[Set[str]] = None) -> List[str]:
     assert _SQLALCHEMY_ENGINE is not None
     assert cluster_name is not None, 'cluster_name cannot be None'
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
         if (_SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
-            rows = session.query(cluster_table).filter(
-                cluster_table.c.name.op('GLOB')(cluster_name)).all()
+            query = session.query(cluster_table.c.name).filter(
+                cluster_table.c.name.op('GLOB')(cluster_name))
         elif (_SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
-            rows = session.query(cluster_table).filter(
+            query = session.query(cluster_table.c.name).filter(
                 cluster_table.c.name.op('SIMILAR TO')(
-                    _glob_to_similar(cluster_name))).all()
+                    _glob_to_similar(cluster_name)))
         else:
             raise ValueError('Unsupported database dialect')
+        if workspaces_filter is not None:
+            query = query.filter(
+                cluster_table.c.workspace.in_(workspaces_filter))
+        rows = query.all()
     return [row.name for row in rows]
 
 
@@ -1076,7 +1147,8 @@ def set_cluster_autostop_value(cluster_name: str, idle_minutes: int,
 def get_cluster_launch_time(cluster_name: str) -> Optional[int]:
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(cluster_table).filter_by(name=cluster_name).first()
+        row = session.query(
+            cluster_table.c.launched_at).filter_by(name=cluster_name).first()
     if row is None or row.launched_at is None:
         return None
     return int(row.launched_at)
@@ -1087,7 +1159,8 @@ def get_cluster_launch_time(cluster_name: str) -> Optional[int]:
 def get_cluster_info(cluster_name: str) -> Optional[Dict[str, Any]]:
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(cluster_table).filter_by(name=cluster_name).first()
+        row = session.query(
+            cluster_table.c.metadata).filter_by(name=cluster_name).first()
     if row is None or row.metadata is None:
         return None
     return json.loads(row.metadata)
@@ -1167,7 +1240,8 @@ def get_cluster_storage_mounts_metadata(
         cluster_name: str) -> Optional[Dict[str, Any]]:
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(cluster_table).filter_by(name=cluster_name).first()
+        row = (session.query(cluster_table.c.storage_mounts_metadata).filter_by(
+            name=cluster_name).first())
     if row is None or row.storage_mounts_metadata is None:
         return None
     return pickle.loads(row.storage_mounts_metadata)
@@ -1196,7 +1270,9 @@ def get_cluster_skylet_ssh_tunnel_metadata(
         cluster_name: str) -> Optional[Tuple[int, int]]:
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(cluster_table).filter_by(name=cluster_name).first()
+        row = session.query(
+            cluster_table.c.skylet_ssh_tunnel_metadata).filter_by(
+                name=cluster_name).first()
     if row is None or row.skylet_ssh_tunnel_metadata is None:
         return None
     return pickle.loads(row.skylet_ssh_tunnel_metadata)
@@ -1230,7 +1306,7 @@ def _get_cluster_usage_intervals(
     if cluster_hash is None:
         return None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(cluster_history_table).filter_by(
+        row = session.query(cluster_history_table.c.usage_intervals).filter_by(
             cluster_hash=cluster_hash).first()
     if row is None or row.usage_intervals is None:
         return None
@@ -1305,7 +1381,8 @@ def set_owner_identity_for_cluster(cluster_name: str,
 def _get_hash_for_existing_cluster(cluster_name: str) -> Optional[str]:
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(cluster_table).filter_by(name=cluster_name).first()
+        row = (session.query(
+            cluster_table.c.cluster_hash).filter_by(name=cluster_name).first())
     if row is None or row.cluster_hash is None:
         return None
     return row.cluster_hash
@@ -1317,8 +1394,10 @@ def get_launched_resources_from_cluster_hash(
         cluster_hash: str) -> Optional[Tuple[int, Any]]:
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(cluster_history_table).filter_by(
-            cluster_hash=cluster_hash).first()
+        row = session.query(
+            cluster_history_table.c.num_nodes,
+            cluster_history_table.c.launched_resources).filter_by(
+                cluster_hash=cluster_hash).first()
     if row is None:
         return None
     num_nodes = row.num_nodes
@@ -1362,17 +1441,56 @@ def _load_storage_mounts_metadata(
 @metrics_lib.time_me
 @context_utils.cancellation_guard
 def get_cluster_from_name(
-        cluster_name: Optional[str]) -> Optional[Dict[str, Any]]:
+        cluster_name: Optional[str],
+        *,
+        include_user_info: bool = True,
+        summary_response: bool = False) -> Optional[Dict[str, Any]]:
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(cluster_table).filter_by(name=cluster_name).first()
+        if summary_response:
+            query = session.query(
+                cluster_table.c.name, cluster_table.c.launched_at,
+                cluster_table.c.handle, cluster_table.c.last_use,
+                cluster_table.c.status, cluster_table.c.autostop,
+                cluster_table.c.to_down, cluster_table.c.owner,
+                cluster_table.c.metadata, cluster_table.c.cluster_hash,
+                cluster_table.c.storage_mounts_metadata,
+                cluster_table.c.cluster_ever_up,
+                cluster_table.c.status_updated_at, cluster_table.c.user_hash,
+                cluster_table.c.config_hash, cluster_table.c.workspace,
+                cluster_table.c.is_managed)
+        else:
+            query = session.query(
+                cluster_table.c.name,
+                cluster_table.c.launched_at,
+                cluster_table.c.handle,
+                cluster_table.c.last_use,
+                cluster_table.c.status,
+                cluster_table.c.autostop,
+                cluster_table.c.to_down,
+                cluster_table.c.owner,
+                cluster_table.c.metadata,
+                cluster_table.c.cluster_hash,
+                cluster_table.c.storage_mounts_metadata,
+                cluster_table.c.cluster_ever_up,
+                cluster_table.c.status_updated_at,
+                cluster_table.c.user_hash,
+                cluster_table.c.config_hash,
+                cluster_table.c.workspace,
+                cluster_table.c.is_managed,
+                # extra fields compared to above query
+                cluster_table.c.last_creation_yaml,
+                cluster_table.c.last_creation_command)
+        row = query.filter_by(name=cluster_name).first()
     if row is None:
         return None
-    user_hash = _get_user_hash_or_current_user(row.user_hash)
-    user = get_user(user_hash)
-    user_name = user.name if user is not None else None
-    last_event = get_last_cluster_event(
-        row.cluster_hash, event_type=ClusterEventType.STATUS_CHANGE)
+    if include_user_info:
+        user_hash = _get_user_hash_or_current_user(row.user_hash)
+        user = get_user(user_hash)
+        user_name = user.name if user is not None else None
+    if not summary_response:
+        last_event = get_last_cluster_event(
+            row.cluster_hash, event_type=ClusterEventType.STATUS_CHANGE)
     # TODO: use namedtuple instead of dict
     record = {
         'name': row.name,
@@ -1389,26 +1507,43 @@ def get_cluster_from_name(
             row.storage_mounts_metadata),
         'cluster_ever_up': bool(row.cluster_ever_up),
         'status_updated_at': row.status_updated_at,
-        'user_hash': user_hash,
-        'user_name': user_name,
-        'config_hash': row.config_hash,
         'workspace': row.workspace,
-        'last_creation_yaml': row.last_creation_yaml,
-        'last_creation_command': row.last_creation_command,
         'is_managed': bool(row.is_managed),
-        'last_event': last_event,
+        'config_hash': row.config_hash,
     }
+    if not summary_response:
+        record['last_creation_yaml'] = row.last_creation_yaml
+        record['last_creation_command'] = row.last_creation_command
+        record['last_event'] = last_event
+    if include_user_info:
+        record['user_hash'] = user_hash
+        record['user_name'] = user_name
 
     return record
 
 
+@_init_db
+@metrics_lib.time_me
+@context_utils.cancellation_guard
+def cluster_with_name_exists(cluster_name: str) -> bool:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        row = session.query(
+            cluster_table.c.name).filter_by(name=cluster_name).first()
+    if row is None:
+        return False
+    return True
+
+
 @_init_db
 @metrics_lib.time_me
 def get_clusters(
     *,  # keyword only separator
     exclude_managed_clusters: bool = False,
-    workspaces_filter: Optional[Set[str]] = None,
+    workspaces_filter: Optional[Dict[str, Any]] = None,
     user_hashes_filter: Optional[Set[str]] = None,
+    cluster_names: Optional[List[str]] = None,
+    summary_response: bool = False,
 ) -> List[Dict[str, Any]]:
     """Get clusters from the database.
 
@@ -1419,13 +1554,48 @@ def get_clusters(
             that has workspace field set to one of the values.
         user_hashes_filter: If specified, only include clusters
             that has user_hash field set to one of the values.
+        cluster_names: If specified, only include clusters
+            that has name field set to one of the values.
     """
     # is a cluster has a null user_hash,
     # we treat it as belonging to the current user.
     current_user_hash = common_utils.get_user_hash()
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        query = session.query(cluster_table)
+        if summary_response:
+            query = session.query(
+                cluster_table.c.name, cluster_table.c.launched_at,
+                cluster_table.c.handle, cluster_table.c.last_use,
+                cluster_table.c.status, cluster_table.c.autostop,
+                cluster_table.c.to_down, cluster_table.c.owner,
+                cluster_table.c.metadata, cluster_table.c.cluster_hash,
+                cluster_table.c.storage_mounts_metadata,
+                cluster_table.c.cluster_ever_up,
+                cluster_table.c.status_updated_at, cluster_table.c.user_hash,
+                cluster_table.c.config_hash, cluster_table.c.workspace,
+                cluster_table.c.is_managed)
+        else:
+            query = session.query(
+                cluster_table.c.name,
+                cluster_table.c.launched_at,
+                cluster_table.c.handle,
+                cluster_table.c.last_use,
+                cluster_table.c.status,
+                cluster_table.c.autostop,
+                cluster_table.c.to_down,
+                cluster_table.c.owner,
+                cluster_table.c.metadata,
+                cluster_table.c.cluster_hash,
+                cluster_table.c.storage_mounts_metadata,
+                cluster_table.c.cluster_ever_up,
+                cluster_table.c.status_updated_at,
+                cluster_table.c.user_hash,
+                cluster_table.c.config_hash,
+                cluster_table.c.workspace,
+                cluster_table.c.is_managed,
+                # extra fields compared to above query
+                cluster_table.c.last_creation_yaml,
+                cluster_table.c.last_creation_command)
         if exclude_managed_clusters:
             query = query.filter(cluster_table.c.is_managed == int(False))
         if workspaces_filter is not None:
@@ -1437,11 +1607,13 @@ def get_clusters(
                 # If current_user_hash is in user_hashes_filter, we include
                 # clusters that have a null user_hash.
                 query = query.filter(
-                    cluster_table.c.user_hash.in_(user_hashes_filter) |
-                    (cluster_table.c.user_hash is None))
+                    (cluster_table.c.user_hash.in_(user_hashes_filter) |
+                     (cluster_table.c.user_hash is None)))
             else:
                 query = query.filter(
                     cluster_table.c.user_hash.in_(user_hashes_filter))
+        if cluster_names is not None:
+            query = query.filter(cluster_table.c.name.in_(cluster_names))
         query = query.order_by(sqlalchemy.desc(cluster_table.c.launched_at))
         rows = query.all()
     records = []
@@ -1459,15 +1631,15 @@ def get_clusters(
 
     # get last cluster event for each row
     cluster_hashes = set(row_to_user_hash.keys())
-    last_cluster_event_dict = _get_last_cluster_event_multiple(
-        cluster_hashes, ClusterEventType.STATUS_CHANGE)
+    if not summary_response:
+        last_cluster_event_dict = _get_last_cluster_event_multiple(
+            cluster_hashes, ClusterEventType.STATUS_CHANGE)
 
     # get user for each row
     for row in rows:
         user_hash = row_to_user_hash[row.cluster_hash]
         user = user_hash_to_user.get(user_hash, None)
         user_name = user.name if user is not None else None
-        last_event = last_cluster_event_dict.get(row.cluster_hash, None)
         # TODO: use namedtuple instead of dict
         record = {
             'name': row.name,
@@ -1486,18 +1658,32 @@ def get_clusters(
             'status_updated_at': row.status_updated_at,
             'user_hash': user_hash,
             'user_name': user_name,
-            'config_hash': row.config_hash,
             'workspace': row.workspace,
-            'last_creation_yaml': row.last_creation_yaml,
-            'last_creation_command': row.last_creation_command,
             'is_managed': bool(row.is_managed),
-            'last_event': last_event,
+            'config_hash': row.config_hash,
         }
+        if not summary_response:
+            record['last_creation_yaml'] = row.last_creation_yaml
+            record['last_creation_command'] = row.last_creation_command
+            record['last_event'] = last_cluster_event_dict.get(
+                row.cluster_hash, None)
 
         records.append(record)
     return records
 
 
+@_init_db
+@metrics_lib.time_me
+def get_cluster_names(exclude_managed_clusters: bool = False,) -> List[str]:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        query = session.query(cluster_table.c.name)
+        if exclude_managed_clusters:
+            query = query.filter(cluster_table.c.is_managed == int(False))
+        rows = query.all()
+    return [row[0] for row in rows]
+
+
 @_init_db
 @metrics_lib.time_me
 def get_clusters_from_history(

sky/jobs/client/sdk.py

@@ -383,15 +383,24 @@ def dashboard() -> None:
 @server_common.check_server_healthy_or_start
 @versions.minimal_api_version(12)
 def pool_apply(
-    task: Union['sky.Task', 'sky.Dag'],
+    task: Optional[Union['sky.Task', 'sky.Dag']],
     pool_name: str,
     mode: 'serve_utils.UpdateMode',
+    workers: Optional[int] = None,
     # Internal only:
     # pylint: disable=invalid-name
     _need_confirmation: bool = False
 ) -> server_common.RequestId[None]:
     """Apply a config to a pool."""
+    remote_api_version = versions.get_remote_api_version()
+    if (workers is not None and
+        (remote_api_version is None or remote_api_version < 19)):
+        raise click.UsageError('Updating the number of workers in a pool is '
+                               'not supported in your API server. Please '
+                               'upgrade to a newer API server to use this '
+                               'feature.')
     return impl.apply(task,
+                      workers,
                       pool_name,
                       mode,
                       pool=True,

sky/jobs/controller.py

@@ -3,6 +3,7 @@
 import asyncio
 import logging
 import os
+import pathlib
 import resource
 import shutil
 import sys
@@ -56,6 +57,7 @@ async def create_background_task(coro: typing.Coroutine) -> None:
     async with _background_tasks_lock:
         task = asyncio.create_task(coro)
         _background_tasks.add(task)
+        # TODO(cooperc): Discard needs a lock?
         task.add_done_callback(_background_tasks.discard)
 
 
@@ -896,6 +898,9 @@ class Controller:
             # some data here.
             raise error
 
+    # Use context.contextual to enable per-job output redirection and env var
+    # isolation.
+    @context.contextual
     async def run_job_loop(self,
                            job_id: int,
                            dag_yaml: str,
@@ -904,13 +909,9 @@ class Controller:
                            env_file_path: Optional[str] = None,
                            pool: Optional[str] = None):
         """Background task that runs the job loop."""
-        # Replace os.environ with ContextualEnviron to enable per-job
-        # environment isolation. This allows each job to have its own
-        # environment variables without affecting other jobs or the main
-        # process.
-        context.initialize()
         ctx = context.get()
-        ctx.redirect_log(log_file)  # type: ignore
+        assert ctx is not None, 'Context is not initialized'
+        ctx.redirect_log(pathlib.Path(log_file))
 
         # Load and apply environment variables from the job's environment file
         if env_file_path and os.path.exists(env_file_path):
@@ -921,7 +922,6 @@ class Controller:
                                 f'{list(env_vars.keys())}')
 
                 # Apply environment variables to the job's context
-                ctx = context.get()
                 if ctx is not None:
                     for key, value in env_vars.items():
                         if value is not None:

sky/jobs/server/core.py

@@ -281,8 +281,7 @@ def launch(
         # Check whether cached jobs controller cluster is accessible
         cluster_name = (
             controller_utils.Controllers.JOBS_CONTROLLER.value.cluster_name)
-        record = global_user_state.get_cluster_from_name(cluster_name)
-        if record is not None:
+        if global_user_state.cluster_with_name_exists(cluster_name):
             # there is a cached jobs controller cluster
             try:
                 # TODO: do something with returned status?
@@ -959,9 +958,10 @@ def pool_apply(
     task: 'sky.Task',
     pool_name: str,
     mode: serve_utils.UpdateMode = serve_utils.DEFAULT_UPDATE_MODE,
+    workers: Optional[int] = None,
 ) -> None:
     """Apply a config to a pool."""
-    return impl.apply(task, pool_name, mode, pool=True)
+    return impl.apply(task, workers, pool_name, mode, pool=True)
 
 
 @usage_lib.entrypoint

sky/jobs/server/server.py

@@ -94,23 +94,27 @@ 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
+    request_task = executor.prepare_request(
         request_id=request.state.request_id,
         request_name='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
-                                                       )
+    if schedule_type == api_requests.ScheduleType.LONG:
+        executor.schedule_prepared_request(request_task)
+    else:
+        # 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)
 
     return stream_utils.stream_response(
         request_id=request_task.request_id,

sky/jobs/utils.py

@@ -156,7 +156,7 @@ def _validate_consolidation_mode_config(
     if current_is_consolidation_mode:
         controller_cn = (
             controller_utils.Controllers.JOBS_CONTROLLER.value.cluster_name)
-        if global_user_state.get_cluster_from_name(controller_cn) is not None:
+        if global_user_state.cluster_with_name_exists(controller_cn):
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.InconsistentConsolidationModeError(
                     f'{colorama.Fore.RED}Consolidation mode for jobs is '

sky/logs/agent.py

@@ -34,7 +34,8 @@ class FluentbitAgent(LoggingAgent):
     def get_setup_command(self,
                           cluster_name: resources_utils.ClusterName) -> str:
         install_cmd = (
-            'if ! command -v fluent-bit >/dev/null 2>&1; then '
+            # pylint: disable=line-too-long
+            'if ! command -v fluent-bit >/dev/null 2>&1 && [ ! -f /opt/fluent-bit/bin/fluent-bit ]; then '
             'sudo apt-get update; sudo apt-get install -y gnupg; '
             # pylint: disable=line-too-long
             'sudo sh -c \'curl https://packages.fluentbit.io/fluentbit.key | gpg --dearmor > /usr/share/keyrings/fluentbit-keyring.gpg\'; '
@@ -51,14 +52,32 @@ class FluentbitAgent(LoggingAgent):
         cfg_path = os.path.join(constants.LOGGING_CONFIG_DIR, 'fluentbit.yaml')
         config_cmd = (f'mkdir -p {constants.LOGGING_CONFIG_DIR} && '
                       f'echo {shlex.quote(cfg)} > {cfg_path}')
+        kill_prior_cmd = (
+            'if [ -f "/tmp/fluentbit.pid" ]; then '
+            # pylint: disable=line-too-long
+            'echo "Killing prior fluent-bit process $(cat /tmp/fluentbit.pid)"; '
+            'kill "$(cat /tmp/fluentbit.pid)" || true; '
+            'fi')
         start_cmd = ('nohup $(command -v fluent-bit || '
                      'echo "/opt/fluent-bit/bin/fluent-bit") '
-                     f'-c {cfg_path} > /tmp/fluentbit.log 2>&1 &')
-        return f'set -e; {install_cmd}; {config_cmd}; {start_cmd}'
+                     f'-c {cfg_path} > /tmp/fluentbit.log 2>&1 & '
+                     'echo $! > /tmp/fluentbit.pid')
+        return ('set -e; '
+                f'{install_cmd}; '
+                f'{config_cmd}; '
+                f'{kill_prior_cmd}; '
+                f'{start_cmd}')
 
     def fluentbit_config(self,
                          cluster_name: resources_utils.ClusterName) -> str:
         cfg_dict = {
+            'parsers': [{
+                'name': 'sky-ray-parser',
+                'format': 'regex',
+                # pylint: disable=line-too-long
+                'regex': r'(?:\x1b\[[\d;]+m)?\((?<worker_name>[^,]+)(?:,\s*rank=(?<rank>\d+))?(?:,\s*pid=(?<pid>\d+))(?:,\s*ip=(?<ip>[\d.]+))?\)(?:\x1b\[[\d;]+m)?\s*(?<log_line>.*)',
+                'types': 'rank:integer pid:integer',
+            }],
             'pipeline': {
                 'inputs': [{
                     'name': 'tail',
@@ -70,6 +89,14 @@ class FluentbitAgent(LoggingAgent):
                     # right after the job completion.
                     'refresh_interval': 1,
                 }],
+                'filters': [{
+                    'name': 'parser',
+                    'match': '*',
+                    'key_name': 'log',
+                    'parser': 'sky-ray-parser',
+                    'preserve_key': 'on',  # preserve field for backwards compat
+                    'reserve_data': 'on',
+                }],
                 'outputs': [self.fluentbit_output_config(cluster_name)],
             }
         }