单机器并发运行cuml K-Means.fit_predict性能骤降问题求助

性能骤降的核心原因 #

1. CUDA上下文切换开销：每个KMeans实例初始化时都会创建CUDA上下文，并发运行时多个实例频繁切换上下文，再加上你的函数中每个实例要执行18次KMeans（每次n_init=10），上下文切换的累积开销被大幅放大，直接拖慢整体速度。
2. GPU计算资源竞争：RTX3090的SM（流式多处理器）数量有限，并发的KMeans实例会抢占SM资源，导致每个实例的计算任务被频繁打断，无法持续利用GPU的并行能力。
3. 隐式数据转换开销：每次调用fit_predict时，pandas DataFrame会被隐式转换为cuDF格式，并发时主机到GPU的数据传输会排队，累积延迟不可忽视。

具体优化建议 #

1. 显式预转换数据集为cuDF #

提前将pandas DataFrame转换为cuDF格式，避免每次fit_predict时重复执行数据拷贝和格式转换：

import cudf
import numpy as np
from cuml.cluster import KMeans
from cuml.metrics.cluster import silhouette_samples

# 预转换数据集，仅执行一次
cudf_df = cudf.from_pandas(df)

def max_silhouette_score(cudf_df):
    sil_scores = []
    test_range = range(2,20)
    for k in test_range:
        kmeans = KMeans(n_clusters=k, n_init=10)
        predictions = kmeans.fit_predict(cudf_df)
        # 直接用cuDF数据计算轮廓系数，避免来回转换
        sil = silhouette_samples(cudf_df, predictions)
        sil_scores.append(float(sil.mean()))
    max_sil_idx = np.argmax(sil_scores)
    max_sil = sil_scores[max_sil_idx]
    max_sil_k = list(test_range)[max_sil_idx]
    return max_sil, max_sil_k

2. 控制GPU资源的并发使用 #

• 设置n_streams参数：允许多个KMeans实例共享GPU流，减少上下文切换开销：

  kmeans = KMeans(n_clusters=k, n_init=10, n_streams=2)
  

• 调整n_init参数：如果业务对聚类稳定性要求可放宽，适当降低n_init的值（比如从10降到5），减少每个KMeans实例的计算量，降低并发时的资源竞争。

3. 使用进程池而非线程池管理并发 #

Python的GIL会限制线程的并行能力，对于GPU任务，使用multiprocessing进程池可以让每个进程拥有独立的CUDA上下文，减少上下文切换的冲突：

from multiprocessing import Pool

# 假设你有多个待处理的数据集
datasets = [cudf.from_pandas(df1), cudf.from_pandas(df2)]

if __name__ == "__main__":
    with Pool(processes=2) as pool:
        results = pool.map(max_silhouette_score, datasets)

4. 启用CUDA统一内存 #

通过设置cuML的内存分配器为统一内存，减少主机与GPU之间的数据拷贝次数：

import cuml
cuml.set_allocator("managed")

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