跳转至

utils.py — Rollout 通用工具

文件概述

提供推理引擎的通用工具函数,包括模型长度获取、HTTP 服务器启动和异步迭代器适配。

核心函数

get_max_position_embeddings

def get_max_position_embeddings(hf_config):
    """从 HuggingFace 配置中获取最大位置嵌入长度

    不同模型使用不同的属性名:
    - max_position_embeddings
    - max_seq_len
    - max_sequence_length
    """

run_uvicorn

async def run_uvicorn(app, args, host):
    """启动 Uvicorn HTTP 服务器

    自动查找可用端口,返回 (port, task)
    用于 vLLM/SGLang/TRT-LLM 的 HTTP 推理服务
    """

内部使用 _UvicornServerAutoPort 自动端口发现:

class _UvicornServerAutoPort:
    """自动寻找可用端口的 Uvicorn 服务器"""
    def __init__(self, app, host, port_start=8000, max_retries=100):
        ...

ensure_async_iterator

async def ensure_async_iterator(iterable):
    """将同步生成器转换为异步迭代器

    用于权重传输场景: 训练引擎产生同步的权重生成器,
    但传输层需要异步消费。
    """
    if hasattr(iterable, '__aiter__'):
        async for item in iterable:
            yield item
    else:
        for item in iterable:
            yield item

与其他模块的关系

  • 被所有推理引擎服务器(vLLM/SGLang/TRT-LLM)使用
  • run_uvicorn 是 HTTP 推理服务的启动入口

小结

这些工具函数为推理引擎的 HTTP 服务提供了基础设施支持。