[core][taskEvents out of GCS][10/n] Fix bugs found after migration

[karticam]

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

To test the entire feature, I took changes from PR 2/n to 9/n, changed the flags so that task_event_buffer to GCS is stopped, ray_task_event_recorder is used, task events are sent to task events head via aggregator agent, and state APIs and ray timeline consume task events from task event head.

essentially these flags were configured:
enable_ray_event: true
enable_ray_task_event_recorder: true
enable_task_events_to_dashboard_head: true
enable_core_worker_task_event_to_gcs: false

Some issues in tests were found. This PR resolves those issues, so that premerge tests pass even even after we switch the flags ON default

[gemini-code-assist]

Code Review

This pull request migrates task events from GCS to the dashboard head, introducing a new TaskEventsHead subprocess module to store and query task events in memory, and a TaskEventManager to handle background reconciliation against GCS worker and job events. The state API and timeline queries are updated to read from this new dashboard-head store when enabled. Feedback on the changes suggests several robustness improvements: handling missing worker info and wrapping gRPC calls in try-except blocks in TaskEventManager, implementing retry loops with delays for pubsub subscriptions to avoid CPU-burning loops, using an insertion-ordered dict for FIFO eviction of dropped task attempts in TaskEventStorage, and wrapping the entire request handler in TaskEventsHead to prevent unhandled 500 errors.

[gemini-code-assist]

If self._get_worker_info(worker_id) returns None (or raises an exception), the tasks of the dead worker are currently ignored and left in the RUNNING state indefinitely. To ensure robust fault tolerance, we should fallback to a default WorkerTableData with a timestamp so that the tasks are still failed gracefully.

Suggested change

	worker_table_data, _ = await asyncio.gather(
	self._get_worker_info(worker_id),
	asyncio.sleep(_MARK_FAILED_ON_WORKER_DEAD_DELAY_S),
	)
	if worker_table_data is None:
	logger.warning(
	f"No worker info found for dead worker {worker_id.hex()}; its tasks "
	"cannot be marked as failed."
	)
	return
	logger.debug(
	f"Marking all running tasks of worker {worker_id.hex()} as failed."
	)
	self._store.mark_tasks_failed_on_worker_dead(worker_id, worker_table_data)
	worker_table_data, _ = await asyncio.gather(
	self._get_worker_info(worker_id),
	asyncio.sleep(_MARK_FAILED_ON_WORKER_DEAD_DELAY_S),
	)
	if worker_table_data is None:
	import time
	logger.warning(
	f"No worker info found for dead worker {worker_id.hex()}; "
	"falling back to generic failure marking."
	)
	worker_table_data = gcs_pb2.WorkerTableData()
	worker_table_data.end_time_ms = int(time.time() * 1000)
	worker_table_data.exit_detail = "Worker died (detailed exit info unavailable)"
	logger.debug(
	f"Marking all running tasks of worker {worker_id.hex()} as failed."
	)
	self._store.mark_tasks_failed_on_worker_dead(worker_id, worker_table_data)

[gemini-code-assist]

Wrap the gRPC call to GetWorkerInfo in a try-except block to handle transient network or gRPC errors gracefully. Otherwise, any gRPC exception will propagate and crash the _on_worker_dead task, preventing task failure marking.

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._gcs_aio_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._gcs_aio_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.warning(f"Failed to get worker info for {worker_id.hex()}: {e}")
	return None

[gemini-code-assist]

If subscriber.subscribe() fails (e.g., during GCS startup or transient network issues), the subscription task will terminate permanently and never retry. Additionally, if subscriber.poll() raises an exception, the loop will immediately retry without any delay, potentially causing a tight CPU-burning loop. We should wrap the subscription in a retry loop and add a small sleep on poll exceptions.

    async def _subscribe_for_worker_deaths(self) -> None:
        subscriber = GcsAioWorkerDeltaSubscriber(address=self._gcs_address)
        while True:
            try:
                await subscriber.subscribe()
                break
            except Exception as e:
                logger.warning(f"Failed to subscribe to worker deaths, retrying: {e}")
                await asyncio.sleep(1)
        while True:
            try:
                for _, worker_delta in await subscriber.poll(
                    batch_size=_SUBSCRI_POLL_BATCH_SIZE
                ):
                    self._handle_worker_delta(worker_delta)
            except Exception:
                logger.exception("Failed handling worker-death notifications.")
                await asyncio.sleep(1)

[gemini-code-assist]

If subscriber.subscribe() fails, the subscription task will terminate permanently and never retry. Additionally, if subscriber.poll() raises an exception, the loop will immediately retry without any delay, potentially causing a tight CPU-burning loop. We should wrap the subscription in a retry loop and add a small sleep on poll exceptions.

