mosecorg/mosec

Mosec – 高性能モデルサービングフレームワーク

概要 – Mosec (Model Serving made Efficient in the Cloud) は、Pythonファーストのライブラリであり、学習済みのMLモデルを高速で本番環境対応のHTTPサービスに変換できます。ネットワークI/O、リクエストのバッチ処理、プロセスオーケストレーションといった負荷の高い処理はRustで記述されているため、使い慣れたPythonコードを維持しながら、低レイテンシかつ高スループットな推論を実現します。

主な機能

• 動的バッチ処理 (Dynamic batching) – 受信したリクエストを自動的にバッチ化（サイズと待機時間を設定可能）し、GPUの稼働率を最大化します。
• パイプライン型ワーカー (Pipeline-style workers) – 複数のプロセス（CPUのみ、GPUのみ、I/O）を連結し、前処理、推論、後処理を別々のステージに分割できます。
• Rustベースのサーバー – 非同期I/Oと効率的なCPU使用率に加え、Prometheusメトリクスと正常終了処理を内蔵しています。
• 柔軟なシリアライズ – デフォルトのJSONに加え、画像や埋め込みデータなどのバイナリデータ用に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ペイロードを http://127.0.0.1:8000 にPOSTします。サービスはMsgpack（バイナリフレンドリー）を介して生成されたJPEG画像を返します。Swagger UIは /openapi/swagger/ で利用可能です。

エコシステムと統合

• HuggingFace diffusers、transformers、PyTorch、JAXなどと連携。
• カスタムシリアライズ、共有メモリIPC、リクエストバリデーション、メトリクス拡張用のミックスインを提供。
• 公式DockerイメージとKubernetesデプロイ用のサンプルを提供。
• Modelz (serverless ML)、Tencent Cloud ML Platform、MOSS会話モデルなどのプロジェクトで使用されています。

ヘルプの取得

• Discord: https://discord.gg/Jq5vxuH69W
• GitHub Issues: バグや機能リクエストはチケットを作成してください。
• Documentation: https://mosecorg.github.io/mosec/

引用 研究でMosecを使用する場合は、提供されているソフトウェアリファレンスを引用してください。