构建基于容器的组件

基于容器的组件提供了将任何语言编写的代码集成到您的管道中的灵活性,只要您可以在 Docker 容器中执行该代码即可。

如果您不熟悉 TFX 管道,请了解有关 TFX 管道核心概念的更多信息

创建基于容器的组件

基于容器的组件由容器化的命令行程序支持。如果您已经拥有容器镜像,则可以使用 TFX 通过使用create_container_component 函数来声明输入和输出,从容器镜像中创建组件。函数参数

  • name: 组件的名称。
  • inputs: 一个字典,将输入名称映射到类型。outputs: 一个字典,将输出名称映射到类型 parameters: 一个字典,将参数名称映射到类型。
  • image: 容器镜像名称,以及可选的镜像标签。
  • command: 容器入口点命令行。不在 shell 中执行。命令行可以使用在编译时用输入、输出或参数替换的占位符对象。占位符对象可以从tfx.dsl.component.experimental.placeholders导入。请注意,Jinja 模板不受支持。

返回值: 一个继承自 base_component.BaseComponent 的 Component 类,可以实例化并在管道中使用。

占位符

对于具有输入或输出的组件,command 通常需要具有在运行时被实际数据替换的占位符。为此提供了几个占位符

  • InputValuePlaceholder: 输入工件值的占位符。在运行时,此占位符将被替换为工件值的字符串表示形式。

  • InputUriPlaceholder: 输入工件参数 URI 的占位符。在运行时,此占位符将被替换为输入工件数据的 URI。

  • OutputUriPlaceholder: 输出工件参数 URI 的占位符。在运行时,此占位符将被替换为组件应存储输出工件数据的 URI。

了解有关TFX 组件命令行占位符的更多信息。

基于容器的组件示例

以下是一个非 Python 组件的示例,它可以下载、转换和上传数据。

import tfx.v1 as tfx

grep_component = tfx.dsl.components.create_container_component(
    name='FilterWithGrep',
    inputs={
        'text': tfx.standard_artifacts.ExternalArtifact,
    },
    outputs={
        'filtered_text': tfx.standard_artifacts.ExternalArtifact,
    },
    parameters={
        'pattern': str,
    },
    # The component code uses gsutil to upload the data to Google Cloud Storage, so the
    # container image needs to have gsutil installed and configured.
    image='google/cloud-sdk:278.0.0',
    command=[
        'sh', '-exc',
        '''
          pattern="$1"
          text_uri="$3"/data  # Adding suffix, because currently the URI are "directories". This will be fixed soon.
          text_path=$(mktemp)
          filtered_text_uri="$5"/data  # Adding suffix, because currently the URI are "directories". This will be fixed soon.
          filtered_text_path=$(mktemp)

          # Getting data into the container
          gsutil cp "$text_uri" "$text_path"

          # Running the main code
          grep "$pattern" "$text_path" >"$filtered_text_path"

          # Getting data out of the container
          gsutil cp "$filtered_text_path" "$filtered_text_uri"
        ''',
        '--pattern', tfx.dsl.placeholders.InputValuePlaceholder('pattern'),
        '--text', tfx.dsl.placeholders.InputUriPlaceholder('text'),
        '--filtered-text', tfx.dsl.placeholders.OutputUriPlaceholder('filtered_text'),
    ],
)