Repository navigation
fix: keep the results of finished parallel tool calls on abort - #7992
Open
AkshayShah03 wants to merge 1 commit into
Open
AkshayShah03 wants to merge 1 commit into
AkshayShah03 wants to merge 1 commit into
Conversation
When an invocation was aborted (abort_signal, a client disconnect, or cancelling the task driving run_async) while one call of a parallel batch was still running, the calls that had already finished were sealed as aborted along with it. Their real results and state_delta were dropped, so the model was told a tool that had run did not, and could call it again. On cancellation the batch executor now keeps the merged response event of the plain results that finished, as it already does when a sibling call raises, on the invocation's abort state. Abort sealing records it before building the synthetic responses, so only the calls that were still running are sealed. Fixes google#7958
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
Problem:
When an invocation was aborted while one call of a parallel batch was still running, the calls that had already finished were sealed as aborted along with it. That covers
abort_signal.set(), a client disconnect on/runor/run_sse, and cancelling the task drivingrun_async. Their real results andstate_deltawere dropped, so the model was told a tool that had run did not, and could call it again (for example, send the same email twice)._gather_or_cancelre-raisesCancelledErrorafter cancelling the unfinished calls._execute_prepared_function_callsonly kept finished results forException/NodeInterruptedError(#7428), so on cancellation they were discarded, and_build_abort_eventssealed every call in the batch.Solution:
As suggested on #7958:
_execute_prepared_function_callsalso catchesasyncio.CancelledError. It merges the plain results of the calls that finished, using the same_is_plain_resultfilter and merge as the error path, and stores them on the invocation's_AbortState. Then it re-raises. Live mode is left out, as in the error path.Runner._synthesize_abort_events_if_neededappends those events through_process_and_append_eventbefore calling_build_abort_events. That function only seals calls without a response, so it now covers just the unfinished ones.The kept event is the normal merged response event, so its
state_deltaandartifact_deltaare persisted too. It's also yielded to the caller along with the abort events, as the synthetic responses already are.If the caller closes the stream early (
aclose()/break), nothing is sealed, as before, so nothing is recorded either.Behaviour change: an aborted batch now records finished calls as finished instead of aborted. Per the issue, this is for maintainers to confirm.
Question: the error path still skips aborted invocations (
_keep_completed_responsesreturns early whenis_aborted). So when a sibling raises in the same step an abort lands, the finished call is still sealed as aborted.test_run_async_tool_error_on_abort_leaves_calls_to_abort_sealingasserts that. I left it unchanged to keep this PR to the cancellation case. Should that path record the finished result too, for consistency?Testing Plan
Unit Tests:
New test in
tests/unittests/test_runners.py:test_run_async_abort_keeps_results_of_finished_parallel_calls. It is parametrized over the abort route (abort_signal.set()/task.cancel()) and over a plainLlmAgentor one inside aWorkflow. It checks:state_deltareaches session state;All 4 cases fail on
main.Full unit suite: the same 6 failures as
main(test_google_llmcontinuation-token tests,test_release_dependencies, andtest_app_details, which also fails onmainin a full run). Nothing new.Manual End-to-End (E2E) Tests:
The repro from #7958 (
send_email+ a slowlookup_order, aborted withabort_signal.set()):The same holds when the consumer task is cancelled instead.