[core][taskEvents out of GCS][4/n] Notify task events head on worker death and job finished

[karticam]

This PR builds on top of: #65028 and therefore has some changes from that PR as well.

GcsTaskManager does task state management based on worker death and job finished notifications. Since as part of the project to move task events out of GCS, we will be putting all this logic on dashboard head (specifically, task_events_head), we need these notifications on task_events_head as well.

This PR adds GcsAioWorkerDeltaSubscriber and GcsAioJobSubscriber and makes task_events_head subscribe to GCS for worker delta notifications and job notifications from GCS.

For workers delta notifications, we get notifications only on worker deaths. For jobs, any action is taken only when the job is finished. Right now, we just append them to a list. More logic (similar to what is being done in GcsTaskManager) will be added in subsequent PRs.

[gemini-code-assist]

Code Review

This pull request refactors the event aggregator agent to publish task events over HTTP to a new dashboard-head sink, TaskEventsHead, instead of directly to GCS. It introduces GcsAioWorkerDeltaSubscriber and GcsAioJobSubscriber to allow the dashboard head to subscribe to GCS pubsub for worker-death and job-finished notifications. The review feedback highlights several critical improvements: resolving potential silent termination, infinite busy-looping, and resource leaks in the GCS pubsub subscription loops; preventing unbounded memory growth in the in-memory buffers by setting a maximum deque length; returning a 400 Bad Request instead of a 500 Internal Error for client-side deserialization failures; and ensuring publisher clients are properly closed when the aggregator agent stops.

[gemini-code-assist]

The subscription loop for worker deaths is vulnerable to silent termination if subscriber.subscribe() fails on startup (e.g., if GCS is temporarily unavailable). Additionally, if GCS restarts or failover occurs, the subscriber should be recreated and re-subscribed. Finally, any persistent connection error will cause the loop to spin infinitely and flood the logs without any backoff, and the subscriber is never closed, leading to resource leaks on GCS. Wrapping the subscriber creation and subscription inside the retry loop with a sleep on exception and a finally block to close the subscriber resolves these issues.

    async def _subscribe_for_worker_deaths(self) -> None:
        while True:
            subscriber = None
            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...")
                await asyncio.sleep(1)
            finally:
                if subscriber:
                    try:
                        await subscriber.close()
                    except Exception:
                        pass

[gemini-code-assist]

The subscription loop for finished jobs has the same issues as the worker deaths loop: silent termination if subscribe() fails on startup, no re-subscription on GCS failover/restart, infinite busy-looping on persistent errors, and unclosed subscriber resources. Re-creating the subscriber inside the retry loop with a sleep on exception and closing it in a finally block resolves these issues.

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:
	while True:
	subscriber = None
	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...")
	await asyncio.sleep(1)
	finally:
	if subscriber:
	try:
	await subscriber.close()
	except Exception:
	pass

[gemini-code-assist]

The in-memory buffers self._events, self._dead_workers, and self._finished_jobs are implemented as unbounded collections.deque objects. In large or long-running clusters, these will grow indefinitely and eventually cause the dashboard head to crash with an Out Of Memory (OOM) error. Even as a temporary buffer, it is highly recommended to set a reasonable maxlen (e.g., 100000 for events and 10000 for workers/jobs) to prevent unbounded memory growth.

Suggested change

	self._events = collections.deque()
	self._dead_workers = collections.deque()
	self._finished_jobs = collections.deque()
	self._events = collections.deque(maxlen=100000)
	self._dead_workers = collections.deque(maxlen=10000)
	self._finished_jobs = collections.deque(maxlen=10000)

[gemini-code-assist]

When deserialization of the task events request fails, the API returns HTTPStatusCode.INTERNAL_ERROR (500). Since this is caused by a malformed or invalid payload from the client, it should return HTTPStatusCode.BAD_REQUEST (400) instead. Returning 500 for client-side errors can trigger false alarms in monitoring and alerting systems that track server-side errors.

Suggested change

	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}",
	)
	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.BAD_REQUEST,
	message=f"Failed to deserialize task events request: {e}",
	)

[gemini-code-assist]

The publisher clients (which lazily initialize aiohttp.ClientSession objects) are never closed when the aggregator agent stops or when run exits. This can lead to resource leaks and Unclosed client session warnings. It is recommended to call close() on the publishers in the finally block.

Suggested change

	try:
	await asyncio.gather(
	self._http_endpoint_publisher.run_forever(),
	self._gcs_publisher.run_forever(),
	self._dashboard_head_publisher.run_forever(),
	)
	finally:
	self._executor.shutdown()
	try:
	await asyncio.gather(
	self._http_endpoint_publisher.run_forever(),
	self._dashboard_head_publisher.run_forever(),
	)
	finally:
	await asyncio.gather(
	self._http_endpoint_publisher.close() if hasattr(self._http_endpoint_publisher, "close") else asyncio.sleep(0),
	self._dashboard_head_publisher.close() if hasattr(self._dashboard_head_publisher, "close") else asyncio.sleep(0),
	return_exceptions=True,
	)
	self._executor.shutdown()