デベロッパー

Vertex AI Pipelines 向けサーバーレス Spark コンポーネントの発表

2022年4月21日

Google Cloud Japan Team

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

デベロッパーと ML エンジニアは Spark ML ワークロードの運用化に関するさまざまな課題に直面しています。課題の一つとして、インフラストラクチャに関するものが考えられます。たとえば、インフラストラクチャクラスタをどのように事前にプロビジョニングするか、データの準備やモデルのトレーニング、モデルの評価などの多岐にわたるタスクを妥当な時間で実行するための十分なリソースをどのように確保するかといった課題です。ほかにも、タスクのオーケストレーションやデータ処理に関する課題もあるでしょう。たとえば、モデルトレーニングのタスクを実行しているときに、必ず最新機能を使用できるようにするにはどうすれば良いかといったものです。

このような課題は、Vertex AI Pipelines をサーバーレス Spark ワークフローを実行する Dataproc と組み合わせて使用し、ML ワークフローを自動化することで解決できます。たとえば、こちらの記事では、Spark ML モデルをトレーニングおよびデプロイして、事前にインフラストラクチャをプロビジョニングすることなくリアルタイムに近い予測を行う方法について説明しています。特に、カスタムの Python コンポーネントを使用することで Vertex AI Pipelines から Dataproc ジョブを起動し、オーケストレートする方法を取り上げています。

本日、Vertex AI Pipelines 向け Dataproc サーバーレスコンポーネントを新しくリリースしたことをお知らせいたします。これにより、Spark、Spark SQL、PySpark、Spark ジョブの MLOps が簡素化されます。

以下のコンポーネントが利用可能になりました。

DataprocPySparkBatchOp: PySpark バッチワークロード
DataprocSparkBatchOp: Spark バッチワークロード
DataprocSparkSqlBatchOp: Spark SQL バッチワークロード
DataprocSparkRBatchOp: SparkR バッチワークロード

これらのコンポーネントには、Vertex AI Pipelines と Dataproc サーバーレスを用いて Spark ベースの ML パイプラインを簡単にオーケストレートするためのネイティブ KFP オペレータがあります。自分で行う必要があるのは、前処理、トレーニング、後処理の各ステップのビルドと、KFP パイプラインのコンパイルだけです。その後、Vertex AI Pipeline と Dataproc サーバーレスを使うことで、インフラストラクチャをプロビジョニングおよび管理する必要なく ML ワークフローをスケーラブルかつ再現可能な、信頼性のある方法で実行できます。

これらのコンポーネントの使用方法についての詳細は、以下の開始チュートリアルをご確認ください。

https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage3/get_started_with_dataproc_serverless_pipeline_components.ipynb

その他にも、PySpark と DataprocPySparkBatchOp コンポーネントを使用するエンドツーエンドの例を以下で紹介します。

Vertex AI Pipeline で Pyspark と Dataproc サーバーレスを使ってローン審査モデルをトレーニングする

このセクションでは、Spark MLlib と DataprocPySparkBatchOp コンポーネントを使って、顧客の銀行ローン利用資格を判断する Spark ML パイプラインをビルドする方法について説明します。特に、データの前処理から、顧客のローン審査の合格率を予測するランダムフォレスト分類器のハイパーパラメータ調整までの Spark MLlib パイプラインをカバーします。

パイプラインは次の図のとおりです。

https://storage.googleapis.com/gweb-cloudblog-publish/images/image1_I4ZTwNn.max-1000x1000.png

パイプラインワークフロー（クリックで拡大）

図に示すように、パイプラインには以下の特徴があります。

Cloud Storage にデータを保存する
DataprocPySparkBatchOp を使ってカテゴリ変数と数値変数を代入する
DataprocPySparkBatchOp を使って RandomForestClassifier のトレーニングを行う
カスタムコンポーネントを実行し Vertex AI Pipelines UI のモデル指標を評価して表示する

モデルがパフォーマンス条件（適合率と再現率の曲線、auPR の下側の領域が最小値 0.5 より高い）に一致する場合は、

DataprocPySparkBatchOp を使って RandomForestClassifier のハイパーパラメータ調整を実施する
Vertex AI モデルレジストリにモデルを登録する

DataprocPySparkBatchOp がどのように機能するかわかりやすくするため、トレーニングステップについて説明します。

DataprocPySparkBatchOp コンポーネントを使って PySpark MLlib モデルをトレーニングする

この例では、PySpark と Spark MLlib をパイプラインステップとして使ってランダムフォレスト分類器のトレーニングを行います。DataprocPySparkBatchOp コンポーネントを使用して Dataproc サーバーレスでトレーニングを実行するには、まずトレーニングのスクリプトを作成する必要があります。

以下は GitHub リポジトリで手に入るスクリプトです。

読み込んでいます...

def main(logger, args):
   """
   Main function
   Args:
       logger: logger
       args: args
   Returns:
       None
   """
   train_path = args.train_path
   model_path = args.model_path
   metrics_path = args.metrics_path
   try:
       logger.info('initializing pipeline training.')
       logger.info('start spark session.')
       spark = (SparkSession.builder
                .master("local[*]")
                .appName("spark go live")
                .config('spark.ui.port', '4050')
                .getOrCreate())
       logger.info(f'spark version: {spark.sparkContext.version}')
       logger.info('start building pipeline.')
       preprocessing_stages = build_preprocessing_components()
       feature_engineering_stages = build_feature_engineering_components()
       model_training_stage = build_training_model_component()
       pipeline = build_pipeline(preprocessing_stages, feature_engineering_stages, model_training_stage)
       logger.info(f'load train data from {train_path}.')
       raw_data = (spark.read.format('csv')
                   .option("header", "true")
                   .schema(DATA_SCHEMA)
                   .load(train_path))
       logger.info('fit model pipeline.')
       train, test = raw_data.randomSplit(RANDOM_QUOTAS, seed=RANDOM_SEED)
       pipeline_model = pipeline.fit(train)
       predictions = pipeline_model.transform(test)
       metrics = get_metrics(predictions, TARGET, 'test')
       for m, v in metrics.items():
           print(f'{m}: {v}')
       logger.info(f'load model pipeline in {model_path}.')
       pipeline.write().overwrite().save(model_path)
       if model_path.startswith('gs://'):
           pipeline.write().overwrite().save(model_path)
       else:
           path(model_path).mkdir(parents=True, exist_ok=True)
           pipeline.write().overwrite().save(model_path)
       logger.info(f'upload metrics under {metrics_path}.')
       if metrics_path.startswith('gs://'):
           bucket = urlparse(model_path).netloc
           metrics_file_path = urlparse(metrics_path).path.strip('/')
           write_metrics(bucket, metrics, metrics_file_path)
       else:
           metrics_version_path = path(metrics_path).parents[0]
           metrics_version_path.mkdir(parents=True, exist_ok=True)
           with open(metrics_path, 'w') as json_file:
               json.dump(metrics, json_file)
           json_file.close()
   except RuntimeError as main_error:
       logger.error(main_error)
   else:
       logger.info('model pipeline training successfully completed!')
       return 0

完全なコードはこちらのノートブックをご覧ください。

Spark セッションを初期化したら、スクリプトが Spark ML パイプラインをビルドし、前処理済みデータを読み込み、トレーニングサンプルとテストサンプルを生成します。さらに、モデルのトレーニングを実施し、Cloud Storage バケットにアーティファクトと指標を保存します。

PySpark スクリプトは、使用前に Cloud Storage バケットにアップロードする必要があります。

読み込んでいます...

完全なコードはこちらのノートブックをご覧ください。

スクリプトを Cloud Storage にアップロードしたら、DataprocPySparkBatchOp を使用してトレーニングコンポーネントを定義できるようになります。main_python_file_uri 引数の値は、PySpark スクリプトの Cloud Storage 内の場所にする必要があります。

以下のように定義します

読み込んでいます...

完全なコードはこちらのノートブックをご覧ください。

DataprocPySparkBatchOp には、以下のパラメータの値を指定する必要があります。

project、Dataproc バッチワークロードを実行するプロジェクト
batch_id、バッチジョブに使用する一意の ID
main_python_file_uri、Spark ドライバとして使うメインの Python ファイルの HCFS URI

お使いのワークロードに他のパラメータが必要でも、DataprocPySparkBatchOp コンポーネントはそのパラメータを受け入れることができます。コンポーネントの詳細については、ドキュメントをご覧ください。

最後に、dsl.pipeline デコレータを使って、このコンポーネントをパイプラインの定義の他のタスクに統合します。それから、パイプラインの定義をコンパイルし、Vertex AI SDK を使って実行します。

この実行中に model_traning_op ステップが実行されると、パイプラインがトレーニングワークロードを Dataproc サーバーレスサービスに送信します。ジョブの実行が正常に完了すると、バッチワークロードの詳細を Cloud Console で確認できます。

https://storage.googleapis.com/gweb-cloudblog-publish/images/image6_La8xVWj.max-1300x1300.png

Dataproc UI のトレーニングジョブのバッチビュー（クリックで拡大）

Vertex AI Pipelines 向け Dataproc サーバーレスコンポーネントを使用することで、Spark ジョブを機械学習ワークフローにいかに簡単に統合できるかご理解いただけたと思います。 Dataproc サーバーレスはインフラストラクチャ全体の管理を引き受けます。さらに、ワークロードのトレーニングでは、実行されるときに必要なリソースしか消費されません。

