What is AI Agent Graceful Shutdown?
An AI agent graceful shutdown is the process of safely terminating a long-running AI agent while preserving its internal state, ensuring that no work is lost and the agent can resume from where it left off. Unlike a hard kill (SIGKILL), a graceful shutdown gives the agent time to finish critical operations, persist state to disk, close connections, and clean up resources before exiting.
AI agents—whether they are autonomous coding assistants, chatbot pipelines, or reinforcement learning inference engines—often maintain complex in-memory state including conversation history, partial reasoning traces, tool execution queues, and cached embeddings. If the process is abruptly terminated, all of that state evaporates. A graceful shutdown protocol ensures this state is serialized and saved before the process ends, making restarts seamless and lossless.
Why Graceful Shutdown Matters for AI Agents
Consider these real-world scenarios where graceful shutdown is critical:
- Cloud deployments with auto-scaling: Kubernetes pods can be evicted at any moment due to resource pressure. Without graceful shutdown, an agent mid-way through a complex reasoning chain loses all progress.
- Rolling updates and deployments: When you push a new version of your agent, the old process receives SIGTERM. If it ignores this, the new version starts with no context of ongoing conversations.
- GPU node preemption: In cloud GPU clusters, nodes can be reclaimed with short notice. Agents running on these nodes must checkpoint their state or risk losing expensive computation.
- Long-running reasoning tasks: Agents performing tree-search, planning, or iterative refinement over many minutes need to save partial results so they don't restart from scratch.
- User experience: An agent that forgets everything it was doing after a restart creates a frustrating, unreliable experience.
The cost of not implementing graceful shutdown is wasted compute, lost user context, and degraded reliability. In production, it's not optional—it's a core requirement of any stateful agent system.
Core Components of State Preservation
To build a robust graceful shutdown system, you need to handle these interconnected components:
- Signal trapping: Intercepting SIGTERM, SIGINT, and optionally SIGUSR1 to trigger the shutdown sequence.
- State snapshotting: Serializing the agent's in-memory state to a durable format (JSON, pickle, Parquet, or a database).
- Checkpoint management: Maintaining multiple checkpoints with timestamps to allow rollback and avoid corruption from incomplete writes.
- Resource cleanup: Closing database connections, WebSocket streams, file handles, and HTTP sessions cleanly.
- Health check coordination: Marking the agent as "draining" so load balancers stop routing new requests while existing work finishes.
- Resume logic: On restart, detecting and loading the saved checkpoint to resume exactly where the agent left off.
Implementation Approaches
Approach 1: Signal-Based Shutdown with Atomic Checkpoints
The most common pattern uses Python's signal module to catch termination signals, set a shutdown flag, and let the main event loop handle the rest gracefully. Here's the basic skeleton:
import signal
import sys
import json
import os
from datetime import datetime
from typing import Optional
class GracefulShutdownHandler:
"""Handles graceful shutdown with atomic state persistence."""
def __init__(self, checkpoint_dir: str = "/tmp/agent_checkpoints"):
self.shutdown_requested = False
self.checkpoint_dir = checkpoint_dir
os.makedirs(self.checkpoint_dir, exist_ok=True)
def request_shutdown(self, signum, frame):
"""Signal handler that sets the shutdown flag."""
print(f"[SHUTDOWN] Received signal {signal.Signals(signum).name}. "
f"Completing current work before exit...")
self.shutdown_requested = True
def register(self):
"""Register signal handlers for graceful termination."""
signal.signal(signal.SIGTERM, self.request_shutdown)
signal.signal(signal.SIGINT, self.request_shutdown)
# SIGUSR1 can be used for manual checkpoint triggers
signal.signal(signal.SIGUSR1, self.request_shutdown)
def is_shutdown_requested(self) -> bool:
return self.shutdown_requestedApproach 2: Structured State Serialization
Your agent state should be serializable to a structured format. Avoid raw pickle for production—it's unsafe and brittle across Python versions. Instead, use JSON with explicit schema validation or a binary format like MessagePack for efficiency:
from dataclasses import dataclass, asdict
from typing import List, Dict, Any
import json
import hashlib
from pathlib import Path
@dataclass
class ConversationTurn:
role: str
content: str
timestamp: str
metadata: Dict[str, Any]
@dataclass
class AgentState:
"""Full serializable state of an AI agent."""
session_id: str
conversation_history: List[ConversationTurn]
pending_tool_calls: List[Dict[str, Any]]
reasoning_tree: Optional[Dict[str, Any]] # partial search tree
embeddings_cache_key: str
llm_context_window_used: int
created_at: str
last_updated: str
def to_json(self) -> str:
return json.dumps(asdict(self), indent=2, default=str)
@classmethod
def from_json(cls, data: str) -> 'AgentState':
raw = json.loads(data)
raw['conversation_history'] = [
ConversationTurn(**turn)
for turn in raw['conversation_history']
]
return cls(**raw)
def compute_checksum(self) -> str:
"""SHA256 hash for integrity verification."""
state_bytes = self.to_json().encode('utf-8')
return hashlib.sha256(state_bytes).hexdigest()Approach 3: Atomic File Writes for Corruption Prevention
Never write directly to the final checkpoint file. Always write to a temporary file first, then atomically rename it. This prevents corruption if the process crashes mid-write:
import tempfile
import shutil
class AtomicCheckpointWriter:
"""Writes checkpoints atomically to prevent corruption."""
def __init__(self, checkpoint_dir: str):
self.checkpoint_dir = Path(checkpoint_dir)
self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
def save(self, state: AgentState, checkpoint_name: str) -> str:
"""Save state atomically. Returns the final file path."""
checkpoint_path = self.checkpoint_dir / f"{checkpoint_name}.json"
# Write to a temporary file in the same directory
# (same filesystem ensures atomic rename works)
with tempfile.NamedTemporaryFile(
mode='w',
dir=str(self.checkpoint_dir),
suffix='.tmp',
delete=False
) as tmp_file:
tmp_file.write(state.to_json())
tmp_file.flush()
os.fsync(tmp_file.fileno()) # Force write to disk
tmp_path = tmp_file.name
# Atomic rename (works on same filesystem)
os.rename(tmp_path, str(checkpoint_path))
# Also write a checksum file for integrity verification
checksum_path = self.checkpoint_dir / f"{checkpoint_name}.sha256"
with open(checksum_path, 'w') as f:
f.write(state.compute_checksum())
return str(checkpoint_path)Complete Production-Grade Implementation
Below is a full, production-ready AI agent with graceful shutdown, checkpointing, and resume capability. This example simulates a conversational agent that processes user messages, maintains conversation history, and can safely restart without losing context.
#!/usr/bin/env python3
"""
Production-grade AI Agent with Graceful Shutdown and State Recovery.
This agent:
- Processes a stream of user messages (simulated)
- Maintains conversation history and pending tool calls
- Traps SIGTERM/SIGINT for graceful shutdown
- Checkpoints state atomically every N messages and on shutdown
- Resumes from the latest checkpoint on restart
"""
import signal
import sys
import json
import os
import time
import hashlib
import tempfile
import threading
from datetime import datetime, timezone
from dataclasses import dataclass, field, asdict
from typing import List, Dict, Any, Optional
from pathlib import Path
from enum import Enum
# ─── Data Models ──────────────────────────────────────────────
class AgentStatus(Enum):
RUNNING = "running"
DRAINING = "draining"
SHUTTING_DOWN = "shutting_down"
STOPPED = "stopped"
@dataclass
class ConversationMessage:
role: str
content: str
timestamp: str
message_id: str
@dataclass
class PendingToolCall:
tool_name: str
arguments: Dict[str, Any]
status: str # "queued", "running", "completed"
result: Optional[Dict[str, Any]] = None
@dataclass
class AgentState:
"""Complete serializable agent state."""
session_id: str
status: str
conversation: List[ConversationMessage] = field(default_factory=list)
pending_tool_calls: List[PendingToolCall] = field(default_factory=list)
messages_processed: int = 0
last_user_message: Optional[str] = None
created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
last_updated: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
def touch(self):
"""Update the last_updated timestamp."""
self.last_updated = datetime.now(timezone.utc).isoformat()
def to_json(self) -> str:
return json.dumps(asdict(self), indent=2, default=str, ensure_ascii=False)
@classmethod
def from_json(cls, data: str) -> 'AgentState':
raw = json.loads(data)
raw['conversation'] = [ConversationMessage(**m) for m in raw['conversation']]
raw['pending_tool_calls'] = [PendingToolCall(**tc) for tc in raw['pending_tool_calls']]
return cls(**raw)
def compute_checksum(self) -> str:
return hashlib.sha256(self.to_json().encode('utf-8')).hexdigest()
# ─── Checkpoint Manager ───────────────────────────────────────
class CheckpointManager:
"""Manages atomic checkpoint writes, rotation, and recovery."""
def __init__(self, checkpoint_dir: str = "./agent_checkpoints", max_checkpoints: int = 5):
self.checkpoint_dir = Path(checkpoint_dir)
self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
self.max_checkpoints = max_checkpoints
self._lock = threading.Lock() # Thread-safe writes
def save(self, state: AgentState) -> str:
"""Atomically save state to a timestamped checkpoint."""
state.touch()
timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S_%f")
checkpoint_name = f"checkpoint_{timestamp}"
with self._lock:
# Atomic write
tmp_path = self.checkpoint_dir / f".{checkpoint_name}.tmp"
final_path = self.checkpoint_dir / f"{checkpoint_name}.json"
with open(tmp_path, 'w', encoding='utf-8') as f:
f.write(state.to_json())
f.flush()
os.fsync(f.fileno())
os.rename(str(tmp_path), str(final_path))
# Write checksum
checksum_path = self.checkpoint_dir / f"{checkpoint_name}.sha256"
with open(checksum_path, 'w') as f:
f.write(state.compute_checksum())
# Rotate old checkpoints
self._rotate()
# Update latest symlink for fast resume
latest_symlink = self.checkpoint_dir / "LATEST"
if latest_symlink.is_symlink() or latest_symlink.exists():
latest_symlink.unlink()
latest_symlink.symlink_to(final_path.name)
return str(final_path)
def load_latest(self) -> Optional[AgentState]:
"""Load the most recent checkpoint. Returns None if no checkpoint exists."""
latest_symlink = self.checkpoint_dir / "LATEST"
if not latest_symlink.is_symlink():
# Fall back to finding the most recent checkpoint file
checkpoints = sorted(self.checkpoint_dir.glob("checkpoint_*.json"))
if not checkpoints:
return None
latest_file = checkpoints[-1]
else:
latest_file = self.checkpoint_dir / latest_symlink.readlink().name
# Verify checksum before loading
checksum_path = self.checkpoint_dir / f"{latest_file.stem}.sha256"
if checksum_path.exists():
stored_checksum = checksum_path.read_text().strip()
with open(latest_file, 'r', encoding='utf-8') as f:
data = f.read()
computed = hashlib.sha256(data.encode('utf-8')).hexdigest()
if stored_checksum != computed:
print(f"[WARNING] Checksum mismatch for {latest_file.name}. "
f"Trying previous checkpoint...")
return self._load_previous_checkpoint(latest_file)
with open(latest_file, 'r', encoding='utf-8') as f:
return AgentState.from_json(f.read())
def _load_previous_checkpoint(self, corrupted_file: Path) -> Optional[AgentState]:
"""Fall back to the next most recent valid checkpoint."""
all_checkpoints = sorted(self.checkpoint_dir.glob("checkpoint_*.json"))
for cp in reversed(all_checkpoints):
if cp == corrupted_file:
continue
try:
checksum_path = self.checkpoint_dir / f"{cp.stem}.sha256"
if checksum_path.exists():
stored = checksum_path.read_text().strip()
with open(cp, 'r', encoding='utf-8') as f:
data = f.read()
if hashlib.sha256(data.encode('utf-8')).hexdigest() == stored:
return AgentState.from_json(data)
else:
with open(cp, 'r', encoding='utf-8') as f:
return AgentState.from_json(f.read())
except Exception:
continue
return None
def _rotate(self):
"""Keep only the most recent N checkpoints."""
checkpoints = sorted(self.checkpoint_dir.glob("checkpoint_*.json"))
if len(checkpoints) > self.max_checkpoints:
for old_cp in checkpoints[:-self.max_checkpoints]:
old_cp.unlink(missing_ok=True)
checksum = self.checkpoint_dir / f"{old_cp.stem}.sha256"
checksum.unlink(missing_ok=True)
# ─── AI Agent Core ────────────────────────────────────────────
class ConversationalAgent:
"""AI agent that processes messages with graceful shutdown support."""
def __init__(self, session_id: str, checkpoint_dir: str = "./agent_checkpoints"):
self.session_id = session_id
self.checkpoint_manager = CheckpointManager(checkpoint_dir)
self.shutdown_handler = GracefulShutdownHandler()
# Try to resume from previous checkpoint
resumed_state = self.checkpoint_manager.load_latest()
if resumed_state:
print(f"[RESUME] Restored session {resumed_state.session_id} "
f"with {resumed_state.messages_processed} messages processed")
self.state = resumed_state
self.state.status = AgentStatus.RUNNING.value
else:
print(f"[INIT] Starting new session: {session_id}")
self.state = AgentState(
session_id=session_id,
status=AgentStatus.RUNNING.value
)
self.shutdown_handler.register()
self._running = True
self._checkpoint_interval = 10 # Checkpoint every N messages
def process_message(self, user_message: str) -> Dict[str, Any]:
"""Simulate processing a user message and generating a response."""
# Add user message to conversation
user_msg = ConversationMessage(
role="user",
content=user_message,
timestamp=datetime.now(timezone.utc).isoformat(),
message_id=f"msg_{self.state.messages_processed + 1}"
)
self.state.conversation.append(user_msg)
self.state.last_user_message = user_message
self.state.messages_processed += 1
# Simulate some processing time
time.sleep(0.05)
# Simulate tool call if message contains specific keywords
tool_response = None
if "weather" in user_message.lower():
tool_call = PendingToolCall(
tool_name="get_weather",
arguments={"location": "extracted_from_message"},
status="completed",
result={"temperature": 72, "condition": "sunny"}
)
self.state.pending_tool_calls.append(tool_call)
tool_response = f"Tool result: {tool_call.result}"
# Generate simulated assistant response
assistant_content = f"Processed: '{user_message}'. "
if tool_response:
assistant_content += tool_response
else:
assistant_content += "No tools needed."
assistant_msg = ConversationMessage(
role="assistant",
content=assistant_content,
timestamp=datetime.now(timezone.utc).isoformat(),
message_id=f"msg_{self.state.messages_processed}_resp"
)
self.state.conversation.append(assistant_msg)
return {
"response": assistant_content,
"messages_processed": self.state.messages_processed,
"conversation_length": len(self.state.conversation)
}
def run(self, message_queue: List[str]):
"""Main processing loop with integrated graceful shutdown."""
print(f"[AGENT] Starting processing loop. Session: {self.session_id}")
for i, message in enumerate(message_queue):
# Check for shutdown signal before processing each message
if self.shutdown_handler.is_shutdown_requested():
print(f"[AGENT] Shutdown requested. "
f"Processed {self.state.messages_processed} messages so far.")
self._perform_shutdown()
return
# Process the message
result = self.process_message(message)
print(f"[AGENT] Message {i+1}/{len(message_queue)}: {result['response'][:60]}...")
# Periodic checkpoint
if (self.state.messages_processed % self._checkpoint_interval == 0 and
self.state.messages_processed > 0):
checkpoint_path = self.checkpoint_manager.save(self.state)
print(f"[CHECKPOINT] Saved state to {checkpoint_path} "
f"({self.state.messages_processed} msgs)")
# Simulate that some messages take longer (shutdown during long ops)
if "complex" in message.lower():
print("[AGENT] Processing complex message (long operation)...")
for chunk in range(5):
time.sleep(0.2)
if self.shutdown_handler.is_shutdown_requested():
print("[AGENT] Shutdown during complex operation. "
"Saving partial state...")
self._perform_shutdown()
return
# Normal completion
print(f"[AGENT] Completed all {len(message_queue)} messages.")
self._perform_shutdown(final=True)
def _perform_shutdown(self, final: bool = False):
"""Execute the graceful shutdown sequence."""
self.state.status = AgentStatus.DRAINING.value
print(f"[SHUTDOWN] Saving final state...")
# Complete any in-flight tool calls
for tc in self.state.pending_tool_calls:
if tc.status == "running":
tc.status = "interrupted"
tc.result = {"error": "Shutdown before completion"}
# Save final checkpoint
checkpoint_path = self.checkpoint_manager.save(self.state)
print(f"[SHUTDOWN] Final state saved to {checkpoint_path}")
print(f"[SHUTDOWN] Summary: {self.state.messages_processed} messages, "
f"{len(self.state.conversation)} conversation turns, "
f"{len(self.state.pending_tool_calls)} tool calls")
self.state.status = AgentStatus.STOPPED.value
self._running = False
class GracefulShutdownHandler:
"""Signal-aware shutdown coordinator."""
def __init__(self):
self.shutdown_requested = threading.Event()
self._original_handlers = {}
def _signal_handler(self, signum, frame):
sig_name = signal.Signals(signum).name
print(f"\n[SIGNAL] Received {sig_name}. Initiating graceful shutdown...")
self.shutdown_requested.set()
# For SIGINT (Ctrl+C), allow a second one to force exit
if signum == signal.SIGINT:
signal.signal(signal.SIGINT, self._force_exit)
def _force_exit(self, signum, frame):
print("[SIGNAL] Second SIGINT received. Forcing exit...")
os._exit(1)
def register(self):
"""Register signal handlers."""
self._original_handlers['SIGTERM'] = signal.signal(
signal.SIGTERM, self._signal_handler
)
self._original_handlers['SIGINT'] = signal.signal(
signal.SIGINT, self._signal_handler
)
# SIGUSR1 for manual checkpoint trigger without shutdown
signal.signal(signal.SIGUSR1, lambda s, f: print(
"[SIGNAL] SIGUSR1 received. Checkpoint will be saved at next interval."
))
def is_shutdown_requested(self) -> bool:
return self.shutdown_requested.is_set()
def wait_with_timeout(self, timeout: float) -> bool:
"""Wait up to timeout seconds for shutdown signal. Returns True if signaled."""
return self.shutdown_requested.wait(timeout=timeout)
# ─── Demo Runner ──────────────────────────────────────────────
def demo_graceful_shutdown():
"""Demonstrate the graceful shutdown workflow."""
import random
print("=" * 60)
print("AI AGENT GRACEFUL SHUTDOWN DEMONSTRATION")
print("=" * 60)
# Simulate a queue of user messages
messages = [
"Hello, what can you do?",
"What is the weather in New York?",
"Tell me a joke",
"This is a complex analysis task that requires deep thought",
"Summarize our conversation so far",
"What is the weather in London?",
"Another complex multi-step reasoning problem here",
"Final message before shutdown",
]
agent = ConversationalAgent(
session_id=f"demo_session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
)
# Show resumed state
if agent.state.messages_processed > 0:
print(f"\n[DEMO] Resumed agent with previous state:")
print(f" - Messages already processed: {agent.state.messages_processed}")
print(f" - Conversation turns: {len(agent.state.conversation)}")
print(f" - Last user message: {agent.state.last_user_message}")
print(f"\n[DEMO] Processing {len(messages)} messages...\n")
# Run in a thread so we can simulate external signals
import threading
def signal_after_delay():
"""Simulate Kubernetes sending SIGTERM after some time."""
# Wait a random amount of time to simulate unpredictable shutdown
delay = random.uniform(0.8, 2.5)
time.sleep(delay)
print("\n[ORCHESTRATOR] Simulating container eviction...")
# Send SIGTERM to ourselves
os.kill(os.getpid(), signal.SIGTERM)
signal_thread = threading.Thread(target=signal_after_delay, daemon=True)
signal_thread.start()
try:
agent.run(messages)
except Exception as e:
print(f"[DEMO] Agent exited with: {e}")
print("\n" + "=" * 60)
print("DEMONSTRATION COMPLETE")
print("=" * 60)
# Show that state was persisted
resumed = CheckpointManager().load_latest()
if resumed:
print(f"\n[VERIFY] Checkpoint verification:")
print(f" - Session: {resumed.session_id}")
print(f" - Messages processed: {resumed.messages_processed}")
print(f" - Conversation turns saved: {len(resumed.conversation)}")
print(f" - Pending tool calls saved: {len(resumed.pending_tool_calls)}")
print(f" - Final status: {resumed.status}")
print(f"\n[VERIFY] Agent can resume from this checkpoint on next start.")
if __name__ == "__main__":
demo_graceful_shutdown()Integrating with Container Orchestration
In Kubernetes environments, you need to handle the pod lifecycle correctly. Here's how to wire the agent into a Kubernetes deployment with proper graceful shutdown:
# kubernetes/deployment.yaml
apiVersion: apps/v1
kind: Deployment
metadata:
name: ai-agent
spec:
replicas: 2
template:
spec:
# Give the agent enough time to checkpoint
terminationGracePeriodSeconds: 30
containers:
- name: agent
image: my-ai-agent:latest
# PreStop hook triggers checkpoint before SIGTERM
lifecycle:
preStop:
exec:
command:
- /bin/sh
- -c
- |
# Send SIGUSR1 to trigger immediate checkpoint
kill -USR1 $(pgrep -f 'ai_agent')
sleep 2
# Startup probe waits for checkpoint recovery
startupProbe:
httpGet:
path: /health
port: 8080
initialDelaySeconds: 5
periodSeconds: 2
failureThreshold: 30
# Readiness probe checks agent is fully resumed
readinessProbe:
httpGet:
path: /ready
port: 8080
periodSeconds: 5
env:
- name: CHECKPOINT_DIR
value: "/var/lib/agent/checkpoints"
volumeMounts:
- name: checkpoint-storage
mountPath: "/var/lib/agent/checkpoints"
volumes:
- name: checkpoint-storage
persistentVolumeClaim:
claimName: agent-checkpoints-pvcHealth Check Endpoint for Orchestration
Your agent should expose HTTP endpoints that reflect its internal status so orchestrators can make informed decisions:
from http.server import HTTPServer, BaseHTTPRequestHandler
import json
class HealthCheckHandler(BaseHTTPRequestHandler):
"""HTTP handler that exposes agent health and readiness."""
# This would be set by the agent instance
agent_instance = None
def do_GET(self):
if self.path == '/health':
self._handle_health()
elif self.path == '/ready':
self._handle_readiness()
elif self.path == '/drain':
self._handle_drain_status()
else:
self.send_error(404)
def _handle_health(self):
"""Liveness probe: is the process alive?"""
self.send_response(200)
self.send_header('Content-Type', 'application/json')
self.end_headers()
self.wfile.write(json.dumps({
"alive": True,
"timestamp": datetime.now(timezone.utc).isoformat()
}).encode())
def _handle_readiness(self):
"""Readiness probe: is the agent ready to accept work?"""
if self.agent_instance and self.agent_instance.state.status == AgentStatus.RUNNING.value:
status_code = 200
ready = True
else:
status_code = 503
ready = False
self.send_response(status_code)
self.send_header('Content-Type', 'application/json')
self.end_headers()
self.wfile.write(json.dumps({
"ready": ready,
"status": self.agent_instance.state.status if self.agent_instance else "unknown",
"messages_processed": self.agent_instance.state.messages_processed if self.agent_instance else 0
}).encode())
def _handle_drain_status(self):
"""Report draining progress for graceful shutdown tracking."""
agent = self.agent_instance
if not agent:
self.send_error(404)
return
self.send_response(200)
self.send_header('Content-Type', 'application/json')
self.end_headers()
self.wfile.write(json.dumps({
"draining": agent.state.status == AgentStatus.DRAINING.value,
"pending_tool_calls": len(agent.state.pending_tool_calls),
"unprocessed_queue_size": agent.state.messages_processed,
"estimated_time_remaining_seconds": 5 if agent._running else 0
}).encode())Best Practices for Production Systems
- Set generous termination grace periods: In Kubernetes, set
terminationGracePeriodSecondsto at least 30-60 seconds to give your agent time to flush state. Calculate this based on your worst-case checkpoint write time plus a buffer. - Use atomic writes always: Never write directly to the final file path. Write to a temporary file, fsync it, then rename. This is the single most important rule for preventing corruption.
- Verify checksums on resume: Always validate the integrity of checkpoint files before loading them. A corrupted checkpoint can cause subtle bugs that are hard to diagnose.
- Keep checkpoints small and frequent: Large checkpoints take longer to write. Design your state schema so checkpoints are incremental or at least bounded in size. For very large state, consider append-only logs with periodic snapshots.
- Test shutdown under load: Run chaos engineering experiments: kill your agent process at random points during heavy processing. Verify that it always resumes correctly with no lost data and no duplicate processing.
- Implement a drain mode: When shutdown is requested, stop accepting new work immediately (mark readiness endpoint as not-ready) but continue processing existing work for a bounded time before force-exiting.
- Handle nested shutdowns: If your agent spawns sub-processes or makes calls to external services, propagate the shutdown signal to them and wait for their clean termination.
- Log everything during shutdown: The shutdown sequence is the most failure-prone part of your agent's lifecycle. Emit structured logs at every step so you can debug partial shutdowns later.
- Use a shutdown timeout watchdog: After initiating graceful shutdown, start a timer. If the agent hasn't exited within the grace period, force-exit to avoid hanging indefinitely.
Common Pitfalls and How to Avoid Them
- Pitfall: Ignoring SIGTERM entirely. Many Python scripts don't register signal handlers and simply die on SIGTERM. Always explicitly handle termination signals.
- Pitfall: Writing checkpoints on the main thread while it's blocked. If your agent's main thread is stuck in a blocking operation (like an LLM API call), it can't process the shutdown signal. Use threading or asyncio with cancellation tokens to interrupt blocking operations.
- Pitfall: Checkpoint file corruption from partial writes. If the process crashes mid-write, the checkpoint file is garbage. Atomic writes (write to temp file, rename) eliminate this entire class of bugs.
- Pitfall: Loading a checkpoint that references resources that no longer exist. If your state contains file paths, network addresses, or in-memory object references, validate or reconstruct them on resume.
- Pitfall: Duplicate message processing after restart. If your agent processes a message, crashes before saving the checkpoint, then resumes and processes the same message again, you get duplicates. Use exactly-once semantics: commit the checkpoint after processing, and deduplicate on resume using message IDs.
- Pitfall: Checkpoint bloat over time. Unbounded conversation history or accumulating state can make checkpoints huge. Implement retention policies, snapshot compaction, or move old data to cold storage.
- Pitfall: Not testing the full restart cycle. Many teams implement graceful shutdown but never test the resume path. Always run integration tests that kill the process, restart it, and verify correct behavior.
Handling Blocking Operations During Shutdown
One of the trickiest challenges is handling shutdown when the agent is blocked on an external call. Here's a pattern using threading with a cancellation event:
import threading
import time
class CancellableOperation:
"""Wrapper for operations that can be interrupted during shutdown."""
def __init__(self, shutdown_event: threading.Event):
self.shutdown_event = shutdown_event
def call_llm_with_cancellation(self, prompt: str, timeout: float = 30.0) -> Optional[str]:
"""
Call an external LLM API but cancel if shutdown is requested.
Uses polling to check for shutdown during the operation.
"""
# Simulate an async LLM call with polling
result_container = {"result": None, "complete": False}
def make_call():
"""Simulated blocking LLM call."""
# In reality, this would be an HTTP request to an LLM service
for i in range(int(timeout * 10)): # Poll every 100ms
if self.shutdown_event.is_set():
print("[LLM_CALL] Cancelled due to shutdown")
return
time.sleep(0.1)
result_container["result"] = f"Response to: {prompt}"
result_container["complete"] = True
thread = threading.Thread(target=make_call, daemon=True)
thread.start()
# Wait for completion or shutdown, whichever comes first
while thread.is_alive() and not self.shutdown_event.is_set():