[db] only query for specific fields of requests (#7658)

* only query for specific fields

* fix ut

* apply the same trick to individual request gets

* fix UT

Author: SeungjinYang

sky/backends/backend_utils.py

@@ -3128,12 +3128,12 @@ def refresh_cluster_records() -> None:
     # request info in backend_utils.py.
     # Refactor this to use some other info to
     # determine if a launch is in progress.
-    requests = requests_lib.get_request_tasks(
-        req_filter=requests_lib.RequestTaskFilter(
-            status=[requests_lib.RequestStatus.RUNNING],
-            include_request_names=['sky.launch']))
     cluster_names_with_launch_request = {
-        request.cluster_name for request in requests
+        request.cluster_name for request in requests_lib.get_request_tasks(
+            req_filter=requests_lib.RequestTaskFilter(
+                status=[requests_lib.RequestStatus.RUNNING],
+                include_request_names=['sky.launch'],
+                fields=['cluster_name']))
     }
     cluster_names_without_launch_request = (cluster_names -
                                             cluster_names_with_launch_request)
@@ -3356,13 +3356,13 @@ def _refresh_cluster_record(cluster_name):
     # request info in backend_utils.py.
     # Refactor this to use some other info to
     # determine if a launch is in progress.
-    requests = requests_lib.get_request_tasks(
-        req_filter=requests_lib.RequestTaskFilter(
-            status=[requests_lib.RequestStatus.RUNNING],
-            include_request_names=['sky.launch'],
-            cluster_names=cluster_names))
     cluster_names_with_launch_request = {
-        request.cluster_name for request in requests
+        request.cluster_name for request in requests_lib.get_request_tasks(
+            req_filter=requests_lib.RequestTaskFilter(
+                status=[requests_lib.RequestStatus.RUNNING],
+                include_request_names=['sky.launch'],
+                cluster_names=cluster_names,
+                fields=['cluster_name']))
     }
     # Preserve the index of the cluster name as it appears on "records"
     cluster_names_without_launch_request = [

sky/jobs/server/server.py

@@ -201,10 +201,13 @@ async def pool_tail_logs(
         request_cluster_name=common.JOB_CONTROLLER_NAME,
     )
 
-    request_task = api_requests.get_request(request.state.request_id)
+    request_task = api_requests.get_request(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,
     )

sky/server/requests/executor.py

@@ -214,10 +214,11 @@ def process_request(self, executor: process.BurstableExecutor,
                 time.sleep(0.1)
                 return
             request_id, ignore_return_value, _ = request_element
-            request = api_requests.get_request(request_id)
+            request = api_requests.get_request(request_id, fields=['status'])
             assert request is not None, f'Request with ID {request_id} is None'
             if request.status == api_requests.RequestStatus.CANCELLED:
                 return
+            del request
             logger.info(f'[{self}] Submitting request: {request_id}')
             # Start additional process to run the request, so that it can be
             # cancelled when requested by a user.

sky/server/requests/preconditions.py

@@ -98,13 +98,15 @@ async def _wait(self) -> bool:
                 return False
 
             # Check if the request has been cancelled
-            request = await api_requests.get_request_async(self.request_id)
+            request = await api_requests.get_request_async(self.request_id,
+                                                           fields=['status'])
             if request is None:
                 logger.error(f'Request {self.request_id} not found')
                 return False
             if request.status == api_requests.RequestStatus.CANCELLED:
                 logger.debug(f'Request {self.request_id} cancelled')
                 return False
+            del request
 
             try:
                 met, status_msg = await self.check()
@@ -166,7 +168,10 @@ async def check(self) -> Tuple[bool, Optional[str]]:
                     api_requests.RequestStatus.RUNNING
                 ],
                 include_request_names=['sky.launch', 'sky.start'],
-                cluster_names=[self.cluster_name]))
+                cluster_names=[self.cluster_name],
+                # Only get the request ID to avoid fetching the whole request.
+                # We're only interested in the count, not the whole request.
+                fields=['request_id']))
         if len(requests) == 0:
             # No running or pending tasks, the start process is done.
             return True, None

sky/server/requests/requests.py

@@ -400,7 +400,8 @@ def kill_cluster_requests(cluster_name: str, exclude_request_name: str):
         for request_task in get_request_tasks(req_filter=RequestTaskFilter(
             status=[RequestStatus.PENDING, RequestStatus.RUNNING],
             exclude_request_names=[exclude_request_name],
-            cluster_names=[cluster_name]))
+            cluster_names=[cluster_name],
+            fields=['request_id']))
     ]
     kill_requests(request_ids)
 
@@ -425,7 +426,8 @@ def kill_requests(request_ids: Optional[List[str]] = None,
                 status=[RequestStatus.PENDING, RequestStatus.RUNNING],
                 # Avoid cancelling the cancel request itself.
                 exclude_request_names=['sky.api_cancel'],
-                user_id=user_id))
+                user_id=user_id,
+                fields=['request_id']))
         ]
     cancelled_request_ids = []
     for request_id in request_ids:
@@ -604,30 +606,42 @@ async def update_status_msg_async(request_id: str, status_msg: str) -> None:
             await _add_or_update_request_no_lock_async(request)
 
 
-_get_request_sql = (f'SELECT {", ".join(REQUEST_COLUMNS)} FROM {REQUEST_TABLE} '
-                    'WHERE request_id LIKE ?')
-
-
-def _get_request_no_lock(request_id: str) -> Optional[Request]:
+def _get_request_no_lock(
+        request_id: str,
+        fields: Optional[List[str]] = None) -> Optional[Request]:
     """Get a SkyPilot API request."""
     assert _DB is not None
+    columns_str = ', '.join(REQUEST_COLUMNS)
+    if fields:
+        columns_str = ', '.join(fields)
     with _DB.conn:
         cursor = _DB.conn.cursor()
-        cursor.execute(_get_request_sql, (request_id + '%',))
+        cursor.execute((f'SELECT {columns_str} FROM {REQUEST_TABLE} '
+                        'WHERE request_id LIKE ?'), (request_id + '%',))
         row = cursor.fetchone()
         if row is None:
             return None
+    if fields:
+        row = _update_request_row_fields(row, fields)
     return Request.from_row(row)
 
 
-async def _get_request_no_lock_async(request_id: str) -> Optional[Request]:
+async def _get_request_no_lock_async(
+        request_id: str,
+        fields: Optional[List[str]] = None) -> Optional[Request]:
     """Async version of _get_request_no_lock."""
     assert _DB is not None
-    async with _DB.execute_fetchall_async(_get_request_sql,
-                                          (request_id + '%',)) as rows:
+    columns_str = ', '.join(REQUEST_COLUMNS)
+    if fields:
+        columns_str = ', '.join(fields)
+    async with _DB.execute_fetchall_async(
+        (f'SELECT {columns_str} FROM {REQUEST_TABLE} '
+         'WHERE request_id LIKE ?'), (request_id + '%',)) as rows:
         row = rows[0] if rows else None
         if row is None:
             return None
+    if fields:
+        row = _update_request_row_fields(row, fields)
     return Request.from_row(row)
 
 
@@ -646,20 +660,23 @@ def get_latest_request_id() -> Optional[str]:
 
 @init_db
 @metrics_lib.time_me
-def get_request(request_id: str) -> Optional[Request]:
+def get_request(request_id: str,
+                fields: Optional[List[str]] = None) -> Optional[Request]:
     """Get a SkyPilot API request."""
     with filelock.FileLock(request_lock_path(request_id)):
-        return _get_request_no_lock(request_id)
+        return _get_request_no_lock(request_id, fields)
 
 
 @init_db_async
 @metrics_lib.time_me_async
 @asyncio_utils.shield
-async def get_request_async(request_id: str) -> Optional[Request]:
+async def get_request_async(
+        request_id: str,
+        fields: Optional[List[str]] = None) -> Optional[Request]:
     """Async version of get_request."""
     # TODO(aylei): figure out how to remove FileLock here to avoid the overhead
     async with filelock.AsyncFileLock(request_lock_path(request_id)):
-        return await _get_request_no_lock_async(request_id)
+        return await _get_request_no_lock_async(request_id, fields)
 
 
 class StatusWithMsg(NamedTuple):
