如何在Bash脚本中正确获取Intel MPI进程的Rank值？

我正在开发一款Bash脚本，通过MPI在超级计算机上为每个处理器分配文件集进行压缩。使用Intel 2019模块，通过mpiexec.hydra启动MPI作业，但脚本中rank=$OMPI_COMM_WORLD_RANK无法获取进程Rank值，该环境变量返回空值，导致无法正确分配文件任务。

脚本原代码如下：

#!/bin/bash

# Source directory containing subfolders
SOURCE_DIR="../data"

# Destination directory to store compressed files
DEST_DIR="../zipFiles"

# Create the destination directory if it doesn't exist
mkdir -p "$DEST_DIR"

# Count the number of subfolders
SUBFOLDER_COUNT=$(find "$SOURCE_DIR" -maxdepth 1 -type d -not -path "$SOURCE_DIR" | wc -l)

# Print the number of subdirectories:
echo "The number of subdirectories is $SUBFOLDER_COUNT"

# Set the number of subfolders to compress in each batch
echo "The maximum number of subdirectories per batch is $DirPerBatch"

# COMPRESSED FILE NUMBER
k=1

cd "$SOURCE_DIR"

rank=$OMPI_COMM_WORLD_RANK

# Calculate the start and end index of files for this rank
start=$((rank * DirPerBatch))
end=$((start + DirPerBatch - 1))
echo "DBG:      rank:$rank"
echo "DBG:      sta: $start"
echo "DBG:      end: $end"
# Initialize the 'files' variable
files=""

# Loop to gather files for this rank
for ((i = start; i <= end && i < SUBFOLDER_COUNT; i++)); do
    subfolder_index=$i
    subfolder_name=$(printf "%05d" "$subfolder_index")

    # Check if the subfolder exists before adding it
    if [ -d "$subfolder_name" ]; then
        files+="$subfolder_name "  # Use += to append to 'files'
    fi
done

# Create the compressed file for this rank
tar czf $DEST_DIR/$((rank + 1)).tar.gz $files
if [ $? -eq 0 ]; then
    echo "Rank $rank: Successfully compressed $rank.tar.gz"
else
    echo "Rank $rank: Failed to compress $rank.tar.gz"
fi

注：变量DirPerBatch在作业提交脚本中已声明并导出。

核心问题：环境变量不匹配 #

OMPI_COMM_WORLD_RANK是Open MPI的环境变量，而你使用的是Intel MPI的mpiexec.hydra启动器，对应的Rank环境变量是PMI_RANK。只需将脚本中的Rank获取语句替换为：

rank=$PMI_RANK

优化后的完整脚本 #

为了增强鲁棒性，补充了Rank变量为空时的默认值、空文件集判断等细节：

#!/bin/bash

# Source directory containing subfolders
SOURCE_DIR="../data"

# Destination directory to store compressed files
DEST_DIR="../zipFiles"

# Create the destination directory if it doesn't exist
mkdir -p "$DEST_DIR"

# Count the number of subfolders
SUBFOLDER_COUNT=$(find "$SOURCE_DIR" -maxdepth 1 -type d -not -path "$SOURCE_DIR" | wc -l)

# Print the number of subdirectories:
echo "The number of subdirectories is $SUBFOLDER_COUNT"

# Set the number of subfolders to compress in each batch
echo "The maximum number of subdirectories per batch is $DirPerBatch"

cd "$SOURCE_DIR"

# Intel MPI Hydra 使用 PMI_RANK 获取进程Rank
rank=$PMI_RANK

# 单进程测试时的 fallback，避免变量为空导致计算错误
if [ -z "$rank" ]; then
    rank=0
    echo "Warning: PMI_RANK not detected, defaulting to rank 0"
fi

# Calculate the start and end index of files for this rank
start=$((rank * DirPerBatch))
end=$((start + DirPerBatch - 1))
echo "DBG:      rank:$rank"
echo "DBG:      sta: $start"
echo "DBG:      end: $end"

# Initialize the 'files' variable
files=""

# Loop to gather files for this rank
for ((i = start; i <= end && i < SUBFOLDER_COUNT; i++)); do
    subfolder_index=$i
    subfolder_name=$(printf "%05d" "$subfolder_index")

    # Check if the subfolder exists before adding it
    if [ -d "$subfolder_name" ]; then
        files+="$subfolder_name "
    fi
done

# 仅当有文件需要压缩时执行tar命令，避免空输入报错
if [ -n "$files" ]; then
    tar czf "$DEST_DIR/$((rank + 1)).tar.gz" $files
    if [ $? -eq 0 ]; then
        echo "Rank $rank: Successfully compressed $((rank + 1)).tar.gz"
    else
        echo "Rank $rank: Failed to compress $((rank + 1)).tar.gz"
    fi
else
    echo "Rank $rank: No subfolders assigned for compression"
fi

关键修改点 #

1. 替换Rank环境变量：用PMI_RANK替代OMPI_COMM_WORLD_RANK，适配Intel MPI Hydra启动器
2. 添加Rank变量校验：当PMI_RANK未设置时（比如单进程调试），默认设为0，避免算术错误
3. 空文件集判断：如果当前Rank没有分配到子文件夹，跳过tar命令并给出提示
4. 路径加双引号：给tar的输出路径添加双引号，兼容路径含空格的场景

内容的提问来源于stack exchange，提问作者Ahmed Adam