构建高效可靠的分布式推理系统：深入解析控制器与模型服务的协同工作

小爷毛毛_卓寿杰
159 阅读6分钟

重磅推荐专栏： 《大模型AIGC》 《课程大纲》 《知识星球》

本专栏致力于探索和讨论当今最前沿的技术趋势和应用领域，包括但不限于ChatGPT和Stable Diffusion等。我们将深入研究大型模型的开发和应用，以及与之相关的人工智能生成内容（AIGC）技术。通过深入的技术解析和实践经验分享，旨在帮助读者更好地理解和应用这些领域的最新进展

在现代互联网应用中，随着用户需求的增长和技术的进步，单一服务器已经难以满足大规模并发请求的需求。为了提升系统的性能和可靠性，开发者们越来越多地采用分布式架构。本文将结合具体的代码示例，深入浅出地探讨如何构建一个高效的分布式推理系统，并详细解析其中的关键组件——控制器（Controller）和服务于具体模型处理的Worker节点。

一、引言

分布式推理系统的核心挑战在于如何合理地分配任务给多个工作节点（workers），以确保系统的负载均衡和高效运行。我们的系统通过引入控制器来协调这些工作节点，并根据不同的流量分发策略动态调整任务分配。接下来，我们将详细介绍这个过程的工作原理及其背后的实现细节。

二、控制器（Controller）

功能概述

控制器作为整个系统的中枢神经，负责管理和调度多个分布式的模型服务节点（workers）。它不仅需要为客户端提供可用的工作节点地址，还要确保任务能够被高效地分配到最合适的节点上。此外，控制器还实现了多种流量分发策略，以适应不同应用场景下的需求。

核心组件
  • Controller类：
    • 管理所有注册的工作节点信息，并根据配置选择相应的调度策略。
    • 定期检查工作节点的心跳信号，移除失效或超时的节点。
class Controller:
    def __init__(self, dispatch_method: str):
        self.worker_info = {}
        self.dispatch_method = DispatchMethod.from_str(dispatch_method)
        self.heart_beat_thread = threading.Thread(
            target=heart_beat_controller, args=(self,)
        )
        self.heart_beat_thread.start()
  • 心跳机制：
    • 使用单独的线程定期执行心跳检查，保证系统能够及时响应工作节点的状态变化。
def heart_beat_controller(controller):
    while True:
        time.sleep(CONTROLLER_HEART_BEAT_EXPIRATION)
        controller.remove_stale_workers_by_expiration()
  • API接口：
    • 提供一系列RESTful API用于外部系统与控制器交互，如注册/刷新工作节点、获取模型列表、查询工作节点状态等。
@app.post("/receive_heart_beat")
async def receive_heart_beat(request: Request):
    data = await request.json()
    exist = controller.receive_heart_beat(data)
    return {"exist": exist}
流量分发策略

控制器支持三种不同的流量分发策略：

  • 抽签法（LOTTERY）：基于概率选择工作节点。每个工作节点被选中的概率与其速度成正比，从而优化整体效率。
if self.dispatch_method == DispatchMethod.LOTTERY:
    worker_speeds = np.array(worker_speeds, dtype=np.float32)
    norm = np.sum(worker_speeds)
    if norm < 1e-4:
        return ""
    worker_speeds = worker_speeds / norm
    pt = np.random.choice(np.arange(len(worker_names)), p=worker_speeds)
    worker_name = worker_names[pt]
    return worker_name
  • 最短队列法（SHORTEST_QUEUE）：选择当前队列长度最小的工作节点来执行新任务，有助于减少等待时间。
elif self.dispatch_method == DispatchMethod.SHORTEST_QUEUE:
    min_index = np.argmin(worker_qlen)
    w_name = worker_names[min_index]
    min_qlen = worker_qlen[min_index]
    self.worker_info[w_name].queue_length += 1
    logger.info(f"names: {worker_names}, queue_lens: {worker_qlen}, ret: {w_name}, len:{min_qlen}")
    return w_name
  • 最大剩余容量法（MAX_REMAIN）：选择剩余容量最大的工作节点，即还有最多空闲资源可以用来处理新任务的工作节点。
