PyFlink Next Generation Python Runtime Introduction

1. PyFlink の新機能

PyFlink 1.14 では多くの新機能が追加されました。主に機能、使いやすさ、パフォーマンスの 3 つの側面に分けられます。

機能面では、State TTL 設定が追加されました。1.14 より前は、Python DataStream API や State を操作する一部の実装が行われていましたが、State TTL 設定のサポートは提供されておらず、Python DataStream API のカスタム関数で State 値を自動的にクリアできませんでした。手動での操作が必要であり、ユーザーフレンドリとは言えませんでした。

使いやすさの面では、以下の機能が追加されました。

依存関係管理で tar.gz フォーマットがサポートされました。

プロファイル機能。PyFlink を記述する際に Python カスタム関数を使用しますが、これらの関数のパフォーマンスボトルネックがどこにあるか明確でない場合があります。プロファイル機能により、Python 関数でパフォーマンスボトルネックが発生した際に、プロファイルを使用してボトルネックの具体的な原因を分析し、最適化を行うことができます。

print 関数。1.14 より前は、カスタムログ情報を出力するには Python の logging モジュールを使用する必要がありました。しかし Python ユーザーにとって、print はログ情報の出力により慣れ親しんだ方法です。そのため、1.14 でこの機能が追加されました。

ローカルデバッグモード。1.14 より前は、Python カスタム関数を使用してローカルで PyFlink ジョブを開発する場合、リモートデバッグを使用する必要があり、操作が煩雑で利用のしきい値が高いという課題がありました。1.14 ではこの点が改善され、ローカルで PyFlink ジョブを記述し Python カスタム関数を使用する場合、自動的にローカルデバッグモードに切り替わり、IDE 上で直接 Python カスタム関数をデバッグできます。

パフォーマンス面では、以下の機能が追加されました。

Operator Fusion。この機能は主に Python DataStream API ジョブで複数のオペレータ操作を連続して実行するシナリオを対象としています。たとえば、2 つの .map 操作がある場合、1.14 より前はこれら 2 つの .map はそれぞれ別の Python ワーカーで実行されていましたが、Operator Fusion の実装後は 1 つのオペレータに統合され、単一の Python ワーカーで実行されることで、大幅なパフォーマンスの最適化が実現しました。

State のシリアル化/デシリアル化の最適化。1.14 より前は、State のシリアル化/デシリアル化に Python 組み込みのシリアライザである pickle を使用していました。pickle はさまざまな Python 定義のデータ構造をシリアル化できますが、State の型情報をデータ構造に含める必要があり、シリアル化後のデータサイズが大きくなるという問題がありました。1.14 ではカスタムシリアライザを使用し、各型に専用のシリアライザを割り当てることで最適化され、シリアル化情報のサイズが小さくなりました。

Finish Bundle の最適化。1.14 より前は Finish Bundle は同期操作でしたが、非同期操作に変更され、パフォーマンスが向上し、チェックポイントが完了しないシナリオの解決にも役立ちます。

2. PyFlink ランタイム

図の左側上部にある Python Table API & SQL および DataStream API は、ユーザーに提供される Python API です。ユーザーはこれら 2 つの Python API を通じて PyFlink ジョブを記述し、py4j ライブラリを通じて Python API を Java API に変換することで、Flink の Java API に対応するジョブ記述を行います。

Table と SQL のジョブには追加のオプティマイザーがあり、共通ルールと Python ルールの 2 種類のルールがあります。なぜ Python ルールが必要なのでしょうか。共通ルールは既存の Table/SQL ジョブ全般に有効ですが、Python ルールの最適化は、PyFlink ジョブ内で Python カスタム関数を使用するシナリオに特化しており、対応するオペレータを抽出できるためです。

ジョブ記述後、対応する Python オペレータを含む JobGraph に変換されます。Python オペレータで記述された JobGraph は TM (Runtime) にサブミットされ実行されます。Runtime 内にも Python オペレータが存在します。

図の右側は Python オペレータの各コンポーネントで、PyFlink ランタイムのコア部分を示しています。主に Java オペレータと Python ワーカーの 2 つに分かれます。

