dask_cudf-不遵守rmm配额

问题描述

我是机器学习和使用GPU的新手-出于这个原因，我对RAPID和dask感到兴奋。

我在AWS EC2 p3.8xlarge上运行。在它上面，我正在使用RAPIDs容器运行docker。我正在使用版本0.16。有一个60GB的EBS。

我有一个大约有8000万条记录的数据集。作为csv，它约为27GB，而作为拼花地板（具有少量功能）则为3.4GB（两种情况都在AWS S3上使用）。

尝试使用dask_cudf使用LocalCUDACluster时，我总是遇到崩溃的工人并遇到问题。创建核心转储并继续执行，从而创建了新的工作程序并最终占用了我计算机上的所有存储。

请参见下面的一些示例执行，显示内存增加，不遵守rmm_pool_size并最终崩溃。我尝试了rmm_pool_size的许多值，无论是在总GPU内存的上方还是下方（据我了解，它应该能够溢出到机器内存中）。

我正在使用以下初始代码：

from dask_cuda import LocalCUDACluster
from distributed import Client,LocalCluster
import dask_cudf


cluster = LocalCUDACluster(
    rmm_pool_size="60GB"  # I've tried 64,100,150 etc. No luck
)
# I also tried setting rmm_managed_memory... 
# I kNow there are other parameters (ucx,etc) but don't kNow whether relevant and how to use

client = Client(cluster)

df = dask_cudf.read_parquet("s3://my-bucket/my-parquet-dir/")

我打印内存使用情况：

mem = df.memory_usage().compute()
print(f"total dataset memory: {mem.sum() / 1024**3}GB")

产生

数据集总内存：50.736539436504245GB

然后，执行我的代码（无论是尝试执行一些EDA，运行KNN还是执行几乎所有其他操作，我都会遇到此行为/错误。

我阅读了文档，阅读了许多博客（主要来自RAPIDS），我遍历了笔记本，但仍然无法正常工作。难道我做错了什么？这不适用于我的设置吗？

任何帮助将不胜感激...

解决方法

设置RMM限制时，每个GPU。因此，如果您的目标是60GB，请将RMM设置为15。刚意识到您仅使用4个GPU

dask-distributed rapids