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

Question

我正在使用TFX在顶点AI上构建一个AI管道。我跟随本教程开始，然后我将管道调整到我自己的数据中，它有超过1亿行的时间序列数据。由于内存问题，我的两个组件中途被杀死，所以我只想为这些组件设置内存需求。我使用KubeflowV2DagRunner编排并启动顶点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))

堆栈溢出已经回答了一个相似问题，这使我找到了一种通向在AI平台上设置内存需求的方法，但是这些信任在KubeflowV2DagRunnerConfig中已经不存在了，所以我处于死胡同。

任何帮助都将不胜感激。

编辑**

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

Answer 1

但是根据这个问题，这个特性即将到来。

另一种解决方案是将TFX管道转换为Kubeflow管道。顶点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

此解决方案的另一种选择是使用数据流波束流道，它允许组件通过顶点运行数据流集群。我仍然需要找到一种为自定义组件指定机器类型的方法。

样本光束输入：

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

]

现在，您将迁移到顶点AI。