weiqingy commented on code in PR #926:
URL: https://github.com/apache/flink-agents/pull/926#discussion_r3793394265
##########
python/flink_agents/runtime/flink_runner_context.py:
##########
@@ -241,6 +294,136 @@ def __await__(self) -> Any:
return result
+class _DurableBatchAsyncExecutionResult(AsyncExecutionResult):
+ def __init__(self, ctx: "FlinkRunnerContext", calls: list[DurableCall]) ->
None:
+ self._ctx = ctx
+ self._calls = calls
+
+ def __await__(self) -> Any:
+ plan = self._ctx._prepare_batch_execution(self._calls)
+ parallelism =
self._ctx.config.get(AgentExecutionOptions.TOOL_CALL_PARALLELISM)
+ timeout_ms =
self._ctx.config.get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT_MS)
+ deadline = time.monotonic() + timeout_ms / 1000 if timeout_ms > 0 else
None
+ suppliers = [supplier for _, supplier in plan.suppliers]
+ batch_futures: list[Any | None] = [None] * len(suppliers)
+ try:
+ executed = yield from _execute_sliding_window_batch(
+ self._ctx.executor,
+ suppliers,
+ parallelism,
+ deadline,
+ timeout_ms,
+ batch_futures,
+ plan.submitted,
+ )
+ except _BatchTimeoutError as exception:
+ executed = _collect_sliding_window_outcomes_on_timeout(
+ batch_futures, plan.submitted, exception
+ )
+ return self._ctx._finalize_batch_execution(self._calls, plan, executed)
+
+
+class _BatchTimeoutError(TimeoutError):
+ """Raised when a durable batch exceeds its deadline."""
+
+
+def _execute_sliding_window_batch(
+ executor: ThreadPoolExecutor,
+ suppliers: list[Any],
+ parallelism: int,
+ deadline: float | None,
+ timeout_ms: int,
+ futures: list[Any | None],
+ submitted: list[bool],
+) -> Any:
+ batch_size = len(suppliers)
+ if batch_size == 0:
+ return []
+
+ parallelism_limit = min(max(parallelism, 1), batch_size)
+ next_to_submit = 0
+ completed = 0
+ counted = [False] * batch_size
+
+ def in_flight() -> int:
+ return sum(
+ 1
+ for i in range(next_to_submit)
+ if futures[i] is not None and not futures[i].done()
+ )
+
+ while completed < batch_size:
+ if deadline is not None and time.monotonic() >= deadline:
+ timeout_message = (
+ f"Async durable batch execution timed out after {timeout_ms}
ms"
+ )
+ raise _BatchTimeoutError(timeout_message)
+
+ while next_to_submit < batch_size and in_flight() < parallelism_limit:
+ futures[next_to_submit] =
executor.submit(suppliers[next_to_submit])
+ submitted[next_to_submit] = True
+ next_to_submit += 1
+
+ for i in range(next_to_submit):
+ if not counted[i] and futures[i].done():
+ counted[i] = True
+ completed += 1
+
+ if completed < batch_size:
+ yield
+
+ return _collect_outcomes(futures)
+
+
+def _collect_sliding_window_outcomes_on_timeout(
+ futures: list[Any | None],
+ submitted: list[bool],
+ timeout_exception: BaseException,
+) -> list[Outcome]:
+ outcomes = []
+ for is_submitted, future in zip(submitted, futures, strict=True):
+ if not is_submitted or future is None:
+ outcomes.append(Outcome.failure(timeout_exception))
+ continue
+ if not future.done():
+ future.cancel()
+ if future.done() and not future.cancelled():
Review Comment:
Thanks, this closes the never-submitted case. I think there's one more case
in the same timeout path though, and I'm curious how you'd want to handle it.
`submitted[i]` is set at `:364` the moment `executor.submit(...)` returns at
`:363`, so it really means "handed to the pool" rather than "actually ran". At
the deadline, `future.cancel()` at `:389` succeeds for anything still sitting
in the queue, which stops it from ever running. Line `:390` even spots this
(`future.done() and not future.cancelled()`), but the slot still falls through
to `Outcome.failure(timeout_exception)` at `:396`. Since `submitted[i]` is
`True`, `_finalize_batch_execution:1017` hands it to `finalizeCallAt` at
`:1025` and it gets recorded as FAILED. After a restart,
`_prepare_batch_execution:984-985` replays that failure rather than running the
tool.
This only shows up when `tool-call.batch.timeout.ms` is set, since `:306`
only builds a deadline when it's positive. With it on, you'd get there through
the busy-pool case you describe at `core_options.py:257-264`: the pool is
`num-async-threads` (default `2 x cores`) shared per subtask,
`tool-call.parallelism` defaults to `cores`, so once the pool fills up `submit`
starts queueing, and any queued slot at the deadline lands here.
One option would be to set the flag as the first line inside the supplier
wrapper instead (`:363-364` here,
`java21/ContinuationActionExecutor.java:210-219` on the Java side). That would
keep both languages saying the same thing, since cancelling a
`CompletableFuture` doesn't stop an already-queued supplier from running
(`java21/ContinuationActionExecutor.java:301`). What do you think the flag
should mean, handed to the pool or actually started?
##########
python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py:
##########
@@ -415,3 +594,458 @@ def collect_kwargs(**kwargs: Any) -> dict[str, Any]:
_close_runner_context(ctx)
assert result == {}
+
+
+def
test_flink_runner_context_durable_execute_all_async_runs_calls_in_parallel() ->
None:
+ j_runner_context = _FakeJavaRunnerContext()
+ config = AgentConfiguration(
+ {"tool-call.batch.timeout.ms": -1, "tool-call.parallelism": 3}
+ )
+ ctx = _create_runner_context(j_runner_context, config=config,
executor_workers=3)
+ sleep_seconds = 0.2
+
+ def slow_call(value: str) -> str:
+ time.sleep(sleep_seconds)
+ return value
+
+ try:
+ start = time.perf_counter()
+ outcomes = _run_async(
+ ctx.durable_execute_all_async(
+ [
+ _durable_call(slow_call, "one"),
+ _durable_call(slow_call, "two"),
+ _durable_call(slow_call, "three"),
+ ]
+ )
+ )
+ elapsed = time.perf_counter() - start
+ finally:
+ _close_runner_context(ctx)
+
+ assert [outcome.value for outcome in outcomes] == ["one", "two", "three"]
+ assert elapsed < sleep_seconds * 2.5
Review Comment:
nit: with `sleep_seconds = 0.2` (`:605`), 3 calls and parallelism 3, moving
to `2.5` shifts the threshold from 400 ms to 500 ms. That gives 100 ms more
room for jitter, but it also halves the gap below the fully-serial 600 ms, from
200 ms down to 100 ms. The check also only ever caught fully serial execution,
since 2 parallel plus 1 serial comes in around 400 ms and passes either way.
A `threading.Barrier(3)` awaited inside `slow_call` would make it
deterministic instead. If execution went serial, the first call would block,
the barrier would break, and the existing value check at `:626` would fail on
its own, so the timing assert could go away. This one did go red on `ut-python
[macos]` at the previous head `204eb7a3`. Would a barrier be worth it here?
##########
docs/content/docs/operations/configuration.md:
##########
@@ -131,9 +131,11 @@ Here is the list of all built-in core configuration
options.
| `max-retries` | 3 | int
| Number of retries when using `ErrorHandlingStrategy.RETRY`.
|
| `retry-wait-interval` | 1 | int
| Base wait interval in seconds between retries when using
`ErrorHandlingStrategy.RETRY`. Uses exponential backoff: the actual wait time
for the Nth retry is `retry-wait-interval * 2^(N-1)` seconds. For example, with
default 1s, waits are 1s, 2s, 4s, etc. Retry count and total wait time are
reported in `ChatResponseEvent` and recorded as metrics (`retryCount`,
`retryWaitSec`) under the connection name. |
| `chat.async` | true | boolean
| Whether chat asynchronously for built-in chat action.
|
-| `tool-call.async` | true | boolean
| Whether process tool call for built-in tool call action.
|
+| `tool-call.async` | true | boolean
| Whether the built-in tool-call action runs each tool via durable async
execution.
|
+| `tool-call.parallelism` | os cpu count | int
| In-flight concurrency for tool calls from one `ToolRequestEvent` batch when
`tool-call.async` is enabled. `1` runs tools serially; values greater than `1`
run a parallel durable batch with a sliding window of at most that many
concurrent tool calls. On **Java**, concurrent in-batch execution requires
**JDK 21+** (Continuation API); below JDK 21 the batch still runs but tool
calls execute serially. **Python** uses the shared async `ThreadPoolExecutor`
and runs batches concurrently regardless of JDK version. Increases in-flight
external calls; after failover, unfinished tools may be submitted again —
side-effecting tools should be idempotent or provide a reconciler. {{< hint
warning >}}**Default is parallel** (`os cpu count`). Chat, RAG, and tool
batches share one `num-async-threads` pool **per operator subtask** (all keys
on that subtask). Built-in actions for a single key run one at a time, so chat
and a tool batch on the **same key** do not overlap in the usual chat → tool
path; delay shows up mainly **across keys** on the same subtask. With defaults
(`num-async-threads = 2× cores`, `tool-call.parallelism = cores`), one batch
can use up to half the pool; several busy keys can still saturate it. Lower
this value or increase `num-async-threads` on hot subtasks. {{< /hint >}} |
+| `tool-call.batch.timeout.ms` | -1 (disabled) | long
(milliseconds) | Overall timeout for one parallel tool-call batch.
Non-positive disables it. On timeout, completed slots keep their outcome and
unfinished slots fail. Timeout cancellation is best-effort; external side
effects from unfinished tool calls may still complete, so side-effecting tools
should be idempotent or provide a reconciler. On **Java**, only enforced on
**JDK 21+**; on JDK 11 the batch fallback ignores this setting and runs to
completion serially. **Python** enforces the deadline in the batch await loop. |
Review Comment:
Now that timed-out slots that never got submitted stay PENDING instead of
failing, a few places still describe the old behaviour:
- this row: "On timeout, completed slots keep their outcome and unfinished
slots fail."
- `docs/content/docs/development/tool_use.md:317-318`, the same sentence.
- `AgentExecutionOptions.java:78-79`, "When the deadline elapses, unfinished
slots are failed", which sits right above the new thread-reclamation text at
`:81-87`.
PENDING versus FAILED is what decides whether a timed-out tool runs again
after a restart, so it seems worth keeping these in step with the code. Could
the three be updated together?
nit: the thread-reclamation note you added went into
`AgentExecutionOptions.java:81-87` and `core_options.py:272-280`, but not into
either doc page. The `tool-call.parallelism` row at `:135` keeps its capacity
warning in a `{{< hint warning >}}` block, if you wanted somewhere to put it.
##########
runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java:
##########
@@ -477,7 +599,7 @@ protected <T> Optional<T> tryGetCachedResult(
} else if (resultPayload != null) {
return Optional.of(OBJECT_MAPPER.readValue(resultPayload,
resultClass));
} else {
- return Optional.of(null);
+ return Optional.empty();
Review Comment:
Happy to leave this out of the PR. One thing on the reasoning though, in
case it matters later.
The PENDING check is real (`:789-797`), but it's the only check.
`matchNextOrClearSubsequentCallResult` returns a hit at `:804-806` once
`matches(functionId, argsDigest)` passes and the slot isn't pending, and
nothing looks at the payloads. So a SUCCEEDED slot with a null result payload
still gets through, and reaches `Optional.of(null)` at `:606`.
That's also how a `null` return gets stored. `serializeDurableResult(null)`
returns `null` rather than a JSON `null` (`:691-693`), so a durable call that
legitimately returns `null` is saved as SUCCEEDED with no payload, and NPEs on
replay. `ToolCallAction` is fine, since the per-execution `catch` at
`ToolCallAction.java:187-197` turns it into a tool failure. The case that's
actually exposed is a user action calling `durableExecute` directly.
`main` carries the identical line (`RunnerContextImpl.java:618` there), so
this PR isn't making anything worse, which is why I think deferring is the
right call.
Worth an issue though, so it doesn't get lost. The same class already
handles this case elsewhere: `readTerminalOutcomeAt` at `:566-568` does `if
(callResult.getResultPayload() == null) return Outcome.success(null);`.
--
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]
