mosecorg/mosec

Mosec – 高性能模型服务框架

简介 – Mosec (Model Serving made Efficient in the Cloud) 是一个以 Python 为主的库，让您能将任何训练好的机器学习模型转变为快速且可投入生产的 HTTP 服务。繁重的工作（网络 I/O、请求批处理、进程编排）皆由 Rust 编写，让您在保持熟悉的 Python 代码的同时，获得低延迟、高吞吐量的推理性能。

主要功能

• 动态批处理 (Dynamic batching) – 自动将传入的请求分组为批次（可设置大小与等待时间），以确保 GPU 保持高负载运行。
• 流水线式工作单元 (Pipeline-style workers) – 串联多个进程（仅 CPU、仅 GPU、I/O），让您能将预处理、推理与后处理拆分至不同阶段。
• Rust 驱动的服务器 – 具备异步 I/O 与高效的 CPU 使用率，并内置 Prometheus 指标与优雅关闭功能。
• 灵活的序列化 – 默认使用 JSON，或通过 mixins 支持 Msgpack 等二进制数据格式（如图像或嵌入向量）。
• 云端就绪 – 可直接与容器、Kubernetes 以及官方 Docker 镜像 mosecorg/mosec 搭配使用。
• OpenAPI 生成 – 为任何 Worker 实现自动生成 Swagger 文档。

典型应用场景

• 将大型生成式模型（如 Stable Diffusion）部署为 REST 端点。
• 提供嵌入向量、重排序或其他轻量级推理服务。
• 建立多阶段流水线，将预处理运行在 CPU，模型推理运行在 GPU。
• 为自定义模型提供 OpenAI 兼容的 API。

安装

# From PyPI (Linux/macOS)
pip install -U mosec
# Or via Conda
conda install -c conda-forge mosec
# Or Pixi
pixi add mosec

若要从源码编译，需安装 Rust，然后执行 make package 以生成 wheel 文件。

快速示例 – 部署 Stable Diffusion

from mosec import Server, Worker, get_logger
from mosec.mixin import MsgpackMixin
from diffusers import StableDiffusionPipeline
import torch, io

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 = ["warmup prompt"] * 4   # warm‑up batch

    def forward(self, prompts: list[str]) -> list[memoryview]:
        results = self.pipe(prompts)
        images = []
        for img in results[0]:
            buf = io.BytesIO()
            img.save(buf, format="JPEG")
            images.append(buf.getbuffer())
        return images

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

执行脚本后，将包含 prompt 字段的 JSON 负载 POST 至 http://127.0.0.1:8000。服务将通过 Msgpack（对二进制友好）返回生成的 JPEG 图像。Swagger UI 可于 /openapi/swagger/ 访问。

生态系统与整合

• 支持 HuggingFace diffusers、transformers、PyTorch、JAX 等。
• 提供用于自定义序列化、共享内存 IPC、请求验证与指标扩展的 mixins。
• 提供官方 Docker 镜像与 Kubernetes 部署示例。
• 已被 Modelz (serverless ML)、腾讯云 ML 平台以及 MOSS 对话模型等项目采用。

获取帮助

• Discord: https://discord.gg/Jq5vxuH69W
• GitHub Issues: 若有错误或功能需求，请开启 issue。
• Documentation: https://mosecorg.github.io/mosec/

引用 若您在研究中使用 Mosec，请引用提供的软件参考。