9月17日、Sergio De Simoneが「Microsoft Open-Sources TauGrid to Simplify AI Workload Management on Kubernetes」と題した記事を公開した。AIワークロードをKubernetes上で動かす際に避けられない「接着剤」問題——サブミットスクリプト、キューのラッパー、ヘルスチェック、結果取得処理を個別に組み合わせ続ける運用上の苦労——をシングルHelmインストールで一掃するOSSプラットフォーム「TauGrid」をMicrosoftが公開した。
Kubernetesでの「接着剤」問題を解消する
AIワークロードをKubernetes上で動かそうとすると、プラットフォームチームは複数のOSSプロジェクト、カスタムスクリプト、運用ツールを個別に組み合わせて管理し続けなければならない。サブミットスクリプト、キューのラッパー、ヘルスチェック、結果の取得処理——いわゆる「接着剤」の部分だ。
Microsoftはこのペインポイントへの回答として、TauGridをオープンソースで公開した。Kubernetes上でAIワークロードの管理・スケジューリング・モニタリングを一元化するクラウドネイティブのプラットフォームである。
TauGridの肝は「シングルHelmインストール」にある。従来は個別に構築・維持が必要だったコンポーネント群と、それらの統合部分を、一括でデプロイできる。内部的には以下のコンポーネントで構成される。
研究者はKubernetesの知識なしにtauCLIでジョブを投入でき、プラットフォームチームはワークスペース、キュー、コンピュートプロファイル、ストレージ、ID管理、オブザーバビリティといった高度な機能を活用できる設計になっている。
tau.yamlでジョブを定義する／カバー範囲と現時点での制約
ワークロードの定義はYAMLファイルで行い、tau runコマンドで投入する。コマンドは設定を検証し、Kubernetes JobまたはKubeRay RayJobを生成。Kueueがクォータと優先度に基づいてキューに積む。
以下は、A100 GPU 1枚で動かすPyTorchトレーニングジョブの設定例だ。
schema_version: 1
name: aks-gpu-quickstart
run:
entrypoint: train.py
workload_kind: rayjob
compute:
gpus: 1
workers: 1
cpus: 16
memory: 64Gi
runtime:
image: mcr.microsoft.com/aks/ai-runtime/ray:py3.12-ray2.56.0-cuda13.0
pip:
- torch>=2.4.0
※ ray2.56.0-cuda13.0はコンテナイメージタグとして元記事に掲載されている値をそのまま転記しているが、Rayの現行安定版系列（2.x台）と比較してメジャーバージョン相当の乖離があるため、実際に使用する際はKubeRay公式ドキュメントおよびMCRのイメージカタログで最新の正確なタグを確認することを推奨する。
実行中はワークロードのステータス、ログ、チェックポイントを追跡する。実験の再現性確保と障害診断のために、実験エビデンスを収集・保存する機能も持つ。ジョブが失敗した場合は、チェックポイントから再開できる。
TauGridがカバーするのは、データ準備から分散トレーニング、ファインチューニング、推論までのエンドツーエンドのAIワークロード管理だ。
動作要件はKubernetes 1.30以上のGPUノード付きクラスター、kubectl、Helm 3.0以上。コードベースは主にGoで書かれており、GitHubリポジトリはAzure orgで管理されているが、設計上はKubernetesアグノスティックを志向しており、AKS以外のKubernetes環境での動作も想定されている。
ただし、現時点では開発途上であることを明示しておく必要がある。ロードマップには以下の機能が計画中として列挙されている。
- マルチテナントワークスペース、RBACとクォータ
- PyTorch DDP/FSDP、DeepSpeed、LoRA/QLoRAワークフローのサポート
- データセットライフサイクル管理
- マルチクラスター・マルチクラウド実行
競合との位置づけ
KubernetesベースのAIワークロード管理プラットフォームはTauGrid一択ではない。KubeflowはCNCF（Cloud Native Computing Foundation）のGraduationプロジェクト入りを目指して成熟が進んでおり、「production-readyなMLシステム」として位置づけられている。NVIDIAのRun:AIも同様の領域をカバーする選択肢だ。
TauGridがこれらと比較してどこで優位性を発揮するかは、今後の機能実装と実運用での評価に委ねられる部分が大きい。
詳細はMicrosoft Open-Sources TauGrid to Simplify AI Workload Management on Kubernetesを参照していただきたい。
