AI推理服务的分布式并发控制——大模型请求调度的锁机制选型与实践

AI推理服务的分布式并发控制——大模型请求调度的锁机制选型与实践

一、大模型推理的并发困境——GPU资源有限，但请求无限

AI 推理服务的并发控制是一个容易被工程团队低估的问题。与传统微服务 API 不同，大模型推理具有以下三个显著特征：

这些特征将并发控制从"性能优化"提升到"可用性保障"的层面。本文将分析在 AI 推理场景下，如何利用分布式并发控制机制实现可靠的请求调度，并对比 Redis、ZooKeeper、etcd 三种后端在推理调度场景中的适用性。

二、底层机制与原理深度剖析

flowchart TB
subgraph Inference["AI 推理服务调度层"]
I1["请求接入网关"]
I2["并发控制中间件"]
I3["GPU 工作节点"]
end
subgraph Backend["并发控制后端选型"]
B1["Redis（信号量/令牌桶）<br/>AP 模型 · 低延迟"]
B2["ZooKeeper（顺序节点队列）<br/>CP 模型 · 公平调度"]
B3["etcd（Lease 租约）<br/>CP 模型 · 云原生友好"]
end
subgraph Core["调度核心策略"]
C1["请求排队：FIFO / 优先级队列"]
C2["并发限流：GPU 显存感知的信号量"]
C3["超时保护：动态 TTL 与 Lease 续期"]
C4["故障转移：Worker 心跳与槽位回收"]
end

I1 –> I2 –> I3
I2 –> Backend
Backend –> Core

GPU 显存感知的信号量模型是推理调度区别于传统限流的核心。传统的信号量只关心"并发数"，但推理场景需要同时关注并发数和显存余量。例如，一张 A100 虽然理论上可以加载 4 个 7B 模型实例，但如果某个请求触发了长序列生成（输出 4096 token），其 KV Cache 会额外占用 2~3GB 显存，此时实际可用槽位可能只有 3 个。因此调度层需要维护每个 Worker 上报的实时显存余量，在分配请求时做"并发数 × 显存"的交叉校验。

请求排队的公平性直接影响用户体验。ZooKeeper 的顺序节点机制天然支持 FIFO 公平队列，每个请求在 /inference-queue 下创建顺序节点后，只 Watch 前驱节点。而 Redis 的方案需要额外的 Lua 脚本来实现队列公平性，复杂度更高但延迟更低（单次入队约 1ms vs ZooKeeper 的 5~10ms）。

超时保护需要分两层设计。第一层是用户请求层面，推理超时后返回降级结果或排队提示；第二层是槽位占用层面，如果一个推理 Worker 在处理过程中假死（进程未退出但不响应），需要借助 Lease 机制在超时后自动回收槽位，避免资源永久泄漏。

三、生产级代码实现与最佳实践

3.1 Redis 信号量实现推理并发控制

