diff --git a/python/ray/dashboard/modules/job/job_agent.py b/python/ray/dashboard/modules/job/job_agent.py index 838806cd5c14..c99800c3d864 100644 --- a/python/ray/dashboard/modules/job/job_agent.py +++ b/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(): diff --git a/python/ray/dashboard/modules/job/job_manager.py b/python/ray/dashboard/modules/job/job_manager.py index ec25a504dd51..715ac8cd444a 100644 --- a/python/ray/dashboard/modules/job/job_manager.py +++ b/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: diff --git a/python/ray/dashboard/modules/job/tests/test_job_agent_initialization.py b/python/ray/dashboard/modules/job/tests/test_job_agent_initialization.py new file mode 100644 index 000000000000..167c23078518 --- /dev/null +++ b/python/ray/dashboard/modules/job/tests/test_job_agent_initialization.py @@ -0,0 +1,373 @@ +import asyncio +import threading +import time +from unittest.mock import AsyncMock, MagicMock, call, patch + +import pytest + +import ray.dashboard.optional_utils as optional_utils +from ray.dashboard.modules.job import ( + job_agent as job_agent_module, + job_manager as job_manager_module, +) +from ray.dashboard.modules.job.job_agent import JobAgent +from ray.dashboard.modules.job.job_manager import JobManager +from ray.job_submission import JobStatus + + +def _make_job_agent(*, is_head: bool) -> JobAgent: + dashboard_agent = MagicMock() + dashboard_agent.is_head = is_head + dashboard_agent.gcs_address = "127.0.0.1:6379" + dashboard_agent.gcs_client = MagicMock() + dashboard_agent.log_dir = "/tmp/ray/session_latest/logs" + return JobAgent(dashboard_agent) + + +class _DecoratedHandler: + gcs_address = "127.0.0.1:6379" + + @optional_utils.init_ray_and_catch_exceptions() + async def handle(self): + return "ok" + + +@pytest.mark.asyncio +async def test_head_job_agent_eagerly_initializes_job_manager(): + agent = _make_job_agent(is_head=True) + job_manager = MagicMock() + + with ( + patch.object(optional_utils, "init_ray_connection") as init_ray_connection, + patch.object(job_agent_module, "JobManager", return_value=job_manager), + ): + await agent.run(server=None) + + init_ray_connection.assert_called_once_with("127.0.0.1:6379") + assert agent._job_manager is job_manager + + +@pytest.mark.asyncio +async def test_worker_job_agent_does_not_initialize_job_manager(): + agent = _make_job_agent(is_head=False) + + with ( + patch.object(optional_utils, "init_ray_connection") as init_ray_connection, + patch.object(job_agent_module, "JobManager") as job_manager_cls, + ): + await agent.run(server=None) + + init_ray_connection.assert_not_called() + job_manager_cls.assert_not_called() + assert agent._job_manager is None + + +@pytest.mark.asyncio +async def test_job_agent_retries_ray_connection_failure(): + agent = _make_job_agent(is_head=True) + job_manager = MagicMock() + + with ( + patch.object( + optional_utils, + "init_ray_connection", + side_effect=[ConnectionError("GCS is unavailable"), None], + ) as init_ray_connection, + patch.object(job_agent_module, "JobManager", return_value=job_manager), + patch.object(job_agent_module.random, "uniform", return_value=0), + patch.object( + job_agent_module.asyncio, "sleep", new_callable=AsyncMock + ) as sleep, + ): + await agent.run(server=None) + + assert init_ray_connection.call_count == 2 + sleep.assert_awaited_once_with(30) + assert agent._job_manager is job_manager + + +@pytest.mark.asyncio +async def test_job_agent_retries_job_manager_construction_failure(): + agent = _make_job_agent(is_head=True) + job_manager = MagicMock() + + with ( + patch.object(optional_utils, "init_ray_connection") as init_ray_connection, + patch.object( + job_agent_module, + "JobManager", + side_effect=[RuntimeError("constructor failed"), job_manager], + ) as job_manager_cls, + patch.object(job_agent_module.random, "uniform", return_value=0), + patch.object( + job_agent_module.asyncio, "sleep", new_callable=AsyncMock + ) as sleep, + ): + await agent.run(server=None) + + assert init_ray_connection.call_count == 2 + assert job_manager_cls.call_count == 2 + sleep.assert_awaited_once_with(30) + assert agent._job_manager is job_manager + + +@pytest.mark.asyncio +async def test_job_agent_retry_delay_is_capped(): + agent = _make_job_agent(is_head=True) + failures = [ConnectionError("GCS is unavailable")] * 6 + + with ( + patch.object( + optional_utils, + "init_ray_connection", + side_effect=[*failures, None], + ), + patch.object(job_agent_module, "JobManager", return_value=MagicMock()), + patch.object(job_agent_module.random, "uniform", return_value=0), + patch.object( + job_agent_module.asyncio, "sleep", new_callable=AsyncMock + ) as sleep, + ): + await agent.run(server=None) + + assert sleep.await_args_list == [ + call(30), + call(60), + call(120), + call(240), + call(300), + call(300), + ] + + +@pytest.mark.asyncio +async def test_job_agent_ray_init_does_not_block_event_loop(): + agent = _make_job_agent(is_head=True) + + def slow_init(_): + time.sleep(0.3) + + with ( + patch.object(optional_utils, "init_ray_connection", side_effect=slow_init), + patch.object(job_agent_module, "JobManager", return_value=MagicMock()), + ): + loop = asyncio.get_running_loop() + start = loop.time() + run_task = asyncio.create_task(agent.run(server=None)) + await asyncio.sleep(0.01) + elapsed = loop.time() - start + await run_task + + assert elapsed < 0.2 + + +def test_init_ray_connection_is_idempotent(): + with ( + patch.object(optional_utils.ray, "is_initialized", return_value=True), + patch.object(optional_utils.ray, "init") as ray_init, + ): + optional_utils.init_ray_connection("127.0.0.1:6379") + + ray_init.assert_not_called() + + +def test_init_ray_connection_uses_dashboard_settings(monkeypatch): + monkeypatch.delenv("RAY_gcs_server_request_timeout_seconds", raising=False) + connection_ready = threading.Event() + + with ( + patch.object(optional_utils, "_ray_connection_ready", connection_ready), + patch.object(optional_utils.ray, "is_initialized", return_value=False), + patch.object(optional_utils.ray, "init") as ray_init, + ): + optional_utils.init_ray_connection("127.0.0.1:6379") + + assert optional_utils.os.environ["RAY_gcs_server_request_timeout_seconds"] == str( + optional_utils.dashboard_consts.GCS_RPC_TIMEOUT_SECONDS + ) + ray_init.assert_called_once_with( + address="127.0.0.1:6379", + log_to_driver=False, + configure_logging=False, + namespace=optional_utils.RAY_INTERNAL_DASHBOARD_NAMESPACE, + _skip_env_hook=True, + ) + assert connection_ready.is_set() + + +def test_init_ray_connection_preserves_init_error_when_shutdown_fails(): + connection_ready = threading.Event() + connection_ready.set() + + with ( + patch.object(optional_utils, "_ray_connection_ready", connection_ready), + patch.object(optional_utils.ray, "is_initialized", return_value=False), + patch.object( + optional_utils.ray, + "init", + side_effect=ConnectionError("GCS is unavailable"), + ), + patch.object( + optional_utils.ray, + "shutdown", + side_effect=RuntimeError("shutdown failed"), + ) as ray_shutdown, + ): + with pytest.raises(ConnectionError, match="GCS is unavailable"): + optional_utils.init_ray_connection("127.0.0.1:6379") + + ray_shutdown.assert_called_once_with() + assert not connection_ready.is_set() + + +@pytest.mark.asyncio +async def test_init_decorator_skips_executor_after_connection_is_ready(): + connection_ready = threading.Event() + connection_ready.set() + + with ( + patch.object(optional_utils, "_ray_connection_ready", connection_ready), + patch.object(optional_utils.ray, "is_initialized", return_value=True), + patch.object(optional_utils, "init_ray_connection") as init_ray_connection, + ): + assert await _DecoratedHandler().handle() == "ok" + + init_ray_connection.assert_not_called() + + +@pytest.mark.asyncio +async def test_init_decorator_uses_executor_until_connection_is_ready(): + connection_ready = threading.Event() + + with ( + patch.object(optional_utils, "_ray_connection_ready", connection_ready), + # ray.init() sets this before all of its post-init hooks have completed. + patch.object(optional_utils.ray, "is_initialized", return_value=True), + patch.object(optional_utils, "init_ray_connection") as init_ray_connection, + ): + assert await _DecoratedHandler().handle() == "ok" + + init_ray_connection.assert_called_once_with("127.0.0.1:6379") + + +@pytest.mark.asyncio +async def test_recovery_scan_retries_then_monitors_non_terminal_jobs(): + pending_job = MagicMock(status=JobStatus.PENDING) + finished_job = MagicMock(status=JobStatus.SUCCEEDED) + manager = MagicMock() + manager._job_info_client.get_all_jobs = AsyncMock( + side_effect=[ + RuntimeError("GCS is unavailable"), + {"pending": pending_job, "finished": finished_job}, + ] + ) + manager._recover_running_jobs_event = asyncio.Event() + manager._monitor_job = MagicMock() + + with ( + patch.object( + job_manager_module.asyncio, "sleep", new_callable=AsyncMock + ) as sleep, + patch.object(job_manager_module, "run_background_task") as run_task, + ): + await JobManager._recover_running_jobs(manager) + + assert manager._job_info_client.get_all_jobs.await_args_list == [ + call(timeout=5), + call(timeout=5), + ] + sleep.assert_awaited_once_with(1) + manager._monitor_job.assert_called_once_with("pending") + run_task.assert_called_once_with(manager._monitor_job.return_value) + assert manager._recover_running_jobs_event.is_set() + + +@pytest.mark.asyncio +async def test_recovery_scan_allows_two_gcs_timeout_phases(): + pending_job = MagicMock(status=JobStatus.PENDING) + manager = MagicMock() + + async def get_all_jobs(*, timeout): + assert timeout == 0.05 + # get_all_jobs first lists keys, then fetches job info. Each phase can + # legitimately consume most of the per-RPC timeout. + await asyncio.sleep(0.04) + await asyncio.sleep(0.04) + return {"pending": pending_job} + + manager._job_info_client.get_all_jobs = AsyncMock(side_effect=get_all_jobs) + manager._recover_running_jobs_event = asyncio.Event() + manager._monitor_job = MagicMock() + + with ( + patch.object(job_manager_module, "_RECOVERY_SCAN_GCS_RPC_TIMEOUT_S", 0.05), + patch.object(job_manager_module, "_RECOVERY_SCAN_PER_ATTEMPT_TIMEOUT_S", 1), + patch.object(job_manager_module, "_RECOVERY_SCAN_TOTAL_BUDGET_S", 2), + patch.object(job_manager_module, "run_background_task") as run_task, + ): + await JobManager._recover_running_jobs(manager) + + manager._job_info_client.get_all_jobs.assert_awaited_once_with(timeout=0.05) + manager._monitor_job.assert_called_once_with("pending") + run_task.assert_called_once_with(manager._monitor_job.return_value) + assert manager._recover_running_jobs_event.is_set() + + +@pytest.mark.asyncio +async def test_recovery_scan_exhaustion_still_unblocks_submissions(): + manager = MagicMock() + manager._job_info_client.get_all_jobs = AsyncMock( + side_effect=RuntimeError("GCS is unavailable") + ) + manager._recover_running_jobs_event = asyncio.Event() + manager._monitor_job = MagicMock() + + with patch.object( + job_manager_module.asyncio, "sleep", new_callable=AsyncMock + ) as sleep: + await JobManager._recover_running_jobs(manager) + + assert manager._job_info_client.get_all_jobs.await_count == 5 + assert sleep.await_args_list == [call(1), call(2), call(4), call(8)] + manager._monitor_job.assert_not_called() + assert manager._recover_running_jobs_event.is_set() + + +@pytest.mark.asyncio +async def test_recovery_scan_enforces_total_deadline(): + manager = MagicMock() + release_scan = asyncio.Event() + scan_tasks = [] + + async def get_all_jobs(*, timeout): + scan_tasks.append(asyncio.current_task()) + try: + await release_scan.wait() + except asyncio.CancelledError: + # Simulate a GCS awaitable that does not finish cancellation promptly. + await release_scan.wait() + + manager._job_info_client.get_all_jobs = AsyncMock(side_effect=get_all_jobs) + manager._recover_running_jobs_event = asyncio.Event() + + with patch.object(job_manager_module, "_RECOVERY_SCAN_TOTAL_BUDGET_S", 0.01): + await asyncio.wait_for(JobManager._recover_running_jobs(manager), timeout=1) + + manager._job_info_client.get_all_jobs.assert_awaited_once_with(timeout=5) + assert manager._recover_running_jobs_event.is_set() + release_scan.set() + await asyncio.gather(*scan_tasks) + + +@pytest.mark.asyncio +async def test_submit_job_fails_instead_of_waiting_forever_for_recovery(): + manager = MagicMock() + manager._recover_running_jobs_event = asyncio.Event() + + with patch.object(job_manager_module, "_RECOVERY_SUBMISSION_WAIT_TIMEOUT_S", 0): + with pytest.raises(RuntimeError, match="recovery did not complete"): + await JobManager.submit_job( + manager, + entrypoint="echo hello", + submission_id="submission-id", + ) diff --git a/python/ray/dashboard/optional_utils.py b/python/ray/dashboard/optional_utils.py index d5923499bffd..17e0216c4635 100644 --- a/python/ray/dashboard/optional_utils.py +++ b/python/ray/dashboard/optional_utils.py @@ -9,6 +9,7 @@ import inspect import logging import os +import threading import time import traceback from collections import namedtuple @@ -41,6 +42,10 @@ logger = logging.getLogger(__name__) +_ray_init_lock = threading.Lock() +# ray.is_initialized() becomes true before ray.init() finishes its post-init hooks. +_ray_connection_ready = threading.Event() + DashboardHeadRouteTable = method_route_table_factory() DashboardAgentRouteTable = method_route_table_factory() @@ -209,6 +214,35 @@ async def browser_request_middleware(request, handler): return browser_request_middleware +def init_ray_connection(address: str) -> None: + """Connect to a running Ray cluster if this process is not connected.""" + with _ray_init_lock: + if ray.is_initialized(): + return + + _ray_connection_ready.clear() + logger.info(f"Connecting to ray with address={address}") + os.environ["RAY_gcs_server_request_timeout_seconds"] = str( + dashboard_consts.GCS_RPC_TIMEOUT_SECONDS + ) + try: + ray.init( + address=address, + log_to_driver=False, + configure_logging=False, + namespace=RAY_INTERNAL_DASHBOARD_NAMESPACE, + _skip_env_hook=True, + ) + except Exception: + try: + ray.shutdown() + except Exception: + logger.exception("Failed to clean up after ray.init() failed.") + raise + else: + _ray_connection_ready.set() + + def init_ray_and_catch_exceptions() -> Callable: """Decorator to be used on methods that require being connected to Ray.""" @@ -218,26 +252,11 @@ async def decorator( self: Union[DashboardAgentModule, DashboardHeadModule], *args, **kwargs ): try: - if not ray.is_initialized(): - try: - address = self.gcs_address - logger.info(f"Connecting to ray with address={address}") - # Set the gcs rpc timeout to shorter - os.environ["RAY_gcs_server_request_timeout_seconds"] = str( - dashboard_consts.GCS_RPC_TIMEOUT_SECONDS - ) - # Init ray without logging to driver - # to avoid infinite logging issue. - ray.init( - address=address, - log_to_driver=False, - configure_logging=False, - namespace=RAY_INTERNAL_DASHBOARD_NAMESPACE, - _skip_env_hook=True, - ) - except Exception as e: - ray.shutdown() - raise e from None + if not _ray_connection_ready.is_set() or not ray.is_initialized(): + loop = asyncio.get_running_loop() + await loop.run_in_executor( + None, init_ray_connection, self.gcs_address + ) return await f(self, *args, **kwargs) except Exception as e: logger.exception(f"Unexpected error in handler: {e}") diff --git a/python/ray/tests/test_state_api_data_source_client_get.py b/python/ray/tests/test_state_api_data_source_client_get.py index c25a40ea18e6..078c2099b6ce 100644 --- a/python/ray/tests/test_state_api_data_source_client_get.py +++ b/python/ray/tests/test_state_api_data_source_client_get.py @@ -1352,7 +1352,9 @@ def ready(self): """ result = await client.get_all_worker_info(limit=2) assert len(result.worker_table_data) == 2 - assert result.total == 4 + # 5 = 1 driver + 3 actor workers + at least 1 dashboard agent connection. + # Aggregator routing or an agent restart can leave additional internal records. + assert result.total >= 5 @pytest.mark.asyncio diff --git a/python/ray/tests/test_task_events_2.py b/python/ray/tests/test_task_events_2.py index b1677a60cd1d..b303d758142e 100644 --- a/python/ray/tests/test_task_events_2.py +++ b/python/ray/tests/test_task_events_2.py @@ -194,7 +194,7 @@ def f(): client = StateApiClient() def list_tasks(exclude_driver): - return client.list( + all_tasks = client.list( StateResource.TASKS, # Filter out this driver options=ListApiOptions( @@ -202,6 +202,14 @@ def list_tasks(exclude_driver): ), raise_on_missing_output=True, ) + # Keep only the script driver's job and exclude internal drivers such as the + # dashboard agent connection used for eager JobManager recovery. + script_job_ids = { + task["job_id"] for task in all_tasks if task["type"] == "NORMAL_TASK" + } + if script_job_ids: + all_tasks = [task for task in all_tasks if task["job_id"] in script_job_ids] + return all_tasks # Check driver running def verify():