@@ -919,9 +936,9 @@ def set_request_cancelled(request_id: str) -> None:
 
 @init_db
 @metrics_lib.time_me
-async def _delete_requests(requests: List[Request]):
+async def _delete_requests(request_ids: List[str]):
     """Clean up requests by their IDs."""
-    id_list_str = ','.join(repr(req.request_id) for req in requests)
+    id_list_str = ','.join(repr(request_id) for request_id in request_ids)
     assert _DB is not None
     await _DB.execute_and_commit_async(
         f'DELETE FROM {REQUEST_TABLE} WHERE request_id IN ({id_list_str})')
@@ -949,18 +966,21 @@ async def clean_finished_requests_with_retention(retention_seconds: int,
             req_filter=RequestTaskFilter(status=RequestStatus.finished_status(),
                                          finished_before=time.time() -
                                          retention_seconds,
-                                         limit=batch_size))
+                                         limit=batch_size,
+                                         fields=['request_id']))
         if len(reqs) == 0:
             break
         futs = []
         for req in reqs:
+            # req.log_path is derived from request_id,
+            # so it's ok to just grab the request_id in the above query.
             futs.append(
                 asyncio.create_task(
                     anyio.Path(
                         req.log_path.absolute()).unlink(missing_ok=True)))
         await asyncio.gather(*futs)
 
-        await _delete_requests(reqs)
+        await _delete_requests([req.request_id for req in reqs])
         total_deleted += len(reqs)
         if len(reqs) < batch_size:
             break

sky/server/server.py

