Merge branch 'main' into pr-log-context

Author: rishikeshdevsot

components/job-orchestration/job_orchestration/scheduler/query/query_scheduler.py

@@ -20,6 +20,7 @@
 import concurrent.futures
 import contextlib
 import datetime
+import multiprocessing
 import pathlib
 import sys
 from abc import ABC, abstractmethod
@@ -94,6 +95,8 @@
 
 reducer_connection_queue: asyncio.Queue | None = None
 
+_MULTIPROCESSING_START_METHOD = "spawn"
+
 
 class DispatchExecutor:
     # Globals for dispatch executor pool
@@ -436,6 +439,10 @@ async def handle_cancelling_search_jobs(db_conn_pool) -> None:
             else:
                 logger.error(f"Failed to cancel job {job_id}.")
 
+            # Yield to the event loop between jobs so cancellation batches do not
+            # monopolize the scheduler when many jobs are being processed.
+            await asyncio.sleep(0)
+
 
 def insert_query_tasks_into_db(db_conn, job_id, archive_ids: list[str]) -> list[int]:
     task_ids = []
@@ -1131,6 +1138,11 @@ async def handle_jobs(
 async def main(argv: list[str]) -> int:
     global reducer_connection_queue
 
+    # The scheduler accepts reducer TCP connections in the parent process. Using "spawn" prevents
+    # child processes from inheriting those sockets and keeping cancelled reducer jobs alive after
+    # the parent closes its copy.
+    multiprocessing.set_start_method(_MULTIPROCESSING_START_METHOD)
+
     args_parser = argparse.ArgumentParser(description="Wait for and run query jobs.")
     args_parser.add_argument("--config", "-c", required=True, help="CLP configuration file.")
 

components/job-orchestration/job_orchestration/scheduler/query/reducer_handler.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import enum
 from enum import Enum
 from typing import Any
@@ -110,11 +111,22 @@ async def _send_msg_to_reducer(msg: bytes, writer: asyncio.StreamWriter):
     await writer.drain()
 
 
+async def _clean_up_task(task: asyncio.Task | None) -> None:
+    if task is None:
+        return
+    if not task.done():
+        task.cancel()
+    with contextlib.suppress(asyncio.CancelledError, asyncio.IncompleteReadError):
+        await task
+
+
 async def handle_reducer_connection(
     reader: asyncio.StreamReader,
     writer: asyncio.StreamWriter,
     reducer_connection_queue: asyncio.Queue,
 ):
+    recv_listener_msg_task: asyncio.Task | None = None
+    recv_reducer_msg_task: asyncio.Task | None = None
     try:
         message_bytes = await _recv_msg_from_reducer(reader)
         if message_bytes is None:
@@ -138,10 +150,8 @@ async def handle_reducer_connection(
         # Transition to next state
         """
         current_wait_state: _ReducerHandlerWaitState = _ReducerHandlerWaitState.JOB_CONFIG
-        recv_listener_msg_task: asyncio.Task | None = asyncio.create_task(
-            msg_queues.get_from_listeners()
-        )
-        recv_reducer_msg_task: asyncio.Task | None = asyncio.create_task(reader.readexactly(1))
+        recv_listener_msg_task = asyncio.create_task(msg_queues.get_from_listeners())
+        recv_reducer_msg_task = asyncio.create_task(reader.readexactly(1))
         while True:
             pending = [recv_listener_msg_task, recv_reducer_msg_task]
             done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
@@ -225,5 +235,7 @@ async def handle_reducer_connection(
                 await msg_queues.put_to_listeners(msg)
                 break
     finally:
+        await _clean_up_task(recv_listener_msg_task)
+        await _clean_up_task(recv_reducer_msg_task)
         writer.close()
         await writer.wait_closed()