Java オペレータには多くのコンポーネントがあり、データサービスや State サービス、チェックポイント、ウォーターマーク、State リクエストの処理が含まれます。Python カスタム関数は JVM ベースの Flink の既存アーキテクチャ上で直接実行できないため、Python のランタイムが必要です。オペレータワーカーがこれを解決します。

具体的な方法は次のとおりです。Python プロセスを起動して Python 定義関数を実行し、Java オペレータが上流データを処理してから、特別な処理を経て対応する Python ワーカーに送信します。ここではプロセス間通信方式が使用されており、図中のデータサービスがこれに該当します。State サービスは Python DataStream API の State 操作を対象としています。Python で State を操作すると、データは Python ワーカーから Java オペレータに返され、Java オペレータが State バックエンドにアクセスして対応する State データを取得し、Python ワーカーに送り返します。その後、ユーザーは State の操作結果を利用できます。

上図は PyFlink ランタイムのワークフローです。その中の役割は、Python オペレータ、Python ランナー、バンドルプロセッサ、コーダー、Python オペレーションです。これらの異なる役割はそれぞれ異なる場所で実行されます。Python オペレータと Python ランナーは Java JVM 内で実行され、上流および下流の Java オペレータとの接続を担当します。一方、バンドルプロセッサ、コーダー、Python オペレーションは PVM 内で実行されます。バンドルプロセッサは既存の Apache Beam フレームワークを利用し、プロセス間通信を使用して Java 側からデータを受信します。コーダーは Python 側のカスタムシリアライザです。Java 側から送信されたデータは、Python オペレータを経由して Python ランナーに送られ、Python ランナーでシリアル化された後、プロセス間通信を通じてバンドルプロセッサに送信されます。バンドルプロセッサはコーダーを通じてシリアル化されたバイナリ配列をデシリアル化し、Python オブジェクトを取得します。最後に、デシリアル化された Python パラメーターは Python オペレーションを通じて関数本体の入力パラメーターとして渡され、Python カスタム関数が呼び出されてカスタム結果が得られます。

上記プロセスのボトルネックは主に以下の点に存在します。まず、計算端でユーザー定義関数を呼び出す際、呼び出し前にフレームワークレイヤーに Python で記述されたオーバーヘッドが存在すること。次に、カスタムシリアル化部分で Java 端と Python 端の両方でデータのシリアル化とデシリアル化が必要なこと。3 つ目はプロセス間の通信です。

上記のボトルネックに対して、一連の最適化が行われました。

計算面では、codegen を使用して既存の Python 関数呼び出しのすべての変数を定数に変更し、関数の実行効率を向上させました。さらに、既存の Python オペレーションのすべての実装を Cython に変更することで、Python を C に変換したのと同等の実装方式となり、パフォーマンスが大幅に向上しました。
シリアル化面では、カスタムシリアライザが提供され、すべて純粋な C 実装であり、Python よりも効率的です。
通信面では、まだ最適化されていません。
シリアル化と通信の問題は、本質的に Java と Python の相互呼び出しの問題、すなわち PyFlink のランタイムアーキテクチャをどう最適化するかの問題です。

3. FFI ベースの PEMJA

Java と Python の相互呼び出しは既に一般的な問題であり、多くの実装が既に存在します。

最初はプロセス間通信に基づくソリューション、すなわちネットワーク通信の方式で、以下のものがあります。

ソケット方式は、すべての通信プロトコルを自ら実装するため非常に柔軟ですが、煩雑です。
py4j 方式は、PyFlink と PySpark の両方がクライアント側でのジョブ記述時に py4j を使用しています。
Alink 方式は、ランタイムで py4j を使用し、Python カスタム関数も備えています。gRPC 方式は、既存の gRPC サービスを利用するためカスタムプロトコルが不要で、カスタムサービスとメッセージを備えています。
さらに、共有メモリ方式はもう一つのプロセス間通信ソリューションです。たとえば TensorFlow on Flink は共有メモリを通じて実装されています。PyArrow Plasma もオブジェクトベースの共有メモリストレージです。
上記のソリューションはすべてプロセス間通信を目的としていますが、Python と Java を同一プロセス内で実行できれば、プロセス間通信による問題を完全に排除できるのではないでしょうか。

