问题描述
我的工作流极大地受益于 GPU 加速,但每个任务的内存要求相对较低(2-4 GB)。我正在使用 dask.dataframe
、dask.distributed.Client
和 dask_cuda.LocalCUDACluster
的组合。该过程将从更多的 CUDA 工作器中受益,因此我想将物理 GPU(Nvidia RTX A600、V100、A100)拆分为多个虚拟/逻辑 GPU,以增加我的 dask_cuda LocalCUDACluster
中的工作器数量。我最初的想法是尝试将在 TensorFlow
中创建的 logical_gpus 传递到 LocalCUDACluster
,但我似乎无法将它们传递到集群中。
我在 docker 环境中工作,我想将这些拆分保留在 python 中。理想情况下,此工作流程可以从本地工作站扩展到多节点 MPI 作业,但我不确定这是否可行,我愿意接受任何建议。
from dask.distributed import Client
from dask_cuda import LocalCUDACluster
from dask_cuda.initialize import initialize
import pandas as pd
import dask.dataframe as dd
import time
# fake function
def my_gpu_sim(x):
"""
GPU simulation which is independent of any others (calls a c++ program in real-world,which saves a
file.)
"""
...
return None
# fake data creation
dic = {'random':['apple' for i in range(40)],'main':[i for i in range(40)]}
df = pd.DataFrame.from_dict(dic)
ddf = dd.from_pandas(df,npartitions=4)
# Configurations
protocol = "ucx"
enable_tcp_over_ucx = True
enable_nvlink = True
enable_infiniband = False
initialize(
create_cuda_context=True,enable_tcp_over_ucx=enable_tcp_over_ucx,enable_infiniband=enable_infiniband,enable_nvlink=enable_nvlink,)
cluster = LocalCUDACluster(local_directory="/tmp/USERNAME",protocol=protocol,rmm_pool_size="35GB"
)
client = Client(cluster)
# Simulation
ddf.map_partitions(lambda df: df.apply(lambda x: my_gpu_sim(x.main),axis=1)).compute(scheduler=client)
解决方法
暂无找到可以解决该程序问题的有效方法,小编努力寻找整理中!
如果你已经找到好的解决方法,欢迎将解决方案带上本链接一起发送给小编。
小编邮箱:dio#foxmail.com (将#修改为@)