使用Docker在GPU上运行CUBLAS的详解

引言

随着深度学习和大数据的迅猛发展，GPU的使用越来越普遍。NVIDIA的CUDA库，尤其是CUBLAS（CUDA Basic Linear Algebra Subprograms），在高性能计算中扮演着重要角色。结合Docker的使用，可以有效地部署和管理GPU加速的应用。本文将深入探讨如何在Docker中使用GPU资源以及运行CUBLAS。

为什么使用Docker?

Docker是一种轻量级的虚拟化技术，可以隔离应用及其依赖环境。其优势包括：

  1. 环境一致性：无论在开发、测试还是生产环境，Docker容器能提供一致的运行环境。
  2. 资源隔离：不同的应用可以在独立的容器中运行，互不影响。
  3. 便捷的部署：通过Docker Hub等平台，用户可以快速获取和分发容器镜像。

安装Docker和NVIDIA驱动

在使用Docker管理GPU资源之前，我们需要安装Docker和NVIDIA的驱动程序。

安装NVIDIA驱动

首先，确保你的系统上安装了适合你GPU型号的NVIDIA驱动。可以使用以下命令检查GPU状态：

nvidia-smi

如果成功显示图形信息，则驱动安装成功。

安装Docker

然后，按照以下命令来安装Docker：

# 更新软件包列表
sudo apt-get update

# 安装Docker
sudo apt-get install -y docker.io

# 启动Docker
sudo systemctl start docker

# 设置Docker开机自启
sudo systemctl enable docker

安装NVIDIA Docker支持

为了使Docker容器能够访问GPU，我们需要安装NVIDIA Container Toolkit：

# 添加NVIDIA包存储库
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L  -s -L  | sudo apt-key add -

curl -s -L  -s -L  | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

# 安装nvidia-docker2
sudo apt-get update
sudo apt-get install -y nvidia-docker2

# 重启Docker
sudo systemctl restart docker

创建Docker镜像

我们可以创建一个简单的Dockerfile，来安装CUBLAS库及其依赖。

Dockerfile示例

# 基础镜像使用NVIDIA的CUDA镜像
FROM nvidia/cuda:11.2-cudnn8-runtime-ubuntu20.04

# 安装依赖
RUN apt-get update && \
    apt-get install -y build-essential

# 复制CUBLAS的源代码
COPY ./cublas_project /workspace/cublas_project

# 设置工作目录
WORKDIR /workspace/cublas_project

# 编译CUBLAS示例代码
RUN nvcc -o cublas_example cublas_example.cu -lcublas

在这个示例中，我们使用了NVIDIA官方的CUDA基础镜像，安装了必要的构建工具，并复制了一个示例CUBLAS项目。

运行CUBLAS示例程序

完成镜像创建后，我们可以通过以下命令构建并运行Docker容器：

# 构建Docker镜像
docker build -t cublas-demo .

# 运行Docker容器
docker run --gpus all cublas-demo ./cublas_example

当容器成功运行时，应能看到CUBLAS示例程序的输出。

CUBLAS代码示例

下面是一个简单的CUBLAS示例代码（cublas_example.cu），该示例实现了两个向量的加法：

#include <iostream>
#include <cublas_v2.h>
#include <cuda_runtime.h>

void vectorAdd(const float *A, const float *B, float *C, int N) {
    cublasHandle_t handle;
    cublasCreate(&handle);

    const float alpha = 1.0f;
    cublasSaxpy(handle, N, &alpha, A, 1, C, 1);

    cublasDestroy(handle);
}

int main() {
    int N = 10;
    float A[N], B[N], C[N];

    // 初始化向量
    for(int i = 0; i < N; ++i) {
        A[i] = static_cast<float>(i);
        B[i] = static_cast<float>(i);
        C[i] = 0.0f;
    }

    // 调用CUBLAS
    vectorAdd(A, B, C, N);

    // 输出结果
    for(int i = 0; i < N; ++i) {
        std::cout << "C[" << i << "] = " << C[i] << std::endl;
    }
    
    return 0;
}

关系图

为了更好地理解Docker与GPU及CUBLAS之间的关系，下面是一个简单的ER图示例：

erDiagram
    DOCKER {
        string containerId
        string image
    }
    GPU {
        string gpuId
        string memory
    }
    CUBLAS {
        string version
        string functionality
    }
    
    DOCKER ||--o{ GPU : uses
    DOCKER ||--o{ CUBLAS : runs

结论

通过以上步骤，我们成功地展示了如何在Docker中使用GPU资源并运行CUBLAS库。利用Docker容器化技术，不仅提高了开发效率，还增强了应用的可移植性。GPU加速的线性代数计算能力为深度学习等高性能计算任务提供了强有力的支持。在今后的研究和开发中，人员可以更灵活地使用这项技术，以适应不断变化的数据科学和机器学习环境。