在Slurm中通过Python脚本调用srun的性能考量与最佳实践

#!/bin/bash
#SBATCH --job-name=python_srun_test
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=00:05:00
#SBATCH --output=slurm-%j.out

echo "Starting myscript.sh on $(hostname)"
echo "Slurm job ID: $SLURM_JOB_ID"

# 激活conda环境或设置Python环境（如果需要）
# source /path/to/your/conda/etc/profile.d/conda.sh
# conda activate my_hpc_env

# 执行Python脚本，Python脚本将调用srun
python running.py "Hello HPC!"

echo "myscript.sh finished."

import subprocess
import sys
import os

def run_hpc_workload(message):
    """
    通过srun调用一个简单的并行HPC程序。
    这里以一个简单的MPI程序为例，实际可以是任何并行应用。
    """
    # 假设你有一个名为 'my_mpi_program' 的MPI可执行文件
    # 并且它位于PATH中或者指定了完整路径
    hpc_program = "./my_mpi_program" # 假设my_mpi_program在当前目录

    # 构建srun命令。注意：srun会继承sbatch分配的资源。
    # 这里我们只传递了程序名和参数。
    # 如果需要更精细的srun控制，可以在这里添加更多srun参数，
    # 但通常sbatch的参数已经足够。
    srun_command = [
        "srun",
        hpc_program,
        message # 传递给HPC程序的参数
    ]

    print(f"Python script is about to call srun: {' '.join(srun_command)}")

    try:
        # 使用subprocess.check_call执行srun命令
        # check_call会在命令返回非零退出码时抛出CalledProcessError
        subprocess.check_call(srun_command)
        print("srun command executed successfully.")
    except subprocess.CalledProcessError as e:
        print(f"Error calling srun: {e}", file=sys.stderr)
        sys.exit(e.returncode)
    except FileNotFoundError:
        print(f"Error: '{hpc_program}' not found. Make sure it's compiled and in the correct path.", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    # 获取从bash脚本传递的参数
    if len(sys.argv) > 1:
        param = sys.argv[1]
    else:
        param = "Default Message"

    print(f"Python script running on host: {os.uname().nodename}")
    print(f"Received parameter: {param}")

    run_hpc_workload(param)
    print("Python script finished its orchestration role.")

#include <mpi.h>
#include <stdio.h>
#include <string.h>

int main(int argc, char** argv) {
    MPI_Init(&argc, &argv);

    int world_size;
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);

    int world_rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);

    char processor_name[MPI_MAX_PROCESSOR_NAME];
    int name_len;
    MPI_Get_processor_name(processor_name, &name_len);

    char message[256];
    if (argc > 1) {
        strncpy(message, argv[1], sizeof(message) - 1);
        message[sizeof(message) - 1] = '\0';
    } else {
        strcpy(message, "No message provided.");
    }

    printf("Hello from processor %s, rank %d out of %d processes. Message: %s\n",
           processor_name, world_rank, world_size, message);

    MPI_Finalize();
    return 0;
}