/**
* 基于 Redis 的推理并发控制服务。
* 设计考量：使用 Redis 信号量模式限制单个 Worker 的并发推理数。
* 每个 Worker 启动时注册槽位，请求到达时竞争槽位，推理完成释放。
* 通过 Lua 脚本保证"检查-扣减"的原子性。
*/
@Service
public class InferenceConcurrencyService {

private final StringRedisTemplate redis;
private final String workerId;
private final int maxConcurrency;

/** 槽位获取的 Lua 脚本：原子地检查并扣减可用槽位数 */
private static final String ACQUIRE_SLOT_LUA =
"local used = tonumber(redis.call('GET', KEYS[1]) or '0') " +
"local max = tonumber(ARGV[1]) " +
"if used < max then " +
" redis.call('INCR', KEYS[1]) " +
" return 1 " +
"else " +
" return 0 " +
"end";

/** 槽位释放的 Lua 脚本：原子地减少占用计数（防止减到负数） */
private static final String RELEASE_SLOT_LUA =
"local used = tonumber(redis.call('GET', KEYS[1]) or '0') " +
"if used > 0 then " +
" redis.call('DECR', KEYS[1]) " +
" return 1 " +
"else " +
" return 0 " +
"end";

public InferenceConcurrencyService(StringRedisTemplate redis,
String workerId,
int maxConcurrency) {
this.redis = redis;
this.workerId = workerId;
this.maxConcurrency = maxConcurrency;
}

/**
* 尝试获取推理槽位。
* 设计考量：waitTimeout 不宜过长，推理场景下排队超过一定时间后
* 直接返回"服务繁忙"比让用户无限等待更合理。
*/
public boolean tryAcquireSlot(long waitTimeoutMs) {
String slotKey = "inference:slot:" + workerId;
long deadline = System.currentTimeMillis() + waitTimeoutMs;

while (System.currentTimeMillis() < deadline) {
Long result = redis.execute(
new DefaultRedisScript<>(ACQUIRE_SLOT_LUA, Long.class),
List.of(slotKey),
String.valueOf(maxConcurrency));

if (result != null && result == 1) {
return true;
}

try {
Thread.sleep(50); // 轮询间隔：50ms 在延迟与 CPU 消耗间做平衡
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
return false;
}
}
return false;
}

/**
* 释放推理槽位。
* 必须在 finally 块中调用，确保异常路径也能正确释放。
*/
public void releaseSlot() {
String slotKey = "inference:slot:" + workerId;
redis.execute(
new DefaultRedisScript<>(RELEASE_SLOT_LUA, Long.class),
List.of(slotKey));
}

/**
* 执行带并发控制的推理任务。
*/
public InferenceResult executeWithConcurrencyControl(
String prompt,
long waitTimeoutMs,
Function<String, InferenceResult> inferenceTask) {
if (!tryAcquireSlot(waitTimeoutMs)) {
return InferenceResult.busy("推理服务繁忙，当前排队请求已满");
}

long startTime = System.currentTimeMillis();
try {
return inferenceTask.apply(prompt);
} catch (Exception e) {
return InferenceResult.error("推理失败: " + e.getMessage());
} finally {
releaseSlot();
long costMs = System.currentTimeMillis() – startTime;
if (costMs > 30_000) {
log.warn("推理耗时较长: worker={}, costMs={}, promptLen={}",
workerId, costMs, prompt.length());
}
}
}
}

3.2 ZooKeeper 实现推理请求优先级队列

