如何并行执行自定义Splittable DoFn答案

【问题标题】：How to execute custom Splittable DoFn in parallel如何并行执行自定义Splittable DoFn
【发布时间】：2022-01-04 07:13:16
【问题描述】：

我正在尝试为 Apache Beam 开发一个用 Python 编写的自定义 I/O 连接器。根据official guideline，Splittable DoFn (SDF) 是我的首选框架。

我尝试运行SDF programming guide 中的伪代码，但是我未能并行执行管道。下面是一个工作示例。

虚拟数据

myfile = open('test_beam.txt', 'w')
for i in range(0, 1000):
    myfile.write("%s\n" % i)

myfile.close

管道

确保将DUMMY_FILE 替换为test_beam.txt 的绝对路径。

import argparse
import logging
import os

import apache_beam as beam

from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions


from time import sleep
import random

from apache_beam.io.restriction_trackers import OffsetRange

DUMMY_FILE = absolute_path_to_dummy_data_file

class FileToWordsRestrictionProvider(beam.transforms.core.RestrictionProvider
                                     ):
    def initial_restriction(self, file_name):
        return OffsetRange(0, os.stat(file_name).st_size)

    def create_tracker(self, restriction):
        return beam.io.restriction_trackers.OffsetRestrictionTracker(
            offset_range=self.initial_restriction(file_name=DUMMY_FILE))

    def restriction_size(self, element, restriction):
        return restriction.size()


class FileToWordsFn(beam.DoFn):
    def process(
        self,
        file_name,
        # Alternatively, we can let FileToWordsFn itself inherit from
        # RestrictionProvider, implement the required methods and let
        # tracker=beam.DoFn.RestrictionParam() which will use self as
        # the provider.
            tracker=beam.DoFn.RestrictionParam(FileToWordsRestrictionProvider())):
        with open(file_name) as file_handle:
            file_handle.seek(tracker.current_restriction().start)
            while tracker.try_claim(file_handle.tell()):
                yield read_next_record(file_handle=file_handle)


def read_next_record(file_handle):
    line_number = file_handle.readline()
    logging.info(line_number)
    sleep(random.randint(1, 5))
    logging.info(f'iam done {line_number}')


def run(args, pipeline_args, file_name):
    pipeline_options = PipelineOptions(pipeline_args)

    with beam.Pipeline(options=pipeline_options) as p:
        execute_pipeline(args, p, file_name)


def execute_pipeline(args, p, file_name):
    _ = (
          p |
          'Create' >> beam.Create([file_name]) |
          'Read File' >> beam.ParDo(FileToWordsFn(file_name=file_name))
    )


if __name__ == '__main__':
    logging.getLogger().setLevel(logging.INFO)

    """Build and run the pipeline."""
    parser = argparse.ArgumentParser()
    # to be added later

    args, pipeline_args = parser.parse_known_args()

    file_name = DUMMY_FILE
    run(args, pipeline_args, file_name)

SDF 取自第一个示例 here，但是，我必须修复一些问题（例如，定义 restriction_size 和 () 的小错位）。此外，我在read_next_record 中引入了随机睡眠，以检查管道是否并行执行（显然不是）。

我构建管道的方式可能有错误？我希望使用我的 SDF 作为管道的第一步，但这样做会导致 AttributeError: 'PBegin' object has no attribute 'windowing'。为了规避这个问题，我遵循了这个post 并添加了一个包含输入file_name 的PCollection。

在管道中并行执行 SDF 的正确方法是什么？

【问题讨论】：

标签： python apache-beam

【解决方案1】：

光束DoFns（包括SplittableDoFns）对输入PCollection进行操作。对于SplittableDoFn，输入通常是源配置的PCollection（例如，输入文件）。当执行SplittableDoFn 时，Beam runner 甚至可以通过使用RestrictionTracker 隔离部分输入读取来并行执行单个输入元素。因此，对于一个文件，这意味着您可能有并行运行的工作人员从同一个文件中读取数据，但偏移量不同。

所以你的实现看起来是正确的，应该已经促进了 Beam runner 的并行执行。

【讨论】：