使用 Redis 键空间通知和租约对 GPU 工作者池进行负载均衡

GPU 池不同于无状态 Web 服务器池，像后者一样对其进行负载均衡会浪费你所拥有的最昂贵的硬件。运行繁重推理作业的 GPU 会被完全占用。向它发送第二个作业时，它不会共享资源，而是会将作业放入队列，于是，一个请求需要排在别人长达数分钟的工作之后等待，而此时另一块 GPU 却处于空闲状态。轮询（Round-robin）和最少连接（least-connections）这些源自 Web 负载均衡的本能反应，在这里都会失效，因为它们不知道“繁忙”对于 GPU 意味着什么。本文将基于 Redis 的两个原生功能——一个用于预留 GPU 的租约（lease）和一个在 GPU 释放时立即发出的键空间通知（keyspace notification）——来构建一个能感知繁忙状态的调度器，并将探讨那些促使我们采用此设计的故障模式。

为什么普通负载均衡在 GPU 上会失效

普通负载均衡器所基于的假设并不适用于 GPU 工作。Web 服务器处理许多并发请求，每个请求都成本低、耗时短，因此将请求均匀地分散到各个服务器上接近最优解，并且过时的负载视图几乎无关紧要。而 GPU 一次处理一个重型任务，每个任务都成本高、耗时长，因此唯一重要的是某个特定的 GPU 当前是否空闲，而一个关于其状态的过时视图，则意味着“立即开始”与“长达数分钟的停滞”之间的天壤之别。

轮询（Round-robin）算法完全忽略了繁忙状态，会“愉快地”将任务 N+1 分配给仍在处理任务 N 的 GPU。最少连接（Least-connections）算法更接近正确答案，但仍然是错误的，因为它计算的是连接数，而不是 GPU 的占用情况，而一个正在进行推理的 GPU 可能持有零个连接，但它却完全不可用。你需要的是一个调度器，其事实的唯一依据是“这个特定的 GPU 是否空闲”，并且这个状态在改变的瞬间就得到更新，同时带有硬性预留机制，以防止两个请求同时声明占用同一个空闲的 GPU。Redis 可以低成本地为你提供这两个部分。

租约即预留

核心状态是每个 GPU 的租约：一个 Redis 键，当 GPU 繁忙时存在，当其空闲时消失。申领 GPU 就是获取其租约，并且申领必须是原子操作，这样两个争夺同一个空闲 GPU 的请求就不会都成功。

// Claim GPU n only if its lease is absent. SET NX is atomic, so exactly one
// caller wins the race. The TTL is a safety net, covered below.
ok, err := rdb.SetNX(ctx, "gpu:lease:"+gpuID, workerID, leaseTTL).Result()
if err != nil {
	return "", err
}
if !ok {
	return "", errBusy // someone else holds this GPU
}

SetNX 是一个命令中的比较并交换：仅当键尚不存在时才设置该键。获胜的 GPU 获得租约并开始其作业。其他所有请求都会看到 ok 返回 false，然后继续尝试另一个 GPU。作业完成后，工作单元会删除租约，GPU 再次变为空闲。如果仅此而已，你仍然需要轮询来查找空闲的 GPU，而轮询正是我们想要避免的延迟和浪费的源头。第二个原语消除了轮询。

键空间通知将“空闲”转变为事件

每当键发生变化时，Redis 都可以发布一个事件，包括键过期或被删除时。该功能就是键空间通知，它让等待中的请求能够在 GPU 空闲出来的瞬间得知，而无需通过计时器轮询。你只需在服务器上为你关心的事件类别（通用键事件和过期事件）启用一次该功能，然后订阅租约键的模式即可。

// Subscribe to deletions and expirations of GPU lease keys. Each message
// means a specific GPU just became free.
pubsub := rdb.PSubscribe(ctx,
	"__keyevent@0__:del",
	"__keyevent@0__:expired",
)
for msg := range pubsub.Channel() {
	freedGPU := parseGPUID(msg.Payload) // the lease key that just vanished
	scheduler.notifyFree(freedGPU)
}

现在，整个流程是端到端事件驱动的。一个请求会尝试使用 SetNX 来声明任何一个空闲的 GPU。如果所有租约都已被持有，它会等待，而不是空转。任何地方的任何作业完成并删除其租约，或者租约过期时，Redis 就会发布该事件，一个等待中的请求会被唤醒，并尝试声明新空闲出来的 GPU。昂贵的硬件从繁忙到重新分配，所花费的时间仅为一条 Redis 消息的传输时间，在稳定状态下无需轮询。

TTL 所针对的故障模式

