da-daken commented on code in PR #926:
URL: https://github.com/apache/flink-agents/pull/926#discussion_r3864600818
##########
python/flink_agents/runtime/flink_runner_context.py:
##########
@@ -241,6 +293,158 @@ 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)
+ started: list[bool] = [False] * len(suppliers)
+ try:
+ executed = yield from _execute_sliding_window_batch(
+ self._ctx.executor,
+ suppliers,
+ parallelism,
+ deadline,
+ timeout_ms,
+ batch_futures,
+ started,
+ )
+ except _BatchTimeoutError as exception:
+ executed = _collect_sliding_window_outcomes_on_timeout(
+ batch_futures, started, exception
+ )
+ return self._ctx._finalize_batch_execution(
+ self._calls, plan, started, executed
+ )
+
+
+class _BatchTimeoutError(TimeoutError):
+ """Raised when a durable batch exceeds its deadline."""
+
+
+def _mark_started_on_run(
+ supplier: Callable[[], Any], started: list[bool], index: int
+) -> Callable[[], Any]:
+ """Wrap a supplier so ``started[index]`` flips only when the worker truly
runs.
+
+ A task queued in a saturated pool but cancelled before it executes keeps
+ ``started[index] == False``, so it is treated as never-run and stays
+ re-executable on recovery instead of being recorded as a timeout failure.
+ """
+
+ def _run() -> Any:
+ started[index] = True
+ return supplier()
+
+ return _run
+
+
+def _execute_sliding_window_batch(
+ executor: ThreadPoolExecutor,
+ suppliers: list[Any],
+ parallelism: int,
+ deadline: float | None,
+ timeout_ms: int,
+ futures: list[Any | None],
+ started: 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:
+ index = next_to_submit
+ futures[index] = executor.submit(
+ _mark_started_on_run(suppliers[index], started, index)
+ )
+ 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],
+ started: list[bool],
+ timeout_exception: BaseException,
+) -> list[Outcome]:
+ outcomes = []
+ for is_started, future in zip(started, futures, strict=True):
+ if not is_started or future is None:
Review Comment:
Good catch! Similarly, I observed that under OpenJDK, when a worker dequeues
a task, AsyncSupply.run() checks whether d.result == null. After cancel is
invoked, tasks already queued will not be executed. This behavior does not
affect other JDK versions either, so I also removed that check on the Java
side, allowing execution to fall through to the cancel method instead.
--
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]
