ml-cellar: Minimal MLOpsを回すためのモデル管理CLIを公開した

Summary

「Minimal MLOpsを回すためのモデル管理CLI」として、ml-cellar をRustで作成しOSSにした
Minimal MLOps Checklist を支援するソフトウェアのうち、モデル管理とモデルServingを担当している

–

背景

MLモデルの運用には、学習コードそのもの以上に「学習済モデルなどの成果物の管理」が効いてきます。どの実験の重みがどこにあるのか、どうやって再現するのか、誰がどのモデルを使うべきなのか。どのモデルを元にファインチューニングを行ったのか。このMLモデルの運用に対して、“MLOps (Machine Learning Operations)” という概念が生まれました。機械学習モデルを「作る」だけでなく、「安全に・安定して・継続的に運用して価値を出し続ける」ための仕組み・文化・技術の総称です。 DevOps（開発と運用の統合）を、機械学習特有の課題（データ、再現性、劣化、監視など）に合わせて拡張したものです。

一方でこの"MLOps"を構築するには沢山のMLOpsエンジニアが必要になり、多大なるコストが発生します。 MLOpsはフルスタックに寄せるほど「正しい」方向に進みます。実験管理、メタデータ、モデルレジストリ、アーティファクトストア、サービング、監視、などなど。ただしその全体像はコストと運用負荷を強く要求します。 BigTechでは巨大なMLOps基盤が整備されている（と思っている）のですが、すべての会社、特に小さなスタートアップや研究チームなどではMLOps基盤を持てるほどのコストを支払えないことが多いです。

“現実的"には、

研究寄りで、実験結果（重みや設定）さえ共有できれば十分
受託/プロダクトで、複数案件に対して fine-tune したモデルを配布したい
まずは「モデルの置き場」と「最低限のルール」だけ欲しい

といったMLOpsへの「ライトな需要」も多いと考えていました。

その「ライトな需要」のチェックリストとして Minimal MLOps Checklist を書いたのですが、今回はMinimal MLOpsを支援するようなソフトウェアを作ってみました。 MLOpsの文脈では全自動化を目指すものが多いですが、コストを下げるため ml-cellarではCLIの手動操作にすることで、最低限の運用がすぐに回ることを重視しました。

既存のサービスとしては、wandb (Weights & Biases) などがあります。 wandbでもArtifactsの管理は可能ではあるのですが、ml-cellarは「モデルを保管して、いつでも使える状態にしたい」というものを重要視しています。「大量の実験に埋もれたモデルを探す」のではなく、「結果の良かったモデルだけを保存して誰でも使える状態にする」方針にしています。使い分けとして実験管理はwandbで雑にアップロードしていって個人責任で管理、モデルの保存と提供はml-cellarでチーム責任で丁寧に管理していく、というのを想定しています。またml-cellarは学習コードに手を加える必要が無いのもメリットで、wandbに比べてすぐにモデル管理を始めることができます。

–

`ml-cellar` について

自分が作った ml-cellar ではMLflowのようなサービスを提供するのではなく、Git LFS を土台にした最小限のモデルレジストリを簡単に扱えるようにするRust製CLIとして提供することにしました。 MLflowのような"全部入り"ではなく、本当に必要な機能に絞った “妥協のMLOps” を目指しています。 Git LFSを用いることで構成を把握するための学習コストが低く、技術負債になりにくい構成を目指しました。また基本的に gitコマンドと git lfsコマンドをラッピングしているため、ml-cellar 無しでも運用は可能です（ml-cellarを使ったら楽になるというだけ）。そのため、ベンダーロックインになりづらい構成にでき、運用もロバストにできると考えています。

ml-cellar においては、モデル成果物を「ワインセラー」のように rack（棚）とbin（瓶）で管理します。以下がrackの例です。