理想路径中存在一个漏洞，而这个漏洞最要紧，因为它会搁置你最昂贵的资源。假设一个工作进程申请了一个 GPU，启动了一个作业，然后崩溃了，或者其网络发生了分区，或者其进程在运行中被终止。它永远不会执行到删除租约的那行代码。租约键将永远留在 Redis 中，那个 GPU 将被永久标记为繁忙，并且永远不会为它触发键空间事件，因为一个已不存在的进程不会去删除任何键。一次崩溃就会将一个 GPU 从资源池中永久移除。

这就是为什么租约要带有 TTL（即上述 SetNX 调用中的 leaseTTL）。TTL 并非为正常情况而设，正常情况下工作进程会在租约到期前就完成任务并删除租约。它是针对异常情况的恢复机制。如果持有者死掉，租约会自行过期，Redis 会触发过期键空间事件，GPU 会在无人干预的情况下自动重新加入资源池。这里的租约与之前讨论的单运行器作业中的租约含义相同：一个由存活的持有者续订、由死掉的持有者放弃的预留。

这意味着一个长作业必须在 TTL 到期前续订其租约，这完全就像一个心跳，否则它将在运行中途因恢复机制而失去其 GPU。要么将 TTL 设置得比预期的最长作业时间更长一些，留出足够余量；要么对于那些无法限定时长的作业，周期性地对租约进行心跳续期。这其中的权衡大家都很熟悉：短 TTL 能快速恢复崩溃的 GPU，但有强行中止慢速作业的风险；长 TTL 对慢速作业是安全的，但会让崩溃的 GPU 被搁置更长时间。应该根据你的作业时长分布来选择 TTL，而不是使用默认值。

申领与等待之间的竞态

一个能感知繁忙状态的调度器存在一个微妙的顺序错误，这个问题值得专门指出，因为它很容易写出来，却又很难复现。设想一个请求检查了所有 GPU，发现它们都处于繁忙状态，然后订阅了空闲事件源。如果在检查和订阅之间的微小间隙中有一个 GPU 空闲出来，那么事件就会在订阅生效前触发，该请求永远不会听到这个事件，于是它会一直等待下一个空闲事件，而这可能需要几分钟之久。GPU 是空闲的，请求在等待，而彼此都不知道对方的存在。

修复方法是颠倒顺序，并在订阅后重新检查。首先订阅空闲事件源，然后尝试申领一个 GPU。如果申领成功，就取消订阅并运行。如果因为所有资源都繁忙而申领失败，就等待你已经订阅的事件源。关键在于，当事件到达时，要再次尝试申领，而不是想当然地认为空闲出来的 GPU 就是你的，因为可能已有另一个等待者捷足先登了。订阅，然后检查，然后等待并重试。这个顺序弥补了那个间隙，而“唤醒后重试”机制则处理了存在多个等待者时，空闲的 GPU 在你获取它之前就被别人申领的情况。

陈旧状态与协调过程

即使有 TTL，分布式状态也会发生漂移。一个租约可能在其工作单元（worker）被短暂分区时被删除，而该工作单元仍然认为自己持有 GPU。如果一个订阅者在过期事件触发的瞬间正在重新连接，那么该事件就可能会被错过。仅仅建立在事件之上的事件驱动系统，会随着时间的推移累积这些微小的差异。解决方法是进行周期性的协调：以一个较慢的间隔，独立地询问每个 GPU 工作单元它实际上在做什么，并修正 Redis 以匹配实际情况。在通常情况下，事件能保持系统的快速和实时。从长远来看，协调过程通过捕获事件错过的漂移来保持系统的正确性。由事件提供快速路径，由协调提供正确性保障，这是任何事件驱动调度器的一个良好默认模式。

当更简单的设计就足够时

当 GPU 是你稀缺而昂贵的资源，并且任务足够长，以至于将一个任务路由到繁忙的 GPU 上是一个代价高昂的错误时，这套机制就有了用武之地。如果你的任务又短又便宜，一个任何空闲工作单元都可以从中拉取任务的普通队列会更简单，效果也一样好，因为当每个任务只需片刻时间时，一个稍有瑕疵的分配所带来的成本可以忽略不计。如果你有一个托管的推理服务为你处理放置问题，那就使用它并跳过所有这些，因为你付费正是为了解决这个问题。

当三件事同时成立时，才需要构建这种租约和通知调度器：GPU 足够昂贵，以至于空闲时间是真正的损失；任务足够长，以至于错误的分配会使请求在一段有意义的时间内被搁置；以及你正在自己运行资源池，而不是租用托管端点。在这些条件下，用于原子预留的 Redis 租约和用于即时交接的 keyspace notification 为你提供了一个无需轮询即可保持昂贵硬件繁忙的调度器，而 TTL 加上一个协调过程则在工作单元死亡和事件丢失时保持其诚实可靠。