Dask cuDF执行ETL过滤操作时触发GPU内存溢出问题求助

• 硬件环境：搭载15GiB显存的Tesla T4 GPU
• 业务场景：运行多步ETL处理流程，近期绝大多数过滤、转换类操作（含少量shuffle步骤）均运行失败
• 待处理数据集：共20个单文件大小为500MB的parquet文件
• 问题背景：已持续排查dask_cudf内存管理问题较长时间，始终未定位根因，本次以触发OOM的过滤步骤为例说明问题，该类过滤操作理论内存开销极低。

集群初始化配置与对应报错 #

集群初始化代码如下：

CUDA_VISIBLE_DEVICES = os.environ.get("CUDA_VISIBLE_DEVICES", "0")

cluster = LocalCUDACluster(
#     rmm_pool_size=get_rmm_size(0.6 * device_mem_size()),
    CUDA_VISIBLE_DEVICES=CUDA_VISIBLE_DEVICES,
    local_directory=os.path.join(WORKING_DIR, "dask-space"),
    device_memory_limit=parse_bytes("12GB")
)
client = Client(cluster)
client

rmm_pool_size参数配置不同时触发的报错存在差异：

• 传入rmm_pool_size参数时：报错提示超出最大池上限
• 不传入rmm_pool_size参数时：抛出如下显存分配错误

MemoryError: std::bad_alloc: CUDA error at: ../include/rmm/mr/device/cuda_memory_resource.hpp:70: cudaErrorMemoryAllocation out of memory

过滤逻辑实现 #

过滤规则为判断数据列col的取值是否存在于包含约8万个有效值的集合中，实现代码如下：

def remove_invalid_values_filter_factory(valid_value_set_or_series):
    def f(df):
        mask = df['col'].isin(valid_value_set_or_series)
        return df.loc[mask]
    return f

# 从其他文件加载有效值
valid_values_info_df = pd.read_csv(...)
# 传入的Series大小仅约1MiB
keep_known_values_only = remove_invalid_values_filter_factory(valid_values_info_df['values'])
# 已尝试传入Python set形式的有效值，两种传参方式均触发报错
# keep_known_values_only = remove_invalid_values_filter_factory(set(valid_values_info_df['values']))

已验证两种有效值传入形式（pandas Series、Python set）均会触发OOM，其中使用的Series大小仅1MiB，不存在参数本身占用过高显存的问题。

报错触发场景 #

读取parquet数据、执行过滤并写出结果的环节触发报错，对应执行代码如下：

%%time
# 报错发生在该处理步骤执行过程中
keep_known_values_only(
    dask_cudf.read_parquet(...)
).to_parquet(...)

诉求 #

已按规范配置LocalCUDACluster与Dask Client，目前怀疑集群配置存在问题，需要可落地的解决方案，修复该问题后需支撑后续更高内存开销的操作正常运行。

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