Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
58 changes: 37 additions & 21 deletions cmd/ateapi/internal/controlapi/functional_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -293,7 +293,7 @@ func setupTest(t *testing.T, ns string) *testContext {

ctx, cancel := context.WithCancel(context.Background())

syncer := NewWorkerPoolSyncer(persistence, workerInformer, workerPoolLister)
syncer := NewWorkerPoolSyncer(persistence, k8sClient, workerInformer, workerPoolLister)
syncer.Start(ctx)

workerFactory.Start(ctx.Done())
Expand Down Expand Up @@ -2446,35 +2446,46 @@ func TestResumeActor_DanglingWorker(t *testing.T) {

deleteWorkerPod(t, tc, ns, "worker-a")

// 5. Wait for the syncer to remove worker A from the store. The actor is
// RESUMING, not RUNNING, so the syncer leaves it bound to the dead worker.
if err := wait.PollUntilContextTimeout(context.Background(), 100*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
_, gerr := tc.persistence.GetWorker(ctx, ns, "pool1", "worker-a")
return gerr != nil, nil
}); err != nil {
t.Fatalf("worker-a not removed from store: %v", err)
}

// 6. Create Worker Pod B
createWorkerPod(t, tc, ns, "worker-b", "node1", "pool1")

// 7. Configure fake Atelet to SUCCEED on Restore
tc.fakeAtelet.FailRestore = nil
tc.fakeAtelet.RestoreCalled = false // reset

// 8. Call ResumeActor again -> Expect success and picking Worker B!
// 8. Call ResumeActor again -> Expect failure: the resume workflow refuses
// to reassign a RESUMING actor whose recorded worker is gone, even though
// a free worker is available.
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: name},
})
if err != nil {
t.Fatalf("ResumeActor failed on retry: %v", err)
if err == nil {
t.Fatalf("expected ResumeActor retry to fail for actor bound to a deleted worker")
}

if !tc.fakeAtelet.RestoreCalled {
t.Errorf("expected Restore to be called on retry")
if tc.fakeAtelet.RestoreCalled {
t.Errorf("expected Restore not to be called on retry")
}

// Verify actor state is RUNNING with worker B assigned
// Verify the actor is still RESUMING and pinned to the dead worker A.
actor, err = tc.persistence.GetActor(context.Background(), testAtespace, name)
if err != nil {
t.Fatalf("failed to get actor from store: %v", err)
}
if actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING {
t.Errorf("expected status RUNNING, got %v", actor.GetStatus())
if actor.GetStatus() != ateapipb.Actor_STATUS_RESUMING {
t.Errorf("expected status RESUMING, got %v", actor.GetStatus())
}
if actor.GetAteomPodName() != "worker-b" {
t.Errorf("expected worker-b assigned, got %v", actor.GetAteomPodName())
if actor.GetAteomPodName() != "worker-a" {
t.Errorf("expected worker-a still assigned, got %v", actor.GetAteomPodName())
}
}

Expand Down Expand Up @@ -2508,29 +2519,34 @@ func TestSuspendActor_DanglingWorker(t *testing.T) {

deleteWorkerPod(t, tc, ns, "worker-1")

// 3. Call SuspendActor -> Should succeed (our fix skips missing pod execution)
actors, _, _ := tc.persistence.ListActors(context.Background(), testAtespace, maxPageSize, "")
t.Logf("Actors in Redis before Suspend: %d", len(actors))
for _, a := range actors {
t.Logf(" Actor: %s/%s/%s", a.GetActorTemplateNamespace(), a.GetActorTemplateName(), a.GetMetadata().GetName())
// 3. The syncer crashes a RUNNING actor whose worker pod vanished.
if err := wait.PollUntilContextTimeout(context.Background(), 100*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
a, gerr := tc.persistence.GetActor(ctx, testAtespace, name)
if gerr != nil {
return false, gerr
}
return a.GetStatus() == ateapipb.Actor_STATUS_CRASHED, nil
}); err != nil {
t.Fatalf("actor not marked CRASHED after worker deletion: %v", err)
}

// 4. SuspendActor now fails: there is no live state left to snapshot.
_, err = tc.client.SuspendActor(context.Background(), &ateapipb.SuspendActorRequest{
Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: name},
})
if err != nil {
t.Fatalf("SuspendActor failed: %v", err)
if err == nil {
t.Fatalf("expected SuspendActor to fail for a crashed actor")
}

// 4. Verify it becomes SUSPENDED in Redis
// 5. Verify the crash stuck and the dead worker binding is cleared.
getResp, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{
Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: name},
})
if err != nil {
t.Fatalf("GetActor failed: %v", err)
}
if getResp.GetStatus() != ateapipb.Actor_STATUS_SUSPENDED {
t.Errorf("expected status SUSPENDED, got %v", getResp.GetStatus())
if getResp.GetStatus() != ateapipb.Actor_STATUS_CRASHED {
t.Errorf("expected status CRASHED, got %v", getResp.GetStatus())
}
if getResp.GetAteomPodNamespace() != "" {
t.Errorf("expected ateom_pod_namespace to be empty, got %v", getResp.GetAteomPodNamespace())
Expand Down
110 changes: 85 additions & 25 deletions cmd/ateapi/internal/controlapi/syncer.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,21 +27,26 @@ import (
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
corev1 "k8s.io/api/core/v1"
apierrors "k8s.io/apimachinery/pkg/api/errors"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/client-go/kubernetes"
"k8s.io/client-go/tools/cache"
)

// WorkerPoolSyncer reconciles the state of worker pods from Kubernetes Informer
// into the store.
type WorkerPoolSyncer struct {
persistence store.Interface
kubeClient kubernetes.Interface
workerInformer cache.SharedIndexInformer
workerPoolLister listersv1alpha1.WorkerPoolLister
}

// NewWorkerPoolSyncer creates a new WorkerPoolSyncer.
func NewWorkerPoolSyncer(persistence store.Interface, workerInformer cache.SharedIndexInformer, workerPoolLister listersv1alpha1.WorkerPoolLister) *WorkerPoolSyncer {
func NewWorkerPoolSyncer(persistence store.Interface, kubeClient kubernetes.Interface, workerInformer cache.SharedIndexInformer, workerPoolLister listersv1alpha1.WorkerPoolLister) *WorkerPoolSyncer {
return &WorkerPoolSyncer{
persistence: persistence,
kubeClient: kubeClient,
workerInformer: workerInformer,
workerPoolLister: workerPoolLister,
}
Expand Down Expand Up @@ -75,8 +80,8 @@ func (s *WorkerPoolSyncer) Start(ctx context.Context) {
return
}
slog.InfoContext(ctx, "Syncer: removing worker from store", slog.String("worker", pod.Namespace+"/"+pod.Name))
if err := s.releaseActorOnDeadWorker(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name); err != nil {
slog.ErrorContext(ctx, "Failed to release actor bound to deleted worker", slog.Any("err", err))
if err := s.markActorCrashed(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name); err != nil {
slog.ErrorContext(ctx, "Failed to crash actor bound to deleted worker", slog.Any("err", err))
}
err := s.persistence.DeleteWorker(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name)
if err != nil {
Expand Down Expand Up @@ -107,7 +112,7 @@ func (s *WorkerPoolSyncer) syncWorkerToStore(ctx context.Context, pod *corev1.Po

if pod.DeletionTimestamp != nil {
slog.InfoContext(ctx, "Syncer: removing worker from store (pod deleting)", slog.String("worker", pod.Namespace+"/"+pod.Name))
if err := s.releaseActorOnDeadWorker(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name); err != nil {
if err := s.markActorCrashed(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name); err != nil {
slog.ErrorContext(ctx, "Failed to release actor bound to soft-deleting worker", slog.Any("err", err))
}
err := s.persistence.DeleteWorker(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name)
Expand All @@ -117,6 +122,12 @@ func (s *WorkerPoolSyncer) syncWorkerToStore(ctx context.Context, pod *corev1.Po
return
}

// If either the ateom server itself or the actor inside it terminated, crash the actor running on the worker.
if ateomTerminated(pod) {
s.handleTerminatedAteom(ctx, pod)
return
}

poolName := pod.Labels[workerPodLabel]
pool, err := s.workerPoolLister.WorkerPools(pod.Namespace).Get(poolName)
if err != nil {
Expand Down Expand Up @@ -185,21 +196,58 @@ func isWorkerEligible(pod *corev1.Pod) bool {
return pod.Status.PodIP != ""
}

// releaseActorOnDeadWorker resets the actor bound to a vanishing worker
// pod back to STATUS_SUSPENDED so the next request reassigns it.
//
// UpdateActor uses optimistic version checking. A concurrent SuspendActor
// or ResumeActor wins; we drop this attempt silently.
//
// Best-effort only. The caller always proceeds to DeleteWorker after this
// returns, so any non-contention failure leaves the actor stranded
// (STATUS_RUNNING, pointer at a pod that no longer exists). Recovery
// then needs a manual SuspendActor.
//
// The long-term fix is a finalizer-based controller that holds the pod
// in Terminating state until the actor is gracefully suspended. Tracked
// in https://github.com/agent-substrate/substrate/issues/23.
func (s *WorkerPoolSyncer) releaseActorOnDeadWorker(ctx context.Context, namespace, pool, podName string) error {
const ateomContainerName = "ateom"

func ateomTerminated(pod *corev1.Pod) bool {
for i := range pod.Status.ContainerStatuses {
cs := &pod.Status.ContainerStatuses[i]
if cs.Name != "ateom" {
continue
}
if cs.LastTerminationState.Terminated != nil || cs.State.Terminated != nil {
return true
}
}
return false
}

func (s *WorkerPoolSyncer) handleTerminatedAteom(ctx context.Context, pod *corev1.Pod) {
slog.InfoContext(ctx, "Syncer: worker ateom terminated; crashing actor and recycling pod",
slog.String("worker", pod.Namespace+"/"+pod.Name))
pool := pod.Labels[workerPodLabel]
worker, err := s.persistence.GetWorker(ctx, pod.Namespace, pool, pod.Name)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To avoid workers being re-used before we have a chance to clean them up, I wonder if it's worth doing something similar to how I'm going to handle graceful termination.

  1. Mark the worker as 'CRASHED' (I'll use 'DRAINING' for graceful termination).
  2. In the ateom, be able to detect that it crashed and reject any further suspend / resume attempts.

This mitigates the race conditions around state detection and clean up.

switch {
case errors.Is(err, store.ErrNotFound):
// Worker already deleted from store, nothing to do.
slog.WarnContext(ctx, "worker is already deleted from store")
case err != nil:
slog.ErrorContext(ctx, "Failed to get worker for terminated ateom; keeping pod so the resync retries", slog.Any("err", err))
return
case worker.Assignment == nil:
// If the worker is not assigned to any actor, delete the worker from pool.
if err := s.persistence.DeleteWorker(ctx, pod.Namespace, pool, pod.Name); err != nil {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this is a new issue, so I don't think we need to put a fix in this PR. I believe there are some race conditions around DeleteWorker. However unlikely, it's possible a different actor gets assigned to this worker after we check the state. So we could potentially be deleting a worker that has an actor assigned. We already check the expected version during updates, but maybe we need to do this for deletes as well?

slog.ErrorContext(ctx, "Failed to delete idle crashed worker from store; keeping pod so the resync retries", slog.Any("err", err))
return
}
default:
// Crash the actor first, if the ateapi server is restarted. If we deleted the worker first, and ate apiserver restarted, we wouldn't be able to crash the Actor anymore.
if err := s.markActorCrashed(ctx, pod.Namespace, pool, pod.Name); err != nil {
slog.ErrorContext(ctx, "Failed to crash actor bound to crashed worker; keeping pod so the resync retries", slog.Any("err", err))
return
}
if err := s.persistence.DeleteWorker(ctx, pod.Namespace, pool, pod.Name); err != nil {
slog.ErrorContext(ctx, "Failed to delete crashed worker from store", slog.Any("err", err))
}
}
// Deleting the pod, because the actor network may not have been cleaned up from the Ateom Pod's net NS.
if err := s.kubeClient.CoreV1().Pods(pod.Namespace).Delete(ctx, pod.Name, metav1.DeleteOptions{

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe I'm totally misunderstanding something, if the ateom crashed wouldn't that automatically rollout the pod because we're using a Deployment under the hood?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is because the ateom container's crash will be detected by kubelet, which will restart the ateom container in the same pod.
You are right that once we delete this pod, the Deployment will automatically create a new one.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤦

And the worry is that there will be dangling networking and/or other resources if we don't manually kill the pod?

Preconditions: metav1.NewUIDPreconditions(string(pod.UID)),
}); err != nil && !apierrors.IsNotFound(err) {
slog.ErrorContext(ctx, "Failed to delete crashed worker pod", slog.Any("err", err))
}
}

func (s *WorkerPoolSyncer) markActorCrashed(ctx context.Context, namespace, pool, podName string) error {
worker, err := s.persistence.GetWorker(ctx, namespace, pool, podName)
if err != nil {
if errors.Is(err, store.ErrNotFound) {
Expand All @@ -210,24 +258,36 @@ func (s *WorkerPoolSyncer) releaseActorOnDeadWorker(ctx context.Context, namespa
if worker.Assignment == nil {
return nil
}
actor, err := s.persistence.GetActor(ctx, worker.Assignment.Actor.Atespace, worker.Assignment.Actor.Name)
atespace := worker.GetAssignment().GetActor().GetAtespace()
actorName := worker.GetAssignment().GetActor().GetName()
if atespace == "" || actorName == "" {
slog.WarnContext(ctx, "cannot release actor on worker, since either atespace or actor name is empty", slog.String("atespace", atespace), slog.String("actor name", actorName))
return nil
}
actor, err := s.persistence.GetActor(ctx, atespace, actorName)
if err != nil {
if errors.Is(err, store.ErrNotFound) {
return nil
}
return err
}
// Skip if a concurrent SuspendActor already cleared the pointer.
if actor.GetAteomPodNamespace() != namespace || actor.GetAteomPodName() != podName {
// Only crash the actor if it's still assigned to the same worker.
if actor.GetAteomPodNamespace() == "" || actor.GetAteomPodName() == "" || actor.GetAteomPodUid() == "" {
slog.WarnContext(ctx, "cannot release actor on worker", slog.String("ateom pod namesace", actor.GetAteomPodNamespace()), slog.String("ateom name", actor.GetAteomPodName()), slog.String("ateom pod uid", actor.GetAteomPodUid()))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Could we clarify this warning? My understanding isn't that we can't release the actor, it's that the actor is no longer assigned to this worker.

Maybe something along the lines:
"Actor is no longer assigned to worker, skipping crash."

return nil
}
if actor.Status != ateapipb.Actor_STATUS_CRASHED {
actor.Status = ateapipb.Actor_STATUS_SUSPENDED
if actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING {
slog.WarnContext(ctx, "actor is not in running state")
return nil
}
if actor.GetAteomPodNamespace() != worker.GetWorkerNamespace() || actor.GetAteomPodName() != worker.GetWorkerPod() || actor.GetAteomPodUid() != worker.GetWorkerPodUid() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: This is a similar class of warning / error as above, where the actor was suspended and no longer assigned to the worker. We should likely log in both cases or neither.

return nil
}
actor.Status = ateapipb.Actor_STATUS_CRASHED
actor.AteomPodNamespace = ""
actor.AteomPodName = ""
actor.AteomPodIp = ""
actor.InProgressSnapshot = ""
actor.AteomPodUid = ""
actor.WorkerPoolName = ""
if _, err := s.persistence.UpdateActor(ctx, actor, actor.GetMetadata().GetVersion()); err != nil && !errors.Is(err, store.ErrPersistenceRetry) {
return err
Expand Down
Loading
Loading