Repository navigation
feat(server): release parked interrupted tasks after a configurable timeout #1338
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -213,6 +213,7 @@ async def _process_event_inner(self, event: Event) -> None: | |
| 'Consumer[%s]: Request completed', self.active_task._task_id | ||
| ) | ||
| self.active_task._request_lock.release() | ||
| self.active_task._schedule_idle_release() | ||
| elif isinstance(event, _RequestStarted): | ||
| logger.debug( | ||
| 'Consumer[%s]: Request started', self.active_task._task_id | ||
|
|
@@ -441,6 +442,7 @@ def __init__( # noqa: PLR0913 | |
| push_sender: PushNotificationSender | None = None, | ||
| on_cleanup: Callable[[ActiveTask], None] | None = None, | ||
| event_stream: TaskEventStream | None = None, | ||
| interrupted_task_timeout: float | None = None, | ||
| ) -> None: | ||
| """Initializes the ActiveTask. | ||
|
|
||
|
|
@@ -454,6 +456,8 @@ def __init__( # noqa: PLR0913 | |
| the task from the ActiveTaskRegistry. | ||
| event_stream: Optional cross-replica stream; applied events are | ||
| published to it so other replicas can observe them. | ||
| interrupted_task_timeout: Optional seconds after which an idle | ||
| interrupted task is finished; ``None`` disables it. | ||
| """ | ||
| # --- Core Dependencies --- | ||
| self._agent_executor = agent_executor | ||
|
|
@@ -466,6 +470,9 @@ def __init__( # noqa: PLR0913 | |
| self._push_sender = push_sender | ||
| self._on_cleanup = on_cleanup | ||
| self._event_stream = event_stream | ||
| self._interrupted_task_timeout = interrupted_task_timeout | ||
| # Pending timer that releases a task parked in an interrupted state. | ||
| self._idle_release_task: asyncio.Task[None] | None = None | ||
|
|
||
| # --- Synchronization Primitives --- | ||
| # `_lock` protects structural lifecycle changes: start(), subscribe() counting, | ||
|
|
@@ -502,6 +509,11 @@ def task_id(self) -> str: | |
| """The ID of the task.""" | ||
| return self._task_id | ||
|
|
||
| @property | ||
| def is_finished(self) -> bool: | ||
| """Whether the task's active lifecycle has ended.""" | ||
| return self._is_finished.is_set() | ||
|
|
||
| async def enqueue_request( | ||
| self, request_context: RequestContext | ||
| ) -> uuid.UUID: | ||
|
|
@@ -692,6 +704,7 @@ async def _run_consumer(self) -> None: | |
| await EventConsumer(self).run() | ||
| finally: | ||
| self._is_finished.set() | ||
| self._cancel_idle_release() | ||
| self._request_queue.shutdown(immediate=True) | ||
| await self._event_queue_agent.close(immediate=True) | ||
|
|
||
|
|
@@ -734,6 +747,7 @@ async def subscribe( | |
| f'Task {self._task_id} is already completed.' | ||
| ) | ||
| self._reference_count += 1 | ||
| self._cancel_idle_release() | ||
| logger.debug( | ||
| 'Subscribe[%s]: Subscribers count: %d', | ||
| self._task_id, | ||
|
|
@@ -814,6 +828,7 @@ async def subscribe( | |
| self._reference_count -= 1 | ||
| # Evaluate if this was the last subscriber on a finished task. | ||
| await self._maybe_cleanup() | ||
| self._schedule_idle_release() | ||
|
|
||
| async def cancel(self, call_context: ServerCallContext) -> Task: | ||
| """Cancels the running active task. | ||
|
|
@@ -940,6 +955,7 @@ async def aclose(self) -> None: | |
| # tasks, because their teardown re-acquires it. | ||
| async with self._lock: | ||
| self._is_finished.set() | ||
| self._cancel_idle_release() | ||
| background_tasks = [ | ||
| task | ||
| for task in (self._producer_task, self._consumer_task) | ||
|
|
@@ -956,6 +972,45 @@ async def aclose(self) -> None: | |
| if isinstance(result, Exception): | ||
| logger.error('Error during aclose', exc_info=result) | ||
|
|
||
| def _schedule_idle_release(self) -> None: | ||
| """(Re)starts the timer that releases a parked interrupted task.""" | ||
| if self._interrupted_task_timeout is None or self._is_finished.is_set(): | ||
| return | ||
| self._cancel_idle_release() | ||
| self._idle_release_task = asyncio.create_task( | ||
| self._release_when_idle(self._interrupted_task_timeout), | ||
| name=f'idle-release:{self._task_id}', | ||
| ) | ||
|
|
||
| def _cancel_idle_release(self) -> None: | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please add docstring |
||
| if self._idle_release_task is not None: | ||
| self._idle_release_task.cancel() | ||
| self._idle_release_task = None | ||
|
|
||
| async def _release_when_idle(self, timeout: float) -> None: | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please introduce try/catch to handle |
||
| """Finishes the task if it is still parked after `timeout` seconds.""" | ||
| await asyncio.sleep(timeout) | ||
| # Read before `_lock` so the checks and release below cannot suspend. | ||
| task = await self._task_manager.get_task() | ||
| async with self._lock: | ||
| if ( | ||
| task is None | ||
| or task.status.state not in INTERRUPTED_TASK_STATES | ||
| or self._is_finished.is_set() | ||
| # The consumer holds one reference; more means subscribers. | ||
| or self._reference_count != 1 | ||
| or self._request_lock.locked() | ||
| or not self._request_queue.empty() | ||
| ): | ||
| return | ||
| logger.debug( | ||
| 'ActiveTask[%s]: Releasing task parked in %s', | ||
| self._task_id, | ||
| TaskState.Name(task.status.state), | ||
| ) | ||
| self._is_finished.set() | ||
| self._request_queue.shutdown(immediate=True) | ||
|
|
||
| async def _maybe_cleanup(self) -> None: | ||
| """Triggers cleanup if task is finished and has no subscribers. | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -96,6 +96,10 @@ | |
| with `DefaultRequestHandler` but is not used: v2 delegates event streaming | ||
| to an in-memory `ActiveTaskRegistry`. Passing a non-``None`` value emits a | ||
| `DeprecationWarning` and logs a warning. | ||
|
|
||
| ``interrupted_task_timeout`` releases the `ActiveTask` of a task idle in | ||
| ``input-required`` or ``auth-required`` after that many seconds; a | ||
| follow-up recreates it from the store. ``None`` (default) never releases. | ||
| """ | ||
|
|
||
| _background_tasks: set[asyncio.Task] | ||
|
|
@@ -107,17 +111,18 @@ | |
| agent_card: AgentCard, | ||
| queue_manager: Any | ||
| | None = None, # Accepted for signature compat; ignored in v2 (warns) | ||
| push_config_store: PushNotificationConfigStore | None = None, | ||
| push_sender: PushNotificationSender | None = None, | ||
| request_context_builder: RequestContextBuilder | None = None, | ||
| extended_agent_card: AgentCard | None = None, | ||
| extended_card_modifier: Callable[ | ||
| [AgentCard, ServerCallContext], Awaitable[AgentCard] | ||
| ] | ||
| | None = None, | ||
| push_url_validator: Callable[[str], Awaitable[bool]] | None = None, | ||
| validate_input_modes: bool = False, | ||
| event_stream: TaskEventStream | None = None, | ||
| interrupted_task_timeout: float | None = None, | ||
| ) -> None: | ||
| if queue_manager is not None: | ||
| message = ( | ||
|
|
@@ -172,6 +177,7 @@ | |
| task_store=self._versioned_store, | ||
| push_sender=self._push_sender, | ||
| event_stream=self._event_stream, | ||
| interrupted_task_timeout=interrupted_task_timeout, | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please validate that timeout value>0 |
||
| ) | ||
| self._background_tasks = set() | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should be moved above async lock