Suggested change

	async def _subscribe_for_finished_jobs(self) -> None:
	subscriber = GcsAioJobSubscriber(address=self._gcs_address)
	await subscriber.subscribe()
	while True:
	try:
	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.")
	async def _subscribe_for_finished_jobs(self) -> None:
	subscriber = GcsAioJobSubscriber(address=self._gcs_address)
	while True:
	try:
	await subscriber.subscribe()
	break
	except Exception as e:
	logger.warning(f"Failed to subscribe to job updates, retrying: {e}")
	await asyncio.sleep(1)
	while True:
	try:
	for _, job_data in await subscriber.poll(
	batch_size=_SUBSCRI_POLL_BATCH_SIZE
	):
	self._handle_job_update(job_data)
	except Exception:
	logger.exception("Failed handling job-finished notifications.")
	await asyncio.sleep(1)

[gemini-code-assist]

Using a set for self._dropped_task_attempts results in arbitrary eviction of tracked dropped attempts when the cap is reached. If a recently dropped attempt is arbitrarily evicted, late events for it could be partially re-stored. Using a dict as an ordered set (since Python 3.7+ dicts preserve insertion order) allows us to perform FIFO eviction, ensuring we keep the most recently dropped attempts tracked.

Suggested change

	def __init__(self):
	self._num_profile_events_dropped = 0
	self._num_task_attempts_dropped_tracked = 0
	self._num_dropped_task_attempts_evicted = 0
	self._dropped_task_attempts: Set[TaskAttempt] = set()
	def record_task_attempt_dropped(self, task_attempt: TaskAttempt) -> None:
	self._dropped_task_attempts.add(task_attempt)
	self._num_task_attempts_dropped_tracked = len(self._dropped_task_attempts)
	def __init__(self):
	self._num_profile_events_dropped = 0
	self._num_task_attempts_dropped_tracked = 0
	self._num_dropped_task_attempts_evicted = 0
	self._dropped_task_attempts: Dict[TaskAttempt, None] = {}
	def record_task_attempt_dropped(self, task_attempt: TaskAttempt) -> None:
	self._dropped_task_attempts[task_attempt] = None
	self._num_task_attempts_dropped_tracked = len(self._dropped_task_attempts)

[gemini-code-assist]

Evict the oldest dropped task attempts first using FIFO order by leveraging the insertion-ordered dict keys.

Suggested change

	num_to_evict = min(
	len(self._dropped_task_attempts), num_to_evict + int(0.1 * num_to_evict)
	)
	self._num_task_attempts_dropped_tracked = len(self._dropped_task_attempts)
	if num_to_evict == 0:
	return
	self._num_dropped_task_attempts_evicted += num_to_evict
	to_evict = list(self._dropped_task_attempts)[:num_to_evict]
	self._dropped_task_attempts.difference_update(to_evict)
	self._num_task_attempts_dropped_tracked = len(self._dropped_task_attempts)
	num_to_evict = min(
	len(self._dropped_task_attempts), num_to_evict + int(0.1 * num_to_evict)
	)
	self._num_task_attempts_dropped_tracked = len(self._dropped_task_attempts)
	if num_to_evict == 0:
	return
	self._num_dropped_task_attempts_evicted += num_to_evict
	to_evict = list(self._dropped_task_attempts.keys())[:num_to_evict]
	for k in to_evict:
	del self._dropped_task_attempts[k]
	self._num_task_attempts_dropped_tracked = len(self._dropped_task_attempts)

[gemini-code-assist]

To prevent potential KeyError crashes if the indices somehow get out of sync, use .get() to safely retrieve the tier from self._primary_index before accessing self._tiers.

Suggested change

	def _mark_task_attempt_failed_if_needed(
	self, attempt: TaskAttempt, failed_ts_ns: int, error_info: RayErrorInfo
	) -> None:
	task_event = self._tiers[self._primary_index[attempt]][attempt]
	# Don't fail a task attempt that already reached a terminal state.
	if self._is_task_terminated(task_event):
	return
	task_event.state_updates.state_ts_ns[TaskStatus.FAILED] = failed_ts_ns
	task_event.state_updates.error_info.CopyFrom(error_info)
	def _mark_task_attempt_failed_if_needed(
	self, attempt: TaskAttempt, failed_ts_ns: int, error_info: RayErrorInfo
	) -> None:
	tier = self._primary_index.get(attempt)
	if tier is None:
	return
	task_event = self._tiers[tier][attempt]
	# Don't fail a task attempt that already reached a terminal state.
	if self._is_task_terminated(task_event):
	return
	task_event.state_updates.state_ts_ns[TaskStatus.FAILED] = failed_ts_ns
	task_event.state_updates.error_info.CopyFrom(error_info)

[gemini-code-assist]

Wrap the entire request processing logic (including convert_to_task_events and store updates) in the try-except block. This prevents any unexpected exceptions or parsing/assertion errors from crashing the request handler and returning an unhandled 500 error.

Suggested change

	try:
	add_events_request = self._deserialize_request(body)
	except Exception as e:
	logger.warning(f"Failed to deserialize task events request: {e}")
	return dashboard_optional_utils.rest_response(
	status_code=dashboard_utils.HTTPStatusCode.INTERNAL_ERROR,
	message=f"Failed to deserialize task events request: {e}",
	)
	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)
	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)
	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}",
	)

[cursor]

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

There are 3 total unresolved issues (including 2 from previous reviews).

^{Reviewed by Cursor Bugbot for commit 73a8272. Configure here.}