为流水线步骤指定机器配置

Kubeflow 流水线组件是创建流水线步骤的工厂函数。每个组件都描述了组件的输入、输出和实现。例如，train_op 是以下代码示例中的一个组件。

例如，训练组件可以将 CSV 文件用作输入，并使用它训练模型。通过为流水线步骤设置机器类型参数，您可以管理流水线中每个步骤的要求。如果您有两个训练步骤，一个步骤用于训练大型数据文件，第二个步骤用于训练小型数据文件，则可以为第一个任务分配更多内存和 CPU，为第二个任务分配更少的资源。

默认情况下，该组件将作为 Vertex AI 运行CustomJob使用e2-standard-4配备 4 个核心 CPU 和 16GB 内存的机器。 如需详细了解如何选择机器类型中列出的 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”，即 1/1000。 您最多可以指定 96 个 CPU。

  • MEMORY_LIMIT：运算符的最大内存限制。此字符串值可以是数字，也可以是数字后跟“K”（千字节）、“M”（兆字节）或“G”（千兆字节）。 最多支持 624 GB。

  • SELECTOR_CONSTRAINT：每个限制条件都是一个键值对标签。如需让容器能够在节点上运行，该节点必须将每个限制条件作为标签。例如：'cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_T4'

    以下限制条件可用：

    • NVIDIA_H100_80GB
    • NVIDIA_A100_80GB
    • NVIDIA_TESLA_A100 (NVIDIA A100 40 GB)
    • 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、内存和加速器配置，自动从支持的机器类型中选择最接近的匹配项。