実際にこれを試みるライブラリがいくつか存在します。最初のソリューションは Python を Java に変換するものです。たとえば、p2j は Python ソースコードを Java ソースコードに変換し、voc は Python コードを直接 Java バイトコードに変換します。このソリューションの本質は、Python を JVM 上で直接実行できるコードに変換することです。しかしこのソリューションにも多くの欠陥があります。Python は常に発展しており、さまざまな構文があり、Python の構文を Java の対応するオブジェクトにマッピングするのは非常に困難です。結局、異なる言語だからです。

2 つ目のソリューションは、Java ベースの Python インタープリター実装です。1 つは Jython です。Python は実際には C 言語で書かれた Python インタープリターのセットです。C で書かれた Python インタープリターは C 上で実行できるので、Java で実装された Python インタープリターも JVM 上で直接実行できるはずです。もう 1 つのソリューションは GraalVM で、Truffle フレームワークを提供し、さまざまなプログラミング言語が共通の構造を使用できるようにします。この構造は JVM 上で実行できるため、各言語が同一プロセス内で動作できます。

上記のソリューションの前提は Python コードを認識できること、すなわち既存のさまざまな Python コードとの互換性が必要ですが、現在、互換性は解決が困難な問題であり、Python から Java への変換を妨げています。

3 つ目は FFI ベースのソリューションです。

FFI の本質は、ホスト言語がゲスト言語をどのように呼び出すか、すなわち Java と Python の相互呼び出しにあります。対応する具体的な実装方式は多くあります。

Java は JNI (Java Native Interface) を提供しており、Java ユーザーは JNI インターフェイスを通じて C で実装されたライブラリを呼び出すことができ、逆も可能です。このインターフェイスにより、JVM ベンダーは JNI に従って実装を行い、Java と C の相互呼び出しを実現します。

Python/C API も同様です。Python は C で実装されたインタープリターであるため、Python コードから C のライブラリを呼び出すことをサポートし、逆も可能です。

Cython は、ソースコードを他の言語が理解できるコードに変換するツールを提供します。たとえば、Python コードを非常に効率的な C 言語コードに変換し、CPython インタープリターに組み込んで直接実行できるため、非常に効率的です。

Ctypes は、C ライブラリをラップすることで Python から C ライブラリを効率的に呼び出せるようにします。

上記の FFI ベースのソリューションの核心は C です。C のブリッジがあれば、Java で書かれたコードは JNI インターフェイスを通じて C を呼び出し、C は CPython API インターフェイスを呼び出すことで、最終的に Java と Python が同一スレッド内で動作することを実現できます。これが PEMJA の全体的な設計思想です。プロセス間通信の問題を解決し、Python 自体が提供する Python/C API を使用しているため互換性の問題がなく、Java によるインタープリター実装の欠陥も克服しています。

JPype が解決するのは Python から Java を呼び出す問題であり、Java から Python を呼び出すことはサポートしていないため、このシナリオには適用できません。

JEP は Java から Python を呼び出すことを実現していますが、具体的な実装に多くの制限があります。まず、ソースコードからしかインストールできず、環境に対する要件が非常に高く、CPython の .source ファイルに依存するため、クロスプラットフォームでのインストールと使用に非常に不利です。JEP の起動エントリーは必ず JEP プログラムである必要があり、クラスライブラリを動的にロードする必要があり、環境変数に事前に設定する必要があるため、サードパーティのミドルウェアプラグインとして別のアーキテクチャで実行するのに不利です。さらにパフォーマンスの課題もあり、Python の GIL 問題を十分に克服できていないため、効率的とは言えません。

PEMJA は上記の問題を基本的に克服し、Java と Python の相互呼び出しをより良く実現します。

上図はいくつかのフレームワークのパフォーマンス比較です。ここでは比較的標準的でシンプルな String upper 関数を使用しています。主にフレームワークレイヤーのオーバーヘッドを比較しており、カスタム関数のパフォーマンスではないため、最もシンプルな関数を使用しています。同時に、既存のさまざまな関数で最もよく使われるデータ構造が文字列であることを考慮し、ここでは文字列を使用しています。

ここでは、これら 4 つのインタープリターで 100 バイトと 1000 バイトのパフォーマンスを比較しています。Jython は期待されたほどの効率ではなく、4 つの実装の中で最もパフォーマンスが低いことがわかります。JEP のパフォーマンスは PEMJA に遠くおよびません。PEMJA は 100 バイトの場合、純粋な Java 実装の約 40% のパフォーマンスですが、1000 バイトの場合、実際には純粋な Java 実装のパフォーマンスを上回ります。

この現象をどう説明するのでしょうか。String upper 自体は Java 実装ですが、Python では C 実装です。関数自体の実行効率は Java より高く、フレームワークのオーバーヘッドが十分に小さいことと相まって、全体のパフォーマンスが Java を上回ります。つまり、一部のシナリオでは Python UDF のパフォーマンスが Java UDF を上回る可能性があるということです。

現在多くのユーザーが Python UDF の代わりに Java UDF を使用している主な理由は、Python UDF のパフォーマンスが Java よりもはるかに劣っているからです。しかし、Java のパフォーマンスが必ずしも Python より優れているわけではない場合、Python にも利点があります。Python はスクリプト言語であり、記述がより便利です。

Java のデーモンスレッドは、PEMJA と対応する Python PVM での初期化と最終的な破棄、およびリソースの作成と解放を担当します。ユーザーは Java で PEMJA インスタンスを使用し、インスタンスは PEMJA 内の対応するインスタンスにマッピングされ、各インスタンスが Python サブインタープリターを作成します。Python サブインタープリターは、グローバル Python インタープリターと比較してより小さな概念で、GIL を制御できます。独自の独立したヒープ空間を持ち、名前空間の分離を実現できます。各スレッドは Python サブインタープリターに対応し、対応する PVM で独自の Python 関数を実行できます。

4. PyFlink Runtime 2.0

PyFlink Runtime 2.0 は PEMJA に基づいています。

上図の左側は PyFlink 1.0 のアーキテクチャです。中には 2 つのプロセスがあり、1 つは Java プロセス、もう 1 つは Python プロセスです。データインタラクションはデータサービスと State サービスを通じて、プロセス IPC 通信で実現されています。

PEMJA を使用すると、データサービスと State サービスを PEMJA ライブラリに置き換えることができ、左側の JVM と右側の PVM を同一プロセス内で実行できるため、IPC プロセス間通信の問題を完全に解決できます。

上図は既存の PyFlink UDF、PEMJA ベースの PyFlink UDF、および Java UDF のパフォーマンス比較です。String upper 関数を使用して 100 バイトと 1000 バイトのパフォーマンスを比較しています。100 バイトの場合、PEMJA 上の UDF 実装は Java UDF のパフォーマンスの約 50% に達しています。1000 バイトの場合、PEMJA 上の UDF のパフォーマンスは Java UDF を上回ります。これはカスタム関数の実装にもよりますが、PEMJA フレームワークの高性能さを示しています。

5. 今後の取り組み

今後、PEMJA フレームワークはオープンソース化されます (2022 年 1 月 14 日に公式オープンソース化)。これは汎用ソリューションであり、PyFlink だけでなく、Java と Python のさまざまな相互呼び出し方式にも利用できるため、PEMJA フレームワークとして独立したオープンソースプロジェクトにします。初期バージョンでは Java から Python 関数の呼び出しのみをサポートし、その後 Python から Java 関数の呼び出しもサポートする予定です。Python で記述された Python DataStream API の State 呼び出し機能は、Python から Java 関数の呼び出しに依存しているためです。さらに、PEMJA で NumPy のネイティブデータ構造をサポートする予定です。このサポートが実装されれば、pandas UDF が利用可能になり、パフォーマンスは飛躍的に向上します。