HeterogeneousCore/MPICore/scripts/cmsenv_mpirun - master - githubmirror/cmssw

HeterogeneousCore/MPICore/scripts/cmsenv_mpirun
118 строк4 KB
Mario Gonzalez
introduced cmsenv_ssh, modified cmsenv_mpirun
18 дек 2025, 15:57
18 дек 2025, 15:570edf2f9
#! /bin/bash


function usage() {
    echo "Usage: $0 [-v, --verbose] [-h, --help] [-b, --release-base <path>] [--] <mpirun options>"
    echo
    echo "Options:"
    echo "  -v, --verbose              Make this script verbose"
    echo "  -h, --help                 Show this help message"
    echo "  -b, --release-base <path>  Remote CMSSW release base path (default: local CMSSW_BASE, CMSSW_RELEASE_BASE or CMSSW_FULL_RELEASE_BASE)"
    echo "  --                         Separates cmsenv_mpirun options from mpirun options"
    exit 0
}

VERBOSE=false
REMOTE_RELEASE_BASE=""
ARG_COUNT=$#

# Parse any arguments specific to this script
while [[ $# -gt 0 ]]; do
    case "$1" in
        -h|--help)
            usage
            ;;
        -v|--verbose)
            VERBOSE=true
            shift
            ;;
        -b|--release-base)
            REMOTE_RELEASE_BASE="$2"
            shift 2
            ;;
        --)
            shift
            break
            ;;
        *)
            if [[ $((ARG_COUNT - $#)) -gt 0 ]]; then
                # if cmsenv_mpirun specific options have been passed, throw an error (-- is missing)
                echo "Error: unrecognised cmsenv_mpirun option '$1'. Did you forget '--' between the cmsenv_mpirun options and the mpirun options?"
                usage
            else
                # no cmsenv_mpirun-specific options passed, "--" is not required
                break
            fi
            ;;
    esac
done

# Remaining arguments are mpirun options
MPIRUN_PREDICATE="$@"


# Set RELEASE_BASE to a path where cmsenv_ssh is found
RELEASE_BASE=
for VAR in CMSSW_BASE CMSSW_RELEASE_BASE CMSSW_FULL_RELEASE_BASE; do
  TEST_PATH=${!VAR}
  if [ -x "$TEST_PATH/bin/$SCRAM_ARCH/cmsenv_ssh" ]; then
    RELEASE_BASE=$TEST_PATH
    break
  fi
done
if [ -z "$RELEASE_BASE" ]; then
  echo "Error: cannot find cmsenv_ssh in the \$PATH."
  exit 1
fi

# if REMOTE_RELEASE_BASE is not set, default to local RELEASE_BASE
if [[ -z "$REMOTE_RELEASE_BASE" ]]; then
    REMOTE_RELEASE_BASE=$RELEASE_BASE
fi

# arguments to cmsenv_ssh. The -- at the end is needed to separate them from 
# subsequent mpirun arguments
SSH_WRAPPER_ARGS=" --release-base $REMOTE_RELEASE_BASE --"


# Get MPI implementation and major version number
MPIRUN_VERSION="$(mpirun --version)"
SSH_WRAPPER_OPTION=""
SSH_WRAPPER_ARGS_OPTION=""
if [[ "$MPIRUN_VERSION" == *"Open MPI"* ]]; then
    MAJOR_VERSION="$(echo "$MPIRUN_VERSION" | awk '/Open MPI/ {print substr($4,1,1)}')"

    if [[ "$MAJOR_VERSION" == 5 ]]; then
        SSH_WRAPPER_OPTION="--prtemca plm_ssh_agent"
        SSH_WRAPPER_ARGS_OPTION="--prtemca plm_ssh_args '${SSH_WRAPPER_ARGS}'"

    elif [[ "$MAJOR_VERSION" == 4 ]]; then
        SSH_WRAPPER_OPTION="--mca plm_rsh_agent"
        SSH_WRAPPER_ARGS_OPTION="--mca plm_rsh_args '${SSH_WRAPPER_ARGS}'"
    else
        echo "Error: unsupported Open MPI major version $MAJOR_VERSION."
        exit 1
    fi
    
elif [[ "$MPIRUN_VERSION" == *"HYDRA"* ]]; then
    SSH_WRAPPER_OPTION="-launcher-exec"
    SSH_WRAPPER_ARGS_OPTION="-env HYDRA_LAUNCHER_EXTRA_ARGS='${SSH_WRAPPER_ARGS}'"
fi

if [[ -z "$SSH_WRAPPER_OPTION" ]]; then
    echo "Error: could not determine MPI implementation from 'mpirun --version' output:"
    echo "$MPIRUN_VERSION"
    exit 1
fi


MPI_CMD="mpirun $SSH_WRAPPER_OPTION $RELEASE_BASE/bin/$SCRAM_ARCH/cmsenv_ssh $SSH_WRAPPER_ARGS_OPTION $MPIRUN_PREDICATE"

if [[ "$VERBOSE" == true ]]; then
    echo "--------------------------------------------------------------------------------"
    echo "MPI command:"
    echo "$MPI_CMD"
    echo "--------------------------------------------------------------------------------"
fi

eval $MPI_CMD