同様のアプローチを使って、他のタスクをパイプラインに統合できます。下のコードは、完全な機械学習ワークフローを実行するパイプラインの定義の例です。このワークフローには、データの前処理およびデータセットの作成、モデルのトレーニング、モデルの評価、ハイパーパラメータの調整、Vertex AI モデルレジストリへのモデルのアップロードなどが含まれます。パイプラインは Dataproc サーバーレスで PySpark ワークロードを実行するための DataprocPySparkBatchOp と、Google Cloud パイプラインコンポーネント SDK の一部である他のコンポーネントを使用します。

読み込んでいます...

@dsl.pipeline(name="dataproc-pyspark-preprocessing",
             description="")
def pipeline(
   preprocessing_batch_id: str = PREPROCESSING_BATCH_ID,
   preprocessing_main_python_file_uri: str = PREPROCESSING_PYTHON_FILE_URI,
   train_data_path: str = FEATURES_TRAIN_URI,
   preprocessed_data_path: str = PROCESSED_DATA_URI,
   dataset_name: str = DATASET_NAME,
   dataset_uri: str = GCS_PREPROCESSED_URI,
   training_batch_id: str = TRAINING_BATCH_ID,
   training_main_python_file_uri: str = TRAINING_PYTHON_FILE_URI,
   train_path: str = PROCESSED_DATA_URI,
   model_path: str = MODEL_URI,
   metrics_path: str = METRICS_URI,
   threshold:float = AUPR_THRESHOLD,
   hpt_batch_id: str = HPT_TRAINING_BATCH_ID,
   hpt_main_python_file_uri: str = HPT_PYTHON_FILE_URI,
   hpt_model_path: str = HPT_MODEL_URI,
   hpt_metrics_path: str = HPT_METRICS_URI,
   custom_container_image: str = RUNTIME_CONTAINER_IMAGE,
   model_name: str = MODEL_NAME,
   project_id: str = PROJECT_ID,
   location: str = REGION,
):
   from google_cloud_pipeline_components.experimental.dataproc import DataprocPySparkBatchOp
   #build preprocessed data args
   build_preprocessing_args_op = build_preprocessing_args(
       train_data_path=train_data_path,
       processed_data_path=preprocessed_data_path
   )
   # preprocess data
   data_preprocessing_op = DataprocPySparkBatchOp(
       project=project_id,
       location=location,
       container_image=custom_container_image,
       batch_id=preprocessing_batch_id,
       main_python_file_uri=preprocessing_main_python_file_uri,
       args=build_preprocessing_args_op.output
   ).after(build_preprocessing_args_op)
   # create dataset
   create_dataset_op = vertex_ai_components.TabularDatasetCreateOp(
       display_name=dataset_name,
       gcs_source=dataset_uri,
       project=project_id,
       location=location,
       ).after(data_preprocessing_op)
  
   # build training data args
   build_training_args_op = build_training_args(
       dataset_uri = create_dataset_op.output,
       train_path=train_path,
       model_path=model_path,
       metrics_path=metrics_path,
   ).after(create_dataset_op)
   # training model
   model_training_op = DataprocPySparkBatchOp(
       project=project_id,
       location=location,
       container_image=custom_container_image,
       batch_id=training_batch_id,
       main_python_file_uri=training_main_python_file_uri,
       args=build_training_args_op.output,
   ).after(build_training_args_op)
   evaluate_model_op = evaluate_model(
       metrics_uri=metrics_path
       ).after(model_traning_op)
  
   # evaluate condition
   with dsl.Condition(evaluate_model_op.outputs['threshold_metric'] >= threshold, name=AUPR_HYPERTUNE_CONDITION):
     build_hpt_args_op = build_training_args(
       dataset_uri = create_dataset_op.output,
       train_path=train_path,
       model_path=hpt_model_path,
       metrics_path=hpt_metrics_path,
   ).after(evaluate_model_op)
     # hypertuning model
     hyperparameter_tuning_op = DataprocPySparkBatchOp(
         project=project_id,
         location=location,
         container_image=custom_container_image,
         batch_id=hpt_batch_id,
         main_python_file_uri=hpt_main_python_file_uri,
         args=build_hpt_args_op.output,
     ).after(model_traning_op)
     # upload model
     register_model(artifact_uri=hpt_model_path).after(hyperparameter_tuning_op)

完全なコードはこちらのノートブックをご覧ください

Vertex AI Pipelines によって、実行中の機械学習ワークフローが可視化されます。次の図は、上のパイプラインの定義を使って実行したエンドツーエンドのワークフローの全体図です。

https://storage.googleapis.com/gweb-cloudblog-publish/original_images/image5_z8lTS4c.gif

Vertex AI Pipelines UI のパイプライン（クリックで拡大）

まとめ

今回のブログ記事では、Vertex AI Pipelines 向けの新しい Dataproc コンポーネントが利用可能になったことを発表しました。また、DataprocPySparkBatchOp を使ってデータの前処理、ローン審査用の PySpark MLlib モデルのトレーニングとハイパーパラメータ調整を行うエンドツーエンドの例を示しました。