@@ -1567,11 +1567,14 @@ async def stream(
     polling_interval = stream_utils.DEFAULT_POLL_INTERVAL
     # Original plain text streaming logic
     if request_id is not None:
-        request_task = await requests_lib.get_request_async(request_id)
+        request_task = await requests_lib.get_request_async(
+            request_id, fields=['request_id', 'schedule_type'])
         if request_task is None:
             print(f'No task with request ID {request_id}')
             raise fastapi.HTTPException(
                 status_code=404, detail=f'Request {request_id!r} not found')
+        # req.log_path is derived from request_id,
+        # so it's ok to just grab the request_id in the above query.
         log_path_to_stream = request_task.log_path
         if not log_path_to_stream.exists():
             # The log file might be deleted by the request GC daemon but the
@@ -1581,6 +1584,7 @@ async def stream(
                 detail=f'Log of request {request_id!r} has been deleted')
         if request_task.schedule_type == requests_lib.ScheduleType.LONG:
             polling_interval = stream_utils.LONG_REQUEST_POLL_INTERVAL
+        del request_task
     else:
         assert log_path is not None, (request_id, log_path)
         if log_path == constants.API_SERVER_LOGS:

sky/server/stream_utils.py

@@ -68,7 +68,14 @@ async def log_streamer(
     if request_id is not None:
         status_msg = rich_utils.EncodedStatusMessage(
             f'[dim]Checking request: {request_id}[/dim]')
-        request_task = await requests_lib.get_request_async(request_id)
+        request_task = await requests_lib.get_request_async(request_id,
+                                                            fields=[
+                                                                'request_id',
+                                                                'name',
+                                                                'schedule_type',
+                                                                'status',
+                                                                'status_msg'
+                                                            ])
 
         if request_task is None:
             raise fastapi.HTTPException(
@@ -89,14 +96,15 @@ async def log_streamer(
                        f'scheduled: {request_id}')
         req_status = request_task.status
         req_msg = request_task.status_msg
+        del request_task
         # Slowly back off the database polling up to every 1 second, to avoid
         # overloading the CPU and DB.
         backoff = common_utils.Backoff(initial_backoff=polling_interval,
                                        max_backoff_factor=10,
                                        multiplier=1.2)
         while req_status < requests_lib.RequestStatus.RUNNING:
             if req_msg is not None:
-                waiting_msg = request_task.status_msg
+                waiting_msg = req_msg
             if show_request_waiting_spinner:
                 yield status_msg.update(f'[dim]{waiting_msg}[/dim]')
             elif plain_logs and waiting_msg != last_waiting_msg:
@@ -197,7 +205,7 @@ async def flush_buffer() -> AsyncGenerator[str, None]:
                     if (req_status.status ==
                             requests_lib.RequestStatus.CANCELLED):
                         request_task = await requests_lib.get_request_async(
-                            request_id)
+                            request_id, fields=['name', 'should_retry'])
                         if request_task.should_retry:
                             buffer.append(
                                 message_utils.encode_payload(
@@ -206,6 +214,7 @@ async def flush_buffer() -> AsyncGenerator[str, None]:
                             buffer.append(
                                 f'{request_task.name!r} request {request_id}'
                                 ' cancelled\n')
+                        del request_task
                     break
             if not follow:
                 # The below checks (cluster status, heartbeat) are not needed

sky/server/uvicorn.py

@@ -46,11 +46,11 @@
 
 # TODO(aylei): use decorator to register requests that need to be proactively
 # cancelled instead of hardcoding here.
-_RETRIABLE_REQUEST_NAMES = [
+_RETRIABLE_REQUEST_NAMES = {
     'sky.logs',
     'sky.jobs.logs',
     'sky.serve.logs',
-]
+}
 
 
 def add_timestamp_prefix_for_server_logs() -> None:
@@ -151,37 +151,38 @@ def _wait_requests(self) -> None:
                 requests_lib.RequestStatus.PENDING,
                 requests_lib.RequestStatus.RUNNING,
             ]
-            reqs = requests_lib.get_request_tasks(
-                req_filter=requests_lib.RequestTaskFilter(status=statuses))
-            if not reqs:
+            requests = [(request_task.request_id, request_task.name)
+                        for request_task in requests_lib.get_request_tasks(
+                            req_filter=requests_lib.RequestTaskFilter(
+                                status=statuses, fields=['request_id', 'name']))
+                       ]
+            if not requests:
                 break
-            logger.info(f'{len(reqs)} on-going requests '
+            logger.info(f'{len(requests)} on-going requests '
                         'found, waiting for them to finish...')
             # Proactively cancel internal requests and logs requests since
             # they can run for infinite time.
-            internal_request_ids = [
+            internal_request_ids = {
                 d.id for d in daemons.INTERNAL_REQUEST_DAEMONS
-            ]
+            }
             if time.time() - start_time > _WAIT_REQUESTS_TIMEOUT_SECONDS:
                 logger.warning('Timeout waiting for on-going requests to '
                                'finish, cancelling all on-going requests.')
-                for req in reqs:
-                    self.interrupt_request_for_retry(req.request_id)
+                for request_id, _ in requests:
+                    self.interrupt_request_for_retry(request_id)
                 break
             interrupted = 0
-            for req in reqs:
-                if req.request_id in internal_request_ids:
-                    self.interrupt_request_for_retry(req.request_id)
-                    interrupted += 1
-                elif req.name in _RETRIABLE_REQUEST_NAMES:
-                    self.interrupt_request_for_retry(req.request_id)
+            for request_id, name in requests:
+                if (name in _RETRIABLE_REQUEST_NAMES or
+                        request_id in internal_request_ids):
+                    self.interrupt_request_for_retry(request_id)
                     interrupted += 1
                 # TODO(aylei): interrupt pending requests to accelerate the
                 # shutdown.
             # If some requests are not interrupted, wait for them to finish,
             # otherwise we just check again immediately to accelerate the
             # shutdown process.
-            if interrupted < len(reqs):
+            if interrupted < len(requests):
                 time.sleep(_WAIT_REQUESTS_INTERVAL_SECONDS)
 
     def interrupt_request_for_retry(self, request_id: str) -> None:

tests/test_api.py

@@ -58,7 +58,7 @@ def __init__(self):
                 self.schedule_type = requests_lib.ScheduleType.LONG
                 self.status_msg = None
 
-        async def mock_get_request(request_id):
+        async def mock_get_request(request_id, fields):
             return MockRequest()
 
         async def mock_get_request_status(request_id):
@@ -150,7 +150,7 @@ def __init__(self):
                 self.schedule_type = requests_lib.ScheduleType.LONG
                 self.status_msg = None
 
-        async def mock_get_request(request_id):
+        async def mock_get_request(request_id, fields):
             return MockRequest()
 
         async def mock_get_request_status(request_id):