[Jobs] Eagerly initialize JobManager on head startup

python/ray/dashboard/modules/job/job_agent.py

@@ -1,6 +1,8 @@
+import asyncio
 import dataclasses
 import json
 import logging
+import random
 import traceback
 
 import aiohttp
@@ -23,6 +25,10 @@
 routes = optional_utils.DashboardAgentRouteTable
 logger = logging.getLogger(__name__)
 
+_INIT_RETRY_BASE_SECONDS = 30
+_INIT_RETRY_MAX_SECONDS = 300
+_INIT_RETRY_JITTER = 0.2
+
 
 class JobAgent(dashboard_utils.DashboardAgentModule):
     def __init__(self, dashboard_agent):
@@ -203,7 +209,37 @@ def get_job_manager(self):
         return self._job_manager
 
     async def run(self, server):
-        pass
+        if not self._dashboard_agent.is_head:
+            return
+
+        retry_delay_s = _INIT_RETRY_BASE_SECONDS
+        attempt = 0
+        while True:
+            try:
+                loop = asyncio.get_running_loop()
+                await loop.run_in_executor(
+                    None, optional_utils.init_ray_connection, self.gcs_address
+                )
+                self.get_job_manager()
+                logger.info(
+                    "Initialized JobManager on the head node and scheduled "
+                    "submission job recovery."
+                )
+                return
+            except Exception:
+                attempt += 1
+                delay_s = retry_delay_s * (
+                    1 + random.uniform(-_INIT_RETRY_JITTER, _INIT_RETRY_JITTER)
+                )
+                logger.warning(
+                    "Failed to initialize JobManager on the head node "
+                    "(attempt %d); retrying in %.1f seconds.",
+                    attempt,
+                    delay_s,
+                    exc_info=True,
+                )
+                await asyncio.sleep(delay_s)
+                retry_delay_s = min(retry_delay_s * 2, _INIT_RETRY_MAX_SECONDS)
 
     @staticmethod
     def is_minimal_module():

python/ray/dashboard/modules/job/job_manager.py

@@ -39,6 +39,21 @@
 
 logger = logging.getLogger(__name__)
 
+_RECOVERY_SCAN_MAX_ATTEMPTS = 5
+_RECOVERY_SCAN_BASE_DELAY_S = 1
+_RECOVERY_SCAN_MAX_DELAY_S = 30
+_RECOVERY_SCAN_GCS_RPC_TIMEOUT_S = 5
+# get_all_jobs lists keys, then fetches job info in a second GCS phase.
+_RECOVERY_SCAN_PER_ATTEMPT_TIMEOUT_S = 2 * _RECOVERY_SCAN_GCS_RPC_TIMEOUT_S + 5
+_RECOVERY_SCAN_TOTAL_BUDGET_S = 45
+_RECOVERY_SUBMISSION_WAIT_TIMEOUT_S = 60
+_RECOVERY_SCAN_TASKS = set()
+
+
+def _consume_recovery_scan_result(task: asyncio.Task) -> None:
+    if not task.cancelled():
+        task.exception()
+
 
 def generate_job_id() -> str:
     """Returns a job_id of the form 'raysubmit_XYZ'.
@@ -111,7 +126,77 @@ async def _recover_running_jobs(self):
         Each will be added to self._running_jobs and reconciled.
         """
         try:
-            all_jobs = await self._job_info_client.get_all_jobs()
+            loop = asyncio.get_running_loop()
+            recovery_deadline = loop.time() + _RECOVERY_SCAN_TOTAL_BUDGET_S
+            for attempt in range(1, _RECOVERY_SCAN_MAX_ATTEMPTS + 1):
+                remaining_s = recovery_deadline - loop.time()
+                if remaining_s <= 0:
+                    logger.error(
+                        "Submission job recovery exceeded its %.1f second budget. "
+                        "Existing non-terminal jobs will not be monitored until "
+                        "the Dashboard Agent restarts.",
+                        _RECOVERY_SCAN_TOTAL_BUDGET_S,
+                    )
+                    return
+                try:
+                    scan_task = asyncio.create_task(
+                        self._job_info_client.get_all_jobs(
+                            timeout=_RECOVERY_SCAN_GCS_RPC_TIMEOUT_S
+                        )
+                    )
+                    _RECOVERY_SCAN_TASKS.add(scan_task)
+                    scan_task.add_done_callback(_RECOVERY_SCAN_TASKS.discard)
+                    scan_task.add_done_callback(_consume_recovery_scan_result)
+                    try:
+                        done, _ = await asyncio.wait(
+                            [scan_task],
+                            timeout=min(
+                                _RECOVERY_SCAN_PER_ATTEMPT_TIMEOUT_S, remaining_s
+                            ),
+                        )
+                    except asyncio.CancelledError:
+                        scan_task.cancel()
+                        raise
+                    if not done:
+                        scan_task.cancel()
+                        raise asyncio.TimeoutError
+                    all_jobs = scan_task.result()
+                    break
+                except Exception:
+                    if attempt == _RECOVERY_SCAN_MAX_ATTEMPTS:
+                        logger.error(
+                            "Failed to fetch submission jobs for recovery after "
+                            "%d attempts. Existing non-terminal jobs will not be "
+                            "monitored until the Dashboard Agent restarts.",
+                            _RECOVERY_SCAN_MAX_ATTEMPTS,
+                            exc_info=True,
+                        )
+                        return
+
+                    delay_s = min(
+                        _RECOVERY_SCAN_BASE_DELAY_S * (2 ** (attempt - 1)),
+                        _RECOVERY_SCAN_MAX_DELAY_S,
+                    )
+                    remaining_s = recovery_deadline - loop.time()
+                    if remaining_s <= delay_s:
+                        logger.error(
+                            "Submission job recovery could not retry within its "
+                            "%.1f second budget. Existing non-terminal jobs will "
+                            "not be monitored until the Dashboard Agent restarts.",
+                            _RECOVERY_SCAN_TOTAL_BUDGET_S,
+                            exc_info=True,
+                        )
+                        return
+                    logger.warning(
+                        "Failed to fetch submission jobs for recovery "
+                        "(attempt %d/%d); retrying in %.1f seconds.",
+                        attempt,
+                        _RECOVERY_SCAN_MAX_ATTEMPTS,
+                        delay_s,
+                        exc_info=True,
+                    )
+                    await asyncio.sleep(delay_s)
+
             for job_id, job_info in all_jobs.items():
                 if not job_info.status.is_terminal():
                     run_background_task(self._monitor_job(job_id))
@@ -530,7 +615,17 @@ async def submit_job(
 
         # Wait for `_recover_running_jobs` to run before accepting submissions to
         # avoid duplicate monitoring of the same job.
-        await self._recover_running_jobs_event.wait()
+        try:
+            await asyncio.wait_for(
+                self._recover_running_jobs_event.wait(),
+                timeout=_RECOVERY_SUBMISSION_WAIT_TIMEOUT_S,
+            )
+        except asyncio.TimeoutError:
+            raise RuntimeError(
+                "Submission job recovery did not complete within "
+                f"{_RECOVERY_SUBMISSION_WAIT_TIMEOUT_S} seconds. Check the Dashboard "
+                "Agent logs for recovery failures."
+            )
 
         logger.info(f"Starting job with submission_id: {submission_id}")
         if entrypoint_label_selector: