让云上的模型服务更高效。

简介

Mosec 是一个高性能且灵活的模型服务（model serving）框架，用于构建支持机器学习模型的后端服务和微服务。它弥合了你刚训练好的任意机器学习模型与高效在线服务 API 之间的鸿沟。

• 高性能：Web 层和任务协调使用 Rust 🦀 构建，在异步 I/O（async I/O）驱动下，不仅提供极快的速度，还能高效利用 CPU 资源
• 易于使用：用户接口完全基于 Python 🐍，用户可以以与离线测试相同的代码，以不依赖具体机器学习框架（ML framework-agnostic）的方式部署模型
• 动态批处理（Dynamic batching）：聚合来自不同用户的请求进行批量推理，并将结果分发回各自用户
• 流水线阶段（Pipelined stages）：为流水线的不同阶段启动多个进程，以处理混合了 CPU/GPU/I/O 的工作负载
• 云原生友好（Cloud friendly）：专为云环境设计，内置模型预热（warmup）、优雅关闭（graceful shutdown）以及 Prometheus 监控指标，可轻松通过 Kubernetes 或其他容器编排系统管理
• 专注做好一件事：专注于在线服务部分，让用户能更专注于模型优化和业务逻辑

安装

Mosec 要求 Python 3.7 或更高版本。可通过以下命令在 Linux 或 macOS 上安装最新的 PyPI 包：

pip install -U mosec
# 或使用 conda 安装
conda install conda-forge::mosec
# 或使用 pixi 安装
pixi add mosec

若需从源码构建，请先安装 Rust，然后运行以下命令：

make package

你将在 dist 文件夹中获得一个 mosec 的 wheel 文件。

使用示例

我们演示如何使用 Mosec 轻松将一个预训练的 Stable Diffusion 模型部署为服务。你需要先安装前置依赖 diffusers 和 transformers：

pip install --upgrade diffusers[torch] transformers

编写服务端

from io import BytesIO
from typing import List

import torch  # type: ignore
from diffusers import StableDiffusionPipeline  # type: ignore

from mosec import Server, Worker, get_logger
from mosec.mixin import MsgpackMixin

logger = get_logger()

class StableDiffusion(MsgpackMixin, Worker):
    def __init__(self):
        self.pipe = StableDiffusionPipeline.from_pretrained(
            "sd-legacy/stable-diffusion-v1-5", torch_dtype=torch.float16
        )
        self.pipe.enable_model_cpu_offload()
        self.example = ["useless example prompt"] * 4  # 预热 (batch_size=4)

    def forward(self, data: List[str]) -> List[memoryview]:
        logger.debug("generate images for %s", data)
        res = self.pipe(data)
        logger.debug("NSFW: %s", res[1])
        images = []
        for img in res[0]:
            dummy_file = BytesIO()
            img.save(dummy_file, format="JPEG")
            images.append(dummy_file.getbuffer())
        return images

if __name__ == "__main__":
    server = Server()
    # 1) `num` 指定将启动多少个进程并行运行。
    # 2) 当 `max_batch_size` 设置为大于 1 的值时，`forward` 函数中的输入数据将是一个列表（批处理）；
    #    否则，输入为单个元素。
    server.append_worker(StableDiffusion, num=1, max_batch_size=4, max_wait_time=10)
    server.run()

运行服务端

python examples/stable_diffusion/server.py --help

python examples/stable_diffusion/server.py --log-level debug --timeout 30000

python examples/stable_diffusion/client.py --prompt "a cute cat playing with a red ball" --output cat.jpg --port 8000

curl http://127.0.0.1:8000/metrics

示例

更多开箱即用的示例可以在 示例 章节中找到，包括：

• Pipeline（流水线）：一个简单的 echo 演示，甚至不包含任何机器学习模型。
• 请求验证（Request validation）：使用类型注解验证请求并生成 OpenAPI 文档。
• 多路由（Multiple route）：在一个服务中部署多个模型。
• Embedding 服务：兼容 OpenAI 的 embedding 服务。
• 重排序服务（Reranking service）：根据查询对一组段落进行重排序。
• 共享内存 IPC（Shared memory IPC）：使用共享内存进行进程间通信。
• 自定义 GPU 分配（Customized GPU allocation）：部署多个副本，每个副本使用不同的 GPU。
• 自定义指标（Customized metrics）：记录您自己的监控指标。
• Jax 即时编译推理（Jax jitted inference）：通过即时编译（just-in-time compilation）加速推理。
• 压缩（Compression）：启用请求/响应压缩。
• PyTorch 深度学习模型：
  • 情感分析（sentiment analysis）：推断句子的情感倾向。
  • 图像识别（image recognition）：对给定图像进行分类。
  • Stable Diffusion：基于文本生成图像，并使用 msgpack 序列化。

配置

• 动态批处理（Dynamic batching）
  • max_batch_size 和 max_wait_time（毫秒） 在调用 append_worker 时配置。
  • 确保使用 max_batch_size 进行推理不会导致 GPU 内存溢出（out-of-memory）。
  • 通常，max_wait_time 应小于批处理推理所需时间。
  • 启用后，当累积请求数达到 max_batch_size 或等待时间达到 max_wait_time 时，将触发一次批处理。在高流量场景下，该功能可显著提升服务性能。
• 其他配置项请参考 参数文档。

部署

• 如果您需要一个已预装 mosec 的 GPU 基础镜像，可以使用官方镜像 mosecorg/mosec。对于更复杂的使用场景，请参考 envd。
• 本服务无需 Gunicorn 或 NGINX，但在必要时仍可配合 Ingress 控制器使用。
• 本服务应作为容器中的 PID 1 进程运行，因为它会管理多个子进程。如果您需要在单个容器中运行多个进程，则需使用进程管理器，例如 Supervisor 或 Horust。
• 请务必收集 指标（metrics）：
  • mosec_service_batch_size_bucket：显示批处理大小的分布。
  • mosec_service_batch_duration_second_bucket：显示每个连接在各阶段的动态批处理耗时（从接收到第一个任务开始计算）。
  • mosec_service_process_duration_second_bucket：显示每个连接在各阶段的处理耗时（包含 IPC 时间，但不包含 mosec_service_batch_duration_second_bucket 所涵盖的时间）。
  • mosec_service_remaining_task：显示当前正在处理的任务数量。
  • mosec_service_throughput：显示服务吞吐量。
• 请通过 SIGINT（CTRL+C）或 SIGTERM（kill {PID}）信号停止服务，因为服务内置了优雅关闭（graceful shutdown）逻辑。

性能调优

• 为您的推理服务找出最佳的 max_batch_size 和 max_wait_time。指标会显示实际批处理大小和批处理耗时的直方图，这些是调整这两个参数的关键依据。
• 尝试将整个推理过程拆分为独立的 CPU 和 GPU 阶段（参考 DistilBERT）。不同阶段将以数据流水线（data pipeline） 方式运行，从而保持 GPU 持续忙碌。
• 您还可以调整每个阶段的工作进程（worker）数量。例如，如果您的流水线包含一个用于预处理的 CPU 阶段和一个用于模型推理的 GPU 阶段，增加 CPU 阶段的 worker 数量有助于为 GPU 阶段生成更多可批处理的数据；增加 GPU 阶段的 worker 数量则能更充分地利用 GPU 内存和计算能力。这两种方式都可能提高 GPU 利用率，从而提升服务吞吐量。
• 对于多阶段服务，请注意：数据在不同阶段之间传递时会通过 serialize_ipc/deserialize_ipc 方法进行序列化/反序列化，因此过大的数据可能导致整个流水线变慢。默认情况下，序列化后的数据通过 Rust 传递到下一阶段，您可以启用共享内存以潜在降低延迟（参考 RedisShmIPCMixin）。
• 您应选择合适的 serialize/deserialize 方法，用于解析用户请求和编码响应。默认两者均使用 JSON，但 JSON 对图像和 embedding 等数据支持不佳。您可以选择更快且二进制兼容的 msgpack（参考 Stable Diffusion）。
• 配置 OpenBLAS 或 MKL 的线程数。它们可能无法自动选择当前 Python 进程所使用的最合适 CPU 核心。您可以通过 env 为每个 worker 单独配置（参考 自定义 GPU 分配）。
• 从客户端启用 HTTP/2。自 v0.8.8 起，mosec 会自动适配用户使用的协议（如 HTTP/2）。

用户

以下是一些正在使用 Mosec 的公司和个人用户：

• Modelz：面向机器学习推理的 Serverless 平台。
• MOSS：一个开源的类 ChatGPT 对话语言模型。
• 腾讯云（TencentCloud）：腾讯云机器学习平台，使用 Mosec 作为核心推理服务器框架。
• TensorChord：云原生 AI 基础设施公司。
• OAT：用于在线大语言模型（LLM）对齐的奖励模型服务。

引用（Citation）

如果您在研究中发现本软件有用，请考虑引用以下内容：

@software{yang2021mosec,
  title = {{MOSEC: Model Serving made Efficient in the Cloud}},
  author = {Yang, Keming and Liu, Zichen and Cheng, Philip},
  howpublished = {https://github.com/mosecorg/mosec},
  year = {2021}
}

贡献（Contributing）

我们欢迎任何形式的贡献。您可以通过 提交 Issue 或在 Discord 上与我们讨论来提供反馈。您也可以直接贡献代码并提交 Pull Request！

开始开发时，您可以使用 envd 创建一个隔离且干净的 Python 和 Rust 环境。更多信息请参阅 envd 文档 或 build.envd 文件。

mosec 快速上手指南

环境准备

• 操作系统：Linux 或 macOS
• Python 版本：3.7 或更高
• 可选依赖（如运行 Stable Diffusion 示例）：

  pip install --upgrade diffusers[torch] transformers
  

💡 提示：国内用户建议配置 PyPI 镜像源（如清华源）以加速安装：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -U mosec

安装步骤

通过 PyPI 安装（推荐）

pip install -U mosec

通过 conda 安装

conda install conda-forge::mosec

通过 pixi 安装

pixi add mosec

如需从源码构建，请先安装 Rust，然后执行：

make package

构建好的 wheel 文件将位于 dist/ 目录下。

基本使用

以下是一个使用 mosec 部署 Stable Diffusion 模型的最简示例：

1. 编写服务代码（server.py）

from io import BytesIO
from typing import List

import torch
from diffusers import StableDiffusionPipeline

from mosec import Server, Worker, get_logger
from mosec.mixin import MsgpackMixin

logger = get_logger()

class StableDiffusion(MsgpackMixin, Worker):
    def __init__(self):
        self.pipe = StableDiffusionPipeline.from_pretrained(
            "sd-legacy/stable-diffusion-v1-5", torch_dtype=torch.float16
        )
        self.pipe.enable_model_cpu_offload()
        self.example = ["useless example prompt"] * 4  # 模型预热

    def forward(self, data: List[str]) -> List[memoryview]:
        logger.debug("generate images for %s", data)
        res = self.pipe(data)
        images = []
        for img in res[0]:
            dummy_file = BytesIO()
            img.save(dummy_file, format="JPEG")
            images.append(dummy_file.getbuffer())
        return images

if __name__ == "__main__":
    server = Server()
    server.append_worker(StableDiffusion, num=1, max_batch_size=4, max_wait_time=10)
    server.run()

2. 启动服务

python server.py --log-level debug --timeout 30000

服务启动后，可通过以下方式访问：

• API 文档：访问 http://127.0.0.1:8000/openapi/swagger/
• 指标监控：curl http://127.0.0.1:8000/metrics

3. 调用示例（客户端）

python examples/stable_diffusion/client.py --prompt "a cute cat playing with a red ball" --output cat.jpg --port 8000

执行后将在当前目录生成 cat.jpg 图像文件。

更多完整示例请参考官方文档：mosec Examples