[core][taskEvents out of GCS][8/n] Reroute state head to task events head for state APIs.

[karticam]

Part of the effort to move task events out of GCS.
This PR builds on top of #65158

#65158 adds logic on task events head to serve queries for task events. This PR adds the logic to route the task related state APIs (ray list tasks) to task events head rather than GCS.

Changes done:

1. Adds a flag RAY_task_events_read_from_dashboard_head which determines whether state head sends the request to GCS or dashboard head (specifically task events head).
2. If the request is send to dashboard head, it makes an HTTP client to task events head. Since both of them are subprocesses on the same node, they just communicate through unix sockets without the need for any auth.
3. The HTTP client is made lazily on the first request.

[gemini-code-assist]

Code Review

This pull request introduces an in-memory task event storage and query system on the dashboard head, allowing task queries to be served directly from the dashboard instead of GCS. It implements pubsub subscribers for worker and job updates, a publisher client, and a dedicated TaskEventsHead module with garbage-collection policies. The review feedback is highly constructive and identifies several critical areas for improvement: a bug in the query logic that incorrectly reverses priority tiers during truncation, fragile pubsub subscriptions that lack startup retry logic, missing exception handling for gRPC calls and event conversion, and silent background task failures due to unhandled exceptions in done callbacks.

[gemini-code-assist]

Reversing the concatenated candidates list reverses the priority tiers, placing Tier 2 (running/failed tasks) at the end of the list. This causes running tasks to be truncated first when a limit is applied. We should instead iterate over candidates directly after updating get_all_task_events to return newest-first within each tier.

Suggested change

	# Iterate newest-first so a limit keeps the most recent events.
	for task_event in reversed(candidates):
	# Iterate newest-first so a limit keeps the most recent events.
	for task_event in candidates:

[gemini-code-assist]

Update get_all_task_events to return newest-first within each tier. This ensures that when candidates are processed, the most recent events within each priority tier are kept, while preserving the priority tier order (Tier 2 first, then Tier 1, then Tier 0).

Suggested change

	def get_all_task_events(self) -> List[gcs_pb2.TaskEvents]:
	"""All stored task events, higher-priority tiers first, oldest-first within a tier."""
	result: List[gcs_pb2.TaskEvents] = []
	for tier in range(self._gc_policy.MAX_PRIORITY - 1, -1, -1):
	result.extend(self._tiers[tier].values())
	return result
	def get_all_task_events(self) -> List[gcs_pb2.TaskEvents]:
	"""All stored task events, higher-priority tiers first, newest-first within a tier."""
	result: List[gcs_pb2.TaskEvents] = []
	for tier in range(self._gc_policy.MAX_PRIORITY - 1, -1, -1):
	result.extend(reversed(self._tiers[tier].values()))
	return result

[gemini-code-assist]

Calling subscriber.subscribe() outside the while True loop makes the subscription fragile. If GCS is temporarily unavailable during startup, the subscription will fail, the background task will terminate, and the dashboard head will never receive notifications for the lifetime of the process. Wrapping the subscription and the loop in a retry block with a delay ensures robustness.

    async def _subscribe_for_worker_deaths(self) -> None:
        while True:
            try:
                subscriber = GcsAioWorkerDeltaSubscriber(address=self.gcs_address)
                await subscriber.subscribe()
                while True:
                    for _, worker_delta in await subscriber.poll(
                        batch_size=_SUBSCRIBER_POLL_BATCH_SIZE
                    ):
                        self._handle_worker_delta(worker_delta)
            except Exception:
                logger.exception("Failed handling worker-death notifications. Retrying in 5 seconds...")
                await asyncio.sleep(5)

    async def _subscribe_for_finished_jobs(self) -> None:
        while True:
            try:
                subscriber = GcsAioJobSubscriber(address=self.gcs_address)
                await subscriber.subscribe()
                while True:
                    for _, job_data in await subscriber.poll(
                        batch_size=_SUBSCRIBER_POLL_BATCH_SIZE
                    ):
                        self._handle_job_update(job_data)
            except Exception:
                logger.exception("Failed handling job-finished notifications. Retrying in 5 seconds...")
                await asyncio.sleep(5)

