1

我正在使用 TFX 在 Vertex AI 上构建 AI 管道。我已经按照本教程开始，然后我将管道调整为我自己的数据，其中包含超过 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))

Stack Overflow 上已经回答了一个类似的问题，这让我找到了一种在 AI Platform 中设置内存要求的方法，但是这些配置在 中不再存在KubeflowV2DagRunnerConfig，所以我陷入了死胡同。

任何帮助将非常感激。

** 编辑 **
我们将组件定义为带有@component装饰器的 python 函数，因此它们中的大多数是自定义组件。对于培训组件，我知道您可以使用本教程tfx.Trainer中解释的类指定机器类型，尽管我的问题是针对不进行任何培训的自定义组件。

4

2 回答 2

4

事实证明你现在不能，但根据这个问题，这个功能即将到来。

另一种解决方案是将 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(...)
于 2021-07-02T17:17:12.623 回答
0

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

样品光束输入：

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

]

到目前为止，您将迁移到 Vertex AI

于 2021-12-20T06:20:57.213 回答