AI 推理服务优雅下线：Drain 状态与等待中的请求处理
AI 推理服务优雅下线：Drain 状态与等待中的请求处理
一、K8s 直接杀 Pod，推理请求全丢了，用户看到 502，老板脸绿了
推理服务和普通 Web 服务不一样的地方在于：一个请求可能跑几十秒甚至几分钟。K8s 默认的terminationGracePeriodSeconds: 30根本不够用。滚动更新时，旧 Pod 被 SIGTERM 杀死，正在跑的推理请求直接中断，用户再重试就报错。你以为是模型崩了，结果只是 Pod 被杀得太快。
优雅下线要做四件事：
- 摘流量—— 不要再给这个 Pod 发新请求
- 等现有的—— 正在跑的请求让它跑完（合理超时内）
- 释放资源—— 优雅关闭连接池、数据库连接、GPU 上下文
- 记录状态—— 把未完成的请求状态持久化，别丢了
二、优雅下线的关键时序与状态转移
核心思路：先在 Service 层面摘掉 Endpoint，再给 Pod 发 SIGTERM。这样新请求不会路由到这个 Pod，只有已经在处理的请求需要等。
三、生产级优雅下线实现
Go 推理服务优雅退出
package main import ( "context" "net/http" "os" "os/signal" "sync" "sync/atomic" "syscall" "time" ) // InferenceServer 带优雅下线能力的推理服务 type InferenceServer struct { server *http.Server activeRequests int64 // 当前正在处理的请求数 draining int32 // 是否进入 draining 模式 wg sync.WaitGroup } func NewInferenceServer(addr string) *InferenceServer { mux := http.NewServeMux() s := &InferenceServer{ server: &http.Server{Addr: addr, Handler: mux}, } mux.HandleFunc("/v1/infer", s.handleInfer) mux.HandleFunc("/healthz", s.handleHealthz) mux.HandleFunc("/readyz", s.handleReadyz) return s } func (s *InferenceServer) handleReadyz(w http.ResponseWriter, r *http.Request) { if atomic.LoadInt32(&s.draining) == 1 { w.WriteHeader(http.StatusServiceUnavailable) w.Write([]byte("draining")) return } // 还检查模型是否加载完毕 w.WriteHeader(http.StatusOK) w.Write([]byte("ok")) } func (s *InferenceServer) handleHealthz(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) w.Write([]byte("ok")) } func (s *InferenceServer) handleInfer(w http.ResponseWriter, r *http.Request) { // 进入 draining 模式后拒绝新请求 if atomic.LoadInt32(&s.draining) == 1 { w.Header().Set("Retry-After", "5") w.WriteHeader(http.StatusServiceUnavailable) w.Write([]byte(`{"error":"server is draining, please retry"}`)) return } atomic.AddInt64(&s.activeRequests, 1) defer atomic.AddInt64(&s.activeRequests, -1) // 模拟推理处理（实际场景是 GPU 推理） ctx := r.Context() select { case <-time.After(30 * time.Second): w.Write([]byte(`{"result":"inference done"}`)) case <-ctx.Done(): // 客户端断开了，不用等了 return } } func (s *InferenceServer) WaitForDrain(timeout time.Duration) error { deadline := time.Now().Add(timeout) ticker := time.NewTicker(500 * time.Millisecond) defer ticker.Stop() for time.Now().Before(deadline) { active := atomic.LoadInt64(&s.activeRequests) if active == 0 { return nil } <-ticker.C } return context.DeadlineExceeded } func (s *InferenceServer) Shutdown(ctx context.Context) error { // Step 1: 标记为 draining，拒绝新请求 atomic.StoreInt32(&s.draining, 1) // Step 2: 等待当前 in-flight 请求处理完 if err := s.WaitForDrain(25 * time.Second); err != nil { log.Printf("drain timeout, %d requests still active", atomic.LoadInt64(&s.activeRequests)) } // Step 3: 关闭 HTTP Server（给剩余连接 5s 缓冲） shutdownCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() return s.server.Shutdown(shutdownCtx) } func (s *InferenceServer) Run() error { sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT) go func() { <-sigCh // 收到终止信号后启动优雅关闭 log.Println("received shutdown signal, start draining...") if err := s.Shutdown(context.Background()); err != nil { log.Printf("shutdown error: %v", err) } }() return s.server.ListenAndServe() }K8s Deployment 配置
apiVersion: apps/v1 kind: Deployment metadata: name: inference-server spec: replicas: 3 strategy: type: RollingUpdate rollingUpdate: maxSurge: 1 maxUnavailable: 0 # 关键：不允许不健康 Pod 存在 template: spec: # 关键参数：给推理请求充足的时间完成 terminationGracePeriodSeconds: 120 # 2 分钟 containers: - name: inference image: inference-server:v2.3.0 ports: - containerPort: 8080 lifecycle: preStop: exec: command: - /bin/sh - -c - | # 等 K8s 网络规则传播（10s 足够） sleep 10 # 调用自身优雅关闭接口 curl -X POST http://localhost:8080/admin/drain # 给足够时间处理 in-flight 请求 sleep 105 readinessProbe: httpGet: path: /readyz port: 8080 periodSeconds: 5 failureThreshold: 2 livenessProbe: httpGet: path: /healthz port: 8080 initialDelaySeconds: 30 periodSeconds: 10 resources: limits: nvidia.com/gpu: 1推理请求端重试策略
import httpx import asyncio from typing import Optional class InferenceClient: """带优雅重试的推理客户端""" def __init__(self, base_url: str, max_retries: int = 3): self.base_url = base_url self.max_retries = max_retries self.client = httpx.AsyncClient(timeout=120.0) async def infer(self, prompt: str) -> Optional[dict]: for attempt in range(self.max_retries): try: resp = await self.client.post( f"{self.base_url}/v1/infer", json={"prompt": prompt} ) if resp.status_code == 200: return resp.json() # 503 + Retry-After：服务在 draining if resp.status_code == 503: retry_after = int( resp.headers.get("Retry-After", "5") ) await asyncio.sleep(retry_after) continue resp.raise_for_status() except httpx.TimeoutException: # 推理超时，可能正在被 kill if attempt < self.max_retries - 1: await asyncio.sleep(2 ** attempt) continue raise except httpx.ConnectError: # 连接失败，Pod 可能已经退了 await asyncio.sleep(1) continue return None四、边界分析与架构权衡
关键边界
| 边界场景 | 风险 | 缓解措施 |
|---|---|---|
| 推理请求超过 terminationGracePeriod | 被 SIGKILL 强制终止 | 合理设置 MaxInferTime，超过则返回部分结果 |
| GPU 显存中的 KV Cache 丢失 | 请求需要重新计算 | 支持断点续推（Checkpoint） |
| HPA 缩容过多 Pod | 剩余 Pod 过载 | 限制缩容速率（scaleDown.stabilizationWindowSeconds: 300） |
| preStop hook 执行失败 | Pod 直接 SIGTERM | Hook 失败不影响容器停止 |
| Client 侧没有重试 | 用户直接看到错误 | Sidecar/API Gateway 层做透明重试 |
| Graceful shutdown 触发 OOM | 等待期间内存增长 | 限制 in-flight 请求上限（队列满即 503） |
权衡决策
Drain 等待时间 vs 可用性：Drain 时间越长，请求成功率高，但发布变慢。生产环境建议：推理 P50 延迟 × 2 + 10s 缓冲。例如 P50=20s，则 terminationGracePeriod 设为 50s。
是否强制中断：超过 Drain 时间的请求，是记日志然后 kill，还是拼命等？推荐 kill 后依赖客户端重试。推理请求应设计为幂等（相同 prompt 返回相同结果或会话 ID），重试没副作用。
sidecar Drain 顺序：如果 Pod 有 Envoy/Istio sidecar，容器关闭顺序很关键。Sidecar 要比主容器晚退出，否则主容器发给 sidecar 的请求会失败。Istio 1.15+ 支持EXIT_ON_ZERO_ACTIVE_CONNECTIONS。
五、总结
优雅下线看起来是运维问题，实际上是架构问题。你的服务如果不支持重试、不支持幂等、不支持部分结果返回，优雅下线做得再好也没用。
三个必须做：
- K8s 层：preStop hook + terminationGracePeriodSeconds > 推理最大耗时
- 应用层：Drain 标记 + in-flight 计数 + 503 + Retry-After
- 客户端：指数退避重试 + 连接池自动发现新 Endpoint
一个自检问题：你的推理服务，从发 SIGTERM 到完全停止，需要多少秒？如果答案是"不知道"，这篇文章值得收藏。