已经会创建任务,为什么还需要限流

创建 100 个 Task,不代表下游服务就适合同时处理 100 个请求。Semaphore(n) 给一段工作 n 个许可,进入前获取,离开时归还。

import asyncio

async def generate(rollout_id, semaphore, state):
    async with semaphore:
        state["active"] += 1
        state["peak"] = max(state["peak"], state["active"])
        try:
            await asyncio.sleep(0)
            return rollout_id
        finally:
            state["active"] -= 1

async def exercise():
    semaphore = asyncio.Semaphore(2)
    state = {"active": 0, "peak": 0}
    results = await asyncio.gather(*(generate(i, semaphore, state) for i in range(5)))
    print("results", results)
    print("peak", state["peak"])

if __name__ == "__main__":
    asyncio.run(exercise())

下载并运行本节 example.py。只使用 Python 标准库;推荐 Python 3.12,最低 Python 3.11。

运行结果:

results [0, 1, 2, 3, 4]
peak 2

这段独立实验把生成逻辑替换为一次异步等待;peak 记录许可区域里的最大任务数。async with semaphore 让进入和离开配对,异常退出也会执行释放。

从方法层面看,进入时相当于 await semaphore.acquire(),退出时执行 semaphore.release()。acquire() 在没有许可时让当前任务等待;release() 归还许可,使等待者有机会继续。async with 自动把成功进入与退出配对,避免异常路径漏掉释放。它不是创建或启动任务的接口。

在 Slime 中找到许可的作用范围

GenerateState 创建信号量;generate_and_rm 在调用生成器的区域使用它。不要只读创建时的数字:许可从哪一行开始持有、在哪一行释放,才决定限制了什么工作。

展开 Slime 源码:GenerateState.__init__

GenerateState.__init__
slime/rollout/sglang_rollout.py · L88–L101
快照 2f2318653f6f794dddd321eff7c9d4b7b174643f· 查看固定版本源码


    def __init__(self, args: Namespace, *, concurrency: int | None = None) -> None:
        # persistent state for the generation process
        self.args = args
        self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True)
        self.processor = load_processor(args.hf_checkpoint, trust_remote_code=True)

        self.semaphore = asyncio.Semaphore(
            concurrency if concurrency is not None else args.sglang_server_concurrency * get_rollout_num_engines(args)
        )
        self.generation_pacer = AsyncPacer()
        self.sampling_params: dict[str, Any] = dict(
            temperature=args.rollout_temperature,
            top_p=args.rollout_top_p,
展开 Slime 源码:generate_and_rm

generate_and_rm
slime/rollout/sglang_rollout.py · L325–L349
快照 2f2318653f6f794dddd321eff7c9d4b7b174643f· 查看固定版本源码

            }
        ]

    # generate
    async with state.semaphore:
        await state.generation_pacer.wait()
        if state.aborted:
            sample.status = Sample.Status.ABORTED
            return sample

        with state.dp_rank_context() as _:
            custom_func_path = sample.generate_function_path or args.custom_generate_function_path
            generate_func = load_function(custom_func_path) if custom_func_path is not None else generate

            if custom_func_path is not None and "evaluation" in inspect.signature(generate_func).parameters:
                generate_call = generate_func(args, sample, sampling_params, evaluation=evaluation)
            else:
                generate_call = generate_func(args, sample, sampling_params)

            if getattr(generate_func, "abort_mode", None) == "request":
                sample = await _run_request_abortable_generate(state, sample, generate_call)
            else:
                sample = await _run_server_abort_generate(state, generate_call)

    sample = inherit_queue_context(queue_input, sample)

generation_pacer 是另一种安排启动节奏的工具。它不等同于信号量,也不能单独保证在途请求数不超上限,第 13 节再解释。

两个常见误解

  • 限制许可数,不会阻止调用方提前创建很多正在等待的 Task;任务对象本身仍会占用内存。
  • 许可只约束 async with 内的区域。区域外的评分或结果处理,不会自动被同一个上限限制。
复述检查

把许可从 2 改成 1,实验的 peak 变成 1;任务总数依然是 5。许可数与工作总数是两回事。

语义核对:Semaphore。

本节对应的全部源码位置

下面把同类用法在 Slime 中的每个扫描位置接回本节讲解;点击文件名可查看完整源码快照。重复出现的 API 共用解释,具体输入与清理由所在函数决定。

展开本节 11 条源码记录
记录 / 讲解源码位置原代码
occ-000004examples/coding_agent_rl/generate.py:92
<module>
_BOOT_SEM = asyncio.Semaphore(CONFIG.boot_concurrency)
occ-000206examples/retool/tool_sandbox.py:38
<module>
SEMAPHORE = asyncio.Semaphore(TOOL_CONFIGS["tool_concurrency"])
occ-000223examples/search-r1/generate_with_search.py:41
<module>
SEMAPHORE = asyncio.Semaphore(SEARCH_R1_CONFIGS["search_concurrency"])
occ-000239examples/search-r1/google_search_server.py:42
<module>
async def fetch(session: aiohttp.ClientSession, url: str, semaphore: asyncio.Semaphore) -> str:
occ-000247examples/search-r1/google_search_server.py:64
fetch_all
semaphore = asyncio.Semaphore(limit)
occ-000837slime/data/transport.py:375
run_rollout_io
_async_limits[key] = asyncio.Semaphore(getattr(args, "rollout_io_concurrency", 4))
occ-001117slime/rollout/sglang_rollout.py:95
GenerateState.__init__
self.semaphore = asyncio.Semaphore(
occ-001998tests/test_score_centering_transport.py:191
test_evaluation_preserves_training_score_centering
semaphore=asyncio.Semaphore(1),
occ-002121tests/test_streaming_rollout.py:129
_generation_state
semaphore=asyncio.Semaphore(8),
occ-002184tests/test_streaming_rollout.py:507
test_partial_abort_buffers_and_resumes_only_aborted_siblings.exercise
state.semaphore = asyncio.Semaphore(2)
occ-002188tests/test_streaming_rollout.py:535
test_generation_pacing_preserves_concurrency_cancellation_and_reuse.exercise
state.semaphore = asyncio.Semaphore(128)

专题首页 · 上一节 · 下一节:07 · Queue:把生产速度和消费速度接起来 · 源码与教程对照 · 全部用法