DirectRunner运行含Splittable DoFn的流式管线报StageNode时间戳错误

问题现象 #

运行基于Splittable DoFn（SDF）生成器的简单流式管线时，Python版本的DirectRunner在SDF生成恰好1条输入后抛出错误：
AssertionError: wrong timestamp for StageNode<inputs=['ref_PCollection_PCollection_3_split'],side_inputs=[].

最小复现代码 #

使用apache_beam.transforms.periodicsequence.ImpulseSeqGenDoFn作为SDF实现：

pipeline_options = PipelineOptions(pipeline_args, streaming=True, save_main_session=True)

start = Timestamp.now()
stop = start + Duration(seconds=(24 * 60 * 60))
interval = 5.0

pipeline = Pipeline(options=pipeline_options)
(pipeline
 | "Create"    >> Create([(start, stop, interval)])
 | "Generator" >> ParDo(ImpulseSeqGenDoFn())
 | "Print"     >> ParDo(lambda x: logging.info(f"GEN: {x}")))

pipeline.run().wait_until_finish()

使用python3 generator.py --runner DirectRunner执行即可复现。

报错栈信息 #

...
INFO:root:GEN: 1632697296.421499
Traceback (most recent call last):
  File "generator.py", line 157, in <module>
    r = pipeline.run()
  File "/usr/local/lib/python3.6/dist-packages/apache_beam/pipeline.py", line 564, in run
    return self.runner.run_pipeline(self, self._options)
  File "/usr/local/lib/python3.6/dist-packages/apache_beam/runners/direct/direct_runner.py", line 131, in run_pipeline
    return runner.run_pipeline(pipeline, options)
  File "/usr/local/lib/python3.6/dist-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 196, in run_pipeline
    pipeline.to_runner_api(default_environment=self._default_environment))
  File "/usr/local/lib/python3.6/dist-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 206, in run_via_runner_api
    return self.run_stages(stage_context, stages)
  File "/usr/local/lib/python3.6/dist-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 385, in run_stages
    runner_execution_context, bundle_context_manager)
  File "/usr/local/lib/python3.6/dist-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 668, in _run_stage
    bundle_context_manager.stage.name))
AssertionError: wrong timestamp for StageNode<inputs=['ref_PCollection_PCollection_3_split'],side_inputs=[]. 

相关排查结论 #

• 完全相同的管线在Dataflow runner上运行符合预期，可按指定间隔生成消息。
• 报错对应的断言逻辑位于/apache_beam/runners/portability/fn_api_runner/fn_runner.py:668，与watermark（水位线）相关：

assert (runner_execution_context.watermark_manager.get_stage_node(
            bundle_context_manager.stage.name).output_watermark()
                < timestamp.MAX_TIMESTAMP), (
            'wrong timestamp for %s. '
            % runner_execution_context.watermark_manager.get_stage_node(
            bundle_context_manager.stage.name))

• 本次使用的ImpulseSeqGenDoFn没有显式操作水位线，也没有被@beam.DoFn.unbounded_per_element()装饰，可能并非专为流式管线设计，将streaming设为False运行也会触发相同错误。
• 当在管线末尾追加一个扩展为iobase.Sink的步骤后，该问题消失，SDF可按正确间隔生成输出，可正常运行的管线如下：

pipeline_options = PipelineOptions(pipeline_args, streaming=True, save_main_session=True)

start = Timestamp.now()
stop = start + Duration(seconds=(24 * 60 * 60))
interval = 5.0

pipeline = Pipeline(options=pipeline_options)
(pipeline
 | "Create"    >> Create([(start, stop, interval)])
 | "Generator" >> ParDo(ImpulseSeqGenDoFn())
 | "Print"     >> ParDo(lambda x: logging.info(f"GEN: {x}"))
 | 'Write'     >> WriteToText("generated.out"))

pipeline.run().wait_until_finish()

待确认疑问 #

该问题属于DirectRunner的缺陷，还是构建基于SDF的流式管线的方式存在错误？

内容的提问来源于stack exchange，提问作者Sha