GPU作业开发

步骤一：开发并上传作业

1. 编写基于CUDA访问GPU的示例程序。

import torch
import numpy as np
from pyspark.sql import SparkSession

def check_executor_gpu_env(records):
    #将数据转化为torch类型
    data = torch.from_numpy(np.array(records)) 
    try:
        if torch.cuda.is_available():
            device = torch.device("cuda")
        else:
            raise Exception('CUDA is not available in executor')
        #将数据写入GPU
        on_gpu = data.to(device) 
        print("gpu device -- {}, sum of data -- {}".format(on_gpu.device, torch.sum(on_gpu)))
    except Exception as e:
        print(e)

if __name__ == "__main__":
    if torch.cuda.is_available():
        print('CUDA is available in driver')
    else:
        print('CUDA is not available in driver')
    spark = SparkSession.builder.appName("TEST_GPU_ENV").getOrCreate()
    rdd = spark.sparkContext.parallelize(range(100), 1)
    rdd.foreachPartition(check_executor_gpu_env)
    spark.stop()

2. 将示例程序上传至HDFS。具体操作，请参见通过控制台上传文件。

步骤二：打包并上传CUDA运行环境

打包项目依赖的Python环境和第三方类库。推荐使用Conda或Virtualenv将依赖类库打包为tar包，此处以Conda为例。

1. 在Linux环境中创建3.7~3.10版本的Python环境，本文以3.10.10版本为例。

conda create --name gpu_env python=3.10.10
conda activate gpu_env

2. 安装GPU相关的依赖包。

conda install pytorch=1.11.0 cudatoolkit=11.3 -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch
conda install cudnn
conda install numpy

3. 打包Conda环境，并命名为gpu_env.tar.gz。

conda install conda-pack=0.7.0 -c conda-forge
conda-pack -f -o gpu_env.tar.gz

4. 将gpu_env.tar.gz上传至HDFS 。具体操作，请参见通过控制台上传文件。

如果您想通过Virtualenv打包项目依赖，请参见Python Package Management。

步骤三：提交GPU作业

Lindorm计算引擎支持以下两种方式提交并管理作业。

• 通过Lindorm控制台提交作业，具体操作，请参见通过控制台管理作业。

• 通过DMS提交作业，具体操作，请参见通过DMS管理作业。

提交作业时，请配置与GPU相关的Spark参数：