用 Redis 心跳自数 Pod,实现多 Pod 并发限流(不查 K8s API)

多个 Pod 一起从 Redis 队列取任务跑(典型如一批 GPU worker),想加个开关:当"正在忙的 Pod"占到"总 Pod 数"的 80% 时,空闲 Pod 就先别再领新任务,留 20% 余量,避免整个机群被压满。

直觉上第一步是"先搞清楚现在有几个 Pod"——于是去查 Kubernetes API 或 Prometheus。但这条路会卡在三个坑上:要 RBAC 权限(容器默认 ServiceAccount 通常 403)、副本会动态扩缩(查到的值要不断刷新)、多集群端点不统一(自建 Prometheus、云托管 TMP/ARMS 各不相同、还各自要 token)。

本文给一个绕开这些的做法:让每个 Pod 用 Redis 心跳"自己数自己",纯 Redis 实现 Pod 计数 + 80% 并发刹车,完全不碰 Kubernetes API——零权限、跨集群一致、副本怎么变都自动跟上。

环境

  • Kubernetes(任意发行版/任意云厂商,方案与厂商无关)
  • 一批同一 Deployment 的消费 Pod,均从同一个 Redis 队列取任务
  • Redis(所有 Pod 本就连着它取任务)
  • 示例语言 Python 3,redis-py 客户端(其它语言同理,核心是几条 Redis 命令)
  • 目标:全机群"忙碌 Pod / 总 Pod > 80%"时刹车

应用场景:为什么要"留 20% 余量"

一批 GPU worker Pod 从 Redis 队列取任务跑重活(模型推理、音视频处理、渲染等),队列经常堆积、任务来得很快。如果放任所有 Pod 一直满载消费,会有几个问题:

  • 突发/高优任务抢不到机器:整个 GPU 机群被普通任务占满时,临时插队的高优任务、失败重试的任务只能干等,没有机器接住;
  • 满载下更容易出事:显存吃紧、资源争抢、处理变慢甚至 OOM,反而拖垮整体吞吐;
  • 没有缓冲:扩容还没起来、或个别 Pod 故障时,没有富余产能顶上。

所以想要一个自我节流的开关:让消费端在"忙碌 Pod 达到总数的 80%"时主动刹车,始终留约 20% 的机器空着,随时能接住高优/突发任务,也给系统留缓冲。

注意这个 80% 是相对"当前实际 Pod 数"的动态比例:副本会手动扩缩(甚至跨多个云的集群),5 个 Pod 时余量是 1 个、20 个 Pod 时余量是 4 个——所以分母 N 必须实时拿到,不能写死。这正是下面要解决的核心难点。

一、为什么不去查 Kubernetes "我有几个 Pod"

要算"忙碌占比",分母是"总 Pod 数 N"。直接问 Kubernetes 拿 N,有三个现实障碍:

  1. 要权限:Pod 里用自带 ServiceAccount 调 API(GET .../deployments 或 list pods),默认 RBAC 大概率返回 403 Forbidden,得让运维单独授一个只读 Role。
  2. 副本动态变:spec.replicas 会扩缩(手动或 HPA),查到的 N 得周期刷新,不能查一次写死。
  3. 多集群不统一:换到不同云,集群内不一定有自建 Prometheus;云托管的 Prometheus(腾讯 TMP、阿里 ARMS)端点和鉴权各不相同,有的还得带 token、Pod 不一定能直连。

这些都是"向外部系统要数据"带来的耦合。换个思路:N 这个数,其实 Pod 们自己就能数出来,不用问外人。

二、核心思路:用 Redis 心跳让 Pod 自己数自己

打个比方:想知道班里几个人,不用找老师要花名册(Kubernetes),而是让每个人每隔一会儿在黑板上写一下自己的名字(黑板=Redis)。数黑板上的名字就是当前人数;有人走了,过一会儿名字被擦掉;新来的写上就被数到。

