Skip to content

Commit 2592bbb

Browse files
author
Jacob Alber
committed
[BREAKING] refactor: Decouple Checkpointing and Execution APIs
With this change, Checkpointing becomes an property of an IWorkflowExecutionEnvironment. This lets environments that are tightly-coupled to their CheckpointManager avoid needing to present APIs that would not work (e.g. taking in an InMemory CheckpointManager for Durable Tasks, for example)
1 parent 1b87a07 commit 2592bbb

29 files changed

Lines changed: 279 additions & 294 deletions

‎dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndRehydrate/Program.cs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,10 @@ private static async Task Main()
3232
var checkpoints = new List<CheckpointInfo>();
3333

3434
// Execute the workflow and save checkpoints
35-
await using Checkpointed<StreamingRun> checkpointedRun = await InProcessExecution
35+
await using StreamingRun checkpointedRun = await InProcessExecution
3636
.StreamAsync(workflow, NumberSignal.Init, checkpointManager);
3737

38-
await foreach (WorkflowEvent evt in checkpointedRun.Run.WatchStreamAsync())
38+
await foreach (WorkflowEvent evt in checkpointedRun.WatchStreamAsync())
3939
{
4040
if (evt is ExecutorCompletedEvent executorCompletedEvt)
4141
{
@@ -72,10 +72,10 @@ private static async Task Main()
7272
Console.WriteLine($"\n\nHydrating a new workflow instance from the {CheckpointIndex + 1}th checkpoint.");
7373
CheckpointInfo savedCheckpoint = checkpoints[CheckpointIndex];
7474

75-
await using Checkpointed<StreamingRun> newCheckpointedRun =
75+
await using StreamingRun newCheckpointedRun =
7676
await InProcessExecution.ResumeStreamAsync(newWorkflow, savedCheckpoint, checkpointManager);
7777

78-
await foreach (WorkflowEvent evt in newCheckpointedRun.Run.WatchStreamAsync())
78+
await foreach (WorkflowEvent evt in newCheckpointedRun.WatchStreamAsync())
7979
{
8080
if (evt is ExecutorCompletedEvent executorCompletedEvt)
8181
{

‎dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointAndResume/Program.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,10 @@ private static async Task Main()
3131
var checkpoints = new List<CheckpointInfo>();
3232

3333
// Execute the workflow and save checkpoints
34-
await using Checkpointed<StreamingRun> checkpointedRun = await InProcessExecution
34+
await using StreamingRun checkpointedRun = await InProcessExecution
3535
.StreamAsync(workflow, NumberSignal.Init, checkpointManager)
3636
;
37-
await foreach (WorkflowEvent evt in checkpointedRun.Run.WatchStreamAsync())
37+
await foreach (WorkflowEvent evt in checkpointedRun.WatchStreamAsync())
3838
{
3939
if (evt is ExecutorCompletedEvent executorCompletedEvt)
4040
{
@@ -71,7 +71,7 @@ private static async Task Main()
7171
CheckpointInfo savedCheckpoint = checkpoints[CheckpointIndex];
7272
// Note that we are restoring the state directly to the same run instance.
7373
await checkpointedRun.RestoreCheckpointAsync(savedCheckpoint, CancellationToken.None);
74-
await foreach (WorkflowEvent evt in checkpointedRun.Run.WatchStreamAsync())
74+
await foreach (WorkflowEvent evt in checkpointedRun.WatchStreamAsync())
7575
{
7676
if (evt is ExecutorCompletedEvent executorCompletedEvt)
7777
{

‎dotnet/samples/GettingStarted/Workflows/Checkpoint/CheckpointWithHumanInTheLoop/Program.cs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,17 +34,17 @@ private static async Task Main()
3434
var checkpoints = new List<CheckpointInfo>();
3535

3636
// Execute the workflow and save checkpoints
37-
await using Checkpointed<StreamingRun> checkpointedRun = await InProcessExecution
37+
await using StreamingRun checkpointedRun = await InProcessExecution
3838
.StreamAsync(workflow, new SignalWithNumber(NumberSignal.Init), checkpointManager)
3939
;
40-
await foreach (WorkflowEvent evt in checkpointedRun.Run.WatchStreamAsync())
40+
await foreach (WorkflowEvent evt in checkpointedRun.WatchStreamAsync())
4141
{
4242
switch (evt)
4343
{
4444
case RequestInfoEvent requestInputEvt:
4545
// Handle `RequestInfoEvent` from the workflow
4646
ExternalResponse response = HandleExternalRequest(requestInputEvt.Request);
47-
await checkpointedRun.Run.SendResponseAsync(response);
47+
await checkpointedRun.SendResponseAsync(response);
4848
break;
4949
case ExecutorCompletedEvent executorCompletedEvt:
5050
Console.WriteLine($"* Executor {executorCompletedEvt.ExecutorId} completed.");
@@ -77,14 +77,14 @@ private static async Task Main()
7777
CheckpointInfo savedCheckpoint = checkpoints[CheckpointIndex];
7878
// Note that we are restoring the state directly to the same run instance.
7979
await checkpointedRun.RestoreCheckpointAsync(savedCheckpoint, CancellationToken.None);
80-
await foreach (WorkflowEvent evt in checkpointedRun.Run.WatchStreamAsync())
80+
await foreach (WorkflowEvent evt in checkpointedRun.WatchStreamAsync())
8181
{
8282
switch (evt)
8383
{
8484
case RequestInfoEvent requestInputEvt:
8585
// Handle `RequestInfoEvent` from the workflow
8686
ExternalResponse response = HandleExternalRequest(requestInputEvt.Request);
87-
await checkpointedRun.Run.SendResponseAsync(response);
87+
await checkpointedRun.SendResponseAsync(response);
8888
break;
8989
case ExecutorCompletedEvent executorCompletedEvt:
9090
Console.WriteLine($"* Executor {executorCompletedEvt.ExecutorId} completed.");

‎dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

3+
using System.Collections.Generic;
34
using System.Text.Json;
45
using System.Threading.Tasks;
56
using Microsoft.Agents.AI.Workflows.Checkpointing;
@@ -54,4 +55,7 @@ ValueTask<CheckpointInfo> ICheckpointManager.CommitCheckpointAsync(string runId,
5455

5556
ValueTask<Checkpoint> ICheckpointManager.LookupCheckpointAsync(string runId, CheckpointInfo checkpointInfo)
5657
=> this._impl.LookupCheckpointAsync(runId, checkpointInfo);
58+
59+
ValueTask<IEnumerable<CheckpointInfo>> ICheckpointManager.RetrieveIndexAsync(string runId, CheckpointInfo? withParent)
60+
=> this._impl.RetrieveIndexAsync(runId, withParent);
5761
}

‎dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointableRunBase.cs‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System;
4+
using System.Collections.Generic;
5+
using System.Threading;
6+
using System.Threading.Tasks;
7+
using Microsoft.Agents.AI.Workflows.Checkpointing;
8+
9+
namespace Microsoft.Agents.AI.Workflows;
10+
11+
/// <summary>
12+
/// Represents a base object for a workflow run that may support checkpointing.
13+
/// </summary>
14+
public abstract class CheckpointableRunBase
15+
{
16+
// TODO: Rename Context?
17+
private readonly ICheckpointingHandle _checkpointingHandle;
18+
19+
internal CheckpointableRunBase(ICheckpointingHandle checkpointingHandle)
20+
{
21+
this._checkpointingHandle = checkpointingHandle;
22+
}
23+
24+
/// <inheritdoc cref="ICheckpointingHandle.IsCheckpointingEnabled"/>
25+
public bool IsCheckpointingEnabled => this._checkpointingHandle.IsCheckpointingEnabled;
26+
27+
/// <inheritdoc cref="ICheckpointingHandle.Checkpoints"/>
28+
public IReadOnlyList<CheckpointInfo> Checkpoints => this._checkpointingHandle.Checkpoints ?? [];
29+
30+
/// <summary>
31+
/// Gets the most recent checkpoint information.
32+
/// </summary>
33+
public CheckpointInfo? LastCheckpoint
34+
{
35+
get
36+
{
37+
if (!this.IsCheckpointingEnabled)
38+
{
39+
return null;
40+
}
41+
42+
var checkpoints = this.Checkpoints;
43+
return checkpoints.Count > 0 ? checkpoints[checkpoints.Count - 1] : null;
44+
}
45+
}
46+
47+
/// <inheritdoc cref="ICheckpointingHandle.RestoreCheckpointAsync"/>
48+
public ValueTask RestoreCheckpointAsync(CheckpointInfo checkpointInfo, CancellationToken cancellationToken = default)
49+
=> this._checkpointingHandle.RestoreCheckpointAsync(checkpointInfo, cancellationToken);
50+
}

‎dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointed.cs‎

Lines changed: 0 additions & 66 deletions
This file was deleted.

‎dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CheckpointManagerImpl.cs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

3+
using System.Collections.Generic;
34
using System.Threading.Tasks;
45

56
namespace Microsoft.Agents.AI.Workflows.Checkpointing;
@@ -27,4 +28,7 @@ public async ValueTask<Checkpoint> LookupCheckpointAsync(string runId, Checkpoin
2728
TStoreObject result = await this._store.RetrieveCheckpointAsync(runId, checkpointInfo).ConfigureAwait(false);
2829
return this._marshaller.Marshal<Checkpoint>(result);
2930
}
31+
32+
public ValueTask<IEnumerable<CheckpointInfo>> RetrieveIndexAsync(string runId, CheckpointInfo? withParent = null)
33+
=> this._store.RetrieveIndexAsync(runId, withParent);
3034
}

‎dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/ICheckpointManager.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,4 +27,16 @@ internal interface ICheckpointManager
2727
/// cref="Checkpoint"/> associated with the specified <paramref name="checkpointInfo"/>.</returns>
2828
/// <exception cref="KeyNotFoundException">Thrown if the checkpoint is not found.</exception>
2929
ValueTask<Checkpoint> LookupCheckpointAsync(string runId, CheckpointInfo checkpointInfo);
30+
31+
/// <summary>
32+
/// Asynchronously retrieves the collection of checkpoint information for the specified run identifier, optionally
33+
/// filtered by a parent checkpoint.
34+
/// </summary>
35+
/// <param name="runId">The unique identifier of the run for which to retrieve checkpoint information. Cannot be null or empty.</param>
36+
/// <param name="withParent">An optional parent checkpoint to filter the results. If specified, only checkpoints with the given parent are
37+
/// returned; otherwise, all checkpoints for the run are included.</param>
38+
/// <returns>A value task representing the asynchronous operation. The result contains a collection of <see
39+
/// cref="CheckpointInfo"/> objects associated with the specified run. The collection is empty if no checkpoints are
40+
/// found.</returns>
41+
ValueTask<IEnumerable<CheckpointInfo>> RetrieveIndexAsync(string runId, CheckpointInfo? withParent = null);
3042
}

‎dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/ICheckpointingHandle.cs‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,21 @@ namespace Microsoft.Agents.AI.Workflows.Checkpointing;
88

99
internal interface ICheckpointingHandle
1010
{
11-
// TODO: Convert this to a multi-timeline (e.g.: Live timeline + forks for orphaned checkpoints due to timetravel)
11+
/// <summary>
12+
/// Gets a value indicating whether checkpointing is enabled for the current operation or process.
13+
/// </summary>
14+
bool IsCheckpointingEnabled { get; }
15+
16+
/// <summary>
17+
/// Gets a read-only list of checkpoint information associated with the current context.
18+
/// </summary>
1219
IReadOnlyList<CheckpointInfo> Checkpoints { get; }
1320

21+
/// <summary>
22+
/// Restores the system state from the specified checkpoint asynchronously.
23+
/// </summary>
24+
/// <param name="checkpointInfo">The checkpoint information that identifies the state to restore. Cannot be null.</param>
25+
/// <param name="cancellationToken">A cancellation token that can be used to cancel the restore operation.</param>
26+
/// <returns>A <see cref="ValueTask"/> that represents the asynchronous restore operation.</returns>
1427
ValueTask RestoreCheckpointAsync(CheckpointInfo checkpointInfo, CancellationToken cancellationToken = default);
1528
}

‎dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/InMemoryCheckpointManager.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,4 +60,7 @@ public ValueTask<Checkpoint> LookupCheckpointAsync(string runId, CheckpointInfo
6060

6161
public bool TryGetLastCheckpoint(string runId, [NotNullWhen(true)] out CheckpointInfo? checkpoint)
6262
=> this.GetRunStore(runId).TryGetLastCheckpointInfo(out checkpoint);
63+
64+
public ValueTask<IEnumerable<CheckpointInfo>> RetrieveIndexAsync(string runId, CheckpointInfo? withParent = null)
65+
=> new(this.GetRunStore(runId).CheckpointIndex.AsReadOnly());
6366
}

0 commit comments

Comments
 (0)