已经会创建任务,为什么还需要限流
创建 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_rmslime/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-000004 | examples/coding_agent_rl/generate.py:92 <module> | _BOOT_SEM = asyncio.Semaphore(CONFIG.boot_concurrency) |
| occ-000206 | examples/retool/tool_sandbox.py:38 <module> | SEMAPHORE = asyncio.Semaphore(TOOL_CONFIGS["tool_concurrency"]) |
| occ-000223 | examples/search-r1/generate_with_search.py:41 <module> | SEMAPHORE = asyncio.Semaphore(SEARCH_R1_CONFIGS["search_concurrency"]) |
| occ-000239 | examples/search-r1/google_search_server.py:42 <module> | async def fetch(session: aiohttp.ClientSession, url: str, semaphore: asyncio.Semaphore) -> str: |
| occ-000247 | examples/search-r1/google_search_server.py:64 fetch_all | semaphore = asyncio.Semaphore(limit) |
| occ-000837 | slime/data/transport.py:375 run_rollout_io | _async_limits[key] = asyncio.Semaphore(getattr(args, "rollout_io_concurrency", 4)) |
| occ-001117 | slime/rollout/sglang_rollout.py:95 GenerateState.__init__ | self.semaphore = asyncio.Semaphore( |
| occ-001998 | tests/test_score_centering_transport.py:191 test_evaluation_preserves_training_score_centering | semaphore=asyncio.Semaphore(1), |
| occ-002121 | tests/test_streaming_rollout.py:129 _generation_state | semaphore=asyncio.Semaphore(8), |
| occ-002184 | tests/test_streaming_rollout.py:507 test_partial_abort_buffers_and_resumes_only_aborted_siblings.exercise | state.semaphore = asyncio.Semaphore(2) |
| occ-002188 | tests/test_streaming_rollout.py:535 test_generation_pacing_preserves_concurrency_cancellation_and_reuse.exercise | state.semaphore = asyncio.Semaphore(128) |
专题首页 · 上一节 · 下一节:07 · Queue:把生产速度和消费速度接起来 · 源码与教程对照 · 全部用法