/**
* 基于 ZooKeeper 的优先级推理队列。
* 设计考量：利用 ZooKeeper 的顺序节点机制构建请求队列，
* 通过节点名称前缀区分优先级（high/normal/low），
* 高优先级请求的节点序号始终排在普通请求之前。
*/
@Service
public class InferencePriorityQueue implements DisposableBean {

private final CuratorFramework client;
private static final String QUEUE_ROOT = "/inference-queue";
private static final String SLOT_ROOT = "/inference-slots";

public InferencePriorityQueue(CuratorFramework client) {
this.client = client;
try {
// 确保根路径存在
if (client.checkExists().forPath(QUEUE_ROOT) == null) {
client.create().creatingParentsIfNeeded()
.forPath(QUEUE_ROOT);
}
if (client.checkExists().forPath(SLOT_ROOT) == null) {
client.create().creatingParentsIfNeeded()
.forPath(SLOT_ROOT);
}
} catch (Exception e) {
throw new RuntimeException("初始化推理队列失败", e);
}
}

/**
* 入队一个推理请求。
*
* @param priority 优先级：high 请求在 normal 之前处理
* @param requestId 请求唯一标识
* @return 队列序号
*/
public int enqueue(Priority priority, String requestId) throws Exception {
String path = QUEUE_ROOT + "/" + priority.name().toLowerCase() + "-";
String createdPath = client.create()
.withMode(CreateMode.EPHEMERAL_SEQUENTIAL)
.forPath(path, requestId.getBytes(StandardCharsets.UTF_8));

// 从路径中提取序号
String seqPart = createdPath.substring(createdPath.lastIndexOf("-") + 1);
return Integer.parseInt(seqPart);
}

/**
* 监听队列并消费下一个推理任务。
* 设计考量：高优先级队列优先处理，同一优先级内 FIFO。
* 使用 Watch 前驱节点的方式避免惊群效应。
*/
public String waitAndDequeue(long timeoutMs) throws Exception {
long deadline = System.currentTimeMillis() + timeoutMs;

// 先检查高优队列，再检查普通队列
Priority[] priorities = {Priority.HIGH, Priority.NORMAL, Priority.LOW};

while (System.currentTimeMillis() < deadline) {
for (Priority priority : priorities) {
String prefix = QUEUE_ROOT + "/" + priority.name().toLowerCase() + "-";
List<String> children = client.getChildren()
.forPath(QUEUE_ROOT);

// 找出当前优先级下序号最小的节点
Optional<String> candidate = children.stream()
.filter(c -> c.startsWith(priority.name().toLowerCase()))
.sorted()
.findFirst();

if (candidate.isPresent()) {
String fullPath = QUEUE_ROOT + "/" + candidate.get();
byte[] data = client.getData().forPath(fullPath);
// 消费后删除节点
client.delete().forPath(fullPath);
return new String(data, StandardCharsets.UTF_8);
}
}
Thread.sleep(100);
}
throw new TimeoutException("推理队列等待超时");
}

@Override
public void destroy() throws Exception {
if (client != null) {
client.close();
}
}

public enum Priority { HIGH, NORMAL, LOW }
}

3.3 GPU 显存感知的动态槽位管理

/**
* GPU Worker 的显存感知槽位注册。
* 设计考量：每个推理 Worker 启动时通过 etcd Lease 注册自身信息，
* 包括最大并发数、当前显存余量、模型名称。
* 调度器根据这些信息动态分配请求，而非使用固定并发数。
*/
@Service
public class GpuWorkerRegistry {

private final Client etcdClient;
private final String workerId;
private long leaseId;
private ScheduledExecutorService keepAliveExecutor;

public GpuWorkerRegistry(Client etcdClient, String workerId) {
this.etcdClient = etcdClient;
this.workerId = workerId;
}

/**
* 注册 Worker 并持续上报心跳和显存信息。
* 设计考量：Lease TTL 设置为 10s，每 3s 续约一次，
* 确保 Worker 宕机后最多 10s 内槽位被回收。
*/
public void register(GpuInfo gpuInfo) throws Exception {
// 创建 Lease
leaseId = etcdClient.getLeaseClient().grant(10)
.get(5, TimeUnit.SECONDS)
.getID();

// 上报 Worker 信息
String key = "/gpu-workers/" + workerId;
String value = JsonUtils.toJson(gpuInfo);

etcdClient.getKVClient().put(
ByteSequence.from(key, StandardCharsets.UTF_8),
ByteSequence.from(value, StandardCharsets.UTF_8),
PutOption.newBuilder().withLeaseId(leaseId).build())
.get(5, TimeUnit.SECONDS);

// 启动续约线程
keepAliveExecutor = Executors.newSingleThreadScheduledExecutor();
keepAliveExecutor.scheduleAtFixedRate(() -> {
try {
// 每次续约时更新显存信息
GpuInfo latest = GpuMonitor.collect();
String updatedKey = "/gpu-workers/" + workerId;
String updatedValue = JsonUtils.toJson(latest);

etcdClient.getKVClient().put(
ByteSequence.from(updatedKey, StandardCharsets.UTF_8),
ByteSequence.from(updatedValue, StandardCharsets.UTF_8),
PutOption.newBuilder().withLeaseId(leaseId).build())
.get(3, TimeUnit.SECONDS);
} catch (Exception e) {
log.warn("Worker 续约失败: workerId={}", workerId, e);
}
}, 3, 3, TimeUnit.SECONDS);

log.info("GPU Worker 注册成功: workerId={}, maxConcurrency={}, vramFree={}MB",
workerId, gpuInfo.maxConcurrency(), gpuInfo.vramFreeMB());
}

/**
* 优雅下线：撤销 Lease，释放槽位。
*/
@PreDestroy
public void deregister() {
try {
if (keepAliveExecutor != null) {
keepAliveExecutor.shutdown();
}
if (leaseId > 0) {
etcdClient.getLeaseClient().revoke(leaseId)
.get(3, TimeUnit.SECONDS);
}
log.info("GPU Worker 已注销: workerId={}", workerId);
} catch (Exception e) {
log.error("Worker 注销失败: workerId={}", workerId, e);
}
}

/**
* GPU 信息快照，用于调度决策。
*/
public record GpuInfo(
String modelName,
int maxConcurrency,
long vramTotalMB,
long vramFreeMB,
int currentRequests,
long timestamp) {}
}