落到 Pod:每个 Pod 本来就连 Redis,那就让它周期性在 Redis 里"打卡报到"。数打卡的=当前 Pod 数。维护两本"打卡本":

  • 存活本:每个 Pod 不管忙不忙,从启动到退出,每隔约 10 秒打一次卡 → 数它=总 Pod 数 N;
  • 忙碌本:Pod 正在跑任务期间打卡(并持续刷新) → 数它=忙碌数 busy

开关:消费前算 busy / N,超过 80% 就刹车。副本手动扩缩会自动反映——新 Pod 一启动就进 N,走掉的 Pod 心跳停了自动掉出 N。

方案结构:两个 Redis 打卡本 + 80% 刹车判断

下面这张动图把"一次消费"的判断走一遍:① 打卡存活 → ② 取任务 → ③ 算 busy/N → ④ 超 80% 就刹车(任务放回),否则标忙处理。

一次消费的判断流程

三、数据结构:两个 Redis 有序集合

用 Redis 的有序集合(sorted set)当打卡本,member = Pod 名,score = 最后打卡时间戳:

  • gpu:pods:alive —— 存活打卡本
  • gpu:pods:busy —— 忙碌打卡本

为什么用"时间戳当 score":数数之前先 ZREMRANGEBYSCORE 把"超过 TTL 没打卡的"删掉(判定为死),再 ZCARD 数剩下的。这样崩溃/被缩掉的 Pod 自动从计数里消失,比"INCR/DECR 计数器"安全——计数器只要有一个 Pod 崩在中间没减回来,就永久泄漏。

唯一性靠两层咬合保证:Kubernetes 的 Pod 名天生唯一(部署名-RS哈希-随机后缀),Redis 集合 member 也唯一(同名反复 ZADD 只更新它的 score、不会变两条)。所以 ZCARD 就是不重复的 Pod 数。

四、完整实现

# gpu_throttle.py —— Redis 自数 Pod + 80% 并发刹车(零 K8s 依赖)
import os, time, socket, threading, signal, redis

