Skip to content

Some improvements to the way we allocate and track snapshot ids #784

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Julian Gutierrez Oschmann (juli4n) wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Some improvements to the way we allocate and track snapshot ids #784

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
225 changes: 112 additions & 113 deletions benchmarking/locust/common/ateapi_pb2.py
View file
Open in desktop

Large diffs are not rendered by default.

6 changes: 3 additions & 3 deletions benchmarking/locust/common/ateapi_pb2_grpc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ def __init__(self, channel):
self.UpdateActor = channel.unary_unary(
'/ateapi.Control/UpdateActor',
request_serializer=ateapi__pb2.UpdateActorRequest.SerializeToString,
response_deserializer=ateapi__pb2.UpdateActorResponse.FromString,
response_deserializer=ateapi__pb2.Actor.FromString,
_registered_method=True)
self.SuspendActor = channel.unary_unary(
'/ateapi.Control/SuspendActor',
Expand Down Expand Up @@ -289,7 +289,7 @@ def add_ControlServicer_to_server(servicer, server):
'UpdateActor': grpc.unary_unary_rpc_method_handler(
servicer.UpdateActor,
request_deserializer=ateapi__pb2.UpdateActorRequest.FromString,
response_serializer=ateapi__pb2.UpdateActorResponse.SerializeToString,
response_serializer=ateapi__pb2.Actor.SerializeToString,
),
'SuspendActor': grpc.unary_unary_rpc_method_handler(
servicer.SuspendActor,
Expand Down Expand Up @@ -448,7 +448,7 @@ def UpdateActor(request,
target,
'/ateapi.Control/UpdateActor',
ateapi__pb2.UpdateActorRequest.SerializeToString,
ateapi__pb2.UpdateActorResponse.FromString,
ateapi__pb2.Actor.FromString,
options,
channel_credentials,
insecure,
Expand Down
21 changes: 10 additions & 11 deletions cmd/ateapi/internal/controlapi/actor_snapshot.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ func (s *Service) GetActorSnapshot(ctx context.Context, req *ateapipb.GetActorSn
if err := validateActorSnapshotRef(req.GetSnapshot(), "snapshot"); err != nil {
return nil, err
}
snapshot, _, _, _, err := s.getActorSnapshot(ctx, req.GetSnapshot())
snapshot, _, _, err := s.getActorSnapshot(ctx, req.GetSnapshot())
if errors.Is(err, store.ErrNotFound) {
return nil, status.Error(codes.NotFound, "ActorSnapshot not found")
}
Expand Down Expand Up @@ -128,7 +128,7 @@ func (s *Service) UpdateActorSnapshotTag(ctx context.Context, req *ateapipb.Upda
}
in := req.GetTag()
atespace, name := in.GetMetadata().GetAtespace(), in.GetMetadata().GetName()
_, _, current, err := s.persistence.GetActorSnapshotByTag(ctx, atespace, name)
_, current, err := s.persistence.GetActorSnapshotByTag(ctx, atespace, name)
if errors.Is(err, store.ErrNotFound) {
return nil, status.Errorf(codes.NotFound, "ActorSnapshot tag %s/%s not found", atespace, name)
}
Expand Down Expand Up @@ -202,29 +202,28 @@ func (s *Service) DeleteActorSnapshotTag(ctx context.Context, req *ateapipb.Dele
return tag, nil
}

func (s *Service) getActorSnapshot(ctx context.Context, ref *ateapipb.ActorSnapshotRef) (*ateapipb.ActorSnapshot, string, *ateapipb.ObjectRef, *ateapipb.ActorSnapshotTag, error) {
func (s *Service) getActorSnapshot(ctx context.Context, ref *ateapipb.ActorSnapshotRef) (*ateapipb.ActorSnapshot, *ateapipb.ObjectRef, *ateapipb.ActorSnapshotTag, error) {
var snapshot *ateapipb.ActorSnapshot
var tag *ateapipb.ActorSnapshotTag
var location string
var err error
switch ref.GetReference().(type) {
case *ateapipb.ActorSnapshotRef_Snapshot:
canonical := ref.GetSnapshot()
snapshot, location, err = s.persistence.GetActorSnapshot(ctx, canonical.GetAtespace(), canonical.GetName())
snapshot, err = s.persistence.GetActorSnapshot(ctx, canonical.GetAtespace(), canonical.GetName())
case *ateapipb.ActorSnapshotRef_Tag:
snapshot, location, tag, err = s.persistence.GetActorSnapshotByTag(ctx, ref.GetTag().GetAtespace(), ref.GetTag().GetName())
snapshot, tag, err = s.persistence.GetActorSnapshotByTag(ctx, ref.GetTag().GetAtespace(), ref.GetTag().GetName())
default:
return nil, "", nil, nil, store.ErrNotFound
return nil, nil, nil, store.ErrNotFound
}
if err != nil {
return nil, "", nil, nil, err
return nil, nil, nil, err
}
canonical := &ateapipb.ObjectRef{Atespace: snapshot.GetMetadata().GetAtespace(), Name: snapshot.GetMetadata().GetName()}
return snapshot, location, canonical, tag, nil
return snapshot, canonical, tag, nil
}

func (s *Service) lockActorSnapshot(ctx context.Context, ref *ateapipb.ActorSnapshotRef) (*store.Lock, *ateapipb.ActorSnapshot, *ateapipb.ObjectRef, *ateapipb.ActorSnapshotTag, error) {
_, _, canonical, _, err := s.getActorSnapshot(ctx, ref)
_, canonical, _, err := s.getActorSnapshot(ctx, ref)
if errors.Is(err, store.ErrNotFound) {
return nil, nil, nil, nil, status.Error(codes.NotFound, "ActorSnapshot not found")
}
Expand All @@ -238,7 +237,7 @@ func (s *Service) lockActorSnapshot(ctx context.Context, ref *ateapipb.ActorSnap
if err != nil {
return nil, nil, nil, nil, fmt.Errorf("while locking actor snapshot: %w", err)
}
snapshot, _, lockedCanonical, tag, err := s.getActorSnapshot(lock.Context(), ref)
snapshot, lockedCanonical, tag, err := s.getActorSnapshot(lock.Context(), ref)
if err != nil || canonical.GetAtespace() != lockedCanonical.GetAtespace() || canonical.GetName() != lockedCanonical.GetName() {
lock.Close()
if errors.Is(err, store.ErrNotFound) {
Expand Down
5 changes: 3 additions & 2 deletions cmd/ateapi/internal/controlapi/crash.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,8 +77,9 @@ func crashActor(ctx context.Context, st store.Interface, actorRef resources.Acto

actor.Status = ateapipb.Actor_STATUS_CRASHED

// InProgressSnapshot is kept for debugging; failed workflow
// steps must never promote it to an ActorSnapshot.
// InProgressSnapshotName and InProgressLocalSnapshotName are kept for
// debugging; failed workflow steps must never promote either of them to an
// ActorSnapshot or to LocalSnapshotInfo.
actor.WorkerAssignment = nil

_, err = st.UpdateActor(ctx, actor, actor.GetMetadata().GetVersion())
Expand Down
12 changes: 6 additions & 6 deletions cmd/ateapi/internal/controlapi/crash_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ func seedActor(t *testing.T, ctx context.Context, st store.Interface, actorRef r
WorkerPodUid: "uid",
WorkerPodIp: "1.2.3.4",
},
InProgressSnapshot: "gs://snapshots/actor-1/reserved",
InProgressSnapshotName: "reserved-snapshot",
}); err != nil {
t.Fatalf("seed actor: %v", err)
}
Expand Down Expand Up @@ -85,9 +85,9 @@ func seedWorker(t *testing.T, ctx context.Context, st store.Interface, actorRef
func seedUnboundActor(t *testing.T, ctx context.Context, st store.Interface, actorRef resources.ActorRef) {
t.Helper()
if _, err := st.CreateActor(ctx, &ateapipb.Actor{
Metadata: &ateapipb.ResourceMetadata{Name: actorRef.Name, Atespace: actorRef.Atespace},
Status: ateapipb.Actor_STATUS_RUNNING,
InProgressSnapshot: "gs://snapshots/actor-1/reserved",
Metadata: &ateapipb.ResourceMetadata{Name: actorRef.Name, Atespace: actorRef.Atespace},
Status: ateapipb.Actor_STATUS_RUNNING,
InProgressSnapshotName: "reserved-snapshot",
}); err != nil {
t.Fatalf("seed unbound actor: %v", err)
}
Expand All @@ -105,8 +105,8 @@ func assertCrashed(t *testing.T, ctx context.Context, st store.Interface, actorR
t.Errorf("status = %v, want %v", got.GetStatus(), ateapipb.Actor_STATUS_CRASHED)
}
// Keep the snapshot uri for debugging.
if got.GetInProgressSnapshot() == "" {
t.Error(`InProgressSnapshot = "", want preserved`)
if got.GetInProgressSnapshotName() == "" {
t.Error(`InProgressSnapshotName = "", want preserved`)
}
if got.GetWorkerAssignment() != nil {
t.Errorf("WorkerAssignment = %v, want cleared", got.GetWorkerAssignment())
Expand Down
6 changes: 4 additions & 2 deletions cmd/ateapi/internal/controlapi/create_actor_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,8 @@ func TestCreateActor_RejectsDifferentTemplateForDataSnapshot(t *testing.T) {
SourceActor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: "source"},
ActorTemplateUid: string(tmpl.GetUID()),
ContentScope: ateapipb.SnapshotContentScope_SNAPSHOT_CONTENT_SCOPE_DATA,
}, "gs://snapshots/data")
SnapshotUri: "gs://snapshots/snapshots/" + testAtespace + "/data-snapshot",
})
if err != nil {
t.Fatalf("CreateActorSnapshot: %v", err)
}
Expand Down Expand Up @@ -237,7 +238,8 @@ func TestCreateActor_RejectsSnapshotWithExternalVolumes(t *testing.T) {
snapshot, err := tc.persistence.CreateActorSnapshot(context.Background(), &ateapipb.ActorSnapshot{
Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "external-volume-snapshot"},
ActorTemplateUid: string(template.GetUID()),
}, "gs://snapshots/external-volume")
SnapshotUri: "gs://snapshots/snapshots/" + testAtespace + "/external-volume-snapshot",
})
if err != nil {
t.Fatalf("CreateActorSnapshot: %v", err)
}
Expand Down
19 changes: 11 additions & 8 deletions cmd/ateapi/internal/controlapi/functional_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -445,8 +445,9 @@ func createAtespace(t *testing.T, tc *testContext, name string) {
func createActorSnapshot(t *testing.T, tc *testContext, name string) *ateapipb.ObjectRef {
t.Helper()
if _, err := tc.persistence.CreateActorSnapshot(context.Background(), &ateapipb.ActorSnapshot{
Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: name},
}, "gs://my-bucket/"+name); err != nil {
Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: name},
SnapshotUri: "gs://my-bucket/snapshots/" + testAtespace + "/" + name,
}); err != nil {
t.Fatalf("CreateActorSnapshot(%s) failed: %v", name, err)
}
return &ateapipb.ObjectRef{Atespace: testAtespace, Name: name}
Expand Down Expand Up @@ -533,7 +534,8 @@ func createTemplateWithContainersAndVolumes(t *testing.T, tc *testContext, ns st
ActorTemplateName: createdTemplate.GetName(),
ActorTemplateUid: string(createdTemplate.GetUID()),
ContentScope: ateapipb.SnapshotContentScope_SNAPSHOT_CONTENT_SCOPE_FULL,
}, "gs://my-bucket/my-folder"); err != nil {
SnapshotUri: "gs://fake-fake-fake/snapshots/" + resources.GoldenActorAtespace + "/" + goldenSnapshot,
}); err != nil {
t.Fatalf("failed to create golden ActorSnapshot: %v", err)
}
createdTemplate.Status = atev1alpha1.ActorTemplateStatus{
Expand Down Expand Up @@ -2304,7 +2306,6 @@ func TestPauseActor(t *testing.T) {
ActorTemplateName: "tmpl1",
Status: ateapipb.Actor_STATUS_PAUSED,
LocalSnapshotInfo: &ateapipb.LocalSnapshotInfo{
SnapshotPrefix: name,
NodeVmsWithLocalSnapshots: []string{"node1"},
},
}
Expand All @@ -2314,12 +2315,14 @@ func TestPauseActor(t *testing.T) {
ignoreUID,
ignoreVersion,
ignoreTimestamps,
protocmp.FilterField(&ateapipb.LocalSnapshotInfo{}, "snapshot_prefix", cmp.Comparer(func(x, y string) bool {
return strings.HasPrefix(y, x)
})),
protocmp.IgnoreFields(&ateapipb.WorkerAssignment{}, "worker_pod_uid"),
protocmp.IgnoreFields(&ateapipb.LocalSnapshotInfo{}, "snapshot_name"),
); diff != "" {
t.Errorf("GetActor response mismatch (-want +got):\n%s", diff)
}
if getResp.GetLocalSnapshotInfo().GetSnapshotName() == "" {
t.Error("LocalSnapshotInfo.SnapshotName is empty, want the name the pause checkpointed under")
}
}

// TestUpdateActor_Success verifies UpdateActor replaces the actor's
Expand Down Expand Up @@ -2516,7 +2519,7 @@ func TestUpdateActorSnapshotTag_Success(t *testing.T) {
t.Errorf("UpdateActorSnapshotTag response mismatch (-want +got):\n%s", diff)
}

_, _, storedTag, err := tc.persistence.GetActorSnapshotByTag(ctx, testAtespace, tagName)
_, storedTag, err := tc.persistence.GetActorSnapshotByTag(ctx, testAtespace, tagName)
if err != nil {
t.Fatalf("GetActorSnapshotByTag failed: %v", err)
}
Expand Down
5 changes: 4 additions & 1 deletion cmd/ateapi/internal/controlapi/syncer.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -385,7 +385,10 @@ func (s *WorkerPoolSyncer) releaseActorOnDeadWorker(ctx context.Context, namespa

actor.Status = ateapipb.Actor_STATUS_CRASHED
actor.WorkerAssignment = nil
actor.InProgressSnapshot = ""
// Both in-progress checkpoints die with the worker: the durable one was
// never uploaded, the local one lived on the node that went away.
actor.InProgressSnapshotName = ""
actor.InProgressLocalSnapshotName = ""

_, err = s.persistence.UpdateActor(ctx, actor, actor.GetMetadata().GetVersion())

Expand Down
9 changes: 6 additions & 3 deletions cmd/ateapi/internal/controlapi/syncer_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -262,8 +262,11 @@ func TestSyncer_DeleteBoundWorker_ClearsActor(t *testing.T) {
WorkerAssignment: &ateapipb.WorkerAssignment{
WorkerNamespace: ns, WorkerPool: pool, WorkerPod: pod, WorkerPodIp: ip,
},
InProgressSnapshot: "gs://snapshots/partial",
LatestSnapshot: &ateapipb.ObjectRef{Atespace: "team-orphan", Name: "last"},
// Both in-progress checkpoints are set so the assertion below covers the
// shared crash path, which cannot know which workflow was in flight.
InProgressSnapshotName: "partial-snapshot",
InProgressLocalSnapshotName: "partial-local-snapshot",
LatestSnapshot: &ateapipb.ObjectRef{Atespace: "team-orphan", Name: "last"},
})
if err != nil {
t.Fatalf("create actor: %v", err)
Expand Down Expand Up @@ -297,7 +300,7 @@ func TestSyncer_DeleteBoundWorker_ClearsActor(t *testing.T) {
}); err != nil {
t.Fatalf("actor not reset to CRASHED: %v", err)
}
if got.GetWorkerAssignment() != nil || got.InProgressSnapshot != "" {
if got.GetWorkerAssignment() != nil || got.InProgressSnapshotName != "" || got.InProgressLocalSnapshotName != "" {
t.Errorf("bind fields not cleared: %+v", got)
}
if got.GetLatestSnapshot().GetName() == "" {
Expand Down
16 changes: 7 additions & 9 deletions cmd/ateapi/internal/controlapi/workflow_pause.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,9 @@ package controlapi

import (
"context"
"crypto/rand"
"errors"
"fmt"
"log/slog"
"time"

"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
"github.com/agent-substrate/substrate/internal/ateattr"
Expand Down Expand Up @@ -95,7 +93,7 @@ func (s *MarkPausingStep) CheckPrerequisite(ctx context.Context, input *PauseInp
}
func (s *MarkPausingStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error {
state.Actor.Status = ateapipb.Actor_STATUS_PAUSING
state.Actor.InProgressSnapshot = fmt.Sprintf("%s-%s-%s", state.Actor.GetMetadata().GetName(), time.Now().Format(time.RFC3339), rand.Text())
state.Actor.InProgressLocalSnapshotName = resources.NewSnapshotName()
updatedActor, err := s.store.UpdateActor(ctx, state.Actor, state.Actor.GetMetadata().GetVersion())
if err != nil {
return err
Expand Down Expand Up @@ -135,7 +133,7 @@ func (s *CallAteletPauseStep) Execute(ctx context.Context, input *PauseInput, st
ateletConn, err := s.dialer.DialForWorker(assignment.GetWorkerNamespace(), assignment.GetWorkerPod())
if err != nil {
if errors.Is(err, ErrWorkerPodNotFound) {
slog.ErrorContext(ctx, "Worker pod gone before checkpoint, crashing actor", "namespace", assignment.GetWorkerNamespace(), "pod", assignment.GetWorkerPod(), "in_progress_snapshot", state.Actor.GetInProgressSnapshot())
slog.ErrorContext(ctx, "Worker pod gone before checkpoint, crashing actor", "namespace", assignment.GetWorkerNamespace(), "pod", assignment.GetWorkerPod(), "in_progress_local_snapshot_name", state.Actor.GetInProgressLocalSnapshotName())
if err := crashActor(ctx, s.store, input.ActorRef, ateattr.OperationPause, ateattr.ReasonWorkerPodGone); err != nil {
slog.ErrorContext(ctx, "Failed to crash actor", slog.String("err", err.Error()))
}
Expand Down Expand Up @@ -163,7 +161,7 @@ func (s *CallAteletPauseStep) Execute(ctx context.Context, input *PauseInput, st
Type: ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL,
Config: &ateletpb.CheckpointRequest_LocalConfig{
LocalConfig: &ateletpb.LocalCheckpointConfiguration{
SnapshotPrefix: state.Actor.InProgressSnapshot,
SnapshotName: state.Actor.InProgressLocalSnapshotName,
},
},
Scope: toAteletSnapshotScope(state.ActorTemplate.Spec.SnapshotsConfig.OnPause),
Expand Down Expand Up @@ -264,16 +262,16 @@ func (s *FinalizePausedStep) Execute(ctx context.Context, input *PauseInput, sta
slog.ErrorContext(ctx, "Node name not found during finalize pause, crashing actor", slog.Any("actor", input.ActorRef))
latestActor.Status = ateapipb.Actor_STATUS_CRASHED
}
// TODO(dberkov) - what if InProgressSnapshot is empty? That shouldn't be possible.
if latestActor.InProgressSnapshot != "" {
// TODO(dberkov) - what if InProgressLocalSnapshotName is empty? That shouldn't be possible.
if latestActor.InProgressLocalSnapshotName != "" {
localInfo := &ateapipb.LocalSnapshotInfo{
SnapshotPrefix: latestActor.InProgressSnapshot,
SnapshotName: latestActor.InProgressLocalSnapshotName,
}
if latestActor.Status != ateapipb.Actor_STATUS_CRASHED {
localInfo.NodeVmsWithLocalSnapshots = []string{nodeName}
}
latestActor.LocalSnapshotInfo = localInfo
latestActor.InProgressSnapshot = ""
latestActor.InProgressLocalSnapshotName = ""
}
sandboxClass := ""
if worker != nil {
Expand Down
10 changes: 5 additions & 5 deletions cmd/ateapi/internal/controlapi/workflow_pause_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ func TestFinalizePausedStep_WorkerGone(t *testing.T) {
WorkerPool: "pool1",
WorkerPod: "worker-pod-1",
},
InProgressSnapshot: "snap-prefix",
InProgressLocalSnapshotName: "local-snap-1",
}
if _, err := st.CreateActor(ctx, actor); err != nil {
t.Fatalf("CreateActor: %v", err)
Expand Down Expand Up @@ -238,8 +238,8 @@ func TestCallAteletPauseStep_DanglingWorkerDoesNotRecordPhantomSnapshot(t *testi
WorkerPool: "pool",
WorkerPod: "pod-gone",
},
InProgressSnapshot: "actor-1-never-written",
LatestSnapshot: tt.prevSnapshot,
InProgressLocalSnapshotName: "actor-1-never-written",
LatestSnapshot: tt.prevSnapshot,
}
created, err := persistence.CreateActor(ctx, actor)
if err != nil {
Expand All @@ -259,8 +259,8 @@ func TestCallAteletPauseStep_DanglingWorkerDoesNotRecordPhantomSnapshot(t *testi
if stored.GetStatus() != ateapipb.Actor_STATUS_CRASHED {
t.Errorf("status = %v, want CRASHED", stored.GetStatus())
}
if got := stored.GetInProgressSnapshot(); got != "actor-1-never-written" {
t.Errorf("InProgressSnapshot = %q, want preserved for debugging", got)
if got := stored.GetInProgressLocalSnapshotName(); got != "actor-1-never-written" {
t.Errorf("InProgressLocalSnapshotName = %q, want preserved for debugging", got)
}
if tt.prevSnapshot == nil {
if stored.GetLatestSnapshot() != nil {
Expand Down
Loading
Loading