python - 如何将大型 python 模型应用于 pyspark-dataframe？

Question

我有:

• 包含一些数据(特征)的大型数据框(parquet 格式，100.000.000 行，4.5TB 大小)
• 几个巨大的 ML 模型(每个模型需要 5-15GB 的内存)
• Spark 集群 (AWS EMR)，典型节点配置为 8 个 CPU、32 个 RAM，可以根据需要进行更改。

我想使用 PySpark 应用它们，但我总是会遇到一些连线错误，例如:

• OOM
• 随机超时(节点不返回任何结果)-> 节点被 YARN 管理器杀死

我通常使用如下代码

def apply_model(partition):
    model = load(...)  # load model only when apply this function to avoid serialization issue
    for row in partition:
        yield model.infer(row)

或

def apply_model(partition):
    model = load(...)  # load model only when apply this function to 
    yield from model.infer(partition)

并应用它使用

df.select(...).rdd.mapPartitions(apply_model)

由于序列化原因，我无法广播模型。

问题 - 如何将基于 python/任何非 jvm 的大型模型应用于 spark dataframe 并避免 spark 异常？

Answer 1

以下是一些有助于提高工作绩效的额外建议:

• 我要做的第一个改变是减小分区大小。如果我现在理解正确的话，你输入的数据是 4.5TB。这意味着如果您有 1000 个分区，那么您最终将在每个执行程序上为每个分区发送 4.5GB!这个尺寸算是相当large ，相反，我会尝试将分区大小保持在 250-500MB 之间。在您的情况下，这大约意味着 ~10000 (4.5TB/500MB) 个分区。

• 通过添加更多执行程序来提高并行度。这将提高数据水平locality从而减少执行时间。理想情况下，每个执行程序应该有 5 个核心，每个集群节点有两个执行程序(如果可能)。每个执行程序的最大核心数不应高于 5，因为这会导致 I/O 瓶颈(当/如果使用磁盘存储时)。

• 至于内存，我认为@rluta 的建议绰绰有余。一般来说，执行程序内存的值太大会对 Java GC 时间产生负面影响，因此 10GB 的上限应该是 spark.executor.memory 的理想值。