diff --git a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java index a0b9269e4..32f248d8b 100644 --- a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java +++ b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java @@ -20,6 +20,8 @@ import org.apache.flink.agents.api.configuration.ConfigOption; +import java.time.Duration; + public class AgentExecutionOptions { public static final ConfigOption ERROR_HANDLING_STRATEGY = new ConfigOption<>( @@ -42,9 +44,43 @@ public class AgentExecutionOptions { public static final ConfigOption 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 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+). + * + *

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 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. + * + *

Separate from {@link #NUM_ASYNC_THREADS} so a large tool batch does not exhaust the global + * async pool. + */ + public static final ConfigOption 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. + * + *

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 TOOL_CALL_BATCH_TIMEOUT = + new ConfigOption<>("tool-call.batch.timeout", Duration.class, Duration.ofMillis(-1)); + public static final ConfigOption RAG_ASYNC = new ConfigOption<>("rag.async", Boolean.class, true); diff --git a/api/src/main/java/org/apache/flink/agents/api/context/Outcome.java b/api/src/main/java/org/apache/flink/agents/api/context/Outcome.java new file mode 100644 index 000000000..0b16ee92c --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/context/Outcome.java @@ -0,0 +1,53 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.api.context; + +/** Result of a durable batch execution slot. */ +public class Outcome { + private final T value; + private final Exception error; + + private Outcome(T value, Exception error) { + this.value = value; + this.error = error; + } + + public static Outcome success(T value) { + return new Outcome<>(value, null); + } + + public static Outcome failure(Exception error) { + return new Outcome<>(null, error); + } + + public boolean isSuccess() { + return error == null; + } + + public boolean isFailure() { + return error != null; + } + + public T getValue() { + return value; + } + + public Exception getError() { + return error; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java b/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java index c3e5d19ba..a6d816300 100644 --- a/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java +++ b/api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java @@ -24,6 +24,7 @@ import org.apache.flink.agents.api.resource.Resource; import org.apache.flink.agents.api.resource.ResourceType; +import java.util.List; import java.util.Map; /** @@ -146,6 +147,21 @@ public interface RunnerContext { */ T durableExecuteAsync(DurableCallable callable) throws Exception; + /** + * Executes multiple durable callables as one batch and returns outcomes in input order. + * + *

On JDK 21+, implementations may submit uncached calls concurrently and yield the current + * action execution until the batch completes. On JDK < 21, this falls back to serial durable + * execution. + * + *

The callable list must be deterministic across recovery: same order and same {@link + * DurableCallable#getId()} values. + * + *

Access to memory and sendEvent are prohibited within the callables. + */ + List> durableExecuteAllAsync(List> callables) + throws Exception; + /** Clean up the resource. */ void close() throws Exception; } diff --git a/docs/content/docs/operations/configuration.md b/docs/content/docs/operations/configuration.md index 9ba9e200e..f1f70a7e3 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -131,7 +131,10 @@ 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.parallel` | true | boolean | When `tool-call.async` is also true (JDK 21+), run multiple tool calls from one `ToolRequestEvent` as one parallel durable batch. Increases in-flight external calls; after failover, unfinished tools may be submitted again — side-effecting tools should be idempotent or provide a reconciler. Set to `false` for serial tool execution. | +| `tool-call.num-async-threads` | os cpu count * 2 | int | Dedicated thread pool size for tool-call async / parallel batch execution. Separate from `num-async-threads` so a large tool batch does not exhaust the global async pool. | +| `tool-call.batch.timeout` | -1ms (disabled) | Duration | 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. | | `rag.async` | true | boolean | Whether retrieve context asynchronously for built-in context retrieval action. | | `num-async-threads` | os cpu count * 2 | int | The thread pool size for async executor. | | `job-identifier` | none | String | The unique identifier of job, remaining consistent after restoring from a savepoint. If not set, uses flink job id. | diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java index c8df332f5..ac9ed762d 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionAgent.java @@ -23,13 +23,33 @@ import org.apache.flink.agents.api.OutputEvent; import org.apache.flink.agents.api.agents.Agent; import org.apache.flink.agents.api.annotation.Action; +import org.apache.flink.agents.api.annotation.ChatModelConnection; +import org.apache.flink.agents.api.annotation.ChatModelSetup; +import org.apache.flink.agents.api.annotation.ToolParam; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.model.BaseChatModelConnection; +import org.apache.flink.agents.api.chat.model.BaseChatModelSetup; import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ChatRequestEvent; +import org.apache.flink.agents.api.event.ChatResponseEvent; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.tools.ToolType; import org.apache.flink.api.java.functions.KeySelector; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + /** * Agent definition for testing async execution functionality. * @@ -71,6 +91,141 @@ public Integer getKey(AsyncRequest request) { } } + /** Chat connection that emits one tool request containing multiple slow tool calls. */ + public static class ToolBatchChatConnection extends BaseChatModelConnection { + public ToolBatchChatConnection( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + } + + @Override + public ChatMessage chat( + List messages, List tools, Map modelParams) { + ChatMessage lastMessage = messages.get(messages.size() - 1); + if (lastMessage.getRole() == MessageRole.TOOL) { + return new ChatMessage(MessageRole.ASSISTANT, lastMessage.getContent()); + } + + String requestId = lastMessage.getContent(); + return new ChatMessage( + MessageRole.ASSISTANT, + "", + List.of( + toolCall("call-1", requestId, 1), + toolCall("call-2", requestId, 2), + toolCall("call-3", requestId, 3))); + } + } + + /** Chat model setup bound to the timed tool. */ + public static class ToolBatchChatModel extends BaseChatModelSetup { + public ToolBatchChatModel(ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + } + + @Override + public Map getParameters() { + return new HashMap<>(); + } + } + + private static Map toolCall(String id, String requestId, int index) { + return Map.of( + "id", + id, + "type", + "function", + "function", + Map.of( + "name", + "timed_tool", + "arguments", + Map.of("request_id", requestId, "call_index", index))); + } + + /** Agent that requests one chat turn that produces multiple slow tool calls. */ + public static class ToolBatchAgent extends Agent { + public ToolBatchAgent(int sleepTimeMs) {} + + @ChatModelConnection + public static ResourceDescriptor toolBatchChatConnection() { + return ResourceDescriptor.Builder.newBuilder(ToolBatchChatConnection.class.getName()) + .build(); + } + + @ChatModelSetup + public static ResourceDescriptor toolBatchChatModel() { + return ResourceDescriptor.Builder.newBuilder(ToolBatchChatModel.class.getName()) + .addInitialArgument("connection", "toolBatchChatConnection") + .addInitialArgument("model", "test-model") + .addInitialArgument("tools", List.of("timed_tool")) + .build(); + } + + @org.apache.flink.agents.api.annotation.Tool( + description = "Records timing for a slow tool call.") + public static String timed_tool( + @ToolParam(name = "request_id") String requestId, + @ToolParam(name = "call_index") String callIndex) { + long start = System.currentTimeMillis(); + try { + Thread.sleep(500); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + long end = System.currentTimeMillis(); + return String.format( + "request=%s,call=%s,start=%d,end=%d", requestId, callIndex, start, end); + } + + @Action(EventType.InputEvent) + public static void requestTools(Event event, RunnerContext ctx) { + InputEvent inputEvent = InputEvent.fromEvent(event); + AsyncRequest request = (AsyncRequest) inputEvent.getInput(); + ctx.sendEvent( + new ChatRequestEvent( + "toolBatchChatModel", + List.of( + new ChatMessage( + MessageRole.USER, String.valueOf(request.id))))); + } + + @Action(EventType.ChatResponseEvent) + public static void emitToolTimings(Event event, RunnerContext ctx) { + ChatResponseEvent responseEvent = ChatResponseEvent.fromEvent(event); + ctx.sendEvent(new OutputEvent(responseEvent.getResponse().getContent())); + } + } + + /** Tool that records its execution time range in the response result. */ + public static class TimedTool extends Tool { + private final int sleepTimeMs; + + public TimedTool(int sleepTimeMs) { + super(new ToolMetadata("timed_tool", "Records timing for a slow tool call.", "{}")); + this.sleepTimeMs = sleepTimeMs; + } + + @Override + public ToolType getToolType() { + return ToolType.FUNCTION; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + int callIndex = parameters.getParameter("call_index", Integer.class); + long start = System.currentTimeMillis(); + try { + Thread.sleep(sleepTimeMs); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + long end = System.currentTimeMillis(); + return ToolResponse.success( + String.format("call=%d,start=%d,end=%d", callIndex, start, end)); + } + } + /** Custom event type for internal agent communication. */ public static class AsyncProcessedEvent extends Event { public static final String EVENT_TYPE = "AsyncProcessedEvent"; diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java index 271b067fa..1c4a62afe 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/AsyncExecutionTest.java @@ -18,6 +18,7 @@ package org.apache.flink.agents.integration.test; import org.apache.flink.agents.api.AgentsExecutionEnvironment; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; import org.apache.flink.agents.runtime.async.ContinuationActionExecutor; import org.apache.flink.streaming.api.datastream.DataStream; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; @@ -27,6 +28,8 @@ import java.util.ArrayList; import java.util.List; +import java.util.regex.Matcher; +import java.util.regex.Pattern; /** * End-to-end tests for Java async execution functionality. @@ -333,6 +336,61 @@ public void testAsyncExecutionIsActuallyParallel() throws Exception { System.out.println("=== Test Passed ==="); } + @Test + public void testToolCallBatchExecutionIsActuallyParallel() throws Exception { + boolean continuationSupported = ContinuationActionExecutor.isContinuationSupported(); + int javaVersion = Runtime.version().feature(); + + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + + int sleepTimeMs = 500; + DataStream inputStream = + env.fromElements( + new AsyncExecutionAgent.AsyncRequest(1, "tool-batch", sleepTimeMs)); + + AgentsExecutionEnvironment agentsEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_ASYNC, true); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_PARALLEL, true); + agentsEnv.getConfig().set(AgentExecutionOptions.TOOL_CALL_NUM_ASYNC_THREADS, 3); + + DataStream outputStream = + agentsEnv + .fromDataStream( + inputStream, new AsyncExecutionAgent.AsyncRequestKeySelector()) + .apply(new AsyncExecutionAgent.ToolBatchAgent(sleepTimeMs)) + .toDataStream(); + + CloseableIterator results = outputStream.collectAsync(); + agentsEnv.execute(); + + List executionRanges = new ArrayList<>(); + while (results.hasNext()) { + String output = results.next().toString(); + Matcher matcher = Pattern.compile("start=(\\d+),end=(\\d+)").matcher(output); + while (matcher.find()) { + long start = Long.parseLong(matcher.group(1)); + long end = Long.parseLong(matcher.group(2)); + executionRanges.add(new long[] {start, end}); + } + } + results.close(); + + Assertions.assertEquals(3, executionRanges.size()); + int overlapCount = countOverlaps(executionRanges); + if (continuationSupported && javaVersion >= 21) { + Assertions.assertTrue( + overlapCount >= 2, + "On JDK 21+, tool calls in one ToolRequestEvent should run in parallel."); + } else { + Assertions.assertEquals( + 0, + overlapCount, + "On JDK < 21, tool-call batch execution should use the sequential fallback."); + } + } + /** * Tests that durableExecute (sync) works correctly. * @@ -387,4 +445,18 @@ public void testDurableExecuteSync() throws Exception { "Output should contain processed data: " + output); } } + + private static int countOverlaps(List executionRanges) { + int overlapCount = 0; + for (int i = 0; i < executionRanges.size(); i++) { + for (int j = i + 1; j < executionRanges.size(); j++) { + long[] range1 = executionRanges.get(i); + long[] range2 = executionRanges.get(j); + if (range1[0] < range2[1] && range2[0] < range1[1]) { + overlapCount++; + } + } + } + return overlapCount; + } } diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java index bc82e1b61..412595125 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ToolCallAction.java @@ -22,6 +22,7 @@ import org.apache.flink.agents.api.configuration.ConfigOption; import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; @@ -34,6 +35,7 @@ import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.tools.FunctionTool; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -50,15 +52,37 @@ public static Action getToolCallAction() throws Exception { List.of(ToolRequestEvent.EVENT_TYPE)); } - @SuppressWarnings("unchecked") public static void processToolRequest(Event event, RunnerContext ctx) { ToolRequestEvent toolRequest = ToolRequestEvent.fromEvent(event); boolean toolCallAsync = ctx.getConfig().get(AgentExecutionOptions.TOOL_CALL_ASYNC); + boolean toolCallParallel = ctx.getConfig().get(AgentExecutionOptions.TOOL_CALL_PARALLEL); Map success = new HashMap<>(); Map error = new HashMap<>(); Map responses = new HashMap<>(); Map externalIds = new HashMap<>(); + List executions = + buildToolCallExecutions(toolRequest, ctx, externalIds, success, error, responses); + + if (toolCallAsync && toolCallParallel && executions.size() > 1) { + executeParallel(executions, ctx, success, error, responses); + } else { + executeSequentially(executions, toolCallAsync, ctx, success, error, responses); + } + + ctx.sendEvent( + new ToolResponseEvent(toolRequest.getId(), responses, success, error, externalIds)); + } + + @SuppressWarnings("unchecked") + private static List buildToolCallExecutions( + ToolRequestEvent toolRequest, + RunnerContext ctx, + Map externalIds, + Map success, + Map error, + Map responses) { + List executions = new ArrayList<>(); for (Map toolCall : toolRequest.getToolCalls()) { String id = String.valueOf(toolCall.get("id")); Map function = (Map) toolCall.get("function"); @@ -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 callArguments = mergedArguments; - DurableCallable callable = - new DurableCallable<>() { - @Override - public String getId() { - return "tool-call"; - } - - @Override - public Class 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 callArguments = mergedArguments; + DurableCallable callable = + new DurableCallable<>() { + @Override + public String getId() { + return "tool-call-" + id; + } + + @Override + public Class 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 executions, + RunnerContext ctx, + Map success, + Map error, + Map responses) { + List> callables = new ArrayList<>(executions.size()); + for (ToolCallExecution execution : executions) { + callables.add(execution.callable); + } + try { + List> 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) { + for (ToolCallExecution execution : executions) { + recordExecutionException(execution, e, success, error, responses); + } + } + } + + private static void executeSequentially( + List executions, + boolean toolCallAsync, + RunnerContext ctx, + Map success, + Map error, + Map responses) { + for (ToolCallExecution execution : executions) { + try { + ToolResponse response = + toolCallAsync + ? ctx.durableExecuteAsync(execution.callable) + : ctx.durableExecute(execution.callable); + recordToolResponse(execution.id, response, success, error, responses); + } catch (Exception e) { + recordExecutionException(execution, e, success, error, responses); + } + } + } + + private static void recordOutcome( + ToolCallExecution execution, + Outcome outcome, + Map success, + Map error, + Map responses) { + if (outcome.isFailure()) { + recordExecutionException(execution, outcome.getError(), success, error, responses); + } else { + recordToolResponse(execution.id, outcome.getValue(), success, error, responses); + } + } + + private static void recordInlineResponse( + String id, + ToolResponse response, + String diagnosticError, + Map success, + Map error, + Map responses) { + recordToolResponse(id, response, success, error, responses); + if (diagnosticError != null) { + error.put(id, diagnosticError); + } + } + + private static void recordExecutionException( + ToolCallExecution execution, + Exception exception, + Map success, + Map error, + Map responses) { + success.put(execution.id, false); + responses.put( + execution.id, + ToolResponse.error(String.format("Tool %s execute failed.", execution.name))); + error.put(execution.id, exception.getMessage()); + } + + private static void recordToolResponse( + String id, + ToolResponse response, + Map success, + Map error, + Map responses) { + success.put(id, response.isSuccess()); + responses.put(id, response); + if (!response.isSuccess() && response.getError() != null) { + error.put(id, response.getError()); + } + } + + private static final class ToolCallExecution { + private final String id; + private final String name; + private final DurableCallable callable; + + private ToolCallExecution(String id, String name, DurableCallable callable) { + this.id = id; + this.name = name; + this.callable = callable; } - ctx.sendEvent( - new ToolResponseEvent(toolRequest.getId(), responses, success, error, externalIds)); } private static Map resolveInjectedArguments(Tool tool, RunnerContext ctx) diff --git a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java index 750f12adc..f8accc967 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionTest.java @@ -18,11 +18,13 @@ package org.apache.flink.agents.plan.actions; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; import org.apache.flink.agents.api.annotation.ToolParam; import org.apache.flink.agents.api.configuration.ReadableConfiguration; import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.MemoryRef; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; @@ -46,6 +48,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; import static org.assertj.core.api.Assertions.assertThat; @@ -76,7 +79,19 @@ public ToolResponse call(ToolParameters parameters) { static class FakeRunnerContext implements RunnerContext { private final List sentEvents = new ArrayList<>(); private final AgentConfiguration config = - new AgentConfiguration(Map.of("tenant_id", "tenant-1")); + new AgentConfiguration( + Map.of( + "tenant_id", + "tenant-1", + AgentExecutionOptions.TOOL_CALL_ASYNC.getKey(), + true, + AgentExecutionOptions.TOOL_CALL_PARALLEL.getKey(), + true)); + private final List durableExecuteIds = new ArrayList<>(); + private final List durableExecuteAsyncIds = new ArrayList<>(); + private final List> durableExecuteAllAsyncIds = new ArrayList<>(); + private final AtomicInteger queryOrderCalls = new AtomicInteger(); + private List> durableExecuteAllAsyncOutcomes; private ToolParameterInjection injection = ToolParameterInjection.fromConfig("tenant_id"); private MemoryObject sensoryMemory; private MemoryObject shortTermMemory; @@ -86,6 +101,21 @@ FakeRunnerContext withInjection(ToolParameterInjection injection) { return this; } + FakeRunnerContext withToolCallAsync(boolean enabled) { + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, enabled); + return this; + } + + FakeRunnerContext withToolCallParallel(boolean enabled) { + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, enabled); + return this; + } + + FakeRunnerContext withDurableExecuteAllAsyncOutcomes(List> outcomes) { + this.durableExecuteAllAsyncOutcomes = outcomes; + return this; + } + FakeRunnerContext withSensoryMemory(Map values) { this.sensoryMemory = new FakeMemoryObject(values); return this; @@ -156,14 +186,35 @@ public Object getActionConfigValue(String key) { @Override public T durableExecute(DurableCallable callable) throws Exception { + durableExecuteIds.add(callable.getId()); return callable.call(); } @Override public T durableExecuteAsync(DurableCallable callable) throws Exception { + durableExecuteAsyncIds.add(callable.getId()); return callable.call(); } + @Override + @SuppressWarnings("unchecked") + public List> durableExecuteAllAsync(List> callables) + throws Exception { + List ids = new ArrayList<>(); + for (DurableCallable callable : callables) { + ids.add(callable.getId()); + } + durableExecuteAllAsyncIds.add(ids); + if (durableExecuteAllAsyncOutcomes != null) { + return (List>) (List) durableExecuteAllAsyncOutcomes; + } + List> outcomes = new ArrayList<>(callables.size()); + for (DurableCallable callable : callables) { + outcomes.add(Outcome.success(callable.call())); + } + return outcomes; + } + @Override public void close() {} } @@ -378,20 +429,123 @@ public Resource getResource(String name, ResourceType type) { assertThat(response.getError()).containsEntry("call-1", "missing resource"); } + @Test + void processToolRequestUsesParallelBatchForMultipleTools() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(); + + ToolCallAction.processToolRequest(toolRequest("queryOrder", "call-1", "call-2"), ctx); + + assertThat(ctx.durableExecuteAllAsyncIds) + .containsExactly(List.of("tool-call-call-1", "tool-call-call-2")); + assertThat(ctx.durableExecuteAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteIds).isEmpty(); + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getResponses().get("call-1").getResult()).isEqualTo("tenant-1:order-1"); + assertThat(response.getResponses().get("call-2").getResult()).isEqualTo("tenant-1:order-2"); + } + + @Test + void processToolRequestUsesSerialAsyncWhenParallelDisabled() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext().withToolCallParallel(false); + + ToolCallAction.processToolRequest(toolRequest("queryOrder", "call-1", "call-2"), ctx); + + assertThat(ctx.durableExecuteAllAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteAsyncIds) + .containsExactly("tool-call-call-1", "tool-call-call-2"); + assertThat(ctx.durableExecuteIds).isEmpty(); + } + + @Test + void processToolRequestUsesSyncWhenAsyncDisabled() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext().withToolCallAsync(false); + + ToolCallAction.processToolRequest(toolRequest("queryOrder", "call-1", "call-2"), ctx); + + assertThat(ctx.durableExecuteAllAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteIds).containsExactly("tool-call-call-1", "tool-call-call-2"); + } + + @Test + void processToolRequestDoesNotBatchSingleTool() throws Exception { + FakeRunnerContext ctx = new FakeRunnerContext(); + + ToolCallAction.processToolRequest(toolRequest("queryOrder"), ctx); + + assertThat(ctx.durableExecuteAllAsyncIds).isEmpty(); + assertThat(ctx.durableExecuteAsyncIds).containsExactly("tool-call-call-1"); + assertThat(ctx.durableExecuteIds).isEmpty(); + } + + @Test + void processToolRequestExcludesMissingToolFromParallelBatch() throws Exception { + FakeRunnerContext ctx = + new FakeRunnerContext() { + @Override + public Resource getResource(String name, ResourceType type) throws Exception { + if ("missingTool".equals(name)) { + throw new RuntimeException("missing resource"); + } + return super.getResource(name, type); + } + }; + + ToolCallAction.processToolRequest( + new ToolRequestEvent( + "model", + List.of( + toolCall("missingTool", "missing-call", "order-0"), + toolCall("queryOrder", "call-1", "order-1"), + toolCall("queryOrder", "call-2", "order-2"))), + ctx); + + assertThat(ctx.durableExecuteAllAsyncIds) + .containsExactly(List.of("tool-call-call-1", "tool-call-call-2")); + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("missing-call", false); + assertThat(response.getError()).containsEntry("missing-call", "missing resource"); + } + + @Test + void processToolRequestRecordsOutcomeFailureAsToolError() throws Exception { + FakeRunnerContext ctx = + new FakeRunnerContext() + .withDurableExecuteAllAsyncOutcomes( + List.of( + Outcome.success(ToolResponse.success("ok")), + Outcome.failure(new RuntimeException("boom")))); + + ToolCallAction.processToolRequest(toolRequest("queryOrder", "call-1", "call-2"), ctx); + + ToolResponseEvent response = ToolResponseEvent.fromEvent(ctx.sentEvents.get(0)); + assertThat(response.getSuccess()).containsEntry("call-1", true); + assertThat(response.getResponses().get("call-1").getResult()).isEqualTo("ok"); + assertThat(response.getSuccess()).containsEntry("call-2", false); + assertThat(response.getResponses().get("call-2").getError()) + .isEqualTo("Tool queryOrder execute failed."); + assertThat(response.getError()).containsEntry("call-2", "boom"); + } + private static ToolRequestEvent toolRequest(String toolName) { - return new ToolRequestEvent( - "model", - List.of( - Map.of( - "id", - "call-1", - "type", - "function", - "function", - Map.of( - "name", - toolName, - "arguments", - Map.of("orderId", "order-1"))))); + return new ToolRequestEvent("model", List.of(toolCall(toolName, "call-1", "order-1"))); + } + + private static ToolRequestEvent toolRequest(String toolName, String... ids) { + List> toolCalls = new ArrayList<>(); + for (int i = 0; i < ids.length; i++) { + toolCalls.add(toolCall(toolName, ids[i], "order-" + (i + 1))); + } + return new ToolRequestEvent("model", toolCalls); + } + + private static Map toolCall(String toolName, String id, String orderId) { + return Map.of( + "id", + id, + "type", + "function", + "function", + Map.of("name", toolName, "arguments", Map.of("orderId", orderId))); } } diff --git a/python/flink_agents/api/core_options.py b/python/flink_agents/api/core_options.py index d196777c4..142f0c383 100644 --- a/python/flink_agents/api/core_options.py +++ b/python/flink_agents/api/core_options.py @@ -254,6 +254,24 @@ class AgentExecutionOptions: default=True, ) + TOOL_CALL_PARALLEL = ConfigOption( + key="tool-call.parallel", + config_type=bool, + default=True, + ) + + TOOL_CALL_NUM_ASYNC_THREADS = ConfigOption( + key="tool-call.num-async-threads", + config_type=int, + default=os.cpu_count() * 2, + ) + + TOOL_CALL_BATCH_TIMEOUT = ConfigOption( + key="tool-call.batch.timeout", + config_type=int, + default=-1, + ) + RAG_ASYNC = ConfigOption( key="rag.async", config_type=bool, diff --git a/python/flink_agents/api/runner_context.py b/python/flink_agents/api/runner_context.py index e110bd2f9..70796de79 100644 --- a/python/flink_agents/api/runner_context.py +++ b/python/flink_agents/api/runner_context.py @@ -16,6 +16,7 @@ # limitations under the License. ################################################################################# from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable, Dict from flink_agents.api.configuration import ReadableConfiguration @@ -24,12 +25,49 @@ from flink_agents.api.metric_group import MetricGroup from flink_agents.api.resource import Resource, ResourceType -__all__ = ["AsyncExecutionResult", "RunnerContext"] +__all__ = ["AsyncExecutionResult", "DurableCall", "Outcome", "RunnerContext"] if TYPE_CHECKING: from flink_agents.api.memory_object import MemoryObject +@dataclass(frozen=True) +class DurableCall: + """A deterministic durable call entry for batch execution.""" + + id: str + func: Callable[..., Any] + args: tuple[Any, ...] = () + kwargs: dict[str, Any] | None = None + reconciler: Callable[[], Any] | None = None + + +@dataclass(frozen=True) +class Outcome: + """Result or failure for one durable batch slot.""" + + value: Any = None + error: BaseException | None = None + + @classmethod + def success(cls, value: Any) -> "Outcome": + """Create a successful outcome.""" + return cls(value=value) + + @classmethod + def failure(cls, error: BaseException) -> "Outcome": + """Create a failed outcome.""" + return cls(error=error) + + def is_success(self) -> bool: + """Return whether this outcome is successful.""" + return self.error is None + + def is_failure(self) -> bool: + """Return whether this outcome is failed.""" + return self.error is not None + + class AsyncExecutionResult: """This class wraps an asynchronous task that will be submitted to a thread pool only when awaited. This ensures lazy submission and serial execution semantics. @@ -312,6 +350,18 @@ async def my_action(event, ctx): An awaitable object that yields the function result when awaited. """ + @abstractmethod + def durable_execute_all_async( + self, + callables: list[DurableCall], + ) -> "AsyncExecutionResult": + """Execute multiple durable callables as one asynchronous durable batch. + + The returned awaitable resolves to a list of ``Outcome`` values in the + same order as the input calls. Individual callable exceptions are returned + as failure outcomes instead of failing the whole batch. + """ + @property @abstractmethod def config(self) -> ReadableConfiguration: diff --git a/python/flink_agents/api/tests/test_core_options.py b/python/flink_agents/api/tests/test_core_options.py index f1e773a83..2abbedf92 100644 --- a/python/flink_agents/api/tests/test_core_options.py +++ b/python/flink_agents/api/tests/test_core_options.py @@ -17,6 +17,7 @@ ################################################################################# import ast import importlib +import os import sys import types from pathlib import Path @@ -107,6 +108,20 @@ def test_agent_config_options_are_explicitly_declared() -> None: assert options["EVENT_LOG_LEVEL"].get_default_value() is EventLogLevel.STANDARD +def test_agent_execution_options_include_parallel_tool_call_options() -> None: + from flink_agents.api.core_options import AgentExecutionOptions + + options = _collect_config_options(AgentExecutionOptions) + assert options["TOOL_CALL_ASYNC"].get_key() == "tool-call.async" + assert options["TOOL_CALL_ASYNC"].get_default_value() is True + assert options["TOOL_CALL_PARALLEL"].get_key() == "tool-call.parallel" + assert options["TOOL_CALL_PARALLEL"].get_default_value() is True + assert options["TOOL_CALL_NUM_ASYNC_THREADS"].get_key() == "tool-call.num-async-threads" + assert options["TOOL_CALL_NUM_ASYNC_THREADS"].get_default_value() == os.cpu_count() * 2 + assert options["TOOL_CALL_BATCH_TIMEOUT"].get_key() == "tool-call.batch.timeout" + assert options["TOOL_CALL_BATCH_TIMEOUT"].get_default_value() == -1 + + def test_unknown_agent_config_option_raises_attribute_error() -> None: from flink_agents.api.core_options import AgentConfigOptions diff --git a/python/flink_agents/plan/actions/tool_call_action.py b/python/flink_agents/plan/actions/tool_call_action.py index f605c4807..5651f453f 100644 --- a/python/flink_agents/plan/actions/tool_call_action.py +++ b/python/flink_agents/plan/actions/tool_call_action.py @@ -16,13 +16,14 @@ # limitations under the License. ################################################################################# import logging +from dataclasses import dataclass from flink_agents.api.core_options import AgentExecutionOptions from flink_agents.api.events.event import Event from flink_agents.api.events.tool_event import ToolRequestEvent, ToolResponseEvent from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType -from flink_agents.api.runner_context import RunnerContext +from flink_agents.api.runner_context import DurableCall, Outcome, RunnerContext from flink_agents.api.tools.tool_parameter_injection import ( InjectedArg, ToolParameterSource, @@ -34,10 +35,18 @@ _logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class _ToolCallExecution: + id: str + name: str + durable_call: DurableCall + + async def process_tool_request(event: Event, ctx: RunnerContext) -> None: """Built-in action for processing tool call requests.""" event = ToolRequestEvent.from_event(event) tool_call_async = ctx.config.get(AgentExecutionOptions.TOOL_CALL_ASYNC) + tool_call_parallel = ctx.config.get(AgentExecutionOptions.TOOL_CALL_PARALLEL) if tool_call_async: # To avoid https://github.com/alibaba/pemja/issues/88, we log a message here. @@ -47,11 +56,53 @@ async def process_tool_request(event: Event, ctx: RunnerContext) -> None: success = {} error = {} external_ids = {} + executions = _build_tool_call_executions( + event, + ctx, + responses, + success, + error, + external_ids, + ) + + if tool_call_async and tool_call_parallel and len(executions) > 1: + await _execute_parallel(executions, ctx, responses, success, error) + else: + await _execute_sequentially( + executions, + tool_call_async=tool_call_async, + ctx=ctx, + responses=responses, + success=success, + error=error, + ) + + ctx.send_event( + ToolResponseEvent( + request_id=event.id, + responses=responses, + external_ids=external_ids, + success=success, + error=error, + ) + ) + + +def _build_tool_call_executions( + event: ToolRequestEvent, + ctx: RunnerContext, + responses: dict, + success: dict, + error: dict, + external_ids: dict, +) -> list[_ToolCallExecution]: + executions = [] for tool_call in event.tool_calls: call_id = tool_call["id"] name = tool_call["function"]["name"] kwargs = tool_call["function"]["arguments"] external_id = tool_call.get("original_id") + external_ids[call_id] = external_id try: tool = ctx.get_resource(name, ResourceType.TOOL) @@ -62,37 +113,105 @@ async def process_tool_request(event: Event, ctx: RunnerContext) -> None: responses[call_id] = f"Tool `{name}` does not exist." success[call_id] = False error.setdefault(call_id, f"Tool `{name}` does not exist.") - external_ids[call_id] = external_id continue - else: - try: - call_kwargs = dict(kwargs or {}) - # Framework-owned injected args must win over model-provided values so - # hidden context such as tenant ids cannot be spoofed by tool calls. - call_kwargs.update(_resolve_injected_arguments(tool, ctx)) - if tool_call_async: - response = await ctx.durable_execute_async( - tool.call, **call_kwargs - ) - else: - response = ctx.durable_execute(tool.call, **call_kwargs) - responses[call_id] = response - success[call_id] = True - except Exception as e: - responses[call_id] = f"Tool `{name}` execute failed." - success[call_id] = False - error[call_id] = str(e) - external_ids[call_id] = external_id - ctx.send_event( - ToolResponseEvent( - request_id=event.id, - responses=responses, - external_ids=external_ids, - success=success, - error=error, + try: + call_kwargs = dict(kwargs or {}) + # Framework-owned injected args must win over model-provided values so + # hidden context such as tenant ids cannot be spoofed by tool calls. + call_kwargs.update(_resolve_injected_arguments(tool, ctx)) + except Exception as e: + responses[call_id] = f"Tool `{name}` execute failed." + success[call_id] = False + error[call_id] = str(e) + continue + + executions.append( + _ToolCallExecution( + id=call_id, + name=name, + durable_call=DurableCall( + id=f"tool-call-{call_id}", + func=tool.call, + kwargs=call_kwargs, + ), + ) ) - ) + return executions + + +async def _execute_parallel( + executions: list[_ToolCallExecution], + ctx: RunnerContext, + responses: dict, + success: dict, + error: dict, +) -> None: + try: + outcomes = await ctx.durable_execute_all_async( + [execution.durable_call for execution in executions] + ) + for execution, outcome in zip(executions, outcomes, strict=True): + _record_outcome(execution, outcome, responses, success, error) + except Exception as e: + for execution in executions: + _record_execution_exception(execution, e, responses, success, error) + + +async def _execute_sequentially( + executions: list[_ToolCallExecution], + *, + tool_call_async: bool, + ctx: RunnerContext, + responses: dict, + success: dict, + error: dict, +) -> None: + for execution in executions: + try: + call = execution.durable_call + if tool_call_async: + response = await ctx.durable_execute_async( + call.func, + *call.args, + **(call.kwargs or {}), + ) + else: + response = ctx.durable_execute( + call.func, + *call.args, + **(call.kwargs or {}), + ) + responses[execution.id] = response + success[execution.id] = True + except Exception as e: # noqa: PERF203 + _record_execution_exception(execution, e, responses, success, error) + + +def _record_outcome( + execution: _ToolCallExecution, + outcome: Outcome, + responses: dict, + success: dict, + error: dict, +) -> None: + if outcome.is_failure(): + _record_execution_exception(execution, outcome.error, responses, success, error) + else: + responses[execution.id] = outcome.value + success[execution.id] = True + + +def _record_execution_exception( + execution: _ToolCallExecution, + exception: BaseException, + responses: dict, + success: dict, + error: dict, +) -> None: + responses[execution.id] = f"Tool `{execution.name}` execute failed." + success[execution.id] = False + error[execution.id] = str(exception) def _resolve_injected_arguments(tool: object, ctx: RunnerContext) -> dict: diff --git a/python/flink_agents/plan/tests/actions/test_tool_call_action.py b/python/flink_agents/plan/tests/actions/test_tool_call_action.py index d9690981a..b329e853f 100644 --- a/python/flink_agents/plan/tests/actions/test_tool_call_action.py +++ b/python/flink_agents/plan/tests/actions/test_tool_call_action.py @@ -22,6 +22,7 @@ from flink_agents.api.events.tool_event import ToolRequestEvent, ToolResponseEvent from flink_agents.api.memory_object import MemoryObject from flink_agents.api.resource import ResourceType +from flink_agents.api.runner_context import Outcome from flink_agents.api.tools import InjectedArg from flink_agents.plan.actions.tool_call_action import process_tool_request from flink_agents.plan.configuration import AgentConfiguration @@ -41,29 +42,48 @@ def __init__( sensory_memory: Any | None = None, short_term_memory: Any | None = None, ) -> None: - self.config = config or AgentConfiguration({"tenant_id": "tenant-1"}) - if isinstance(self.config, AgentConfiguration): + if config is None: + self.config = AgentConfiguration({"tenant_id": "tenant-1"}) self.config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, False) + else: + self.config = config self.injected_args = injected_args or { "tenant_id": InjectedArg.from_config("tenant_id") } self.sensory_memory = sensory_memory self.short_term_memory = short_term_memory self.sent_events = [] + self.durable_execute_calls = [] + self.durable_execute_async_calls = [] + self.durable_execute_all_async_calls = [] + self.durable_execute_all_async_outcomes = None def get_resource(self, name: str, type: ResourceType) -> FunctionTool: - assert name == "query_order" assert type == ResourceType.TOOL + if name != "query_order": + msg = f"Tool `{name}` does not exist." + raise ValueError(msg) return FunctionTool( func=PythonFunction.from_callable(query_order), injected_args=self.injected_args, ) - def durable_execute(self, func: Any, **kwargs: Any) -> Any: - return func(**kwargs) + def durable_execute(self, func: Any, *args: Any, **kwargs: Any) -> Any: + self.durable_execute_calls.append((func, args, kwargs)) + return func(*args, **kwargs) - async def durable_execute_async(self, func: Any, **kwargs: Any) -> Any: - return func(**kwargs) + async def durable_execute_async(self, func: Any, *args: Any, **kwargs: Any) -> Any: + self.durable_execute_async_calls.append((func, args, kwargs)) + return func(*args, **kwargs) + + async def durable_execute_all_async(self, callables: list[Any]) -> list[Outcome]: + self.durable_execute_all_async_calls.append(callables) + if self.durable_execute_all_async_outcomes is not None: + return self.durable_execute_all_async_outcomes + return [ + Outcome.success(call.func(*call.args, **(call.kwargs or {}))) + for call in callables + ] def send_event(self, event: Any) -> None: self.sent_events.append(event) @@ -102,7 +122,10 @@ def get_fields(self) -> dict[str, Any]: class _WrongConfig: def get(self, option: Any) -> bool: - assert option == AgentExecutionOptions.TOOL_CALL_ASYNC + assert option in ( + AgentExecutionOptions.TOOL_CALL_ASYNC, + AgentExecutionOptions.TOOL_CALL_PARALLEL, + ) return False @@ -277,14 +300,105 @@ def test_tool_call_action_uses_sync_execution_in_test_context() -> None: assert ctx.config.get(AgentExecutionOptions.TOOL_CALL_ASYNC) is False -def tool_request() -> ToolRequestEvent: +def test_tool_call_action_uses_parallel_batch_for_multiple_tools() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, True) + ctx = _Context(config=config) + + asyncio.run(process_tool_request(tool_request("call-1", "call-2"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.responses == { + "call-1": "tenant-1:order-call-1", + "call-2": "tenant-1:order-call-2", + } + assert response.success == {"call-1": True, "call-2": True} + assert len(ctx.durable_execute_all_async_calls) == 1 + assert [call.id for call in ctx.durable_execute_all_async_calls[0]] == [ + "tool-call-call-1", + "tool-call-call-2", + ] + assert ctx.durable_execute_async_calls == [] + + +def test_tool_call_action_uses_serial_async_when_parallel_disabled() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, False) + ctx = _Context(config=config) + + asyncio.run(process_tool_request(tool_request("call-1", "call-2"), ctx)) + + assert ctx.durable_execute_all_async_calls == [] + assert len(ctx.durable_execute_async_calls) == 2 + + +def test_tool_call_action_does_not_batch_single_tool() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, True) + ctx = _Context(config=config) + + asyncio.run(process_tool_request(tool_request("call-1"), ctx)) + + assert ctx.durable_execute_all_async_calls == [] + assert len(ctx.durable_execute_async_calls) == 1 + + +def test_tool_call_action_excludes_missing_tool_from_parallel_batch() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, True) + ctx = _Context(config=config) + + asyncio.run(process_tool_request(tool_request("call-1", "missing"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.success["call-1"] is True + assert response.success["missing"] is False + assert len(ctx.durable_execute_async_calls) == 1 + assert ctx.durable_execute_all_async_calls == [] + + +def test_tool_call_action_records_parallel_outcome_failure() -> None: + config = AgentConfiguration({"tenant_id": "tenant-1"}) + config.set(AgentExecutionOptions.TOOL_CALL_ASYNC, True) + config.set(AgentExecutionOptions.TOOL_CALL_PARALLEL, True) + ctx = _Context(config=config) + ctx.durable_execute_all_async_outcomes = [ + Outcome.success("ok"), + Outcome.failure(ValueError("boom")), + ] + + asyncio.run(process_tool_request(tool_request("call-1", "call-2"), ctx)) + + response = ToolResponseEvent.from_event(ctx.sent_events[0]) + assert response.responses["call-1"] == "ok" + assert response.success["call-1"] is True + assert response.responses["call-2"] == "Tool `query_order` execute failed." + assert response.success["call-2"] is False + assert response.error["call-2"] == "boom" + + +def tool_request(*call_ids: str) -> ToolRequestEvent: + if not call_ids: + call_ids = ("call-1",) return ToolRequestEvent( model="model", tool_calls=[ { - "id": "call-1", + "id": call_id, "type": "function", - "function": {"name": "query_order", "arguments": {"order_id": "order-1"}}, + "function": { + "name": "missing_tool" if call_id == "missing" else "query_order", + "arguments": { + "order_id": "order-1" + if len(call_ids) == 1 + else f"order-{call_id}" + }, + }, } + for call_id in call_ids ], ) diff --git a/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py b/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py index 0e2762d73..8a7eb712a 100644 --- a/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py +++ b/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py @@ -42,6 +42,7 @@ "java.lang.Float": float, "java.lang.Double": float, "java.util.List": list, + "java.time.Duration": int, } @@ -84,11 +85,20 @@ def _java_type_matches_python(java_type_name: str, python_config_type: type) -> return python_config_type.__name__ == java_simple_name -def normalize_java_default(java_default: Any, python_config_type: type) -> Any: +def _java_duration_to_millis(java_duration: Any) -> int: + return int(java_duration.toMillis()) + + +def normalize_java_default( + java_default: Any, java_type_name: str, python_config_type: type +) -> Any: """Convert a Java default value into a Python-comparable form.""" if java_default is None: return None + if java_type_name == "java.time.Duration": + return _java_duration_to_millis(java_default) + if hasattr(java_default, "name") and callable(java_default.name): enum_name = java_default.name() if isinstance(python_config_type, type) and issubclass(python_config_type, Enum): @@ -122,6 +132,7 @@ def assert_python_option_matches_java( python_default = python_option.get_default_value() java_default = normalize_java_default( java_option.getDefaultValue(), + java_type_name, python_config_type, ) assert python_default == java_default, ( diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index ada5b8e30..b3ff7bea4 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -17,6 +17,7 @@ ################################################################################# import logging import os +import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from functools import partial @@ -26,6 +27,7 @@ from typing_extensions import override from flink_agents.api.configuration import ReadableConfiguration +from flink_agents.api.core_options import AgentExecutionOptions from flink_agents.api.events.event import Event from flink_agents.api.memory.long_term_memory import ( BaseLongTermMemory, @@ -36,6 +38,8 @@ from flink_agents.api.resource import Resource, ResourceType from flink_agents.api.runner_context import ( AsyncExecutionResult, + DurableCall, + Outcome, RunnerContext, ) from flink_agents.runtime.durable_execution import ( @@ -73,6 +77,14 @@ class _ReconcilerExecutionPlan: needs_append_pending: bool = False +@dataclass(frozen=True) +class _BatchExecutionPlan: + outcomes: list[Outcome] + suppliers: list[tuple[int, Callable[[], Any]]] + needs_reservation: bool = False + execution_start: int = -1 + + class _DurableExecutionResult: """Wrapper that holds result and triggers recording when unwrapped.""" @@ -241,6 +253,58 @@ 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) + futures = [ + self._ctx.tool_call_executor.submit(supplier) for _, supplier in plan.suppliers + ] + timeout_ms = self._ctx.config.get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT) + deadline = time.monotonic() + timeout_ms / 1000 if timeout_ms >= 0 else None + while any(not future.done() for future in futures): + if deadline is not None and time.monotonic() >= deadline: + exception = TimeoutError( + f"Async durable batch execution timed out after {timeout_ms} ms" + ) + executed = _collect_outcomes_on_timeout(futures, exception) + return self._ctx._finalize_batch_execution(self._calls, plan, executed) + yield + + executed = _collect_outcomes(futures) + return self._ctx._finalize_batch_execution(self._calls, plan, executed) + + +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( + futures: list[Any], timeout_exception: TimeoutError +) -> list[Outcome]: + outcomes = [] + for future in futures: + 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 + + class FlinkRunnerContext(RunnerContext): """Providing context for agent execution in Flink Environment. @@ -256,6 +320,7 @@ def __init__( j_runner_context: Any, agent_plan_json: str, executor: ThreadPoolExecutor, + tool_call_executor: ThreadPoolExecutor, j_resource_adapter: Any, ) -> None: """Initialize a flink runner context with the given java runner context. @@ -273,7 +338,9 @@ def __init__( self.__agent_plan.resource_providers, self.__agent_plan.config ) self.__resource_cache.set_java_resource_adapter(j_resource_adapter) + self.__config = self.__agent_plan.config self.executor = executor + self.tool_call_executor = tool_call_executor def set_long_term_memory(self, ltm: InternalBaseLongTermMemory) -> None: """Set long term memory instance to this context. @@ -490,6 +557,24 @@ def _serialize_call_payloads( exception_payload = cloudpickle.dumps(exception) if exception else None return result_payload, exception_payload + def _read_call_result_at(self, index: int) -> _PersistedCallResult | None: + current = self._j_runner_context.getCallResultFieldsAt(index) + if current is None: + return None + + function_id, args_digest, status, result_payload, exception_payload = current + return _PersistedCallResult( + function_id=function_id, + args_digest=args_digest, + status=status, + result_payload=bytes(result_payload) + if result_payload is not None + else None, + exception_payload=( + bytes(exception_payload) if exception_payload is not None else None + ), + ) + def _peek_current_call_result(self) -> _PersistedCallResult | None: current = self._j_runner_context.getCurrentCallResultFields() if current is None: @@ -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 = "" + base = self._j_runner_context.getCurrentCallIndex() + outcomes: list[Outcome | None] = [] + suppliers: list[tuple[int, Callable[[], Any]]] = [] + needs_reservation = False + execution_start = -1 + + for index, call in enumerate(calls): + current = self._read_call_result_at(base + index) + if current is None: + needs_reservation = True + if execution_start < 0: + execution_start = index + outcomes.append(None) + suppliers.append((index, self._callable_for_durable_call(call))) + continue + + if not self._call_matches(current, call, args_digest): + self._j_runner_context.clearCallResultsFromAndPersist(base + index) + needs_reservation = True + execution_start = index + outcomes.append(None) + suppliers.append((index, self._callable_for_durable_call(call))) + for remaining_index in range(index + 1, len(calls)): + outcomes.append(None) + suppliers.append( + ( + remaining_index, + self._callable_for_durable_call(calls[remaining_index]), + ) + ) + break + + if current.status == "PENDING": + outcomes.append(None) + suppliers.append( + ( + index, + call.reconciler or self._callable_for_durable_call(call), + ) + ) + else: + outcomes.append(self._read_terminal_outcome(current)) + + if needs_reservation: + function_ids = [call.id for call in calls[execution_start:]] + self._j_runner_context.reservePendingBatch(function_ids, args_digest) + + return _BatchExecutionPlan( + outcomes=outcomes, + suppliers=suppliers, + needs_reservation=needs_reservation, + execution_start=execution_start, + ) + + def _finalize_batch_execution( + self, + calls: list[DurableCall], + plan: _BatchExecutionPlan, + executed: list[Outcome], + ) -> list[Outcome]: + args_digest = "" + base = self._j_runner_context.getCurrentCallIndex() + outcomes = list(plan.outcomes) + for (call_index, _), outcome in zip(plan.suppliers, executed, strict=True): + result_payload, exception_payload = self._serialize_call_payloads( + outcome.value, + outcome.error, + ) + self._j_runner_context.finalizeCallAt( + base + call_index, + calls[call_index].id, + args_digest, + result_payload, + exception_payload, + ) + outcomes[call_index] = outcome + self._j_runner_context.advanceCallIndexBy(len(calls)) + return outcomes + + @override + def durable_execute_all_async( + self, + callables: list[DurableCall], + ) -> AsyncExecutionResult: + return _DurableBatchAsyncExecutionResult(self, callables) + @override def durable_execute( self, @@ -748,6 +937,8 @@ def config(self) -> ReadableConfiguration: ReadableConfiguration The configuration for flink agents. """ + if hasattr(self, "_FlinkRunnerContext__config"): + return self.__config return self.__agent_plan.config @override @@ -766,12 +957,13 @@ def create_flink_runner_context( j_runner_context: Any, agent_plan_json: str, executor: ThreadPoolExecutor, + tool_call_executor: ThreadPoolExecutor, j_resource_adapter: Any, job_identifier: str, ) -> FlinkRunnerContext: """Used to create a FlinkRunnerContext Python object in Pemja environment.""" ctx = FlinkRunnerContext( - j_runner_context, agent_plan_json, executor, j_resource_adapter + j_runner_context, agent_plan_json, executor, tool_call_executor, j_resource_adapter ) ltm = _init_long_term_memory(ctx, job_identifier) if ltm is not None: diff --git a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py index a41a33c7b..59f5283c1 100644 --- a/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py +++ b/python/flink_agents/runtime/tests/test_flink_runner_context_reconcilable.py @@ -15,7 +15,7 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# -import asyncio +import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from typing import Any, Callable @@ -23,6 +23,8 @@ import cloudpickle import pytest +from flink_agents.api.runner_context import DurableCall +from flink_agents.plan.configuration import AgentConfiguration from flink_agents.runtime.durable_execution import ( _compute_args_digest, _compute_function_id, @@ -45,6 +47,22 @@ def __init__(self) -> None: self.current_call_index = 0 self.operations: list[str] = [] + def getCurrentCallIndex(self) -> int: + return self.current_call_index + + def getCallResultFieldsAt(self, index: int) -> list[Any] | None: + self.operations.append(f"read_at:{index}") + if index < len(self.call_results): + current = self.call_results[index] + return [ + current.function_id, + current.args_digest, + current.status, + current.result_payload, + current.exception_payload, + ] + return None + def getCurrentCallResultFields(self) -> list[Any] | None: self.operations.append("peek") if self.current_call_index < len(self.call_results): @@ -128,27 +146,76 @@ def clearCallResultsFromCurrentIndexAndPersist(self) -> None: self.operations.append("clear") self.call_results = self.call_results[: self.current_call_index] + def reservePendingBatch(self, function_ids: list[str], args_digest: str) -> None: + self.operations.append(f"reserve:{len(function_ids)}") + for function_id in function_ids: + self.call_results.append( + _StoredCallResult( + function_id=function_id, + args_digest=args_digest, + status="PENDING", + ) + ) + + def finalizeCallAt( + self, + index: int, + function_id: str, + args_digest: str, + result_payload: bytes | None, + exception_payload: bytes | None, + ) -> None: + self.operations.append(f"finalize_at:{index}") + current = self.call_results[index] + assert current.status == "PENDING" + assert current.function_id == function_id + assert current.args_digest == args_digest + self.call_results[index] = _StoredCallResult( + function_id=function_id, + args_digest=args_digest, + status="FAILED" if exception_payload is not None else "SUCCEEDED", + result_payload=result_payload, + exception_payload=exception_payload, + ) + + def clearCallResultsFromAndPersist(self, index: int) -> None: + self.operations.append(f"clear_at:{index}") + self.call_results = self.call_results[:index] + + def advanceCallIndexBy(self, count: int) -> None: + self.operations.append(f"advance:{count}") + self.current_call_index += count + def _create_runner_context( j_runner_context: _FakeJavaRunnerContext, + config: AgentConfiguration | None = None, + tool_call_workers: int = 2, ) -> FlinkRunnerContext: ctx = FlinkRunnerContext.__new__(FlinkRunnerContext) ctx._j_runner_context = j_runner_context - ctx.executor = ThreadPoolExecutor(max_workers=1) + ctx.executor = ThreadPoolExecutor(max_workers=2) + ctx.tool_call_executor = ThreadPoolExecutor(max_workers=tool_call_workers) ctx._FlinkRunnerContext__agent_plan = None ctx._FlinkRunnerContext__ltm = None + ctx._FlinkRunnerContext__config = config or AgentConfiguration({}) return ctx def _close_runner_context(ctx: FlinkRunnerContext) -> None: ctx.executor.shutdown(wait=True) + ctx.tool_call_executor.shutdown(wait=True) def _run_async(result: Any) -> object: - async def _await_result() -> Any: - return await result - - return asyncio.run(_await_result()) + iterator = result.__await__() + value = None + while True: + try: + iterator.send(value) + value = None + except StopIteration as e: # noqa: PERF203 + return e.value def _preload_pending( @@ -415,3 +482,196 @@ 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": -1}) + ctx = _create_runner_context(j_runner_context, config=config, tool_call_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( + [ + DurableCall("tool-call-1", slow_call, args=("one",)), + DurableCall("tool-call-2", slow_call, args=("two",)), + DurableCall("tool-call-3", slow_call, args=("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 + assert [result.status for result in j_runner_context.call_results] == [ + "SUCCEEDED", + "SUCCEEDED", + "SUCCEEDED", + ] + assert j_runner_context.current_call_index == 3 + + +def test_flink_runner_context_durable_execute_all_async_initial_batch() -> None: + j_runner_context = _FakeJavaRunnerContext() + ctx = _create_runner_context(j_runner_context) + + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [ + DurableCall("tool-call-1", _call_value, args=("one",)), + DurableCall("tool-call-2", _call_value, args=("two",)), + ] + ) + ) + finally: + _close_runner_context(ctx) + + assert [outcome.value for outcome in outcomes] == ["call:one", "call:two"] + assert [result.status for result in j_runner_context.call_results] == [ + "SUCCEEDED", + "SUCCEEDED", + ] + assert j_runner_context.current_call_index == 2 + assert "reserve:2" in j_runner_context.operations + assert "finalize_at:0" in j_runner_context.operations + assert "finalize_at:1" in j_runner_context.operations + + +def test_flink_runner_context_durable_execute_all_async_recovers_partial_batch() -> None: + j_runner_context = _FakeJavaRunnerContext() + j_runner_context.call_results.extend( + [ + _StoredCallResult( + function_id="tool-call-1", + args_digest="", + status="SUCCEEDED", + result_payload=cloudpickle.dumps("cached-one"), + ), + _StoredCallResult( + function_id="tool-call-2", + args_digest="", + status="PENDING", + ), + ] + ) + call_count = 0 + + def tracked_call(value: str) -> str: + nonlocal call_count + call_count += 1 + return _call_value(value) + + ctx = _create_runner_context(j_runner_context) + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [ + DurableCall("tool-call-1", tracked_call, args=("one",)), + DurableCall("tool-call-2", tracked_call, args=("two",)), + ] + ) + ) + finally: + _close_runner_context(ctx) + + assert [outcome.value for outcome in outcomes] == ["cached-one", "call:two"] + assert call_count == 1 + assert j_runner_context.call_results[1].status == "SUCCEEDED" + assert j_runner_context.current_call_index == 2 + + +def test_flink_runner_context_durable_execute_all_async_returns_cached_failure() -> None: + j_runner_context = _FakeJavaRunnerContext() + j_runner_context.call_results.append( + _StoredCallResult( + function_id="tool-call-1", + args_digest="", + status="FAILED", + exception_payload=cloudpickle.dumps(ValueError("cached failure")), + ) + ) + ctx = _create_runner_context(j_runner_context) + + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [DurableCall("tool-call-1", _call_value, args=("one",))] + ) + ) + finally: + _close_runner_context(ctx) + + assert outcomes[0].is_failure() + assert isinstance(outcomes[0].error, ValueError) + assert str(outcomes[0].error) == "cached failure" + assert j_runner_context.current_call_index == 1 + + +def test_flink_runner_context_durable_execute_all_async_collects_failures() -> None: + j_runner_context = _FakeJavaRunnerContext() + ctx = _create_runner_context(j_runner_context) + + def fail_call() -> str: + msg = "failed" + raise ValueError(msg) + + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [ + DurableCall("tool-call-1", _call_value, args=("one",)), + DurableCall("tool-call-2", fail_call), + ] + ) + ) + finally: + _close_runner_context(ctx) + + assert outcomes[0].value == "call:one" + assert outcomes[1].is_failure() + assert isinstance(outcomes[1].error, ValueError) + assert [result.status for result in j_runner_context.call_results] == [ + "SUCCEEDED", + "FAILED", + ] + assert j_runner_context.current_call_index == 2 + + +def test_flink_runner_context_durable_execute_all_async_timeout_keeps_completed_results() -> None: + j_runner_context = _FakeJavaRunnerContext() + config = AgentConfiguration({"tool-call.batch.timeout": 10}) + ctx = _create_runner_context(j_runner_context, config=config) + + def slow_call() -> str: + time.sleep(0.05) + return "slow" + + try: + outcomes = _run_async( + ctx.durable_execute_all_async( + [ + DurableCall("tool-call-1", lambda: "fast"), + DurableCall("tool-call-2", slow_call), + ] + ) + ) + finally: + _close_runner_context(ctx) + + assert outcomes[0].value == "fast" + assert outcomes[1].is_failure() + assert isinstance(outcomes[1].error, TimeoutError) + assert [result.status for result in j_runner_context.call_results] == [ + "SUCCEEDED", + "FAILED", + ] + assert j_runner_context.current_call_index == 2 diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java index 7c1a526d2..54ccbbd0d 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java @@ -17,7 +17,13 @@ */ package org.apache.flink.agents.runtime.async; +import org.apache.flink.agents.api.context.Outcome; + +import java.time.Duration; +import java.util.List; +import java.util.concurrent.Callable; import java.util.function.Supplier; +import java.util.stream.Collectors; /** * Executor for Java actions that supports asynchronous execution. @@ -56,6 +62,30 @@ public T executeAsync(ContinuationContext context, Supplier supplier) { return supplier.get(); } + /** + * Executes all suppliers as one batch. In JDK 11, this falls back to serial execution and + * captures each supplier's success or failure as an {@link Outcome}. + * + * @param context the continuation context + * @param suppliers the suppliers to execute + * @param timeout ignored in the JDK 11 fallback + * @param the result type + * @return outcomes in supplier order + */ + public List> executeAllAsync( + ContinuationContext context, List> suppliers, Duration timeout) { + return suppliers.stream() + .map( + supplier -> { + try { + return Outcome.success(supplier.call()); + } catch (Exception e) { + return Outcome.failure(e); + } + }) + .collect(Collectors.toList()); + } + public void close() {} /** diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java index ae0661ead..f5e0070c1 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImpl.java @@ -17,15 +17,22 @@ */ package org.apache.flink.agents.runtime.context; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.runtime.ResourceCache; +import org.apache.flink.agents.runtime.actionstate.CallResult; import org.apache.flink.agents.runtime.async.ContinuationActionExecutor; import org.apache.flink.agents.runtime.async.ContinuationContext; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; import java.util.concurrent.Callable; import java.util.function.Supplier; +import java.util.stream.Collectors; /** * Java-specific implementation of RunnerContext that includes ContinuationActionExecutor for async @@ -33,6 +40,7 @@ */ public class JavaRunnerContextImpl extends RunnerContextImpl { private final ContinuationActionExecutor continuationExecutor; + private final ContinuationActionExecutor toolCallContinuationExecutor; private ContinuationContext continuationContext; public JavaRunnerContextImpl( @@ -41,9 +49,28 @@ public JavaRunnerContextImpl( AgentPlan agentPlan, ResourceCache resourceCache, String jobIdentifier, - ContinuationActionExecutor continuationExecutor) { + ContinuationActionExecutor continuationExecutor, + ContinuationActionExecutor toolCallContinuationExecutor) { super(agentMetricGroup, mailboxThreadChecker, agentPlan, resourceCache, jobIdentifier); this.continuationExecutor = continuationExecutor; + this.toolCallContinuationExecutor = toolCallContinuationExecutor; + } + + public JavaRunnerContextImpl( + FlinkAgentsMetricGroupImpl agentMetricGroup, + Runnable mailboxThreadChecker, + AgentPlan agentPlan, + ResourceCache resourceCache, + String jobIdentifier, + ContinuationActionExecutor continuationExecutor) { + this( + agentMetricGroup, + mailboxThreadChecker, + agentPlan, + resourceCache, + jobIdentifier, + continuationExecutor, + continuationExecutor); } public ContinuationActionExecutor getContinuationExecutor() { @@ -75,6 +102,159 @@ private T durableExecuteAsyncWithReconcile( callable, reconcileCallable, () -> executeAsyncCallable(callable)); } + @Override + public List> durableExecuteAllAsync(List> callables) + throws Exception { + if (callables.isEmpty()) { + return List.of(); + } + if (durableExecutionContext == null) { + return executeAllWithoutDurableState(callables); + } + + String argsDigest = ""; + int base = durableExecutionContext.getCurrentCallIndex(); + BatchExecutionPlan plan = buildBatchExecutionPlan(callables, base, argsDigest); + + reservePendingBatchIfNeeded(callables, argsDigest, plan); + + List> executed = executeOutcomeSuppliers(plan.suppliers); + finalizeExecutedOutcomes(callables, base, argsDigest, plan, executed); + + advanceCallIndexBy(callables.size()); + return plan.outcomes; + } + + private BatchExecutionPlan buildBatchExecutionPlan( + List> callables, int base, String argsDigest) throws Exception { + BatchExecutionPlan plan = new BatchExecutionPlan<>(callables.size()); + for (int i = 0; i < callables.size(); i++) { + DurableCallable 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 reconcileCallable = callable.reconciler(); + Callable 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 void markReservationStart(BatchExecutionPlan plan, int callIndex) { + plan.needsReservation = true; + if (plan.executionStart < 0) { + plan.executionStart = callIndex; + } + } + + private void addExecutableCall( + BatchExecutionPlan plan, int callIndex, Callable executionCallable) { + plan.outcomes.add(null); + plan.suppliers.add(executionCallable); + plan.executableCallIndexes.add(callIndex); + } + + private void appendRemainingExecutions( + List> callables, BatchExecutionPlan plan, int startIndex) { + for (int i = startIndex; i < callables.size(); i++) { + DurableCallable remaining = callables.get(i); + addExecutableCall(plan, i, remaining::call); + } + } + + private void reservePendingBatchIfNeeded( + List> callables, String argsDigest, BatchExecutionPlan plan) { + if (!plan.needsReservation) { + return; + } + List ids = new ArrayList<>(); + for (DurableCallable callable : + callables.subList(plan.executionStart, callables.size())) { + ids.add(callable.getId()); + } + reservePendingBatch(ids, argsDigest); + } + + private void finalizeExecutedOutcomes( + List> callables, + int base, + String argsDigest, + BatchExecutionPlan plan, + List> executed) + throws Exception { + for (int i = 0; i < executed.size(); i++) { + int callIndex = plan.executableCallIndexes.get(i); + Outcome outcome = executed.get(i); + DurableCallable callable = callables.get(callIndex); + finalizeCallAt( + base + callIndex, + callable.getId(), + argsDigest, + serializeDurableResult(outcome.getValue()), + serializeDurableException(outcome.getError())); + plan.outcomes.set(callIndex, outcome); + } + } + + private static class BatchExecutionPlan { + private final List> outcomes; + private final List> suppliers = new ArrayList<>(); + private final List executableCallIndexes = new ArrayList<>(); + private boolean needsReservation; + private int executionStart = -1; + + private BatchExecutionPlan(int size) { + this.outcomes = new ArrayList<>(size); + } + } + + private List> executeAllWithoutDurableState(List> callables) { + List> suppliers = new ArrayList<>(); + for (DurableCallable callable : callables) { + suppliers.add(callable::call); + } + return executeOutcomeSuppliers(suppliers); + } + + private List> executeOutcomeSuppliers(List> suppliers) { + if (suppliers.isEmpty()) { + return List.of(); + } + if (continuationExecutor == null || continuationContext == null) { + return suppliers.stream() + .map( + supplier -> { + try { + return Outcome.success(supplier.call()); + } catch (Exception e) { + return Outcome.failure(e); + } + }) + .collect(Collectors.toList()); + } + Duration timeout = getConfig().get(AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT); + return toolCallContinuationExecutor.executeAllAsync( + continuationContext, suppliers, timeout); + } + private T executeAsyncCallable(DurableCallable callable) throws Exception { Supplier wrappedSupplier = @@ -103,4 +283,9 @@ private T executeAsyncCallable(DurableCallable callable) throws Exception throw (Exception) e.getCause(); } } + + @Override + public void close() throws Exception { + super.close(); + } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java index 1395bdc56..be36a4e8b 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java @@ -26,6 +26,7 @@ import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.MemoryUpdate; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.memory.BaseLongTermMemory; import org.apache.flink.agents.api.resource.Resource; @@ -59,7 +60,7 @@ */ public class RunnerContextImpl implements RunnerContext { - private static final ObjectMapper OBJECT_MAPPER = + protected static final ObjectMapper OBJECT_MAPPER = new ObjectMapper().registerModule(new JavaTimeModule()); public static class MemoryContext { @@ -265,6 +266,20 @@ public T durableExecuteAsync(DurableCallable callable) throws Exception { return durableExecute(callable); } + @Override + public List> durableExecuteAllAsync(List> callables) + throws Exception { + List> outcomes = new ArrayList<>(callables.size()); + for (DurableCallable callable : callables) { + try { + outcomes.add(Outcome.success(durableExecute(callable))); + } catch (Exception e) { + outcomes.add(Outcome.failure(e)); + } + } + return outcomes; + } + /** * Executes a durable call using the completion-only state machine. * @@ -417,6 +432,13 @@ public void appendPendingCall(String functionId, String argsDigest) { } } + public void reservePendingBatch(List functionIds, String argsDigest) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null && !functionIds.isEmpty()) { + durableExecutionContext.reservePendingBatch(functionIds, argsDigest); + } + } + /** Finalizes the pending durable call slot at the current call index. */ public void finalizeCurrentCall( String functionId, String argsDigest, byte[] resultPayload, byte[] exceptionPayload) { @@ -427,6 +449,26 @@ public void finalizeCurrentCall( } } + public void finalizeCallAt( + int index, + String functionId, + String argsDigest, + byte[] resultPayload, + byte[] exceptionPayload) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null) { + durableExecutionContext.finalizeCallAt( + index, functionId, argsDigest, resultPayload, exceptionPayload); + } + } + + public void advanceCallIndexBy(int count) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null) { + durableExecutionContext.advanceCallIndexBy(count); + } + } + /** * Clears persisted call results from the current call index onward and persists immediately. */ @@ -437,12 +479,23 @@ public void clearCallResultsFromCurrentIndexAndPersist() { } } - /** - * Returns the current durable call result as an array of fields for bridge consumers, or null - * if no persisted slot exists at the current call index. - */ - public Object[] getCurrentCallResultFields() { - CallResult current = getCurrentCallResult(); + public void clearCallResultsFromAndPersist(int index) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null) { + durableExecutionContext.clearCallResultsFromAndPersist(index); + } + } + + public int getCurrentCallIndex() { + mailboxThreadChecker.run(); + if (durableExecutionContext == null) { + return 0; + } + return durableExecutionContext.getCurrentCallIndex(); + } + + public Object[] getCallResultFieldsAt(int index) { + CallResult current = getCallResultAt(index); if (current == null) { return null; } @@ -455,6 +508,40 @@ public Object[] getCurrentCallResultFields() { }; } + /** + * Returns the current durable call result as an array of fields for bridge consumers, or null + * if no persisted slot exists at the current call index. + */ + public Object[] getCurrentCallResultFields() { + if (durableExecutionContext == null) { + return null; + } + return getCallResultFieldsAt(durableExecutionContext.getCurrentCallIndex()); + } + + protected Outcome readTerminalOutcomeAt( + int index, String functionId, String argsDigest, Class resultClass) + throws Exception { + CallResult callResult = getCallResultAt(index); + if (callResult == null || callResult.isPending()) { + throw new IllegalStateException( + String.format( + "Expected a terminal durable call result at index %s for " + + "functionId=%s, argsDigest=%s", + index, functionId, argsDigest)); + } + if (callResult.getExceptionPayload() != null) { + DurableExecutionException exception = + OBJECT_MAPPER.readValue( + callResult.getExceptionPayload(), DurableExecutionException.class); + return Outcome.failure(exception.toException()); + } + if (callResult.getResultPayload() == null) { + return Outcome.success(null); + } + return Outcome.success(OBJECT_MAPPER.readValue(callResult.getResultPayload(), resultClass)); + } + protected CallResult getCurrentCallResult() { mailboxThreadChecker.run(); if (durableExecutionContext != null) { @@ -463,6 +550,14 @@ protected CallResult getCurrentCallResult() { return null; } + protected CallResult getCallResultAt(int index) { + mailboxThreadChecker.run(); + if (durableExecutionContext != null) { + return durableExecutionContext.getCallResultAt(index); + } + return null; + } + protected Optional tryGetCachedResult( String functionId, String argsDigest, Class resultClass) throws Exception { Object[] cached = matchNextOrClearSubsequentCallResult(functionId, argsDigest); @@ -637,8 +732,12 @@ public ActionState getActionState() { * yet have a persisted slot. */ public CallResult getCurrentCallResult() { - if (currentCallIndex < recoveryCallResults.size()) { - return recoveryCallResults.get(currentCallIndex); + return getCallResultAt(currentCallIndex); + } + + public CallResult getCallResultAt(int index) { + if (index < recoveryCallResults.size()) { + return recoveryCallResults.get(index); } return null; } @@ -738,6 +837,15 @@ public void appendPendingCall(String functionId, String argsDigest) { argsDigest); } + public void reservePendingBatch(List functionIds, String argsDigest) { + for (String functionId : functionIds) { + CallResult pending = CallResult.pending(functionId, argsDigest); + actionState.addCallResult(pending); + recoveryCallResults.add(pending); + } + persistActionState(); + } + /** * Replaces the current persisted slot with a terminal call result and advances the current * call index. @@ -747,42 +855,52 @@ public void finalizeCurrentCall( String argsDigest, byte[] resultPayload, byte[] exceptionPayload) { - CallResult current = getCurrentCallResult(); + finalizeCallAt( + currentCallIndex, functionId, argsDigest, resultPayload, exceptionPayload); + currentCallIndex++; + } + + public void finalizeCallAt( + int index, + String functionId, + String argsDigest, + byte[] resultPayload, + byte[] exceptionPayload) { + CallResult current = getCallResultAt(index); if (current == null) { throw new IllegalStateException( String.format( - "Cannot finalize current call at index %s because no persisted " - + "slot exists", - currentCallIndex)); + "Cannot finalize call at index %s because no persisted slot exists", + index)); } if (!current.matches(functionId, argsDigest)) { throw new IllegalStateException( String.format( - "Cannot finalize current call at index %s because the persisted " - + "slot does not match functionId=%s, argsDigest=%s", - currentCallIndex, functionId, argsDigest)); + "Cannot finalize call at index %s because the persisted slot does not match functionId=%s, argsDigest=%s", + index, functionId, argsDigest)); } if (!current.isPending()) { throw new IllegalStateException( String.format( - "Cannot finalize current call at index %s because the persisted " - + "slot is not pending", - currentCallIndex)); + "Cannot finalize call at index %s because the persisted slot is not pending", + index)); } CallResult terminal = new CallResult(functionId, argsDigest, resultPayload, exceptionPayload); - actionState.replaceCallResult(currentCallIndex, terminal); - recoveryCallResults.set(currentCallIndex, terminal); + actionState.replaceCallResult(index, terminal); + recoveryCallResults.set(index, terminal); persistActionState(); LOG.debug( "Finalized and persisted CallResult at index {}: functionId={}, argsDigest={}", - currentCallIndex, + index, functionId, argsDigest); + } - currentCallIndex++; + public void advanceCallIndexBy(int count) { + currentCallIndex += count; } /** @@ -794,12 +912,21 @@ public void clearCallResultsFromCurrentIndexAndPersist() { persistActionState(); } - private void clearCallResultsFromCurrentIndex() { - actionState.clearCallResultsFrom(currentCallIndex); + public void clearCallResultsFromAndPersist(int index) { + clearCallResultsFrom(index); + persistActionState(); + } + + public void clearCallResultsFrom(int index) { + actionState.clearCallResultsFrom(index); recoveryCallResults = new ArrayList<>( recoveryCallResults.subList( - 0, Math.min(currentCallIndex, recoveryCallResults.size()))); + 0, Math.min(index, recoveryCallResults.size()))); + } + + private void clearCallResultsFromCurrentIndex() { + clearCallResultsFrom(currentCallIndex); } private void persistActionState() { diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java index 46d09e0a2..bfa5d7b4d 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java @@ -188,7 +188,10 @@ public void open() throws Exception { // init context manager for runner context creation and memory contexts contextManager = new ActionTaskContextManager( - agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS)); + agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS), + agentPlan + .getConfig() + .get(AgentExecutionOptions.TOOL_CALL_NUM_ASYNC_THREADS)); mailboxProcessor = getMailboxProcessor(); diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java index eba25d5d4..286f038e5 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java @@ -73,13 +73,19 @@ class ActionTaskContextManager implements AutoCloseable { private final Map continuationContexts; private final Map pythonAwaitableRefs; - private ContinuationActionExecutor continuationActionExecutor; + private final ContinuationActionExecutor actionContinuationExecutor; + private final ContinuationActionExecutor toolCallContinuationExecutor; - ActionTaskContextManager(int numAsyncThreads) { + ActionTaskContextManager(int numAsyncThreads, int numToolCallAsyncThreads) { this.actionTaskMemoryContexts = new HashMap<>(); this.continuationContexts = new HashMap<>(); this.pythonAwaitableRefs = new HashMap<>(); - this.continuationActionExecutor = new ContinuationActionExecutor(numAsyncThreads); + this.actionContinuationExecutor = new ContinuationActionExecutor(numAsyncThreads); + this.toolCallContinuationExecutor = new ContinuationActionExecutor(numToolCallAsyncThreads); + } + + ActionTaskContextManager(int numAsyncThreads) { + this(numAsyncThreads, numAsyncThreads); } /** @@ -111,7 +117,7 @@ RunnerContextImpl createOrGetRunnerContext( @Nullable InteranlBaseLongTermMemory longTermMemory) { if (isJava) { if (runnerContext == null) { - if (continuationActionExecutor == null) { + if (actionContinuationExecutor == null || toolCallContinuationExecutor == null) { throw new IllegalStateException( "ContinuationActionExecutor has not been initialized."); } @@ -122,7 +128,8 @@ RunnerContextImpl createOrGetRunnerContext( agentPlan, resourceCache, jobIdentifier, - continuationActionExecutor); + actionContinuationExecutor, + toolCallContinuationExecutor); if (longTermMemory != null) { runnerContext.setLongTermMemory(longTermMemory); } @@ -324,8 +331,11 @@ public void close() throws Exception { runnerContext = null; } } - if (continuationActionExecutor != null) { - continuationActionExecutor.close(); + if (actionContinuationExecutor != null) { + actionContinuationExecutor.close(); + } + if (toolCallContinuationExecutor != null) { + toolCallContinuationExecutor.close(); } } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java index 67c80f38d..29bbbe110 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java @@ -74,6 +74,7 @@ public class PythonActionExecutor { private final JavaResourceAdapter javaResourceAdapter; private final String jobIdentifier; private PyObject pythonAsyncThreadPool; + private PyObject pythonToolCallAsyncThreadPool; private PyObject pythonRunnerContext; public PythonActionExecutor( @@ -102,6 +103,13 @@ public void open() throws Exception { interpreter.invoke( CREATE_ASYNC_THREAD_POOL, agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS)); + pythonToolCallAsyncThreadPool = + (PyObject) + interpreter.invoke( + CREATE_ASYNC_THREAD_POOL, + agentPlan + .getConfig() + .get(AgentExecutionOptions.TOOL_CALL_NUM_ASYNC_THREADS)); pythonRunnerContext = (PyObject) @@ -110,6 +118,7 @@ public void open() throws Exception { runnerContext, new ObjectMapper().writeValueAsString(agentPlan), pythonAsyncThreadPool, + pythonToolCallAsyncThreadPool, javaResourceAdapter, jobIdentifier); } @@ -192,6 +201,9 @@ public void close() throws Exception { if (pythonAsyncThreadPool != null) { interpreter.invoke(CLOSE_ASYNC_THREAD_POOL, pythonAsyncThreadPool); } + if (pythonToolCallAsyncThreadPool != null) { + interpreter.invoke(CLOSE_ASYNC_THREAD_POOL, pythonToolCallAsyncThreadPool); + } if (pythonRunnerContext != null) { try { diff --git a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java index 4dc30d817..284ff441c 100644 --- a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java +++ b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java @@ -17,12 +17,20 @@ */ package org.apache.flink.agents.runtime.async; +import org.apache.flink.agents.api.context.Outcome; + import jdk.internal.vm.Continuation; import jdk.internal.vm.ContinuationScope; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.TimeoutException; import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -58,17 +66,22 @@ public ContinuationActionExecutor(int numAsyncThreads) { * @return true if the action completed, false if waiting for async execution */ public boolean executeAction(ContinuationContext context, Runnable action) { - // Check if we have a pending async Future from previous yield + // Wait while async work is still pending, unless the batch deadline elapsed — then resume + // so executeAllAsync can finalize timed-out slots. + if (context.hasPendingAsync() && !context.isBatchDeadlineElapsed()) { + return false; + } + Future pending = context.getPendingFuture(); if (pending != null) { - if (!pending.isDone()) { - // Async task not done yet, return false to wait - return false; - } - // Async task done, clear the pending future and resume LOG.debug("Async task done..."); context.setPendingFuture(null); } + Future pendingBatch = context.getPendingBatchFuture(); + if (pendingBatch != null) { + LOG.debug("Async batch done..."); + context.setPendingBatchFuture(null); + } Continuation currentContinuation = context.getCurrentContinuation(); if (currentContinuation == null) { @@ -148,6 +161,102 @@ public T executeAsync(ContinuationContext context, Supplier supplier) thr return (T) context.getAsyncResultRef().get(); } + /** + * Executes all suppliers as one async batch and returns one {@link Outcome} per supplier. + * Supplier failures are captured in their own outcome so one failed supplier does not abort the + * whole batch. + * + * @param context the continuation context for this action + * @param suppliers the suppliers to execute + * @param timeout the timeout for the whole batch; null or non-positive means no timeout + * @param the result type + * @return outcomes in supplier order + */ + public List> executeAllAsync( + ContinuationContext context, List> suppliers, Duration timeout) + throws Exception { + context.clearAsyncState(); + + List>> futures = new ArrayList<>(suppliers.size()); + for (Callable supplier : suppliers) { + futures.add( + CompletableFuture.supplyAsync( + () -> { + try { + return Outcome.success(supplier.call()); + } catch (Exception e) { + return Outcome.failure(e); + } + }, + asyncExecutor)); + } + + CompletableFuture barrier = + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])); + long deadlineNanos = getDeadlineNanos(timeout); + context.setPendingBatchFuture(barrier, deadlineNanos); + + while (!barrier.isDone()) { + if (System.nanoTime() >= deadlineNanos) { + TimeoutException exception = + new TimeoutException( + "Async durable batch execution timed out after " + timeout); + barrier.cancel(true); + context.setPendingBatchFuture(null); + return collectBatchOutcomesOnTimeout(futures, exception); + } + Continuation.yield(SCOPE); + } + + context.setPendingBatchFuture(null); + return collectBatchOutcomes(futures); + } + + /** + * Collects per-slot outcomes after the batch barrier completes normally. + * + *

Each supplier already wraps success and failure into an {@link Outcome}, so {@code join()} + * returns that outcome rather than throwing for ordinary tool exceptions. + */ + private static List> collectBatchOutcomes( + List>> futures) { + List> results = new ArrayList<>(futures.size()); + for (CompletableFuture> future : futures) { + results.add(future.join()); + } + return results; + } + + /** + * Collects per-slot outcomes when the batch deadline elapses. + * + *

Completed slots keep their success or failure outcome. Only slots that are still running + * (or become cancelled) are finalized as timeout failures. {@code cancel(true)} is attempted + * only for unfinished futures; a future that completes between the check and cancel stays + * non-cancelled and is collected as a normal outcome. + */ + private static List> collectBatchOutcomesOnTimeout( + List>> futures, TimeoutException timeoutException) { + List> results = new ArrayList<>(futures.size()); + for (CompletableFuture> future : futures) { + if (!future.isDone()) { + future.cancel(true); + } + if (future.isDone() && !future.isCancelled()) { + results.add(future.join()); + } else { + results.add(Outcome.failure(timeoutException)); + } + } + return results; + } + + private long getDeadlineNanos(Duration timeout) { + return timeout == null || timeout.isZero() || timeout.isNegative() + ? Long.MAX_VALUE + : System.nanoTime() + timeout.toNanos(); + } + public void close() { asyncExecutor.shutdownNow(); } diff --git a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java index fbff2b118..0523abb5b 100644 --- a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java +++ b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationContext.java @@ -27,6 +27,10 @@ public class ContinuationContext { private Continuation currentContinuation; private volatile Future pendingFuture; + private volatile Future pendingBatchFuture; + /** {@link Long#MAX_VALUE} means the pending batch has no deadline. */ + private volatile long pendingBatchDeadlineNanos = Long.MAX_VALUE; + private final AtomicReference asyncResult = new AtomicReference<>(); private final AtomicReference asyncException = new AtomicReference<>(); @@ -46,6 +50,40 @@ public void setPendingFuture(Future pendingFuture) { this.pendingFuture = pendingFuture; } + public Future getPendingBatchFuture() { + return pendingBatchFuture; + } + + public void setPendingBatchFuture(Future pendingBatchFuture) { + setPendingBatchFuture(pendingBatchFuture, Long.MAX_VALUE); + } + + /** + * Records a pending batch barrier and its absolute nanoTime deadline. + * + *

{@link #hasPendingAsync()} only reflects whether the barrier future is still running. + * {@link ContinuationActionExecutor#executeAction} consults {@link #isBatchDeadlineElapsed()} + * separately so a timed-out batch can still resume and finalize unfinished slots. + */ + public void setPendingBatchFuture(Future pendingBatchFuture, long deadlineNanos) { + this.pendingBatchFuture = pendingBatchFuture; + this.pendingBatchDeadlineNanos = + pendingBatchFuture == null ? Long.MAX_VALUE : deadlineNanos; + } + + public boolean hasPendingAsync() { + return isPending(pendingFuture) || isPending(pendingBatchFuture); + } + + /** Whether the pending batch deadline has elapsed. No deadline means {@link Long#MAX_VALUE}. */ + public boolean isBatchDeadlineElapsed() { + return System.nanoTime() >= pendingBatchDeadlineNanos; + } + + private static boolean isPending(Future future) { + return future != null && !future.isDone(); + } + public AtomicReference getAsyncResultRef() { return asyncResult; } @@ -56,6 +94,8 @@ public AtomicReference getAsyncExceptionRef() { public void clearAsyncState() { pendingFuture = null; + pendingBatchFuture = null; + pendingBatchDeadlineNanos = Long.MAX_VALUE; asyncResult.set(null); asyncException.set(null); } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java index ff378fb6a..62676e939 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java @@ -20,6 +20,9 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.configuration.Configuration; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.runtime.actionstate.ActionState; @@ -32,6 +35,9 @@ import org.junit.jupiter.api.Test; import java.util.HashMap; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; @@ -195,6 +201,154 @@ void testDurableExecuteAsyncReconcilableReconcileExceptionPersistsFailure() thro assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); } + @Test + void testDurableExecuteAllAsyncInitialBatchPersistsOutcomes() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + JavaRunnerContextImpl context = createContext(new ActionState(null), executor); + TestDurableCallable first = + new TestDurableCallable<>("batch-1", String.class, () -> "one"); + TestDurableCallable second = + new TestDurableCallable<>("batch-2", String.class, () -> "two"); + + List> outcomes = context.durableExecuteAllAsync(List.of(first, second)); + + assertEquals("one", outcomes.get(0).getValue()); + assertEquals("two", outcomes.get(1).getValue()); + assertEquals(1, executor.getExecuteAllAsyncCallCount()); + assertEquals(List.of(2), executor.getExecuteAllAsyncBatchSizes()); + assertEquals(1, first.getCallCount()); + assertEquals(1, second.getCallCount()); + assertEquals(3, persistCallCount.get()); + assertEquals(2, context.getDurableExecutionContext().getCurrentCallIndex()); + List persisted = + context.getDurableExecutionContext().getActionState().getCallResults(); + assertEquals(2, persisted.size()); + assertEquals("batch-1", persisted.get(0).getFunctionId()); + assertTrue(persisted.get(0).isSuccess()); + assertEquals("batch-2", persisted.get(1).getFunctionId()); + assertTrue(persisted.get(1).isSuccess()); + } + + @Test + void testDurableExecuteAllAsyncReconcilesPendingSlot() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + ActionState actionState = new ActionState(null); + actionState.addCallResult(CallResult.pending("batch-1", "")); + JavaRunnerContextImpl context = createContext(actionState, executor); + TestReconcilableCallable callable = + new TestReconcilableCallable<>( + "batch-1", + String.class, + () -> fail("call should not be executed"), + () -> "recovered"); + + List> outcomes = context.durableExecuteAllAsync(List.of(callable)); + + assertEquals("recovered", outcomes.get(0).getValue()); + assertEquals(0, callable.getCallCount()); + assertEquals(1, callable.getReconcileCount()); + assertEquals(1, executor.getExecuteAllAsyncCallCount()); + assertEquals(1, persistCallCount.get()); + assertTrue(actionState.getCallResults().get(0).isSuccess()); + assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); + } + + @Test + void testDurableExecuteAllAsyncRecoversPartialFinalizedBatch() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + ActionState actionState = new ActionState(null); + actionState.addCallResult( + new CallResult("batch-1", "", OBJECT_MAPPER.writeValueAsBytes("cached-one"))); + actionState.addCallResult( + new CallResult("batch-2", "", OBJECT_MAPPER.writeValueAsBytes("cached-two"))); + actionState.addCallResult(CallResult.pending("batch-3", "")); + JavaRunnerContextImpl context = createContext(actionState, executor); + TestDurableCallable first = + new TestDurableCallable<>( + "batch-1", String.class, () -> fail("cached slot should not execute")); + TestDurableCallable second = + new TestDurableCallable<>( + "batch-2", String.class, () -> fail("cached slot should not execute")); + TestDurableCallable third = + new TestDurableCallable<>("batch-3", String.class, () -> "fresh-three"); + + List> outcomes = + context.durableExecuteAllAsync(List.of(first, second, third)); + + assertEquals("cached-one", outcomes.get(0).getValue()); + assertEquals("cached-two", outcomes.get(1).getValue()); + assertEquals("fresh-three", outcomes.get(2).getValue()); + assertEquals(0, first.getCallCount()); + assertEquals(0, second.getCallCount()); + assertEquals(1, third.getCallCount()); + assertEquals(1, executor.getExecuteAllAsyncCallCount()); + assertEquals(List.of(1), executor.getExecuteAllAsyncBatchSizes()); + assertEquals("batch-3", actionState.getCallResults().get(2).getFunctionId()); + assertTrue(actionState.getCallResults().get(2).isSuccess()); + assertEquals(1, persistCallCount.get()); + assertEquals(3, context.getDurableExecutionContext().getCurrentCallIndex()); + } + + @Test + void testDurableExecuteAllAsyncReturnsCachedFailureOutcome() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + ActionState actionState = new ActionState(null); + actionState.addCallResult( + new CallResult( + "batch-1", + "", + null, + OBJECT_MAPPER.writeValueAsBytes( + RunnerContextImpl.DurableExecutionException.fromException( + new IllegalStateException("cached failure"))))); + JavaRunnerContextImpl context = createContext(actionState, executor); + TestDurableCallable callable = + new TestDurableCallable<>( + "batch-1", String.class, () -> fail("cached slot should not execute")); + + List> outcomes = context.durableExecuteAllAsync(List.of(callable)); + + assertTrue(outcomes.get(0).isFailure()); + assertTrue(outcomes.get(0).getError().getMessage().contains("cached failure")); + assertEquals(0, callable.getCallCount()); + assertEquals(0, executor.getExecuteAllAsyncCallCount()); + assertEquals(0, persistCallCount.get()); + assertEquals(1, context.getDurableExecutionContext().getCurrentCallIndex()); + } + + @Test + void testDurableExecuteAllAsyncTimeoutKeepsCompletedOutcomes() throws Exception { + InspectingContinuationActionExecutor executor = new InspectingContinuationActionExecutor(); + executor.setUseTimeoutCollection(true); + JavaRunnerContextImpl context = createContext(new ActionState(null), executor); + ((Configuration) context.getConfig()) + .set( + AgentExecutionOptions.TOOL_CALL_BATCH_TIMEOUT, + java.time.Duration.ofMillis(10)); + TestDurableCallable first = + new TestDurableCallable<>("batch-1", String.class, () -> "fast"); + TestDurableCallable second = + new TestDurableCallable<>( + "batch-2", + String.class, + () -> { + Thread.sleep(100); + return "slow"; + }); + + List> outcomes = context.durableExecuteAllAsync(List.of(first, second)); + + assertEquals("fast", outcomes.get(0).getValue()); + assertTrue(outcomes.get(1).isFailure()); + assertInstanceOf(TimeoutException.class, outcomes.get(1).getError()); + List persisted = + context.getDurableExecutionContext().getActionState().getCallResults(); + assertTrue(persisted.get(0).isSuccess()); + assertTrue(persisted.get(1).isFailure()); + assertEquals(2, context.getDurableExecutionContext().getCurrentCallIndex()); + executor.close(); + } + private JavaRunnerContextImpl createContext( ActionState actionState, ContinuationActionExecutor executor) { JavaRunnerContextImpl context = @@ -225,7 +379,10 @@ private JavaRunnerContextImpl createContext( private static final class InspectingContinuationActionExecutor extends ContinuationActionExecutor { private Runnable beforeExecute; + private boolean useTimeoutCollection; private int executeAsyncCallCount; + private int executeAllAsyncCallCount; + private final List executeAllAsyncBatchSizes = new java.util.ArrayList<>(); private InspectingContinuationActionExecutor() { super(1); @@ -240,12 +397,61 @@ public T executeAsync(ContinuationContext context, Supplier supplier) { return supplier.get(); } + @Override + public List> executeAllAsync( + ContinuationContext context, + List> suppliers, + java.time.Duration timeout) { + executeAllAsyncCallCount++; + executeAllAsyncBatchSizes.add(suppliers.size()); + if (useTimeoutCollection) { + return collectTimedOutOutcomes(suppliers); + } + List> outcomes = new java.util.ArrayList<>(suppliers.size()); + for (Callable supplier : suppliers) { + try { + outcomes.add(Outcome.success(supplier.call())); + } catch (Exception e) { + outcomes.add(Outcome.failure(e)); + } + } + return outcomes; + } + + private List> collectTimedOutOutcomes(List> suppliers) { + List> outcomes = new java.util.ArrayList<>(suppliers.size()); + for (int i = 0; i < suppliers.size(); i++) { + if (i == 0) { + try { + outcomes.add(Outcome.success(suppliers.get(i).call())); + } catch (Exception e) { + outcomes.add(Outcome.failure(e)); + } + } else { + outcomes.add(Outcome.failure(new TimeoutException("batch timeout"))); + } + } + return outcomes; + } + private void setBeforeExecute(Runnable beforeExecute) { this.beforeExecute = beforeExecute; } + private void setUseTimeoutCollection(boolean useTimeoutCollection) { + this.useTimeoutCollection = useTimeoutCollection; + } + private int getExecuteAsyncCallCount() { return executeAsyncCallCount; } + + private int getExecuteAllAsyncCallCount() { + return executeAllAsyncCallCount; + } + + private List getExecuteAllAsyncBatchSizes() { + return executeAllAsyncBatchSizes; + } } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryRefTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryRefTest.java index e16361552..615c6c748 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryRefTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/MemoryRefTest.java @@ -21,6 +21,7 @@ import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.MemoryRef; +import org.apache.flink.agents.api.context.Outcome; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.memory.BaseLongTermMemory; import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; @@ -128,6 +129,16 @@ public T durableExecuteAsync(DurableCallable callable) throws Exception { return callable.call(); } + @Override + public List> durableExecuteAllAsync(List> callables) + throws Exception { + List> outcomes = new ArrayList<>(callables.size()); + for (DurableCallable callable : callables) { + outcomes.add(Outcome.success(callable.call())); + } + return outcomes; + } + @Override public void close() throws Exception {} }