Skip to content

Handle graceful worker termination - #517

Open
sfunkenhauser wants to merge 2 commits into
agent-substrate:mainfrom
sfunkenhauser:worker_graceful_termination
Open

Handle graceful worker termination#517
sfunkenhauser wants to merge 2 commits into
agent-substrate:mainfrom
sfunkenhauser:worker_graceful_termination

Conversation

@sfunkenhauser

@sfunkenhauser sfunkenhauser commented Jul 24, 2026

Copy link
Copy Markdown
Collaborator

Part of #23

The overall plan is:

  1. Trap SIGTERM in the ateom, and forward SIGTERM to all application containers.
  2. Once all application containers exit, ateom will exit itself.
  3. The workerpool syncer will mark workers as 'DRAINING' when a pod deletion is triggered (pod.DeletionTimestamp) to prevent actors from being scheduled. On the pod Deleted event, the worker and actor (if it was assigned) will be cleaned up.

This PR implements 3), there are 2 more PRs coming which will implement 1) and 2) for each ateom (gvisor and uVM).

@sfunkenhauser

Copy link
Copy Markdown
Collaborator Author

@zoez7 Zoe Zhao (zoez7) self-assigned this Jul 24, 2026
@sfunkenhauser
sfunkenhauser force-pushed the worker_graceful_termination branch from 8ae646e to 4cfd15c Compare July 27, 2026 16:21
@sfunkenhauser
sfunkenhauser force-pushed the worker_graceful_termination branch from 4cfd15c to 5d0ac0c Compare July 27, 2026 16:59
// the bound actor here — inside the pod ateom has received SIGTERM and is
// gracefully shutting the actor down. Actor cleanup happens on the Pod
// Deleted event.
if err := s.markWorkerDraining(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name); err != nil {

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.

We might fail to mark the worker as draining due to version mismatch, if it was selected by another actor "won" and marked it as assigned during resuming?
Maybe we need to retry this, and add another step in the ResumeActor workflow to confirm the worker is not in Draining state before returning "success"?

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 will be retried on subsequent informer updates. We could add retry logic to every update call, but I think that will get complicated. I'd rather rely on the informer (and in the future a workqueue) to retry failed operations.

Marking a worker as draining will always be best effort, so the ateom needs to handle the case where a resume/suspend operation lands on a pod that's getting deleted.

ResumeActor already checks in a few places to make sure the worker isn't draining (AssignWorker), but there is always going to be a race between reading worker state and performing some other action (AteletRestore for example).

If we end up calling AteletRestore on a pod that's being deleted, the ateom will return an error and fail the resume workflow.

If the AteletRestore succeeds before the pod is deleted, we should consider the resume successful.

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.

If we end up calling AteletRestore on a pod that's being deleted, the ateom will return an error and fail the resume workflow.

I see, so we'll need to (1) update ateom to reject Restore requests after receiving SIGTERM, and (2) update ateapi to return a retriable error, so client retries the resume operation. Did I understand that correctly?

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.

(1) That's correct.
(2) Ultimately yes, but this is going to take a bit more work. Right now if you resume an actor with status Actor_STATUS_RESUMING, and the worker is in a bad state (NOT_FOUND and I'm adding DRAINING to this as well) the actor will be crashed. I just don't know if it's safe to try and resume the actor on another worker. We need to take a closer look at the different resume workflow states and figure it out. Right now I think it's safer to just crash in this case.

Concretely for this change, I would have the ateom reject the Restore request with a non-retriable error and the resume workflow would crash the actor. Let me know if you think this is being overly cautious.

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.

Re (2) I think this is a good starting point. Pick-another-worker-on-retry is an SLO improvement we can later make.

var workers []*ateapipb.Worker
var pageToken string
for {
wPage, nextToken, err := s.persistence.ListWorkers(ctx, 1000, pageToken)

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.

This will needed to be sharded once we have multiple instances of ate-apiserver.

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.

I'm not sure if the sharding needs to be addressed now, but looks like Julian Gutierrez Oschmann (@juli4n) already has #538 open.

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 just an optimization right? Ultimately I think we should move the syncer out of the ate-apiserver. Maybe we can wait for that decision before optimizing this.

Comment thread cmd/ateapi/internal/controlapi/workflow_resume.go Outdated

@zoez7 Zoe Zhao (zoez7) left a comment

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.

LGTM. Need to resolve conflicts before merging.

return nil, err
}
if !eligible {
if !eligible || worker.GetState() == ateapipb.Worker_STATE_DRAINING {

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.

This should be && worker.GetState() == ateapipb.Worker_STATE_ACTIVE as well?

(It's not backwards compatible, but at this stage I think requiring devs to rebuild ate-system is still ok, soon that may not be the case.)

return fmt.Errorf("while checking worker eligibility: %w", err)
}
if eligible {
if eligible && worker.GetState() == ateapipb.Worker_STATE_ACTIVE {

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.

🤖 should-fix 🟡 – This requires STATE_ACTIVE, but findFreeWorker below only excludes STATE_DRAINING. The two disagree about STATE_UNSPECIFIED, and that is the state every pre-existing worker record has: State is set only on CreateWorker, and syncWorkerToStore's update path never backfills it, so records written before this PR stay unspecified indefinitely. ValidateWorker tolerates that value explicitly "for backward compatibility", so it is an expected state, not a corrupt one.

After an upgrade, this loop then rejects a legacy worker that holds a valid claim for the actor, releases the claim, and re-schedules — and if the pool has no other free worker, the resume fails even though the claimed worker was healthy. Using the same predicate at both sites (!= STATE_DRAINING) would fix it, or backfill State in the update path so legacy records converge on ACTIVE.

return nil
}

actor.Status = ateapipb.Actor_STATUS_CRASHED

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.

🤖 question 🟢 – This is the right end state once ateom suspends the actor on SIGTERM, but steps 1 and 2 are still to come. Until they land, a running actor is no longer marked SUSPENDED when its worker pod goes away, so every routine worker rollout leaves its actors reading CRASHED.

Nothing gates resume on status, so they stay resumable and this is a labeling change rather than data loss — and the old SUSPENDED arguably implied a clean snapshot that was never taken. Is that interim acceptable, or would this transition be better landed alongside step 2?

// pod pointers are cleared.
//
// UpdateActor uses optimistic version checking. A concurrent SuspendActor
// or ResumeActor wins; we drop this attempt silently.

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.

🤖 nit 🟢 – Stale: the ErrPersistenceRetry tolerance is gone, so contention is no longer dropped silently — it propagates and stops reconcileDeadWorker from deleting the worker record. That is deliberate for real errors, but on contention the actor was already updated by whoever won, and the dead worker's record now lingers until the next startup reconcile.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants