パイプライン ステップのマシン構成を指定する

Kubeflow Pipelines コンポーネントは、パイプライン ステップを作成するファクトリ関数です。各コンポーネントでは、コンポーネントの入力、出力、実装を記述します。たとえば、以下のコードサンプルでは、train_op がコンポーネントです。

たとえば、トレーニング コンポーネントで CSV ファイルを入力として受け取り、それを使用してモデルをトレーニングできます。パイプライン ステップにマシンタイプ パラメータを設定することで、パイプラインの各ステップの要件を管理できます。2 つのトレーニング ステップがあり、最初のステップでは巨大なデータファイルを使用し、二番目のステップでは小さなデータファイルを使用する場合は、最初のタスクにより多くのメモリと CPU を割り当て、二つ目のタスクに割り当てるリソースを少なくすることができます。

デフォルトでは、コンポーネントは 4 コアの CPU と 16 GB のメモリを搭載した e2-standard-4 マシンを使用して、Vertex AI CustomJob として動作します。マシンタイプにリストされている Google Cloud 固有のマシンリソースの選択の詳細については、Vertex AI Pipelines を使用して Google Cloud マシンリソースをリクエストするをご覧ください。

次の例は、ステップに CPU、メモリ、GPU 構成を設定する方法を示しています。

from kfp import dsl

@dsl.pipeline(name='custom-container-pipeline')
def pipeline():
  generate = generate_op()
  train = (
    train_op(
      training_data=generate.outputs['training_data'],
      test_data=generate.outputs['test_data'],
      config_file=generate.outputs['config_file'])
    .set_cpu_limit('CPU_LIMIT')
    .set_memory_limit('MEMORY_LIMIT')
    .add_node_selector_constraint(SELECTOR_CONSTRAINT)
    .set_accelerator_limit(ACCELERATOR_LIMIT))

次のように置き換えます。

  • CPU_LIMIT: このオペレータの CPU の上限。この文字列値は、数値（CPU の数を表す整数値）か、数値の後に「m」が続く値になります（m は 1/1000 を意味します）。最大で 96 個の CPU を指定できます。

  • MEMORY_LIMIT: このオペレータのメモリの上限。この文字列値は、数値か、数値の後に「K」（キロバイト）、M（メガバイト）、または「G」（ギガバイト）が後に続く文字列です。最大 624 GB がサポートされます。

  • SELECTOR_CONSTRAINT: 各制約は Key-Value ペアのラベルです。コンテナをノードで実行できるようにするには、ノードで各制約がラベルとして指定されている必要があります。例: 'cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_K80'

    次の制約が利用可能です。

    • NVIDIA_H100_80GB
    • NVIDIA_A100_80GB
    • NVIDIA_TESLA_A100（NVIDIA A100 40GB）
    • NVIDIA_TESLA_K80
    • NVIDIA_TESLA_P4
    • NVIDIA_TESLA_P100
    • NVIDIA_TESLA_T4
    • NVIDIA_TESLA_V100
    • NVIDIA_L4

    • TPU_V2
    • TPU_V3

  • ACCELERATOR_LIMIT: オペレータのアクセラレータ（GPU または TPU）の制限。正の整数を指定できます。使用可能な GPU とその構成方法の詳細については、GPU をご覧ください。使用可能な TPU とその構成方法の詳細については、TPU をご覧ください。

CustomJob が現在サポートしている特定のマシンタイプでは、最大 96 個の CPU と 624 GB のメモリに制限されています。Vertex AI Pipelines は、指定された CPU、メモリ、アクセラレータ構成に基づいて、サポートされているマシンタイプから最も近いマシンタイプを自動的に選択します。