[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

TESTING:
To test correct working of ray task event recorder (and other functionality related to the whole effort of migrating task events out of GCS), I turned changed the apt flags and triggered the CI suite in this PR: #65253

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

[karticam]

done. updated the function and the docstrings

[MengjinYan]

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

[karticam]

yes. its covered here:

ray/python/ray/tests/test_task_events_2.py

Line 800 in 12436ea

def check_file(type, task_name, expected_log, expect_no_end=False):

[MengjinYan]

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

[karticam]

The whole test suite has _SYSTEM_CONFIG, which sets both the delays to 1000ms. If we remove the _system_config overrides, _SYSTEM_CONFIG would be applicable. Since _SYSTEM_CONFIG sets both the delays to 1000ms and system configs take precedence over environment variable settings, removing the custom _system_config would lead to setting both the delays to 1000ms, and we don't want that in the test.

[karticam]

Also added a comment on why we are doing this

[MengjinYan]

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

[karticam]

yeah. we can do another test to exercise rate limiting on the dashboard head as well. will take it as a followup PR.