elif self.dispatch_method == DispatchMethod.MAX_REMAIN:
    max_index = np.argmax(worker_rlen)
    max_rlen = worker_rlen[max_index]
    if max_rlen <= 0:
        return REQUEST_FULL_FLAG
    w_name = worker_names[max_index]
    self.worker_info[w_name].queue_length += 1
    logger.info(f"names: {worker_names}, queue remain lens: {worker_rlen}, ret: {w_name},max remain len:{max_rlen}")
    return w_name

三、模型服务（Model Worker）

功能概述

每个模型服务实例运行特定的机器学习模型，处理来自客户端的推理请求并返回结果。它们与控制器协作，通过HTTP请求向控制器报告自身状态，接收任务分配指令。

核心组件
  • VLLMWorker类：
    • 扩展自基类BaseModelWorker，实现了具体的模型推理逻辑。
    • 包含异步生成流式输出的功能，允许实时推送计算进度给客户端。
class VLLMWorker(BaseModelWorker):
    async def generate_stream(self, params):
        # 实现具体的推理逻辑...
        async for request_output in results_generator:
            # 处理推理结果并返回JSON格式的数据
            yield (json.dumps(ret) + "\0").encode()
  • 异步推理引擎：
    • 使用vllm.engine.async_llm_engine.AsyncLLMEngine作为底层推理引擎，支持高效的并发处理。
engine_args = AsyncEngineArgs.from_cli_args(args)
engine = AsyncLLMEngine.from_engine_args(engine_args)
  • API接口：
    • 提供了一系列API端点，例如生成文本流、获取模型详情、统计token数量等。
@app.post("/worker_generate_stream")
async def api_generate_stream(request: Request):
    params = await request.json()
    await acquire_worker_semaphore()
    request_id = random_uuid()
    params["request_id"] = request_id
    generator = worker.generate_stream(params)
    background_tasks = create_background_tasks(request_id)
    return StreamingResponse(generator, background=background_tasks)

四、系统整合与扩展

系统整合
  • 控制器与模型服务之间的通信：两者之间通过HTTP协议进行交互，控制器负责协调任务分配和状态监控，而模型服务专注于具体的模型推理任务。
  • 负载均衡与容错性：通过合理的流量分发策略以及心跳检测机制，确保即使部分工作节点出现故障，系统仍能保持良好的性能和服务质量。
扩展能力
  • 多模型支持：控制器可以管理多个不同类型的模型服务，只需将新模型的服务地址注册到控制器即可。
  • 灵活配置：用户可以根据实际需求调整调度策略、设置SSL加密传输、限制并发连接数等参数。
  • 易于维护：清晰的模块划分使得系统的各个部分相对独立，便于后期的功能扩展和技术升级。

五、实践应用

为了更好地理解上述理论知识的应用场景，我们可以考虑以下几个方面：

  • 大规模在线推理平台：在一个拥有众多用户的在线推理平台上，控制器可以帮助有效地分配计算资源，提高用户体验。
  • 微服务架构下的模型部署：对于采用微服务架构的应用来说，这种设计模式非常适合部署复杂的AI模型，因为它允许按需扩展工作节点的数量，同时保持系统的稳定性和可维护性。

六、结论

通过对控制器和模型服务的深入分析，我们了解到如何构建一个高效可靠的分布式推理系统。通过合理的流量分发策略、心跳检测机制以及灵活的任务分配方式，该系统不仅具备强大的计算能力和良好的扩展性，还能确保高可用性和稳定性。希望这篇文章能够帮助读者深入了解这一领域的技术细节，并启发更多关于如何构建高效可靠分布式系统的思考。

参考文献

希望通过这篇博客文章，你对分布式推理系统的构建有了更深的理解。如果有任何问题或需要进一步的帮助，请不要犹豫，随时联系我！

avatar
文章
阅读
粉丝