da-daken commented on code in PR #926:
URL: https://github.com/apache/flink-agents/pull/926#discussion_r3864604546
##########
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:
+ outcomes.append(Outcome.failure(timeout_exception))
+ continue
+ if not future.done():
+ future.cancel()
+ if future.done() and not future.cancelled():
+ try:
+ outcomes.append(Outcome.success(future.result()))
+ except Exception as e:
+ outcomes.append(Outcome.failure(e))
+ else:
+ outcomes.append(Outcome.failure(timeout_exception))
+ return outcomes
+
+
+def _collect_outcomes(futures: list[Any]) -> list[Outcome]:
+ outcomes = []
+ for future in futures:
+ try:
+ outcomes.append(Outcome.success(future.result()))
+ except Exception as e: # noqa: PERF203
+ outcomes.append(Outcome.failure(e))
+ return outcomes
+
+
+def _collect_outcomes_on_timeout(
Review Comment:
removed.
--
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]
