From e8f1f4e78512a1efc35715c6c28d2ce896a861c9 Mon Sep 17 00:00:00 2001 From: Jacob Alber Date: Mon, 8 Sep 2025 15:22:41 -0400 Subject: [PATCH] fix: Fix Checkpoint Restore when Rehydrating Run (#642) When checkpointing we did not persist the set of instantiated executors. This means, in turn, when we restore from a checkpoint when using Resume(Stream) rather than restoring a checkpoint in the context of an already existing (Streaming)Run, the executors never got reinstantiated and there were no executors to notify that a state should be loaded. The fix is to ensure we persist the list and reinstantiate the executors on rehydration. * Also adds a rehydration restore test --- .../Execution/RunnerStateData.cs | 3 ++- .../InProc/InProcessRunner.cs | 6 +++--- .../InProc/InProcessRunnerContext.cs | 14 ++++++++++---- .../05_Simple_Workflow_Checkpointing.cs | 16 ++++++++++++++-- .../SampleSmokeTest.cs | 19 +++++++++++++++++++ 5 files changed, 48 insertions(+), 10 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.Workflows/Execution/RunnerStateData.cs b/dotnet/src/Microsoft.Agents.Workflows/Execution/RunnerStateData.cs index 9677ce6e57..877518718a 100644 --- a/dotnet/src/Microsoft.Agents.Workflows/Execution/RunnerStateData.cs +++ b/dotnet/src/Microsoft.Agents.Workflows/Execution/RunnerStateData.cs @@ -5,8 +5,9 @@ using Microsoft.Agents.Workflows.Checkpointing; namespace Microsoft.Agents.Workflows.Execution; -internal class RunnerStateData(Dictionary> queuedMessages, List outstandingRequests) +internal class RunnerStateData(HashSet instantiatedExecutors, Dictionary> queuedMessages, List outstandingRequests) { + public HashSet InstantiatedExecutors { get; } = instantiatedExecutors; public Dictionary> QueuedMessages { get; } = queuedMessages; public List OutstandingRequests { get; } = outstandingRequests; } diff --git a/dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunner.cs b/dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunner.cs index a6d7081c1c..4bf2f50e8c 100644 --- a/dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunner.cs +++ b/dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunner.cs @@ -228,12 +228,12 @@ internal class InProcessRunner : ISuperStepRunner, ICheckpointingRunner this._workflowInfoCache = this.Workflow.ToWorkflowInfo(); } - RunnerStateData runnerData = await this.RunContext.ExportStateAsync().ConfigureAwait(false); Dictionary edgeData = await this.EdgeMap.ExportStateAsync().ConfigureAwait(false); await prepareTask.ConfigureAwait(false); await this.RunContext.StateManager.PublishUpdatesAsync(this.StepTracer).ConfigureAwait(false); + RunnerStateData runnerData = await this.RunContext.ExportStateAsync().ConfigureAwait(false); Dictionary stateData = await this.RunContext.StateManager.ExportStateAsync().ConfigureAwait(false); Checkpoint checkpoint = new(this.StepTracer.StepNumber, this._workflowInfoCache, runnerData, stateData, edgeData); @@ -261,9 +261,9 @@ internal class InProcessRunner : ISuperStepRunner, ICheckpointingRunner } await this.RunContext.StateManager.ImportStateAsync(checkpoint).ConfigureAwait(false); - Task executorNotifyTask = this.RunContext.NotifyCheckpointLoadedAsync(cancellation); - await this.RunContext.ImportStateAsync(checkpoint).ConfigureAwait(false); + + Task executorNotifyTask = this.RunContext.NotifyCheckpointLoadedAsync(cancellation); ValueTask republishRequestsTask = this.RunContext.RepublishUnservicedRequestsAsync(cancellation); await this.EdgeMap.ImportStateAsync(checkpoint).ConfigureAwait(false); diff --git a/dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunnerContext.cs b/dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunnerContext.cs index 3a57dd9302..b2d30f21f4 100644 --- a/dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunnerContext.cs +++ b/dotnet/src/Microsoft.Agents.Workflows/InProc/InProcessRunnerContext.cs @@ -136,8 +136,9 @@ internal class InProcessRunnerContext : IRunnerContext } Dictionary> queuedMessages = this._nextStep.ExportMessages(); - - RunnerStateData result = new(queuedMessages, this._externalRequests.Values.ToList()); + RunnerStateData result = new(instantiatedExecutors: [.. this._executors.Keys], + queuedMessages, + outstandingRequests: [.. this._externalRequests.Values]); return new(result); } @@ -154,7 +155,7 @@ internal class InProcessRunnerContext : IRunnerContext } } - internal ValueTask ImportStateAsync(Checkpoint checkpoint) + internal async ValueTask ImportStateAsync(Checkpoint checkpoint) { if (this.QueuedEvents.Count > 0) { @@ -163,6 +164,11 @@ internal class InProcessRunnerContext : IRunnerContext RunnerStateData importedState = checkpoint.RunnerData; + Task[] executorTasks = importedState.InstantiatedExecutors + .Where(id => !this._executors.ContainsKey(id)) + .Select(id => this.EnsureExecutorAsync(id, tracer: null).AsTask()) + .ToArray(); + this._nextStep = new StepContext(); this._nextStep.ImportMessages(importedState.QueuedMessages); @@ -176,6 +182,6 @@ internal class InProcessRunnerContext : IRunnerContext this._externalRequests[request.RequestId] = request; } - return default; + await Task.WhenAll(executorTasks).ConfigureAwait(false); } } diff --git a/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/Sample/05_Simple_Workflow_Checkpointing.cs b/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/Sample/05_Simple_Workflow_Checkpointing.cs index f0755c064a..7658871f79 100644 --- a/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/Sample/05_Simple_Workflow_Checkpointing.cs +++ b/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/Sample/05_Simple_Workflow_Checkpointing.cs @@ -13,7 +13,7 @@ internal static class Step5EntryPoint { private static CheckpointManager CheckpointManager { get; } = new(); - public static async ValueTask RunAsync(TextWriter writer, Func userGuessCallback) + public static async ValueTask RunAsync(TextWriter writer, Func userGuessCallback, bool rehydrateToRestore = false) { Workflow workflow = Step4EntryPoint.CreateWorkflowInstance(out JudgeExecutor judge); Checkpointed> checkpointed = @@ -30,7 +30,19 @@ internal static class Step5EntryPoint checkpoints.Should().HaveCount(6, "we should have two checkpoints, one for each step"); judge.Tries.Should().Be(2); - await checkpointed.RestoreCheckpointAsync(checkpoints[2], CancellationToken.None).ConfigureAwait(false); + CheckpointInfo targetCheckpoint = checkpoints[2]; + + if (rehydrateToRestore) + { + checkpointed = await InProcessExecution.ResumeStreamAsync(workflow, targetCheckpoint, CheckpointManager, CancellationToken.None) + .ConfigureAwait(false); + handle = checkpointed.Run; + } + else + { + await checkpointed.RestoreCheckpointAsync(checkpoints[2], CancellationToken.None).ConfigureAwait(false); + } + judge.Tries.Should().Be(1); cancellationSource.Dispose(); diff --git a/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/SampleSmokeTest.cs b/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/SampleSmokeTest.cs index 9e3277dfaf..ed1fec105b 100644 --- a/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/SampleSmokeTest.cs +++ b/dotnet/tests/Microsoft.Agents.Workflows.UnitTests/SampleSmokeTest.cs @@ -103,6 +103,25 @@ public class SampleSmokeTest Assert.Equal("You guessed correctly! You Win!", guessResult); } + [Fact] + public async Task Test_RunSample_Step5aAsync() + { + using StringWriter writer = new(); + + VerifyingPlaybackResponder responder = new( + // Iteration 1 + ("Guess the number.", 50), + ("Your guess was too high. Try again.", 23), + + // Iteration 2 + ("Your guess was too high. Try again.", 23), + ("Your guess was too low. Try again.", 42) + ); + + string guessResult = await Step5EntryPoint.RunAsync(writer, userGuessCallback: responder.InvokeNext, rehydrateToRestore: true); + Assert.Equal("You guessed correctly! You Win!", guessResult); + } + [Fact] public async Task Test_RunSample_Step6Async() {