[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]

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]

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)

[karticam]

in same parity as GcsTaskManager

[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.}

[MengjinYan]

Wondering should we add tests to verify the added log task log info conversion?

[karticam]

I moved this piece of code to #65123, and added the test there

Commit: 464a0b5

[MengjinYan]

nit: the src and dst can be confusing. Better to add some descriptions to talk about the conversion is from where to where.

[MengjinYan]

Wondering is the task_log_info field covered in any of the test state api tests?

[MengjinYan]

Nit: if we've already set the env vars through monkeypatch, should we remove the settings to the _system_config?

[MengjinYan]

Could be address in a separate PR but wondering is the behavior worth testing with the new setup?