在 Vertex AI 中为单个 TFX 管道组件指定机器类型答案

【问题标题】：Specify machine type for a single TFX pipeline component in Vertex AI在 Vertex AI 中为单个 TFX 管道组件指定机器类型
【发布时间】：2021-09-02 16:23:28
【问题描述】：

我正在使用 TFX 在 Vertex AI 上构建 AI 管道。我已经关注this tutorial 开始，然后我将管道调整为我自己的数据，其中包含超过 1 亿行时间序列数据。由于内存问题，我的一些组件在中途被杀死，所以我只想为这些组件设置内存要求。我使用KubeflowV2DagRunner 使用以下代码在 Vertex AI 中编排和启动管道：

runner = tfx.orchestration.experimental.KubeflowV2DagRunner(
    config=tfx.orchestration.experimental.KubeflowV2DagRunnerConfig(
        default_image = 'gcr.io/watch-hop/hop-tfx-covid:0.6.2'
    ),
    output_filename=PIPELINE_DEFINITION_FILE)

_ = runner.run(
    create_pipeline(
        pipeline_name=PIPELINE_NAME,
        pipeline_root=PIPELINE_ROOT,
        data_path=DATA_ROOT, metadata_path=METADATA_PATH))

similar question 已在 Stack Overflow 上得到答复，这使我找到了 set memory requirements in AI Platform 的方法，但这些配置在 KubeflowV2DagRunnerConfig 中不再存在，所以我陷入了死胡同。

任何帮助将不胜感激。

** 编辑 **
我们使用 @component 装饰器将我们的组件定义为 python 函数，因此它们中的大多数都是自定义组件。对于训练组件，我知道您可以使用 tfx.Trainer 类指定机器类型，如 this tutorial 中所述，但我的问题是针对未进行任何训练的自定义组件。

【问题讨论】：

标签： tfx google-cloud-ai-platform-pipelines

【解决方案1】：

事实证明你现在不能，但根据这个issue，这个功能即将推出。

另一种解决方案是将您的 TFX 管道转换为 Kubeflow 管道。 Vertex AI 管道支持 kubeflow，您可以使用这些管道在组件级别设置内存和 cpu 约束。

@component // imported from kfp.dsl
def MyComponent(Input[Dataset] input_data):
  // ...

@pipeline // imported from kfp.dsl
def MyPipeline(...):
  component = MyComponent(...)
  component.set_memory_limit('64G') // alternative to set_memory_request(...)

【讨论】：

哦，伙计，今天花了我半天时间 :( 显然 Vertex AI 需要 set_memory_limit 和 set_cpu_limit，但不是 _request 对应项。它在提供带有 ram >= what's in 的机器类型方面仍然表现正确set_memory_limit。因此，对于 Vertex AI 管道，正确的代码应该是：component = MyComponent(...) component.set_memory_limit('64G')
谢谢，我会更新的。

【解决方案2】：

此解决方案的另一种选择是使用数据流光束运行器，它允许组件通过 Vertex 运行数据流集群。我仍在寻找一种方法来为自定义组件指定机器类型

样本光束输入：

BIG_QUERY_WITH_DIRECT_RUNNER_BEAM_PIPELINE_ARGS = [
--project=  GOOGLE_CLOUD_PROJECT,
--temp_location= GCS_LOCAITON,
--runner=DataflowRunner

]

现在您将迁移到 Vertex AI

【讨论】：