da-daken commented on code in PR #926:
URL: https://github.com/apache/flink-agents/pull/926#discussion_r3864602380
##########
runtime/src/test/java/org/apache/flink/agents/runtime/context/JavaRunnerContextImplDurableExecuteAsyncTest.java:
##########
@@ -195,6 +208,357 @@ void
testDurableExecuteAsyncReconcilableReconcileExceptionPersistsFailure() thro
assertEquals(1,
context.getDurableExecutionContext().getCurrentCallIndex());
}
+ @Test
+ void testDurableExecuteAsyncCompletionOnlyReExecutesPendingSlot() throws
Exception {
+ InspectingContinuationActionExecutor executor = new
InspectingContinuationActionExecutor();
+ ActionState actionState = new ActionState(null);
+ actionState.addCallResult(CallResult.pending("tool-call", ""));
+ JavaRunnerContextImpl context = createContext(actionState, executor);
+ TestDurableCallable<String> callable =
+ new TestDurableCallable<>("tool-call", String.class, () ->
"recovered");
+
+ String result = context.durableExecuteAsync(callable);
+
+ assertEquals("recovered", result);
+ assertEquals(1, callable.getCallCount());
+ assertEquals(1, executor.getExecuteAsyncCallCount());
+ assertEquals(1, persistCallCount.get());
+ assertEquals(1,
context.getDurableExecutionContext().getCurrentCallIndex());
+ CallResult persisted =
+
context.getDurableExecutionContext().getActionState().getCallResults().get(0);
+ assertTrue(persisted.isSuccess());
+ }
+
+ @Test
+ void testDurableExecuteAllAsyncInitialBatchPersistsOutcomes() throws
Exception {
+ InspectingContinuationActionExecutor executor = new
InspectingContinuationActionExecutor();
+ JavaRunnerContextImpl context = createContext(new ActionState(null),
executor);
+ TestDurableCallable<String> first =
+ new TestDurableCallable<>("batch-1", String.class, () ->
"one");
+ TestDurableCallable<String> second =
+ new TestDurableCallable<>("batch-2", String.class, () ->
"two");
+
+ List<Outcome<String>> 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<CallResult> 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<String> callable =
+ new TestReconcilableCallable<>(
+ "batch-1",
+ String.class,
+ () -> fail("call should not be executed"),
+ () -> "recovered");
+
+ List<Outcome<String>> 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<String> first =
+ new TestDurableCallable<>(
+ "batch-1", String.class, () -> fail("cached slot
should not execute"));
+ TestDurableCallable<String> second =
+ new TestDurableCallable<>(
+ "batch-2", String.class, () -> fail("cached slot
should not execute"));
+ TestDurableCallable<String> third =
+ new TestDurableCallable<>("batch-3", String.class, () ->
"fresh-three");
+
+ List<Outcome<String>> 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<String> callable =
+ new TestDurableCallable<>(
+ "batch-1", String.class, () -> fail("cached slot
should not execute"));
+
+ List<Outcome<String>> outcomes =
context.durableExecuteAllAsync(List.of(callable));
+
+ assertTrue(outcomes.get(0).isFailure());
+ assertInstanceOf(IllegalStateException.class,
outcomes.get(0).getError());
+ 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 testDurableExecuteAllAsyncReturnsDeserializeFailureAsOutcome() throws
Exception {
+ InspectingContinuationActionExecutor executor = new
InspectingContinuationActionExecutor();
+ ActionState actionState = new ActionState(null);
+ actionState.addCallResult(
+ new CallResult(
+ "batch-1",
+ "",
+
"not-valid-json".getBytes(java.nio.charset.StandardCharsets.UTF_8),
+ null));
+ JavaRunnerContextImpl context = createContext(actionState, executor);
+ TestDurableCallable<String> callable =
+ new TestDurableCallable<>(
+ "batch-1", String.class, () -> fail("cached slot
should not execute"));
+
+ List<Outcome<String>> outcomes =
context.durableExecuteAllAsync(List.of(callable));
+
+ assertTrue(outcomes.get(0).isFailure());
+ assertInstanceOf(JsonProcessingException.class,
outcomes.get(0).getError());
+ assertEquals(0, callable.getCallCount());
+ assertEquals(1,
context.getDurableExecutionContext().getCurrentCallIndex());
+ }
+
+ @Test
+ void testDurableExecuteAllAsyncPassesParallelismFromConfig() throws
Exception {
+ InspectingContinuationActionExecutor executor = new
InspectingContinuationActionExecutor();
+ JavaRunnerContextImpl context = createContext(new ActionState(null),
executor);
+ ((Configuration)
context.getConfig()).set(AgentExecutionOptions.TOOL_CALL_PARALLELISM, 4);
+ TestDurableCallable<String> callable =
+ new TestDurableCallable<>("batch-1", String.class, () -> "ok");
+
+ List<Outcome<String>> outcomes =
context.durableExecuteAllAsync(List.of(callable));
+
+ assertEquals("ok", outcomes.get(0).getValue());
+ assertEquals(4, executor.getLastExecuteAllAsyncMaxParallelism());
+ executor.close();
+ }
+
+ @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_MS, 100L);
+ TestDurableCallable<String> first =
+ new TestDurableCallable<>("batch-1", String.class, () ->
"fast");
+ TestDurableCallable<String> second =
+ new TestDurableCallable<>(
+ "batch-2",
+ String.class,
+ () -> {
+ Thread.sleep(200);
+ return "slow";
+ });
+
+ List<Outcome<String>> 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());
+ assertEquals(Duration.ofMillis(100),
executor.getLastExecuteAllAsyncTimeout());
+ assertEquals(1, first.getCallCount());
+ assertEquals(1, second.getCallCount());
+ List<CallResult> persisted =
+
context.getDurableExecutionContext().getActionState().getCallResults();
+ assertTrue(persisted.get(0).isSuccess());
+ assertTrue(persisted.get(1).isFailure());
+ assertEquals(2,
context.getDurableExecutionContext().getCurrentCallIndex());
+ executor.close();
+ }
+
+ @Test
+ void testDurableExecuteAllAsyncTimeoutLeavesUnsubmittedSlotsPending()
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_MS, 100L);
+ ((Configuration)
context.getConfig()).set(AgentExecutionOptions.TOOL_CALL_PARALLELISM, 2);
+ TestDurableCallable<String> first =
+ new TestDurableCallable<>(
+ "batch-1",
+ String.class,
+ () -> {
+ Thread.sleep(200);
+ return "one";
+ });
+ TestDurableCallable<String> second =
+ new TestDurableCallable<>(
+ "batch-2",
+ String.class,
+ () -> {
+ Thread.sleep(200);
+ return "two";
+ });
+ TestDurableCallable<String> third =
+ new TestDurableCallable<>("batch-3", String.class, () ->
"three");
+ TestDurableCallable<String> fourth =
+ new TestDurableCallable<>("batch-4", String.class, () ->
"four");
+
+ List<Outcome<String>> outcomes =
+ context.durableExecuteAllAsync(List.of(first, second, third,
fourth));
+
+ assertTrue(outcomes.get(0).isFailure());
+ assertTrue(outcomes.get(1).isFailure());
+ assertTrue(outcomes.get(2).isFailure());
+ assertTrue(outcomes.get(3).isFailure());
+ List<CallResult> persisted =
+
context.getDurableExecutionContext().getActionState().getCallResults();
+ assertTrue(persisted.get(0).isFailure());
+ assertTrue(persisted.get(1).isFailure());
+ assertTrue(persisted.get(2).isPending());
+ assertTrue(persisted.get(3).isPending());
+ assertEquals(4,
context.getDurableExecutionContext().getCurrentCallIndex());
+ executor.close();
+ }
+
+ @Test
+ void
testDurableExecuteAllAsyncTimeoutLeavesQueuedButUnstartedSlotsPending() throws
Exception {
Review Comment:
Thanks! Covered by the new
testToolCallBatchTimeoutCancelsQueuedButUnstartedSlots.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
