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

Question

我正在使用TFX在Vertex AI上构建一个AI管道。我已经按照本教程开始了，然后我将管道调整为我自己的数据，该数据具有超过100M行的时间序列数据。由于内存问题，我的几个组件在中途被终止，所以我只想为这些组件设置内存需求。我使用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平台中设置内存需求的方法，但这些配置在KubeflowV2DagRunnerConfig中已经不存在了，所以我走到了死胡同。

任何帮助都将不胜感激。

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

Answer 1

事实证明，目前还不能，但根据这个问题，这个功能即将推出。

另一种解决方案是将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(...)

Answer 2

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

样本光束输入：

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

]

现在您将迁移到Vertex AI