指定管道步驟的機器設定

Kubeflow 管道元件是工廠函式，可建立管道步驟。每個元件都會說明元件的輸入內容、輸出內容和實作方式。舉例來說，train_op 是下列程式碼範例中的元件。

舉例來說，訓練元件可以將 CSV 檔案做為輸入內容，並用來訓練模型。在管道步驟中設定機器類型參數，即可管理管道中每個步驟的需求。假設您有兩個訓練步驟，其中一個步驟會訓練大型資料檔案，另一個步驟則會訓練小型資料檔案，您可以為第一個工作分配更多記憶體和 CPU，為第二個工作分配較少資源。

根據預設，這個元件會以 Vertex AI CustomJob 形式執行，使用 e2-standard-4 機器，具備 4 個 CPU 核心和 16 GB 記憶體。如要進一步瞭解如何選取「機器類型」中列出的其中一種 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_type(ACCELERATOR_TYPE)
    .set_accelerator_limit(ACCELERATOR_LIMIT))

更改下列內容：

  • CPU_LIMIT：這個運算子的 CPU 上限。這個字串值可以是數字 (CPU 數量的整數值)，也可以是數字後接「m」，表示 1/1000。最多可指定 96 個 CPU。

  • MEMORY_LIMIT：這個運算子的記憶體上限。這個字串值可以是數字，也可以是數字後接「K」(KB)、「M」(MB) 或「G」(GB)。最多支援 624 GB。

  • SELECTOR_CONSTRAINT：每個限制都是鍵/值組合標籤。 如要讓容器在節點上執行，節點必須將每個限制條件設為標籤。例如： 'cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_T4'

    可用的限制如下：

    • NVIDIA_GB200+ (包括 GPUDirect-RDMA)
    • NVIDIA_B200* (包括 GPUDirect-RDMA)
    • NVIDIA_H100_MEGA_80GB* (包括 GPUDirect-TCPXO)
    • NVIDIA_H100_80GB
    • NVIDIA_H200_141GB* (包括 GPUDirect-RDMA)
    • NVIDIA_A100_80GB
    • NVIDIA_TESLA_A100 (NVIDIA A100 40GB)
    • NVIDIA_TESLA_P4
    • NVIDIA_TESLA_P100
    • NVIDIA_TESLA_T4
    • NVIDIA_TESLA_V100
    • NVIDIA_L4

    • TPU_V2
    • TPU_V3

  • ACCELERATOR_TYPE：加速器類型。如要進一步瞭解可用的 GPU 和設定方式，請參閱「GPU」。如要進一步瞭解可用的 TPU 類型及設定方式，請參閱 TPU。詳情請參閱 accelerator_type 參數參考資料。

  • ACCELERATOR_LIMIT：運算子的加速器 (GPU 或 TPU) 限制。您可以指定正整數。如要進一步瞭解可用的 GPU 和設定方式，請參閱「GPU」。如要進一步瞭解可用的 TPU 和設定方式，請參閱「TPU」。

CustomJob 支援特定機器類型，最多可使用 96 個 CPU 和 624 GB 記憶體。Vertex AI Pipelines 會根據您指定的 CPU、記憶體和加速器設定，從支援的機器類型中自動選取最接近的類型。