Skip to content

Commit c1b5113

Browse files
authored
Actor Pause/Resume flow (agent-substrate#227)
Implements `PAUSED` state for issue agent-substrate#119. The snapshot files are kept locally on node VM in a separate folder. At resume time, scheduler uses a node VM hint and picks up a worker from the same node where files were stored at suspend time. The local file management solution is temporary and will be replaced once @msau42 introduces a new component that is supposed to manage files on the node VM. - [X] Tests pass - [X] Manual tests with counter demo ``` >kubectl ate create actor my-counter-1 --template ate-demo-counter/counter >curl -X POST -H "Host: my-counter-1.actors.resources.substrate.ate.dev" http://localhost:8000 hello from: 169.254.17.2 | preserved memory count: 1 >curl -X POST -H "Host: my-counter-1.actors.resources.substrate.ate.dev" http://localhost:8000 hello from: 169.254.17.2 | preserved memory count: 2 >kubectl ate pause actor my-counter-1 -o json > curl -X POST -H "Host: my-counter-1.actors.resources.substrate.ate.dev" http://localhost:8000 hello from: 169.254.17.2 | preserved memory count: 3 ``` ### Breaking change This PR introduces a breaking change in the Actor proto. All existing actor needs to be recreated, prior testing PAUSE functionality.
1 parent e09cb61 commit c1b5113

22 files changed

Lines changed: 2195 additions & 350 deletions

File tree

cmd/ateapi/internal/controlapi/functional_test.go

Lines changed: 112 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ package controlapi
1717
import (
1818
"context"
1919
"fmt"
20+
"log"
2021
"net"
2122
"os"
2223
"os/exec"
@@ -60,9 +61,11 @@ var (
6061

6162
func TestMain(m *testing.M) {
6263
cmd := exec.Command("bash", "../../../../hack/run-tool.sh", "setup-envtest", "use", "--print", "path")
64+
var stderr strings.Builder
65+
cmd.Stderr = &stderr
6366
out, err := cmd.Output()
6467
if err != nil {
65-
os.Exit(1)
68+
log.Fatalf("setup-envtest: %v (stderr: %s)", err, stderr.String())
6669
}
6770
binaryAssetsDirectory := strings.TrimSpace(string(out))
6871

@@ -73,19 +76,19 @@ func TestMain(m *testing.M) {
7376

7477
cfg, err = testEnv.Start()
7578
if err != nil {
76-
os.Exit(1)
79+
log.Fatalf("testEnv.Start: %v", err)
7780
}
7881

7982
// Create ate-system namespace
8083
k8sClient, err := kubernetes.NewForConfig(cfg)
8184
if err != nil {
82-
os.Exit(1)
85+
log.Fatalf("kubernetes.NewForConfig: %v", err)
8386
}
8487
_, err = k8sClient.CoreV1().Namespaces().Create(context.Background(), &corev1.Namespace{
8588
ObjectMeta: metav1.ObjectMeta{Name: "ate-system"},
8689
}, metav1.CreateOptions{})
8790
if err != nil && !strings.Contains(err.Error(), "already exists") {
88-
os.Exit(1)
91+
log.Fatalf("create ate-system namespace: %v", err)
8992
}
9093

9194
// Create shared Atelet Pod
@@ -106,14 +109,14 @@ func TestMain(m *testing.M) {
106109
}
107110
createdAtelet, err := k8sClient.CoreV1().Pods("ate-system").Create(context.Background(), ateletPod, metav1.CreateOptions{})
108111
if err != nil && !strings.Contains(err.Error(), "already exists") {
109-
os.Exit(1)
112+
log.Fatalf("create atelet pod: %v", err)
110113
}
111114
if err == nil {
112115
createdAtelet.Status.PodIPs = []corev1.PodIP{{IP: "127.0.0.1"}}
113116
createdAtelet.Status.Phase = corev1.PodRunning
114117
_, err = k8sClient.CoreV1().Pods("ate-system").UpdateStatus(context.Background(), createdAtelet, metav1.UpdateOptions{})
115118
if err != nil {
116-
os.Exit(1)
119+
log.Fatalf("update atelet pod status: %v", err)
117120
}
118121
}
119122

@@ -122,7 +125,7 @@ func TestMain(m *testing.M) {
122125
ateletpb.RegisterAteomHerderServer(ateletGrpcServer, fakeAtelet)
123126
ateletLis, err := net.Listen("tcp", "127.0.0.1:8085")
124127
if err != nil {
125-
os.Exit(1)
128+
log.Fatalf("listen on 127.0.0.1:8085: %v", err)
126129
}
127130
go func() {
128131
if err := ateletGrpcServer.Serve(ateletLis); err != nil {
@@ -136,7 +139,7 @@ func TestMain(m *testing.M) {
136139

137140
err = testEnv.Stop()
138141
if err != nil {
139-
os.Exit(1)
142+
log.Fatalf("testEnv.Stop: %v", err)
140143
}
141144

142145
os.Exit(code)
@@ -868,6 +871,7 @@ func TestResumeActor(t *testing.T) {
868871
ActorTemplate: "tmpl1",
869872
ActorId: id,
870873
Ip: "127.0.0.1",
874+
NodeName: "node1",
871875
}
872876

873877
if diff := cmp.Diff(wantWorker, actorWorker, protocmp.Transform(), protocmp.IgnoreFields(&ateapipb.Worker{}, "version"), protocmp.IgnoreFields(&ateapipb.Worker{}, "worker_pod_uid")); diff != "" {
@@ -1123,14 +1127,112 @@ func TestSuspendActor(t *testing.T) {
11231127
ActorTemplateNamespace: ns,
11241128
ActorTemplateName: "tmpl1",
11251129
Status: ateapipb.Actor_STATUS_SUSPENDED,
1126-
LastSnapshot: fmt.Sprintf("gs://my-bucket/%s/tmpl1/%s/", ns, id),
1130+
LatestSnapshotInfo: &ateapipb.SnapshotInfo{
1131+
Type: ateapipb.SnapshotType_SNAPSHOT_TYPE_EXTERNAL,
1132+
Data: &ateapipb.SnapshotInfo_External{
1133+
External: &ateapipb.ExternalSnapshotInfo{
1134+
SnapshotUriPrefix: fmt.Sprintf("gs://fake-fake-fake/%s/", id),
1135+
},
1136+
},
1137+
},
11271138
},
11281139
}
11291140

1130-
if diff := cmp.Diff(want, getResp, protocmp.Transform(), protocmp.IgnoreFields(&ateapipb.Actor{}, "version", "last_snapshot", "ateom_pod_uid")); diff != "" {
1141+
if diff := cmp.Diff(want, getResp,
1142+
protocmp.Transform(),
1143+
protocmp.IgnoreFields(&ateapipb.Actor{}, "version"),
1144+
protocmp.IgnoreFields(&ateapipb.Actor{}, "ateom_pod_uid"),
1145+
protocmp.FilterField(&ateapipb.ExternalSnapshotInfo{}, "snapshot_uri_prefix", cmp.Comparer(func(x, y string) bool {
1146+
return strings.HasPrefix(y, x)
1147+
})),
1148+
); diff != "" {
11311149
t.Errorf("GetActor response mismatch (-want +got):\n%s", diff)
11321150
}
1151+
}
1152+
1153+
// TestPauseActor tests the full workflow of pausing a running actor.
1154+
// Workflow:
1155+
// 1. Creates a mock ActorTemplate.
1156+
// 2. Creates a mock Atelet Pod on 'node1'.
1157+
// 3. Creates a mock worker Pod on 'node1'.
1158+
// 4. Waits for the WorkerPoolSyncer to mirror the worker to Redis.
1159+
// 5. Creates an actor.
1160+
// 6. Calls ResumeActor to transition it to RUNNING.
1161+
// 7. Calls PauseActor RPC.
1162+
// 8. Verifies that the fake Atelet received the Pause call.
1163+
func TestPauseActor(t *testing.T) {
1164+
ns := namespaceForTest("ns-pause")
1165+
tc := setupTest(t, ns)
1166+
defer tc.cleanup()
1167+
1168+
createTemplate(t, tc, ns)
1169+
1170+
createWorkerPod(t, tc, ns, "worker-1", "node1")
1171+
1172+
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
1173+
ActorTemplateNamespace: ns,
1174+
ActorTemplateName: "tmpl1",
1175+
ActorId: "id1",
1176+
})
1177+
if err != nil {
1178+
t.Fatalf("CreateActor failed: %v", err)
1179+
}
1180+
id := "id1"
1181+
1182+
// Resume first to make it running
1183+
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
1184+
ActorId: id,
1185+
})
1186+
if err != nil {
1187+
t.Fatalf("ResumeActor failed: %v", err)
1188+
}
1189+
1190+
// Pause
1191+
_, err = tc.client.PauseActor(context.Background(), &ateapipb.PauseActorRequest{
1192+
ActorId: id,
1193+
})
1194+
if err != nil {
1195+
t.Fatalf("PauseActor failed: %v", err)
1196+
}
11331197

1198+
if !tc.fakeAtelet.CheckpointCalled {
1199+
t.Errorf("expected atelet Checkpoint to be called")
1200+
}
1201+
1202+
getResp, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{
1203+
ActorId: id,
1204+
})
1205+
if err != nil {
1206+
t.Fatalf("GetActor failed: %v", err)
1207+
}
1208+
want := &ateapipb.GetActorResponse{
1209+
Actor: &ateapipb.Actor{
1210+
ActorId: id,
1211+
ActorTemplateNamespace: ns,
1212+
ActorTemplateName: "tmpl1",
1213+
Status: ateapipb.Actor_STATUS_PAUSED,
1214+
LatestSnapshotInfo: &ateapipb.SnapshotInfo{
1215+
Type: ateapipb.SnapshotType_SNAPSHOT_TYPE_LOCAL,
1216+
Data: &ateapipb.SnapshotInfo_Local{
1217+
Local: &ateapipb.LocalSnapshotInfo{
1218+
SnapshotPrefix: "id1",
1219+
NodeVmsWithLocalSnapshots: []string{"node1"},
1220+
},
1221+
},
1222+
},
1223+
},
1224+
}
1225+
1226+
if diff := cmp.Diff(want, getResp,
1227+
protocmp.Transform(),
1228+
protocmp.IgnoreFields(&ateapipb.Actor{}, "version"),
1229+
protocmp.IgnoreFields(&ateapipb.Actor{}, "ateom_pod_uid"),
1230+
protocmp.FilterField(&ateapipb.LocalSnapshotInfo{}, "snapshot_prefix", cmp.Comparer(func(x, y string) bool {
1231+
return strings.HasPrefix(y, x)
1232+
})),
1233+
); diff != "" {
1234+
t.Errorf("GetActor response mismatch (-want +got):\n%s", diff)
1235+
}
11341236
}
11351237

11361238
// TestValidation tests the negative validation cases for all gRPC methods.
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package controlapi
16+
17+
import (
18+
"context"
19+
"errors"
20+
21+
"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
22+
"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
23+
"google.golang.org/grpc/codes"
24+
"google.golang.org/grpc/status"
25+
)
26+
27+
func (s *Service) PauseActor(ctx context.Context, req *ateapipb.PauseActorRequest) (*ateapipb.PauseActorResponse, error) {
28+
if err := validatePauseActorRequest(req); err != nil {
29+
return nil, err
30+
}
31+
32+
actor, err := s.actorWorkflow.PauseActor(ctx, req.GetActorId())
33+
if err != nil {
34+
if errors.Is(err, store.ErrPersistenceRetry) {
35+
return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
36+
}
37+
if errors.Is(err, store.ErrNotFound) {
38+
return nil, status.Errorf(codes.NotFound, "Actor %s not found", req.GetActorId())
39+
}
40+
return nil, err
41+
}
42+
43+
return &ateapipb.PauseActorResponse{Actor: actor}, nil
44+
}
45+
46+
func validatePauseActorRequest(req *ateapipb.PauseActorRequest) error {
47+
if req.GetActorId() == "" {
48+
return status.Error(codes.InvalidArgument, "id is required")
49+
}
50+
return nil
51+
}

cmd/ateapi/internal/controlapi/syncer.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@ func (s *WorkerPoolSyncer) syncWorkerToStore(ctx context.Context, pod *corev1.Po
120120
WorkerPod: pod.Name,
121121
Ip: pod.Status.PodIP,
122122
WorkerPodUid: string(pod.UID),
123+
NodeName: pod.Spec.NodeName,
123124
})
124125
if err != nil && !errors.Is(err, store.ErrAlreadyExists) {
125126
slog.ErrorContext(ctx, "Failed to create worker in store", slog.Any("err", err))

cmd/ateapi/internal/controlapi/syncer_test.go

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,15 @@ func TestSyncer_DeleteBoundWorker_ClearsActor(t *testing.T) {
183183
ActorId: actorID, ActorTemplateNamespace: ns, ActorTemplateName: "tmpl",
184184
Status: ateapipb.Actor_STATUS_RUNNING,
185185
AteomPodNamespace: ns, AteomPodName: pod, AteomPodIp: ip,
186-
LastSnapshot: "gs://snapshots/last", InProgressSnapshot: "gs://snapshots/partial",
186+
InProgressSnapshot: "gs://snapshots/partial",
187+
LatestSnapshotInfo: &ateapipb.SnapshotInfo{
188+
Type: ateapipb.SnapshotType_SNAPSHOT_TYPE_EXTERNAL,
189+
Data: &ateapipb.SnapshotInfo_External{
190+
External: &ateapipb.ExternalSnapshotInfo{
191+
SnapshotUriPrefix: "gs://snapshots/last",
192+
},
193+
},
194+
},
187195
}); err != nil {
188196
t.Fatalf("create actor: %v", err)
189197
}
@@ -210,7 +218,7 @@ func TestSyncer_DeleteBoundWorker_ClearsActor(t *testing.T) {
210218
if got.AteomPodName != "" || got.AteomPodNamespace != "" || got.AteomPodIp != "" || got.InProgressSnapshot != "" {
211219
t.Errorf("bind fields not cleared: %+v", got)
212220
}
213-
if got.LastSnapshot == "" {
214-
t.Errorf("LastSnapshot must be preserved")
221+
if got.GetLatestSnapshotInfo().GetExternal().SnapshotUriPrefix == "" {
222+
t.Errorf("External SnapshotUriPrefix must be preserved")
215223
}
216224
}

cmd/ateapi/internal/controlapi/workflow.go

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ func (w *ActorWorkflow) SuspendActor(ctx context.Context, id string) (*ateapipb.
170170
state := &SuspendState{}
171171

172172
// Acquire lock and get the timeout context for the workflow
173-
// Lock TTL is 7 seconds, with 2 seconds padding for workflow timeout
173+
// Lock TTL is 30 seconds, with 2 seconds padding for workflow timeout
174174
ctx, releaseLock, err := w.acquireActorLock(ctx, id, 30*time.Second, 2*time.Second)
175175
if err != nil {
176176
return nil, err
@@ -191,6 +191,35 @@ func (w *ActorWorkflow) SuspendActor(ctx context.Context, id string) (*ateapipb.
191191
return state.Actor, nil
192192
}
193193

194+
// PauseActor executes the workflow to pause a running actor. Idempotent.
195+
func (w *ActorWorkflow) PauseActor(ctx context.Context, id string) (*ateapipb.Actor, error) {
196+
input := &PauseInput{
197+
ActorID: id,
198+
}
199+
state := &PauseState{}
200+
201+
// Acquire lock and get the timeout context for the workflow
202+
// Lock TTL is 30 seconds, with 2 seconds padding for workflow timeout
203+
ctx, releaseLock, err := w.acquireActorLock(ctx, id, 30*time.Second, 2*time.Second)
204+
if err != nil {
205+
return nil, err
206+
}
207+
defer releaseLock()
208+
209+
steps := []WorkflowStep[*PauseInput, *PauseState]{
210+
&LoadActorForPauseStep{store: w.store, actorTemplateLister: w.actorTemplateLister},
211+
&MarkPausingStep{store: w.store},
212+
&CallAteletPauseStep{dialer: w.dialer},
213+
&FinalizePausedStep{store: w.store},
214+
}
215+
216+
if err := RunWorkflow(ctx, input, state, steps); err != nil {
217+
return nil, err
218+
}
219+
220+
return state.Actor, nil
221+
}
222+
194223
func (w *ActorWorkflow) acquireActorLock(ctx context.Context, id string, ttl time.Duration, padding time.Duration) (context.Context, func(), error) {
195224
lockKey := "lock:actor:" + id
196225
lockValue := uuid.New().String()

0 commit comments

Comments
 (0)