[gemini-code-assist]

The gRPC call self._worker_info_stub.GetWorkerInfo is not wrapped in a try-except block. If GCS is down or there is a network issue, it will raise a grpc.RpcError, causing the background task _on_worker_dead to crash ungracefully. Wrapping it in a try-except block ensures the background task handles the error gracefully.

Suggested change

	async def _get_worker_info(
	self, worker_id: bytes
	) -> Optional[gcs_pb2.WorkerTableData]:
	if self._worker_info_stub is None:
	self._worker_info_stub = gcs_service_pb2_grpc.WorkerInfoGcsServiceStub(
	self.aiogrpc_gcs_channel
	)
	reply = await self._worker_info_stub.GetWorkerInfo(
	gcs_service_pb2.GetWorkerInfoRequest(worker_id=worker_id)
	)
	if not reply.HasField("worker_table_data"):
	return None
	return reply.worker_table_data
	async def _get_worker_info(
	self, worker_id: bytes
	) -> Optional[gcs_pb2.WorkerTableData]:
	if self._worker_info_stub is None:
	self._worker_info_stub = gcs_service_pb2_grpc.WorkerInfoGcsServiceStub(
	self.aiogrpc_gcs_channel
	)
	try:
	reply = await self._worker_info_stub.GetWorkerInfo(
	gcs_service_pb2.GetWorkerInfoRequest(worker_id=worker_id)
	)
	if not reply.HasField("worker_table_data"):
	return None
	return reply.worker_table_data
	except Exception as e:
	logger.error(f"Failed to get worker info for {worker_id.hex()}: {e}")
	return None

[gemini-code-assist]

convert_to_task_events and the subsequent store updates are called outside the try-except block. If any conversion or processing error occurs, it will result in an unhandled 500 error. Wrapping the entire processing logic in the try-except block ensures robust error handling.

    @routes.post("/api/task_events")
    async def add_task_events(
        self, request: aiohttp.web.Request
    ) -> aiohttp.web.Response:
        body = await request.read()
        try:
            add_events_request = self._deserialize_request(body)
            task_events, dropped_task_attempts = convert_to_task_events(add_events_request)
            self._store.record_data_loss_from_worker(dropped_task_attempts)
            for task_event in task_events:
                self._store.add_or_replace_task_event(task_event)
            logger.debug(
                "Received %d task events (%d attempts stored)",
                len(task_events),
                self._store.num_task_events_stored,
            )
            return dashboard_optional_utils.rest_response(
                status_code=dashboard_utils.HTTPStatusCode.OK,
                message="",
            )
        except Exception as e:
            logger.warning(f"Failed to process task events request: {e}")
            return dashboard_optional_utils.rest_response(
                status_code=dashboard_utils.HTTPStatusCode.INTERNAL_ERROR,
                message=f"Failed to process task events request: {e}",
            )

[gemini-code-assist]

Background tasks are created without any exception handling in their done callbacks. If a background task raises an exception, it will fail silently with a "Task exception was never retrieved" warning. Adding a done callback that logs exceptions ensures background task failures are visible.

Suggested change

	def _spawn(self, coro) -> None:
	task = asyncio.create_task(coro)
	self._background_tasks.add(task)
	task.add_done_callback(self._background_tasks.discard)
	def _spawn(self, coro) -> None:
	task = asyncio.create_task(coro)
	self._background_tasks.add(task)
	def done_callback(t):
	self._background_tasks.discard(t)
	try:
	t.result()
	except asyncio.CancelledError:
	pass
	except Exception:
	logger.exception("Background task failed")
	task.add_done_callback(done_callback)

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 506e9bf. Configure here.}

[sampan-s-nayak]

nit: I feel this comment could be simplified and written more briefly also too much context

[karticam]

done

[sampan-s-nayak]

nit: we can move this check to init (only checked if _READ_TASK_EVENTS_FROM_DASHBOARD_HEAD is true)

[karticam]

done

[sampan-s-nayak]

LGTM, as discussed we can use a common task_event_module across all code paths calling this module (aggregator agent, state manager and other one-off cases)