GPU 指南¶
Mars 可以利用 NVIDIA 显卡执行,但对于 Mars 里的不同模块,需要一些必要的依赖。
安装¶
Mars tensor 依赖于 CuPy,假设你的 CUDA 驱动是 10.1,用如下方式安装 cupy:
pip install cupy-cuda101
参考 安装 cupy 获取更多信息。
Mars DataFrame 依赖于 RAPIDS cuDF,使用 conda 安装 cuDF:
conda install -c rapidsai -c nvidia -c conda-forge \
-c defaults cudf=0.13 python=3.7 cudatoolkit=10.1
参考 安装 cuDF 获取更多信息。
CUDA 上运行 Mars tensor¶
Tensor 通过指定 gpu=True
来指定在 GPU 上创建。创建 Tensor 和 随机抽样 中的方法都支持这个参数。
>>> import mars.tensor as mt
>>> a = mt.random.rand(10, 10, gpu=True) # indicate to create tensor on CUDA
>>> a.sum().execute() # execution will happen on CUDA
记住,创建 tensor 的时候,并没有实际的 GPU 内存分配。当 .execute()
触发的时候,GPU 上才会有真正的内存分配和计算。
对于一个主存上创建的 tensor,调用 .to_gpu()
来指示把数据移至 GPU。
>>> b = mt.random.rand(10, 10) # indicate to create on main memory
>>> b = b.to_gpu() # indicate to move data to GPU memory
>>> b.sum().execute()
调用 .to_cpu()
来指示把数据移到主存。
>>> c = b.to_cpu() # b is allocated on GPU, move back to main memory
>>> c.sum().execute() # execution will happen on CPU
CUDA 上运行 Mars DataFrame¶
Mars 可以直接读取 CSV 文件至显存。
>>> import mars.dataframe as md
>>> df = md.read_csv('data.csv', gpu=True) # indicates to read csv into GPU memory
>>> df.groupby('a').sum().execute() # execution will happen on GPU
对于主存上的 DataFrame,调用 .to_gpu()
指示把数据移至显存。
>>> import mars.tensor as mt
>>> import mars.dataframe as md
>>> df = md.DataFrame(mt.random.rand(10, 10)) # indicate to create on main memory
>>> df = df.to_gpu() # indicate to move data to GPU memory
调用 .to_cpu()
来指示把数据移到主存。
>>> df2 = df.to_cpu() # df is allocated on GPU, move back to main memory
>>> df2.sum().execute() # execution will happen on CPU
单卡¶
本地基于多线程的调度器 能在单卡上正常运行。以上的例子都可以在单卡上工作。
多卡¶
对于 Mars tensor,可以利用单机多卡。
>>> import mars.tensor as mt
>>> t = mt.random.rand(10000, 10000, gpu=True)
>>> t.sum().execute()
如上代码会利用所有可见的显卡来计算。
如果要限制使用的显卡,可以设置环境变量 CUDA_VISIBLE_DEVICES
。
CUDA_VISIBLE_DEVICES=0,3,5 ipython
这会让 ipython 只使用 0、3、5 显卡。因此在 ipython 中跑的 Mars tensor 任务只会使用这些显卡。
对于 Mars DataFrame,单机基于线程的调度器不能利用多卡来计算,在这种情况下,请使用分布式调度器。
分布式¶
Mars scheduler 和 web 的启动命令完全相同,参考 在集群中部署。
Mars worker 只能绑定到一张显卡上,因此如果你想使用多张显卡,启动和显卡个数一样多的 worker。
绑定到某个显卡的 worker 可以按如下命令启动:
mars-worker -a <worker_ip> -p <worker_port> -s <scheduler_ip>:<scheduler_port> --cuda-device 0
这个 worker 会绑定到显卡 0。
参考 启动 Worker 的其他命令 获取更多信息。
一旦一个集群创建好了,可以运行如下代码。
>>> import mars.tensor as mt
>>> from mars.session import new_session
>>> new_session('http://<web_ip>:<web_port>').as_default()
>>> t = mt.random.rand(20, 20, gpu=True)
>>> t.sum().execute() # run on workers which are bind to GPU