DoFn 构造了多少次？答案

【问题标题】：How many times are DoFn's constructed?DoFn 构造了多少次？
【发布时间】：2019-09-09 16:22:51
【问题描述】：

我正在使用 apache beam python SDK 和 Dataflow 编写推理管道，以使用 TensorFlow 模型进行预测。我在 DoFn 中有预测步骤，但我不想每次处理捆绑包时都加载模型，因为这非常昂贵。来自文档here，“如果需要，将在工作人员上创建参数 DoFn 的新实例，并在此实例上调用 DoFn.Setup 方法。这可能是通过反序列化或其他方式。PipelineRunner 可以重用 DoFn多个捆绑包的实例。异常终止（通过抛出异常）的 DoFn 将永远不会被重用。我注意到如果我这样写代码

class StatefulGetEmbeddingsDoFn(beam.DoFn):
    def __init__(self, model_dir):
         self.model = None # initialize
         self.model_dir = model_dir

    def process(self, element):
         if not self.model: # load model if model hasn't been loaded yet
             global i
             i += 1
             logging.info('Getting model: {}'.format(i))
             self.model = Model(saved_model_dir=self.model_dir)


         ids, b64 = element
         embeddings = self.model.predict(b64)

         res = [
            {
                'image': _id,
                'embeddings': embedding.tolist()
            } for _id, embedding in zip(ids, embeddings)
         ]
         return res

似乎每个工人都不止一次地加载了模型（我有一个大约 30-40 台机器的集群）。有没有办法防止模型被多次加载？我原以为这个 DoFn 只能在每台机器上构建一次，但从日志来看，情况似乎并非如此......

【问题讨论】：

标签： python apache-beam dataflow

【解决方案1】：

我知道这是一个较老的问题，但我最初的想法是使用 setup 和 start_bundle 方法。

https://beam.apache.org/releases/pydoc/2.22.0/apache_beam.transforms.core.html#apache_beam.transforms.core.DoFn.setup

【讨论】：