model_registry_repository/ # repository（モデルレジストリと呼称）
  vit-l/ # rack （モデル成果物の種類、ワインで言うところの品種に相当）
    0.1/ # bin （バージョン付きのモデル成果物の集合、ワインでいうところのビンテージ）
      # configファイル、重みファイル、ログなど
      config.yaml
      checkpoint.pth
      result.json
      log.log
    0.2/
      ...
  vit-m/
    ...

このときモデルを「vit-l/0.2」という呼び方で一意に定めることができ、配布や参照が楽になります。配布においても、

ml-cellar clone {your_ml_registry}
ml-cellar download {path to a directory or a file}

でファイルが手に入ります。

またGitLFSのリモートリポジトリとしてGitHubを用いることで、モデルのアップロードに対してPull requestを用いた開発を行うことができます。多くのソフトウェアエンジニアにとって慣れたインターフェイスで「開発」「レビュー」「承認」のフローを行うことができ、エンジニア側の作業コスト、心理的コストを下げることができます。権限の管理もgithubアカウントに結びつけるだけになります。

–

使い方について

詳しくは README 参照
GitHub LFSはストレージ/帯域に上限があるため、超えると追加課金や制限が起きることに注意してください
- Git LFSの上限超過により発生した費用は利用者の責任でお願いします。

リポジトリの作成

Install GitLFS

curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
sudo apt-get update
sudo apt-get install git-lfs
git lfs install

Install Rust

curl https://sh.rustup.rs -sSf | sh

Install ml-cellar.

cargo install ml-cellar

Model registryを作成

mkdir {your_ml_registry}
cd {your_ml_registry}
ml-cellar init

リモートリポジトリ登録

git remote add origin {remote_repository}

.gitattributesでLFS対象を決める

# --- Log ---
*.log   filter=lfs diff=lfs merge=lfs -text

# --- Your data ---
*.db   filter=lfs diff=lfs merge=lfs -text

rack（モデルごとの「棚」）を作成

ml-cellar rack {my_algorithm}

rackの構成

{your_ml_registry}/
  {my_algorithm}/
    README.md
    config.toml
    template.md

config.tomlで「必須成果物」を定義
- required_files：揃っていないといけないファイル
- optional_files：アップロードしても良いファイル（"*” で何でも可）

[artifact]
required_files = ["config.yaml", "checkpoint.pth"]
optional_files = ["*"]

README.md を作成
- モデルの用途、学習データ、注意点、推奨ユースケース、相性の良いタスクなどが書かれていると、後から探す人にとって助かる情報となります

成果物のコミット

モデルをコミットし、PRを出す

# branch作成
git sw -c feat/my_algorithm/0.1

# 成果物が揃っているかチェック
ml-cellar check {my_algorithm}/0.1/

# Commit
git add {path}
git commit -m "{commit message}"

# Remote repository に pushする
ml-cellar push origin HEAD

成果物の配布

clone

ml-cellar clone {your_ml_registry}

clone直後はGit LFSのファイルの実体化がされていない状態なので、実態化をします
- 「必要なものだけ実体化できる」運用は、モデルが増えてきても余計なものを実体化せず、容量を削減することができます

ml-cellar download {path to a directory or a file}

リポジトリの整備

GitHub Git LFS を用いる場合、GitHubの機能をそのまま用いることができます。便利なものとして以下の機能が挙げられます。

Pull Request

Git LFS のリモートリポジトリとして GitHub を利用することで、モデルのアップロードを Pull Request ベースのワークフローで管理できます。多くのソフトウェアエンジニアにとって馴染みのあるインターフェイス上で「開発 → レビュー → 承認」の流れを実現できるため、エンジニアの実作業コストだけでなく心理的な負担も減らせます。

CODEOWNER

また、ラック（rack）/モデル名ごとに CODEOWNERS を設定して、各ラックを保護することも可能です。以下は .github/CODEOWNERS の例です

/vit/           @cv-team
/llm-model/     @llm-team
/**/project_A/  @project_A

プロジェクト用のモデルはプロジェクトを運用している人たちのレビューが入り、基盤モデルはR&Dチームのレビューが入る、といったことが制御できます。

