weiqingy commented on code in PR #926:
URL: https://github.com/apache/flink-agents/pull/926#discussion_r3661644118
##########
api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java:
##########
@@ -42,9 +44,43 @@ public class AgentExecutionOptions {
public static final ConfigOption<Boolean> CHAT_ASYNC =
new ConfigOption<>("chat.async", Boolean.class, true);
+ /** Whether the built-in tool-call action runs each tool via durable async
execution. */
public static final ConfigOption<Boolean> TOOL_CALL_ASYNC =
new ConfigOption<>("tool-call.async", Boolean.class, true);
+ /**
+ * Whether multiple tool calls from one {@code ToolRequestEvent} run as
one parallel durable
+ * batch when {@link #TOOL_CALL_ASYNC} is also enabled (JDK 21+).
+ *
+ * <p>Default is {@code true}. A parallel batch raises the number of
in-flight external calls;
+ * after failover, tools whose results were not yet persisted may be
submitted again.
+ * Side-effecting tools should be idempotent or provide a {@code
reconciler()}. Set to {@code
+ * false} to keep serial async or sync tool execution.
+ */
+ public static final ConfigOption<Boolean> TOOL_CALL_PARALLEL =
+ new ConfigOption<>("tool-call.parallel", Boolean.class, true);
+
+ /**
+ * Size of the dedicated thread pool used for tool-call async and parallel
batch execution.
+ *
+ * <p>Separate from {@link #NUM_ASYNC_THREADS} so a large tool batch does
not exhaust the global
+ * async pool.
+ */
+ public static final ConfigOption<Integer> TOOL_CALL_NUM_ASYNC_THREADS =
+ new ConfigOption<>(
+ "tool-call.num-async-threads",
+ Integer.class,
+ Runtime.getRuntime().availableProcessors() * 2);
+
+ /**
+ * Overall timeout for one parallel tool-call batch.
+ *
+ * <p>Non-positive values disable the timeout. When the deadline elapses,
unfinished slots are
+ * failed; slots that already completed keep their success or failure
outcome.
+ */
+ public static final ConfigOption<Duration> TOOL_CALL_BATCH_TIMEOUT =
+ new ConfigOption<>("tool-call.batch.timeout", Duration.class,
Duration.ofMillis(-1));
Review Comment:
This is the repo's first `ConfigOption<Duration>`, and I do not think it can
be set from any documented route.
`AgentConfiguration.get()` dispatches `isAssignableFrom`, `String`,
`Integer`, `Long`, `Float`, `Double`, `Boolean`, `isEnum()`, then `throw new
ClassCastException`. There is no `Duration` branch, and YAML values arrive as
`String` or `Integer`. So `getConfig().get(TOOL_CALL_BATCH_TIMEOUT)` at
`JavaRunnerContextImpl.java:253` throws for any value a user sets, after `:119`
has already persisted N PENDING slots, and the catch-all at
`ToolCallAction.java:171` then reports every tool in the batch failed.
The programmatic route clears `get()` but not plan serialization:
`AgentPlan.writeObject` uses a bare `new ObjectMapper()` with no
`JavaTimeModule`, and I reproduced `InvalidDefinitionException` on
`java.time.Duration` against the pinned jackson-databind. On Python,
`core_options.py:271` is `config_type=int`, so the documented `30s` raises
`ValueError` mid-action.
CI stays green because
`JavaRunnerContextImplDurableExecuteAsyncTest.java:324-327` sets the option
in-process, where `isAssignableFrom` short-circuits. And
`check_java_python_config_options_parity.py:45` adds `"java.time.Duration":
int` to the type table, so the harness now certifies exactly this pair.
`ConfigurationUtils.convertValue` is already imported in
`AgentConfiguration` for the enum branch and it handles `Duration`; retyping
this option to `Long` millis would instead let the harness stay strict. Do you
have a preference? Asking because the batch deadline is the only bound on a
hung tool wedging the fan-in, and it currently ships both default-off and
unsettable.
##########
runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java:
##########
@@ -240,12 +397,61 @@ public <T> T executeAsync(ContinuationContext context,
Supplier<T> supplier) {
return supplier.get();
}
+ @Override
+ public <T> List<Outcome<T>> executeAllAsync(
+ ContinuationContext context,
+ List<Callable<T>> suppliers,
+ java.time.Duration timeout) {
+ executeAllAsyncCallCount++;
+ executeAllAsyncBatchSizes.add(suppliers.size());
+ if (useTimeoutCollection) {
+ return collectTimedOutOutcomes(suppliers);
+ }
+ List<Outcome<T>> outcomes = new
java.util.ArrayList<>(suppliers.size());
+ for (Callable<T> supplier : suppliers) {
+ try {
+ outcomes.add(Outcome.success(supplier.call()));
+ } catch (Exception e) {
+ outcomes.add(Outcome.failure(e));
+ }
+ }
+ return outcomes;
+ }
+
+ private <T> List<Outcome<T>> collectTimedOutOutcomes(List<Callable<T>>
suppliers) {
Review Comment:
`testToolCallBatchExecutionIsActuallyParallel` covers the real java21 path
well, fallback branch included, which matters because surefire runs the
exploded `target/classes` and so never loads the `META-INF/versions/21`
classes. The deadline leg looks like the remaining gap, and this stub is what
makes it look covered.
`collectTimedOutOutcomes` hard-codes "index 0 succeeds, everything else is a
`TimeoutException`" and never reads the `timeout` argument. So in
`testDurableExecuteAllAsyncTimeoutKeepsCompletedOutcomes` the 10 ms deadline
set at `:324-327` and the `Thread.sleep(100)` in the second supplier are both
inert, supplier 1 is never invoked, and its `getCallCount()` is never asserted.
A deadline computed with the wrong sign, a `cancel` on the wrong future, or a
barrier that never resolves would all still pass. What it does cover, a
timed-out slot persisted FAILED with the cursor still advancing by 2, is worth
keeping.
Now that the e2e path exists, would a batch that overruns a short
`tool-call.batch.timeout` be the cheapest way to get `getDeadlineNanos` and
`collectBatchOutcomesOnTimeout` genuinely exercised? That may have to wait on
the `Duration` conversion on `AgentExecutionOptions.java:82`, since setting the
option is what trips the plan serializer.
##########
plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java:
##########
@@ -79,55 +103,160 @@ public static void processToolRequest(Event event,
RunnerContext ctx) {
diagnosticError = e.getMessage();
}
- if (tool != null) {
- try {
- // Framework-owned injected args must win over
model-provided values so hidden
- // context such as tenant ids cannot be spoofed by a tool
call payload.
- mergedArguments.putAll(resolveInjectedArguments(tool,
ctx));
- ToolResponse response;
- final Tool toolRef = tool;
- final Map<String, Object> callArguments = mergedArguments;
- DurableCallable<ToolResponse> callable =
- new DurableCallable<>() {
- @Override
- public String getId() {
- return "tool-call";
- }
-
- @Override
- public Class<ToolResponse> getResultClass() {
- return ToolResponse.class;
- }
-
- @Override
- public ToolResponse call() throws Exception {
- return toolRef.call(new
ToolParameters(callArguments));
- }
- };
- response =
- toolCallAsync
- ? ctx.durableExecuteAsync(callable)
- : ctx.durableExecute(callable);
- success.put(id, response.isSuccess());
- responses.put(id, response);
- if (!response.isSuccess() && response.getError() != null) {
- error.put(id, response.getError());
- }
- } catch (Exception e) {
- success.put(id, false);
- responses.put(
- id, ToolResponse.error(String.format("Tool %s
execute failed.", name)));
- error.put(id, e.getMessage());
- }
- } else {
- success.put(id, false);
- responses.put(
- id, ToolResponse.error(String.format("Tool %s does not
exist.", name)));
- error.put(id, diagnosticError != null ? diagnosticError :
"Tool does not exist.");
+ if (tool == null) {
+ recordInlineResponse(
+ id,
+ ToolResponse.error(String.format("Tool %s does not
exist.", name)),
+ diagnosticError != null ? diagnosticError : "Tool does
not exist.",
+ success,
+ error,
+ responses);
+ continue;
}
+
+ try {
+ // Framework-owned injected args must win over model-provided
values so hidden
+ // context such as tenant ids cannot be spoofed by a tool call
payload.
+ mergedArguments.putAll(resolveInjectedArguments(tool, ctx));
+ } catch (Exception e) {
+ recordInlineResponse(
+ id,
+ ToolResponse.error(String.format("Tool %s execute
failed.", name)),
+ e.getMessage(),
+ success,
+ error,
+ responses);
+ continue;
+ }
+
+ final Tool toolRef = tool;
+ final Map<String, Object> callArguments = mergedArguments;
+ DurableCallable<ToolResponse> callable =
+ new DurableCallable<>() {
+ @Override
+ public String getId() {
+ return "tool-call-" + id;
+ }
+
+ @Override
+ public Class<ToolResponse> getResultClass() {
+ return ToolResponse.class;
+ }
+
+ @Override
+ public ToolResponse call() throws Exception {
+ return toolRef.call(new
ToolParameters(callArguments));
+ }
+ };
+ executions.add(new ToolCallExecution(id, name, callable));
+ }
+ return executions;
+ }
+
+ private static void executeParallel(
+ List<ToolCallExecution> executions,
+ RunnerContext ctx,
+ Map<String, Boolean> success,
+ Map<String, String> error,
+ Map<String, ToolResponse> responses) {
+ List<DurableCallable<ToolResponse>> callables = new
ArrayList<>(executions.size());
+ for (ToolCallExecution execution : executions) {
+ callables.add(execution.callable);
+ }
+ try {
+ List<Outcome<ToolResponse>> outcomes =
ctx.durableExecuteAllAsync(callables);
+ for (int i = 0; i < outcomes.size(); i++) {
+ recordOutcome(executions.get(i), outcomes.get(i), success,
error, responses);
+ }
+ } catch (Exception e) {
Review Comment:
Any throw out of `durableExecuteAllAsync` after partial progress marks every
tool failed, including slots already durably finalized SUCCEEDED. The action is
then marked completed, so those results never surface on a later run. Lost, not
delayed.
The serial path just below degrades per tool; this one is all-or-nothing.
Routes in besides the `Duration` conversion on `AgentExecutionOptions.java:82`:
a `JsonProcessingException` out of `finalizeExecutedOutcomes`
(`JavaRunnerContextImpl.java:211`), or a deserialization failure in
`readTerminalOutcomeAt`.
Would recording the batch-level exception only against slots with no outcome
yet work here, or is failing the whole set the intent?
##########
python/flink_agents/runtime/flink_runner_context.py:
##########
@@ -636,6 +721,110 @@ def wrapped_func(*a: Any, **kw: Any) -> Any:
return wrapped_func
+ def _call_matches(
+ self, current: _PersistedCallResult, call: DurableCall, args_digest:
str
+ ) -> bool:
+ return current.function_id == call.id and current.args_digest ==
args_digest
+
+ def _read_terminal_outcome(self, current: _PersistedCallResult) -> Outcome:
+ if current.exception_payload is not None:
+ return
Outcome.failure(cloudpickle.loads(current.exception_payload))
+ if current.result_payload is None:
+ return Outcome.success(None)
+ return Outcome.success(cloudpickle.loads(current.result_payload))
+
+ def _callable_for_durable_call(self, call: DurableCall) -> Callable[[],
Any]:
+ kwargs = call.kwargs or {}
+ return partial(call.func, *call.args, **kwargs)
+
+ def _prepare_batch_execution(self, calls: list[DurableCall]) ->
_BatchExecutionPlan:
+ args_digest = ""
Review Comment:
The batch keys slots on `("tool-call-<id>", "")`
(`tool_call_action.py:134`). The Python single-call path keys on
`(_compute_function_id(func), args_digest)` (`:482-483`), and since the
callable is `tool.call`, `_compute_function_id` returns the same string for
every tool. Java has no such split: `RunnerContextImpl.java:604-605` uses
`getId()` plus `""` for the single-call state machine too.
So a job that checkpoints mid-batch and restarts with
`tool-call.parallel=false` replays cleanly on Java, while on Python
`matchNextOrClearSubsequentCallResult` cannot match the persisted key,
truncates the journal, and calls every completed tool again. Same flip that
reaches the serial-path issue on `JavaRunnerContextImpl.java:193`.
The per-call `functionId` rename was headed for its own issue, so not asking
for it here. What caught my eye is that the batch path adopted it while the
single-call path did not, so the two disagree inside Python today. Is
documenting the flip as not recoverable across a restart the right stopgap
until that lands?
##########
runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java:
##########
@@ -75,6 +102,159 @@ private <T> T durableExecuteAsyncWithReconcile(
callable, reconcileCallable, () ->
executeAsyncCallable(callable));
}
+ @Override
+ public <T> List<Outcome<T>>
durableExecuteAllAsync(List<DurableCallable<T>> callables)
+ throws Exception {
+ if (callables.isEmpty()) {
+ return List.of();
+ }
+ if (durableExecutionContext == null) {
+ return executeAllWithoutDurableState(callables);
+ }
+
+ String argsDigest = "";
+ int base = durableExecutionContext.getCurrentCallIndex();
+ BatchExecutionPlan<T> plan = buildBatchExecutionPlan(callables, base,
argsDigest);
+
+ reservePendingBatchIfNeeded(callables, argsDigest, plan);
+
+ List<Outcome<T>> executed = executeOutcomeSuppliers(plan.suppliers);
+ finalizeExecutedOutcomes(callables, base, argsDigest, plan, executed);
+
+ advanceCallIndexBy(callables.size());
+ return plan.outcomes;
+ }
+
+ private <T> BatchExecutionPlan<T> buildBatchExecutionPlan(
+ List<DurableCallable<T>> callables, int base, String argsDigest)
throws Exception {
+ BatchExecutionPlan<T> plan = new
BatchExecutionPlan<>(callables.size());
+ for (int i = 0; i < callables.size(); i++) {
+ DurableCallable<T> callable = callables.get(i);
+ CallResult current = getCallResultAt(base + i);
+ if (current == null) {
+ markReservationStart(plan, i);
+ addExecutableCall(plan, i, callable::call);
+ continue;
+ }
+ if (!current.matches(callable.getId(), argsDigest)) {
+ clearCallResultsFromAndPersist(base + i);
+ plan.needsReservation = true;
+ plan.executionStart = i;
+ addExecutableCall(plan, i, callable::call);
+ appendRemainingExecutions(callables, plan, i + 1);
+ break;
+ }
+ if (current.isPending()) {
+ Callable<T> reconcileCallable = callable.reconciler();
+ Callable<T> executionCallable =
+ reconcileCallable != null ? reconcileCallable :
callable::call;
+ addExecutableCall(plan, i, executionCallable);
+ } else {
+ plan.outcomes.add(
+ readTerminalOutcomeAt(
+ base + i, callable.getId(), argsDigest,
callable.getResultClass()));
+ }
+ }
+ return plan;
+ }
+
+ private <T> void markReservationStart(BatchExecutionPlan<T> plan, int
callIndex) {
+ plan.needsReservation = true;
+ if (plan.executionStart < 0) {
+ plan.executionStart = callIndex;
+ }
+ }
+
+ private <T> void addExecutableCall(
+ BatchExecutionPlan<T> plan, int callIndex, Callable<T>
executionCallable) {
+ plan.outcomes.add(null);
+ plan.suppliers.add(executionCallable);
+ plan.executableCallIndexes.add(callIndex);
+ }
+
+ private <T> void appendRemainingExecutions(
+ List<DurableCallable<T>> callables, BatchExecutionPlan<T> plan,
int startIndex) {
+ for (int i = startIndex; i < callables.size(); i++) {
+ DurableCallable<T> remaining = callables.get(i);
+ addExecutableCall(plan, i, remaining::call);
+ }
+ }
+
+ private <T> void reservePendingBatchIfNeeded(
+ List<DurableCallable<T>> callables, String argsDigest,
BatchExecutionPlan<T> plan) {
+ if (!plan.needsReservation) {
+ return;
+ }
+ List<String> ids = new ArrayList<>();
+ for (DurableCallable<T> callable :
+ callables.subList(plan.executionStart, callables.size())) {
+ ids.add(callable.getId());
+ }
+ reservePendingBatch(ids, argsDigest);
Review Comment:
Reserving PENDING for every call in the batch, non-reconcilable ones
included, is the intended shape, and it carries the guarantee that such a slot
is re-executed on recovery. The batch path honours that at `:147-151`. The
serial path can read those same slots and does not.
Before this PR the two never met. PENDING was written only by
`durableExecuteWithReconcile` (`RunnerContextImpl.java:612`, `:618`) and the
Python bridge's `_prepare_reconciler_execution`, both gated on a non-null
`reconciler()`, and that method gates its own cache read on
`!current.isPending()` at `:622`. The completion-only path never needed that
gate, because non-reconcilable calls never wrote a PENDING slot. Now they do:
`ToolCallAction.java:135` never overrides `reconciler()` and the interface
default is `null`.
`CallResult.matches` compares only `functionId` and `argsDigest`
(`CallResult.java:160-162`), so a reserved slot reads as a cache hit at
`RunnerContextImpl.java:757`, both payloads are null, and `tryGetCachedResult`
throws NPE on `return Optional.of(null)` at `:575`. `ToolCallAction.java:192`
catches it, so the tool is reported failed rather than re-executed.
Two routes onto the serial path after a mid-batch crash, with the reserved
slots still checkpointed. Setting `tool-call.parallel=false` and restarting is
one. The other needs no config change: a trailing tool whose resource fails to
resolve on the recovery run is dropped at `ToolCallAction.java:114`,
`executions.size()` falls to 1, and the `> 1` guard at `:67` routes to serial.
Treating a PENDING slot as a miss in `matchNextOrClearSubsequentCallResult`
/ `tryGetCachedResult` would extend the same re-execute guarantee to that path.
Would that be the right place for it, or is there a reason the serial read
should keep seeing a hit?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
