import os
import socket
import multiprocessing

# 获取当前主机名
host = socket.gethostname()

# 获取当前进程编号
rank = int(os.environ.get("SLURM_PROCID", "0"))

# 获取进程数量
size = int(os.environ.get("SLURM_NTASKS", "1"))

# 计算任务分配
n = 10000000
start = int(n / size) * rank
end = start + int(n / size)

# 执行计算任务
for i in range(start, end):
    # TODO: 在这里编写具体的计算任务代码
    pass

# 输出计算结果
print("Process %d on %s: %d" % (rank, host, result))

# 等待所有进程执行完毕
if rank == 0:
    print("All processes finished.")


#!/bin/bash
#SBATCH --job-name=my_job
#SBATCH --partition=compute
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=00:10:00

# 进入脚本所在的目录
cd /path/to/script

# 加载Python环境
module load python/3.8.5

# 启动MPI进程
srun python my_script.py