またブランチ保護 / Rulesets と組み合わせることで、さらに堅牢な運用にできます。

PR を必須にする
必要な承認数（N 人）を設定する
CODEOWNERS の承認を必須にする

などで、気づかれないままモデルが追加されたり、差し替えられたりすることを防げます。

他にもありますが、詳しくはドキュメントを参照してください

Document整備

実務上全て手作業で評価結果を書くのはコストがかかります。またスプレッドシートなどにまとめていくのも、スプレッドシート上の結果がどのモデルなのかの紐付けが大変です。学習側のコードで、READMEに書きたい情報をjsonにいれておくことで、ml-cellar ではtemplateファイルから自動で評価結果を作成する半手動化がサポートされています。

テンプレートを以下のように作って ml-cellar docs test_registry/test_docs/base/0.1/ のようにコマンドをうつと

### {{version}}

- Write the summary for the model
- Why you updated the model
- What dataset you added
- What you fixed in the algorithm
- What you updated in training parameters

<details>
<summary> Evaluation results </summary>

- Performance
  - Device: {performance.device}
  - Inference time(ms): {performance.time_ms}
  - Training time: {performance.gpu_num} GPU * {performance.batch_size} Batchsize * {performance.training_days} days
- Dataset
  - name: {dataset.name}
  - Version: {dataset.version}
  - Frame number: {dataset.frame_number}
- mAP for test dataset

| Metric         | Value              |
| -------------- | ------------------ |
| mAP (0.5:0.95) | {ap.map.50_95.all} |
| mAP (0.5)      | {ap.map.50.all}    |

- Per-Class AP for test dataset

| Class   | Samples                                  | Accuracy               |
| ------- | ---------------------------------------- | ---------------------- |
| person  | {test_dataset.per_class_samples.person}  | {ap.per_class.person}  |
| bicycle | {test_dataset.per_class_samples.bicycle} | {ap.per_class.bicycle} |
| car     | {test_dataset.per_class_samples.car}     | {ap.per_class.car}     |
| dog     | {test_dataset.per_class_samples.dog}     | {ap.per_class.dog}     |

</details>

以下のようなドキュメントを標準出力に流します。これをREADMEに貼り付けることで低コストでドキュメントを作成することができます

### test_docs/base/0.1

- Write the summary for the model
- Why you updated the model
- What dataset you added
- What you fixed in the algorithm
- What you updated in training parameters

<details>
<summary> Evaluation results </summary>

- Performance
  - Device: RTX4090
  - Inference time(ms): 7.87
  - Training time: 2 GPU * 32 Batchsize * 4.7 days
- Dataset
  - name: myCOCO
  - Version: 2.1.1
  - Frame number: 21321
- mAP for test dataset

| Metric         | Value      |
| -------------- | ---------- |
| mAP (0.5:0.95) | 0.40453625 |
| mAP (0.5)      | 0.598      |

- Per-Class AP for test dataset

| Class   | Samples | Accuracy |
| ------- | ------- | -------- |
| person  | 647     | 0.5517   |
| bicycle | 916     | 0.2899   |
| car     | 663     | 0.3941   |
| dog     | 988     | 0.586    |

詳しくはtemplate.mdを用いたドキュメント化を参照してください

projectの運用

実務で

案件A向けに fine-tune したモデル
案件B向けに別のデータで fine-tune したモデル

を分けたくなることがあるかと思います。 ml-cellar は project-based model registry の考え方も用意し、案件単位の管理に対応できます。詳しくはドキュメントを参照してください

–

今後について

現状の制約として、GitHub LFSは無限でないことが挙げられます。

ストレージ上限
帯域（ダウンロード）上限
超過時の追加課金 or 制限

ml-cellar では Custom Transfer AgentでS3等に逃がす設計にする予定です（まだ未実装です）

現状の運用としては、

LFS対象を本当に必要な成果物に絞る
チェックポイントを “bestだけ” にする（全epoch保存しない）

あたりが現実的な対策になります。