From dbe5a30ad8f05d1183bbb19c03c621da69323345 Mon Sep 17 00:00:00 2001
From: Haven Xia <haoyuxia@google.com>
Date: Wed, 5 Aug 2026 20:18:18 -0700
Subject: [PATCH] atelet: size the shutdown drain to the worker checkpoint
 grace

---
 cmd/atelet/main.go                | 28 +++++++++++++++++++++-------
 manifests/ate-install/atelet.yaml | 11 ++++++-----
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go
index 95a4e16c4..9824af581 100644
--- a/cmd/atelet/main.go
+++ b/cmd/atelet/main.go
@@ -81,7 +81,7 @@ var (
 	logLevelFlag = pflag.String("log-level", "info", "Minimum log level: debug, info, warn, or error.")
 
 	drainDelay   = pflag.Duration("drain-delay", 0, "How long to keep accepting new RPCs after SIGTERM before starting the gRPC drain.")
-	drainTimeout = pflag.Duration("drain-timeout", 5*time.Minute, "Deadline for the graceful gRPC drain on shutdown. In-flight RPCs still running past it are forcefully cancelled.")
+	drainTimeout = pflag.Duration("drain-timeout", 60*time.Minute, "Deadline for the graceful gRPC drain on shutdown. In-flight RPCs still running past it are forcefully cancelled. Sized to the worker pod's 60m termination grace: a Checkpoint/Restore still in flight when atelet is terminated must not be cancelled before the worker's own grace expires.")
 )
 
 func main() {
@@ -248,12 +248,26 @@ func drainOnShutdown(ctx context.Context, srv *grpc.Server, readiness *serverboo
 			srv.GracefulStop()
 			close(drainComplete)
 		}()
-		select {
-		case <-drainComplete:
-			slog.InfoContext(ctx, "Drain completed within deadline")
-		case <-time.After(*drainTimeout):
-			slog.WarnContext(ctx, "Drain deadline exceeded; forcing stop")
-			srv.Stop()
+		// log progress so the pod does not just sit in Terminating silently.
+		start := time.Now()
+		deadline := time.NewTimer(*drainTimeout)
+		defer deadline.Stop()
+		progress := time.NewTicker(time.Minute)
+		defer progress.Stop()
+		for {
+			select {
+			case <-drainComplete:
+				slog.InfoContext(ctx, "Drain completed within deadline")
+				return
+			case <-progress.C:
+				slog.InfoContext(ctx, "Still draining",
+					slog.Duration("elapsed", time.Since(start).Round(time.Second)),
+					slog.Duration("remaining", (*drainTimeout-time.Since(start)).Round(time.Second)))
+			case <-deadline.C:
+				slog.WarnContext(ctx, "Drain deadline exceeded; forcing stop")
+				srv.Stop()
+				return
+			}
 		}
 	}()
 	return done
diff --git a/manifests/ate-install/atelet.yaml b/manifests/ate-install/atelet.yaml
index 9d830d71c..2bdf615b7 100644
--- a/manifests/ate-install/atelet.yaml
+++ b/manifests/ate-install/atelet.yaml
@@ -65,11 +65,12 @@ spec:
     spec:
       serviceAccountName: atelet
       # Budget for the full shutdown sequence: --drain-delay (0 for atelet) +
-      # --drain-timeout (5m). Sized long because Checkpoint/Restore stream
-      # multi-GiB snapshots and must not be force-cancelled mid-upload. The sum
-      # must fit within terminationGracePeriodSeconds, plus slack for the force
+      # --drain-timeout (60m). Sized to the worker pod's hardcoded 60m
+      # termination grace (workerTerminationGracePeriodSeconds): atelet drives
+      # the worker's Checkpoint/Restore, so a terminating atelet must not
+      # force-cancel one before the worker's own grace expires, plus slack for the force
       # Stop() and the tracer/meter flush on exit.
-      terminationGracePeriodSeconds: 330
+      terminationGracePeriodSeconds: 3630
       containers:
       - name: atelet
         image: ko://github.com/agent-substrate/substrate/cmd/atelet
@@ -80,7 +81,7 @@ spec:
         # Graceful shutdown knobs. The sum must fit within
         # terminationGracePeriodSeconds above.
         - --drain-delay=0s
-        - --drain-timeout=5m
+        - --drain-timeout=60m
         # atelet does no mounts, netlink, device, or namespace operations (those
         # live in the ateom worker pod) — it only reads/writes the
         # /var/lib/ateom-gvisor hostPath as root, so it needs no Linux