Dask cuDF 中的 memory 用完

Question

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

我首先设置一个 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 集群，但我两者都有。 另外，直观上过滤操作不应该是 memory 昂贵，所以我不知道该怎么做。 我认为我设置集群的方式有问题，修复它会使 rest 的更多 memory 昂贵的操作希望也能正常工作。

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

Answer 1

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

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

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