Dataproc ML ライブラリを使用して Spark データ パイプラインを Gemini やその他の AI モデルに接続

※この投稿は米国時間 2025 年 10 月 4 日に、Google Cloud blog に投稿されたものの抄訳です。

多くのデータ サイエンス チームは、Dataproc マネージド クラスタで実行される Apache Spark を利用して、大規模なデータ準備を効率的に行っています。これらのチームがデータ パイプラインを ML モデルに直接接続しようとする中、統合プロセスを簡単にするための明確な機会が生まれます。しかし、Vertex AI のモデルを使用して Spark DataFrame で推論を実行するには、通常、カスタム開発が必要になるため、単一のエンドツーエンド ワークフローの構築が複雑化します。

この問題を解決するために、Google は Dataproc の AI / ML 推論を簡素化するように設計された新しいオープンソースの Python ライブラリを開発しています。このライブラリは、Apache Spark ジョブを接続して、モデル推論から始まる一般的な ML フレームワークと Vertex AI の機能を使用します。またオープンソースであるため、アプリケーション コードで直接使用でき、その動作を完全に把握できます。

仕組み

Dataproc ML は、SparkML スタイルのビルダー パターンに従って、Spark ユーザーが使い慣れた感覚で利用できるように構築されています。使用するモデルを構成し、DataFrame で .transform() を呼び出します。一般的な推論のユースケースをいくつか見てみましょう。

1. Gemini モデルを Spark データに適用する

Gemini などの生成 AI モデルを Spark DataFrame の列に適用できます。これは、分類、抽出、要約などのタスクを大規模に行う場合に便利です。この例では、「city」列と「country」列を含む DataFrame を使用し、簡単なプロンプトを指定して Gemini で新しい列を作成します。

PyPi からインストールして、ローカル環境でテストできます。

大規模にデプロイやテストを行うには、Dataproc バージョン 2.3-ml クラスタを作成します。

この例を gemini_spark.py というファイルにコピーします。

ドキュメントで説明されているように、ハンドラはカスタマイズされたオプションをサポートする柔軟性を備えています。

Dataproc クラスタにジョブを送信します。

2. PyTorch モデルと TensorFlow モデルで推論を実行する

このライブラリでは、Gemini エンドポイントを呼び出すだけでなく、Google Cloud Storage から直接読み込まれたモデルファイルを使用して推論を実行することもできます。PyTorchModelHandler（および TensorFlow の同様のハンドラ）を使用して、モデルの重みを読み込み、プリプロセッサを定義し、ワーカーノードで直接推論を実行できます。これは、個別のモデル提供エンドポイントを管理せずに、大規模なバッチ推論を実行する場合に便利です。

パフォーマンスを重視

このライブラリは単なるラッパーではありません。大規模な Dataproc クラスタで推論を実行するように設計されており、推論用にいくつかの最適化が含まれています。

• ベクトル化されたデータ転送: pandas_udf を使用して、Spark と Python ワーカー プロセス間でデータを効率的に移動します。
• 接続の再利用: エンドポイントへの接続はパーティション間で再利用され、オーバーヘッドが削減されます。
• 再試行ロジック: ライブラリは、HTTP 429（リソースが枯渇）などのエラーを指数バックオフと再試行で自動的に処理します。

使ってみる

オープンソース リポジトリを確認し、ドキュメントをお読みいただくことで、今すぐ使用を開始できます。

今後数か月以内に、このライブラリに以下の機能を追加する予定です。

1. Spark Connect のサポート: これにより、BigQuery Studio ノートブック内で上記の機能を使用できるようになります。
2. Vertex AI のインテグレーション: 推論を容易にするため、以下に対して ModelHandler を追加する予定です。
   1. オンライン推論のために Vertex モデル エンドポイントを直接呼び出す
   2. Vertex モデルを参照し、Spark ワーカーにローカライズする
   3. エンベディング モデルを含む、Vertex Model Garden でホストされているモデルを参照する
3. その他の最適化: 推論のランタイムを強化するために、入力データフレームを自動的に再パーティション分割
4. サードパーティ統合: HuggingFace のオープンソース モデルを参照

現在、Google は、Google Compute Engine ML イメージの Dataproc と Apache Spark 向け Google Cloud Serverless ランタイムにこのライブラリをデフォルトで含めるよう取り組んでいます。

皆様がどのようなものを構築されるか楽しみにしております。 また、Spark での AI/ML の利用をさらに簡素化するためのフィードバックや機能リクエストがおありになる場合には、dataproc-feedback@google.com までお問い合わせください。

ー デベロッパー アドボケイト Brad Miro

ー シニア ソフトウェア エンジニア Deependra Patel