Python 多處理池 map - Databricks 中的 Memory 問題

Question

我在 Databricks 環境中運行一個 python 組件，它創建了一組 JSON 條消息，每條 JSON 條消息都使用 Avro 模式編碼。 編碼花費了更長的時間（編碼具有復雜 JSON 結構的 10K 消息需要 8 分鍾），因此我嘗試將多處理與池 map function 一起使用。該過程似乎在第一次執行時運行良好，但是對於后續運行，性能是降級並最終因 oom 錯誤而失敗。 我確保在執行結束時發出 pool.close() 和 pool.join() 但不確定它是否真的釋放了 memory。當我查看 Databricks Ganglia UI 時，它顯示交換 memory 和 CPU 利用率每次運行都在增加。 我還嘗試減少池的數量（驅動程序節點有 8 個內核，因此嘗試使用 6 個和 4 個池）以及 maxtasksperchild=1 但仍然無濟於事。 我想知道我是否做錯了什么。 以下是我現在使用的代碼。 想知道是什么引起了這里的問題。 任何指針/建議表示贊賞。

from multiprocessing import Pool
import multiprocessing
import json
from avro.io import *
import avro.schema
from avro_json_serializer import AvroJsonSerializer, AvroJsonDeserializer
import pyspark.sql.functions as F

def create_json_avro_encoding(row):
  row_dict = row.asDict(True)
  json_data = json.loads(avro_serializer.to_json(row_dict))
  #print(f"JSON created { multiprocessing.current_process().name }")
  return json_data 

avro_schema = avro.schema.SchemaFromJSONData(avro_schema_dict, avro.schema.Names())
avro_serializer = AvroJsonSerializer(avro_schema)
records = df.collect()
pool_cnt = int(multiprocessing.cpu_count()*0.5)
print(f"No of records: {len(records)}")
print(f"starting timestamp {datetime.now().isoformat(sep=' ')}")
with Pool(pool_cnt, maxtasksperchild=1) as pool:
  json_data_ret = pool.map(create_json_avro_encoding, records)
  pool.close()
  pool.join()

Answer 1

你不應該在加入之前關閉游泳池。 事實上，在 with 塊中使用它時根本不應該關閉池，它會在退出 with 塊時自動關閉。