四、边界分析与架构权衡

推理调度后端选型决策矩阵：

Redis 的槽位竞争窗口问题：

在高并发场景下，Redis 信号量虽然能通过 Lua 脚本保证"检查-扣减"原子性，但由于 Redis 的主从复制是异步的，当主节点故障发生切换时，可能出现新主尚未收到最新的槽位计数，导致实际分配的槽位超过 maxConcurrency 的情况。对于推理调度而言，暂时性的超发虽然不会立即造成服务崩溃（GPU 有少量显存余量缓冲），但如果多个 Worker 同时超发，累积的显存占用仍可能触发 OOM。解决思路是对 Worker 侧做兜底保护：Worker 在收到推理请求后，先检查当前显存余量是否足以加载 KV Cache，不足时直接返回拒绝而非强行执行。

ZooKeeper 的适用场景边界：

ZooKeeper 的公平调度和优先级队列能力在推理场景中确实强大，但不是没有代价的。ZooKeeper 集群的写入操作需要过半节点确认，当推理请求 QPS 超过 5000 时，ZooKeeper 的延迟会显著上升。此外，大量临时节点的创建和删除会产生可观的磁盘 I/O 和快照开销。如果推理服务的 QPS 不高（< 1000），且对公平性有明确要求（如付费用户优先、紧急任务插队），ZooKeeper 是合适的选择。如果 QPS 在万级，Redis + 本地队列的组合方案反而更经济。

etcd 的云原生适配优势：

在 Kubernetes 环境中，etcd 天然嵌入控制面，Worker 的注册和心跳完全可以借 K8s 的 Pod 生命周期来管理——Pod 启动时注册，Pod 终止（preStop hook）时主动注销，Lease 仅作为兜底的故障检测机制。这种设计消除了在应用层维护心跳续约线程的复杂性。

推荐选型指南：

• 团队已有 Redis → 使用 Redis 信号量 + Lua 脚本方案，覆盖 80% 的推理调度场景。
• 部署在 Kubernetes 上、QPS 小于 5000 → 使用 etcd 的 Lease 机制做 Worker 注册和槽位回收。
• 对请求公平性有严格要求（VIP 用户插队、任务优先级分级）→ ZooKeeper 顺序节点队列。
• 超大规模推理集群（百卡级以上、QPS 万级）→ 考虑 Redis Cluster + 本地令牌桶的两级限流架构。

五、总结

AI 推理服务的并发控制不是简单的"设置一个最大并发数"，而是一个涉及资源感知、请求调度、故障恢复的系统工程：

在工程实践中，建议先从 Redis 信号量方案起步，满足绝大部分推理调度需求。当业务规模增长到需要优先级调度或 Worker 健康管理时，再按需引入 ZooKeeper 或 etcd。工具服务于场景，场景决定架构。