使用 DataFlow PubSub 到 Cloud Storage 很慢

Question

我正在使用以下示例片段的略微調整版本將 PubSub 消息寫入 GCS：

 class WriteToGCS(beam.DoFn): def __init__(self, output_path, prefix): self.output_path = output_path self.prefix = prefix def process(self, key_value, window=beam.DoFn.WindowParam): """Write messages in a batch to Google Cloud Storage.""" start_date = window.start.to_utc_datetime().date().isoformat() start = window.start.to_utc_datetime().isoformat() end = window.end.to_utc_datetime().isoformat() shard_id, batch = key_value filename = f'{self.output_path}/{start_date}/{start}-{end}/{self.prefix}-{start}-{end}-{shard_id:03d}' with beam.io.gcsio.GcsIO().open(filename=filename, mode="w") as f: for message_body in batch: f.write(f"{message_body},".encode("utf-8"))

然而，它的速度非常慢。 這就是它在圖中的樣子。 有沒有辦法加快這一步？ 訂閱每秒獲得 500 個元素，因此每秒 3-10 個元素跟不上。

管道如下所示：

 class JsonWriter(beam.PTransform): def __init__(self, window_size, path, prefix, num_shards=20): self.window_size = int(window_size) self.path = path self.prefix = prefix self.num_shards = num_shards def expand(self, pcoll): return ( pcoll | "Group into fixed windows" >> beam.WindowInto(window.FixedWindows(self.window_size, 0)) | "Decode windowed elements" >> beam.ParDo(Decode()) | "Group into batches" >> beam.BatchElements() | "Add key" >> beam.WithKeys(lambda _: random.randint(0, int(self.num_shards) - 1)) | "Write to GCS" >> beam.ParDo(WriteToGCS(self.path, self.prefix)) )

Answer 1

通過以下步驟解決了這個問題：

 | "Add Path" >> beam.ParDo(AddFilename(self.path, self.prefix, self.num_shards)) | "Group Paths" >> beam.GroupByKey() | "Join and encode" >> beam.ParDo(ToOneString()) | "Write to GCS" >> beam.ParDo(WriteToGCS())

首先，我添加每個元素的路徑，然后按路徑分組（以防止同時寫入同一個文件/分片）。 然后我連接（新行分隔）每個組（路徑）的 JSON 字符串，然后將整個 string 寫入文件。 這極大地減少了對 GCS 的調用，因為我不會將每個元素一個一個地寫入文件，而是先將它們連接起來，然后再寫入一次。 希望它可以幫助某人。