Dask cuDF 内存不足答案

【问题标题】：Running out of memory in Dask cuDFDask cuDF 内存不足
【发布时间】：2022-07-02 07:36:28
【问题描述】：

最近一段时间，我一直试图在我最近的项目中解决 dask_cudf 中的内存管理问题，但似乎我遗漏了一些东西，我需要你的帮助。我正在开发具有 15 GiB 内存的 Tesla T4 GPU。我有几个 ETL 步骤，但 GPU 最近似乎在大多数步骤上都失败了（其中大多数只是过滤或转换步骤，但很少有旋转洗牌）。我的数据包含大约 20 500MB parquet 文件。对于这个特定的问题，我将提供一段用于过滤的代码，这会使 GPU 由于内存不足而失败。

我首先设置一个 CUDA 集群：

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 参数，错误会有所不同。当提供参数时，我得到超过最大池限制，否则我得到以下错误： MemoryError: std::bad_alloc: CUDA error at: ../include/rmm/mr/device/cuda_memory_resource.hpp:70: cudaErrorMemoryAllocation out of memory

接下来，我创建了一个我打算对数据执行的过滤操作（它围绕检查列中的值是否出现在包含大约 80000 个值的集合中）：

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

# Load valid values from another file
valid_values_info_df = pd.read_csv(...)
# The series is around 1 MiB in size
keep_known_values_only = remove_invalid_values_filter_factory(valid_values_info_df[\'values\'])
# Tried both and both cause the error
# keep_known_values_only = remove_invalid_values_filter_factory(set(valid_values_info_df[\'values\']))

最后我对数据应用这个过滤操作并得到错误：

%%time
# Error occures during this processing step
keep_known_values_only(
    dask_cudf.read_parquet(...)
).to_parquet(...)

我感到完全迷失了，由于在没有 Dask 的情况下使用 cuDF 或未设置 CUDA 集群，我遇到的大多数来源都有此错误，但我两者都有。另外，从直觉上讲，过滤操作不应该占用大量内存，所以我不知道该怎么做。我认为我设置集群的方式有问题，修复它会使其他更昂贵的内存操作也能正常工作。

我会很感激你的帮助，谢谢！

标签： python python-3.x gpu dask rapids

【解决方案1】：

为此，我将使用 dask-sql 来利用它的核心处理能力。

至于 dask_cudf 功能失败，请在 cudf repo 中提出问题，并具有最低可重现性！我们将不胜感激！ :)

你可能不想一起做 dask_cudf 和 RMM，除非你真的必须知道你在做什么（这就像 RAPIDS 超级使用模式，当你需要真正最大化用于算法的 GPU 大小时）。如果您的使用需要这样做（并且在您使用镶木地板文件时它似乎并不在这里，这就是我没有深入研究它的原因），它真的可以提供帮助。

【讨论】：