# ---- 配置(走环境变量) ----
REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
QUEUE     = os.environ.get("TASK_QUEUE", "task_queue")
TTL       = int(os.environ.get("THROTTLE_TTL", "30"))      # 多久没打卡判定为死
RATIO     = float(os.environ.get("THROTTLE_RATIO", "0.8")) # 刹车阈值 80%
HB        = max(2, TTL // 3)                                # 心跳间隔(TTL 的 1/3)
ALIVE, BUSY = "gpu:pods:alive", "gpu:pods:busy"
POD = os.environ.get("POD_NAME") or socket.gethostname()   # 优先用注入的真 Pod 名

r = redis.from_url(REDIS_URL)
_stop = threading.Event()

# ---- 原子占槽位(Lua:清死的→数 N 和 busy→没满就标忙) ----
_ACQUIRE = r.register_script("""
local now=tonumber(ARGV[1]); local ttl=tonumber(ARGV[2])
local ratio=tonumber(ARGV[3]); local pod=ARGV[4]
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now-ttl)   -- 清掉死的 alive
redis.call('ZREMRANGEBYSCORE', KEYS[2], '-inf', now-ttl)   -- 清掉死的 busy
local N = redis.call('ZCARD', KEYS[1]); if N < 1 then N = 1 end
if redis.call('ZSCORE', KEYS[2], pod) then                 -- 已在忙(重入)就续期
  redis.call('ZADD', KEYS[2], now, pod); return 1
end
local busy = redis.call('ZCARD', KEYS[2])
local cap = math.floor(ratio * N); if cap < 1 then cap = 1 end  -- max(1) 防小 N 取整成 0
if busy < cap then
  redis.call('ZADD', KEYS[2], now, pod)                    -- 占到槽位,标忙
  return 1
end
return 0                                                    -- 已达 80%,刹车
""")

# ---- 存活心跳(整个进程生命周期一直打) ----
def _alive_loop():
    while not _stop.wait(HB):
        try: r.zadd(ALIVE, {POD: time.time()})
        except Exception: pass

def start():
    r.zadd(ALIVE, {POD: time.time()})                       # 启动立刻打一次
    threading.Thread(target=_alive_loop, daemon=True).start()

def try_acquire():
    """抢一个消费槽位。占到=True(可消费),没占到=已达 80%。"""
    try:
        return _ACQUIRE(keys=[ALIVE, BUSY],
                        args=[time.time(), TTL, RATIO, POD]) == 1
    except Exception:
        return True   # Redis 异常策略:放行(别让限流器故障把业务卡死);要保守改 False

def release():
    try: r.zrem(BUSY, POD)
    except Exception: pass

def shutdown(*_):
    _stop.set()
    try: r.zrem(ALIVE, POD); r.zrem(BUSY, POD)              # 优雅退出:把自己摘掉
    except Exception: pass
    os._exit(0)

# ---- 处理期间持续续 busy 心跳(防长任务被当过期) ----
def _process_with_hb(task):
    done = threading.Event()
    def renew():
        while not done.wait(HB):
            try: r.zadd(BUSY, {POD: time.time()})
            except Exception: pass
    threading.Thread(target=renew, daemon=True).start()
    try:
        process(task)            # ← 替换成你的真实处理逻辑
    finally:
        done.set()

# ---- 主消费循环 ----
def consume_loop():
    start()
    signal.signal(signal.SIGTERM, shutdown)                 # 缩容/重启时优雅摘掉
    signal.signal(signal.SIGINT, shutdown)
    while not _stop.is_set():
        item = r.blpop(QUEUE, timeout=5)                    # 取任务(队列堆积时立即返回)
        if not item:
            continue
        task = item[1]
        if not try_acquire():                               # 占不到=已 80%,任务放回
            r.lpush(QUEUE, task)                             # 放回队头,不丢
            time.sleep(1); continue
        try:
            _process_with_hb(task)
        finally:
            release()                                       # 干完释放槽位

def process(task):
    raise NotImplementedError("替换成你的处理逻辑")

if __name__ == "__main__":
    consume_loop()

更稳的 Pod 名:socket.gethostname() 在 Kubernetes 默认等于 Pod 名,但若 Pod spec 改了 hostname 就不一定。推荐用 Downward API 注入,代码里优先读 POD_NAME:

# Deployment 的容器 env 加:
- name: POD_NAME
  valueFrom: { fieldRef: { fieldPath: metadata.name } }

五、关键设计点

1. 占槽位为什么要用 Lua 原子化。 多个空闲 Pod 可能同时判断"还没到 80%"然后一起领任务,瞬时越线。把"清死的 + 数 N/busy + 判断 + 标忙"放进一段 Lua,Redis 单线程保证它原子执行,不会一起越过阈值。

2. 过期防泄漏。 数数前先 ZREMRANGEBYSCORE ... now-TTL 删掉过期成员。Pod 崩了心跳停,TTL(默认 30 秒)后自动从两个集合消失,不会被当成"活着/忙碌"。

3. 心跳必须持续续期。 存活心跳由后台线程一直打;忙碌心跳在处理任务期间也要持续 ZADD(代码里 _process_with_hb),否则一个跑很久的任务,Pod 会在 TTL 后被自己当成过期,比例算错。

4. 优雅退出。 收到 SIGTERM(Kubernetes 缩容/滚动更新会发)立刻把自己从两个集合摘掉,N 不用等 TTL 过期就准确。

六、行为与边界

  • 副本手动扩缩自动跟:新 Pod 启动即进 N,缩掉的 Pod TTL 后自动掉出 N——不用关心你是 kubectl scale 还是改 chart 重发;
  • 跨集群一致:只依赖 Redis,任何云/任何集群行为一样;
  • 数的是"实际在跑的 Pod":比"期望副本数"更贴合"别把干活的机器压满"这个真实目标;
  • 滞后约一个心跳(10~30 秒):新 Pod 进 N、缩容 Pod 出 N 有延迟,对软限流无所谓;
  • 滚动更新瞬时偏大:旧 Pod 名停止心跳前、新 Pod 名已出现,中间最多 TTL 秒两者都在,N 短暂 +1,可忽略;
  • 小 N 边界:阈值用 max(1, floor(0.8*N)) 兜底,否则 N 很小时 0.8*N 取整成 0 会永远不消费;
  • Redis 异常策略:限流器查 Redis 失败时,示例选"放行"(别让限流器自身故障把业务全卡死),想保守可改"暂停"。

七、要不要套个框架(Dramatiq / Celery / RQ)

  • 现在是裸 blpop 自己写循环消费 → 上面这套直接用,它本身就是完整的"消费 + 限流",不必引入框架;
  • 已经在用 Dramatiq → 用它自带的 ConcurrentRateLimiter(基于 Redis 的分布式并发限制器),limit 设成本方案心跳数出来的 0.8*N,别另写循环;
  • 用 Celery / RQ → 框架不动,把"心跳 + 80% 判断"作为任务里的一道前置检查塞进去。

只为"80% 刹车"这一件事,裸 Redis 最省、不增加依赖;只有还想要重试/死信/定时这些,才值得上 Dramatiq。

八、相关方案与参考(并发那层有现成库)

本文不是发明新轮子,而是把两个成熟模式组合起来:

  • 分布式信号量(并发计数):"Redis sorted set + 时间戳过期"实现的并发限制,是经典模式——《Redis in Action》第 6 章 "Counting semaphores" 即此。GitHub 上现成可用的库不少:Python 的 bluele/redis-semaphoreredis-rate-limiters,TypeScript 的 swarthy/redis-semaphore,Go 的 amitaifrey/redisemaphore,Java 则有 Redisson 的 RPermitExpirableSemaphore。本文代码里 busy 那半截就是它。
  • Redis 在线计数(presence):用 sorted set + 心跳记录"谁在线",是统计在线实例数的常见做法。本文的"存活打卡本"就是它。

这些信号量库的容量是固定值(创建时传死)。本文要解决的是 容量随 Pod 扩缩动态变、还得不查 Kubernetes 自己数出来——也就是把上面两半拼起来,把"存活心跳数出来的 0.8×N"喂给信号量。这层组合没有现成整包,但拼起来也就这点量。

所以两条路任选:① 用现成信号量库当并发限制器、容量传 0.8×N,N 仍由本文的存活心跳算;② 直接用本文这份(信号量 + presence + 原子 Lua + 崩溃过期都合在一起)。

快速参考

Redis 数据结构

集合 含义 写入 计数
gpu:pods:alive 存活 Pod 每个 Pod 每 ~10s ZADD self now ZCARD(先删过期)= N
gpu:pods:busy 忙碌 Pod 跑任务期间 ZADD self now 持续续期 ZCARD(先删过期)= busy

核心命令

ZADD gpu:pods:alive <now> <pod>                 # 打卡/续期
ZREMRANGEBYSCORE gpu:pods:alive -inf <now-TTL>  # 清掉过期(死)的
ZCARD gpu:pods:alive                            # 当前存活 Pod 数 N
ZREM gpu:pods:busy <pod>                        # 释放忙碌槽位

几条铁律

  • 心跳是周期性续期,不是启动报一次——否则崩溃的 Pod 永远被算成活的;
  • 占槽位用 Lua 原子判断+标忙,避免多个 Pod 一起越过 80%;
  • 处理长任务时忙碌心跳要持续续,否则自己被当过期;
  • Pod 名用 Downward API 注入的 metadata.name 最稳;
  • 阈值 max(1, floor(0.8*N)) 兜底小 N;
  • SIGTERM 时把自己从集合摘掉,N 更及时。

参考链接

posted @ 2026-06-26 10:51  Hello_worlds  阅读(2)  评论(0)    收藏  举报