PySpark处理图像时内存占用异常问题排查求助

问题背景 #

在使用PySpark进行图像聚类、增量PCA建模时，发现内存占用远高于理论估算值：

• 单张700×800×3的uint8图像理论仅占1.68MB，但实际单张内存消耗约10MB
• 调用rdd.take(50)后内存消耗接近800MB，远高于估算的84MB
• 调整spark.driver.memory/spark.executor.memory、调用rdd.unpersist()均无法有效释放内存，导致内核崩溃或连接拒绝

复现代码 #

from pyspark import SparkContext, SparkConf
from psutil import virtual_memory
import numpy as np

conf = SparkConf().set('spark.driver.memory', '2g')
conf = conf.set('spark.executor.memory', '1g')
sc = SparkContext(conf=conf)

print('first:', virtual_memory().free/1000000, 'MB')
rdd = sc.parallelize(np.random.randint(low=0, high=256, size=(100, 700, 800, 3), dtype='uint8'))
print('second:', virtual_memory().free/1000000, 'MB')
vals = rdd.take(50)
print('third:', virtual_memory().free/1000000, 'MB')

内存输出示例 #

first: 2157.563904 MB
second: 1967.8208 MB
third: 1172.652032 MB

核心原因分析 #

1. Spark默认序列化的巨大开销 #

Spark默认使用Java序列化，对numpy数组的处理效率极低：

• numpy的uint8数组在Java中会被拆解为单个Byte对象存储，每个对象仅对象头就占8~16字节（取决于JVM指针压缩配置）
• 单张700×800×3的图像会生成1,680,000个Byte对象，仅这部分就占用13.44MB（1.68e6 × 8字节），加上元数据开销，实际内存占用远超原生数组大小

2. Driver端多份数据副本 #

当前代码先在Driver的Python进程中生成完整的numpy数组，再通过parallelize传入Spark JVM，导致：

• Python进程持有一份原始数组内存
• Spark JVM持有一份序列化后的数组内存
• 调用take()时，JVM需将数据反序列化后再传输回Python进程，生成第三份副本
  三者叠加导致内存消耗呈倍数增长

3. 内存观测的误区 #

用psutil查看的是系统全局空闲内存，但Spark的内存分为：

• Driver的JVM堆内存（由spark.driver.memory控制）
• Driver的Python进程内存
• Executor的JVM内存
  系统内存数值包含了所有进程的内存占用，无法准确反映Spark实际使用的内存，容易造成估算偏差

4. 内存释放的延迟 #

调用rdd.unpersist()仅标记RDD数据为可回收，实际释放需等待JVM的垃圾回收（GC）触发；而Python端的vals变量持有数据时，Python进程的内存也不会自动释放

解决方案 #

1. 启用Kryo高效序列化 #

Kryo序列化对原生数据类型和numpy数组的压缩、序列化效率远高于Java序列化，能大幅降低内存占用：

from pyspark import SparkContext, SparkConf
from psutil import virtual_memory
import numpy as np

conf = SparkConf().set('spark.driver.memory', '2g')
conf = conf.set('spark.executor.memory', '1g')
# 启用Kryo序列化
conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
# 注册numpy数组类型，让Kryo正确处理
conf.registerKryoClasses([np.ndarray])

sc = SparkContext(conf=conf)

2. 避免Driver端生成大数组 #

不要在Driver端生成完整图像数组再parallelize，改为直接在Executor端生成或读取文件：

• 若生成测试数据：用sc.parallelize(range(100)).map(lambda _: np.random.randint(0,256,size=(700,800,3),dtype='uint8'))，让每个Executor分区独立生成图像
• 若读取真实图像：用Spark读取二进制文件或专用图像处理库，避免数据经过Driver端

3. 准确观测内存使用 #

通过Spark UI（默认地址http://localhost:4040）查看：

• Storage标签页：RDD的实际内存占用、分区存储情况
• Executors标签页：Driver和Executor的JVM堆内存使用情况
  这能精准定位内存消耗的来源，避免依赖系统内存数值的误判

4. 正确释放内存 #

• Spark JVM端：调用rdd.unpersist()后，手动触发JVM GC（仅用于调试，生产环境不建议频繁调用）：

  rdd.unpersist()
  sc._jvm.System.gc()
  

• Python端：删除变量并触发Python GC：

  del vals
  import gc
  gc.collect()
  

5. 调整RDD存储级别 #

若不需要全量内存存储，改用序列化后的存储级别，减少内存占用：

from pyspark.storagelevel import StorageLevel

rdd = rdd.persist(StorageLevel.MEMORY_AND_DISK_SER)

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