Python asyncio.Semaphore 控制多模型并发：避免推理请求同时打满显存

多模型推理服务最容易被忽略的风险，是把所有请求都交给 asyncio.gather() 后同时启动。任务数量一多，显存、模型服务连接和上游限流会一起承压。把 asyncio.Semaphore 放在真正的推理调用前，可以把“同时运行多少个模型任务”变成一个明确、可验收的配置。

并发上限应该保护推理临界区，而不是简单限制任务创建；用 async with semaphore 包住 model_infer，再用超时和任务组保证名额最终归还。

为什么 gather 一次启动所有推理任务

asyncio.gather() 负责并发等待，并不替你限制同时执行的协程数量。下面的代码会一次创建所有任务：

import asyncio

async def run_all(prompts):
    return await asyncio.gather(*(model_infer(prompt) for prompt in prompts))

这段写法适合任务数量很小、上游本身有可靠排队能力的情况。对本地多模型服务而言，问题在于任务刚被调度就可能同时占用显存。需要保护的是推理调用本身，而不是列表遍历。

把并发门放在真正的推理调用前

先把并发入口收窄到一个函数。图中只保留正文实际使用的三个节点：asyncio.create_task 负责安排任务，Semaphore 负责放行，model_infer 执行推理。

import asyncio

semaphore = asyncio.Semaphore(2)

async def infer_one(prompt):
    async with semaphore:
        return await model_infer(prompt)

async def run_all(prompts):
    tasks = [asyncio.create_task(infer_one(prompt)) for prompt in prompts]
    return await asyncio.gather(*tasks)

这里可以创建很多等待中的任务，但同时进入 model_infer 的任务最多为 2 个。把信号量放在 infer_one 外层会连排队时间也占住名额，通常会让吞吐下降；放在模型调用内部更容易和显存占用对应。

超时和异常要释放并发名额

不要手写“获取后 try，最后 release”的重复模板，除非确实需要把信号量跨函数传递。官方文档给出的等价语义是 async with semaphore 自动在退出时释放名额；再把单次调用放进 asyncio.timeout，超时也会离开上下文。

async def infer_one(prompt):
    async with semaphore:
        async with asyncio.timeout(20):
            return await model_infer(prompt)

批量任务需要统一收口时，可以使用 TaskGroup。它和 gather 的一个重要差别是：嵌套任务出现异常时，任务组会取消剩余任务。取消传播后，每个任务仍会退出自己的 async with semaphore，并发名额不会永久卡住。

async def run_checked(prompts):
    results = []
    async with asyncio.TaskGroup() as group:
        tasks = [group.create_task(infer_one(prompt)) for prompt in prompts]
    for task in tasks:
        results.append(task.result())
    return results

并发值怎么从小处开始

信号量的初始值不是显存容量的精确换算。模型大小、KV cache、输入长度、量化方式和服务端批处理都会改变实际峰值。先从 1 或 2 开始，观察显存峰值、单请求延迟、超时数和上游拒绝数，再逐步增大。

如果不同模型的显存成本差异很大，共用一个整数信号量只能得到粗粒度保护。可以先按模型组拆分信号量；不要为了追求一个全局数字，把轻量模型和大模型塞进同一条无法解释的排队规则。

几个容易误判的边界

Semaphore 不是速率限制器

它限制的是同时持有名额的任务数，不保证每秒请求数。如果上游要求固定速率，还需要单独的时间窗口或令牌桶策略。

任务创建数量仍然可能过大

信号量能挡住进入临界区的任务，但几万个等待中的 Task 仍会占用内存。输入规模很大时，再加一个有界队列或分批提交。

不要把成功响应当成显存已经释放

推理客户端返回后，服务端可能还在清理缓存。并发值应以监控到的峰值和错误率为依据，而不是只看 Python 协程是否完成。

上线前的核对清单

• 信号量是否紧贴 model_infer，而非包住无关的序列化和落盘。
• 超时、取消、异常路径是否都能离开 async with semaphore。
• 是否记录并发上限、排队时长、推理耗时、超时数和上游拒绝数。
• 输入很大时，是否同时限制了等待中的任务数量。

相关问题

把 Semaphore 改成 BoundedSemaphore 有必要吗？

如果担心代码存在多释放，BoundedSemaphore 会在计数超过初始值时抛出 ValueError，更适合做测试期的保护；正常使用上下文管理器时，两者都应保持成对 acquire/release。

为什么不用 sleep 控制推理并发？

sleep 只能错开启动时间，不能表达“前一个任务是否已经离开显存临界区”。信号量直接绑定资源占用，更容易在完成、超时和取消时得到一致状态。

先让并发上限可观测，再根据真实峰值调整，通常比一开始追求更高吞吐更稳。