Run python on a supercomputer

原创

flavorfan

发布于 2021-01-31 22:05:49

2.1K0

发布于 2021-01-31 22:05:49

引言

有任务需要处理一堆收集来得开源数据集，在服务器单机跑了一天才给结果，多方咨询有HPC可以用，或者叫supercomputer，或者叫计算机集群，大部分的简称grid。看了wiki、confluence，给出一堆链接在脑海中织出密密麻麻的蜘蛛网——无从下手。居然没有use case出发端到端的参考demo，真是无力吐槽。自力更生求助google，youtube，stack overflow，梳理下来，简而言之，可以理解分而治之多线程的多处理核（cpu/gpu）的版本，涉及算力资源调度引入slurm，涉及通讯引入mpi。

题目本来想起中文名，supercomputer翻译过来和python配合，总觉得有点违和。还是起个看过youtube视频的名字。

一个demo

使用随机数解决定量问题的数学方法通常称为蒙特卡洛方法。例如，考虑从圆与内接圆的正方形的面积之比来估计Pi的值的问题。

Serial, Pure python( baseline)

pi_serial.py

# 计算圆中点的个数
def monte_carlo_pi_part(n):
    count=0
    for i in range(n):
        x = random.random()
        y = random.random()

        if x * x + y* y <=1:
            count +=1
    return count

if __name__ ==  '__main__':
    print('Run across a single CPU')
    # Number of points to simulate (50 million)
    n = 5000000000

    count = monte_carlo_pi_part(n)

    print('Estimated value of Pi:: ', count / (n * 1) * 4)

使用的slurm提交脚本

#!/bin/bash
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --time=0-1:00:00

# Load required modules
module load Python/2.7.11-GCC-4.9.2

# Run python script
time python pi_serial.py

numba (加速30x)

仅仅导入numba库，用上@jit简单就能加速30倍。

from numba import jit

@jit #<--magic
def monte_carlo_pi_part(n)
    count=0
    for i in range(n):
        x = random.random()
        y = random.random()

        if x * x + y* y <=1:
            count +=1
    return count

NUMBA + Multiprocess (8 cpu) （再加速6x）

from multiprocessing import Pool
np = 8

# Number of points to simulate (50 million)
n = 5000000000

# Splite up across CPUs
part_count = [int(n /np) for i in range(np) ]

# Run simulation in parallel
pool = Pool(processes=np
count = pool.map(monte_calor_pi_part, part_count)

print('Estimated value of Pi:: ', sum(count) / (n * 1) * 4)

NUMBAA + MPI（100 CPU）(再加速1.5x）

from mpi4py import MPI

# who am I?
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# Number of points to simulate (50 million)
n = 5000000000

if rank == 0:
    master()
else:
    slave()


def master():
    points_per_node = int(n / (size-1))

    for worker in range(1,size):
        comm.send(points_per_node, dest=worker)
    
    received_processes = 0
    count = 0

    # Await results
    while received_processes < (size -1):
        count += comm.recv(source=MPI.ANY_SOURCE)
        received_processes += 1

    print('Estimated value of Pi:: ', count / (n * 1) * 4)

def slave():
    point_to_calculate = comm.rece(source=0)
    count = monte_carlo_pi_part(points_to_calculate)
    comm.send(count, dest=0)

最佳实践

在HPC上使用slurm运行python需要有自己的环境，有两种方式：

1）用pip --user 或者 conda 之类构建隔离的环境；

2）用singularity容器构建环境

推荐使用方式2）。后面就是