Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

EVA-3575 - sbatch wrapper script #170

Merged
merged 3 commits into from
May 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -28,3 +28,4 @@ nb-configuration.xml
.DS_Store
*.log
venv/
.vscode/
109 changes: 109 additions & 0 deletions slurm-scripts/eva-sbatch
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
#!/bin/bash
set -euo pipefail

if [[ $# == 0 ]] || [[ "$1" == "help" ]] || [[ "$1" == "--help" ]]; then
echo "eva-sbatch - EVA's sbatch wrapper"
echo
echo "Usage: eva-sbatch <job_name> [optional sbatch args] [run|dryrun] [code to run]"
echo
echo "Example: eva-sbatch python-version run python --version"
echo
echo "The following defaults are used:"
echo "-t 168:00:00 --mem=6G -J <job_name> -o <job_name>.out -e <job_name>.err"
echo
echo "The entire job script will be written to <job_name>.job, and then run with sbatch unless dryrun is used."
echo
echo "For full sbatch documentation, see https://slurm.schedmd.com/sbatch.html"
exit 0
fi

if [[ $# -lt 3 ]]; then
echo "Usage: eva-sbatch <job_name> [optional sbatch args] [run|dryrun] [code to run]"
echo "Example: eva-sbatch python-version run python --version"
exit 1
fi

# Job name is required to be the first argument
JOB_NAME="$1"
shift

# Defaults
WALL_TIME="168:00:00" # one week
MEMORY="6G"
OUT_FILE="${JOB_NAME}.out"
ERR_FILE="${JOB_NAME}.err"

# Iterate through args until you reach "run" or "dryrun"
# Everything beforehand is an SBATCH directive
OTHER_ARGS=()
while (($#)); do
if [[ "$1" == "run" ]] || [[ "$1" == "dryrun" ]]; then
break

# Overwrite defaults when needed
elif [[ "$1" == -t ]]; then
WALL_TIME="$2"
shift; shift
elif [[ "$1" == --time* ]]; then
WALL_TIME="${1#*=}"
shift

elif [[ "$1" == --mem* ]]; then # note mem doesn't have a short form
MEMORY="${1#*=}"
shift

elif [[ "$1" == -o ]]; then
OUT_FILE="$2"
shift; shift
elif [[ "$1" == --output* ]]; then
OUT_FILE="${1#*=}"
shift

elif [[ "$1" == -e ]]; then
ERR_FILE="$2"
shift; shift
elif [[ "$1" == --error* ]]; then
ERR_FILE="${1#*=}"
shift

# Otherwise save to OTHER_ARGS
else
OTHER_ARGS+=("
#SBATCH $1")
shift
while (($#)); do
if [[ "$1" == -* ]] || [[ "$1" == "run" ]] || [[ "$1" == "dryrun" ]]; then
break
fi
OTHER_ARGS+=("$1")
shift
done
fi

done

# Everything afterwards is the command to submit
RUN_OR_DRYRUN="$1"
shift
COMMAND=$@

# Write the job file
cat << EOF > ${JOB_NAME}.job
#!/bin/bash

#SBATCH -t ${WALL_TIME}
#SBATCH --mem=${MEMORY}
#SBATCH -J ${JOB_NAME}
#SBATCH -o ${OUT_FILE}
#SBATCH -e ${ERR_FILE}
${OTHER_ARGS[*]}

${COMMAND}
EOF

# Run the job if requested
if [[ $RUN_OR_DRYRUN == "run" ]]; then
sbatch ${JOB_NAME}.job
else
echo "Run using: sbatch ${JOB_NAME}.job"
fi
76 changes: 76 additions & 0 deletions slurm-scripts/eva-srun
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
#!/bin/bash
set -euo pipefail

if [[ $# == 0 ]] || [[ "$1" == "help" ]] || [[ "$1" == "--help" ]]; then
echo "eva-srun - EVA's srun wrapper, for interactive jobs"
echo
echo "Usage: eva-srun [optional srun args] [run|dryrun] [code to run]"
echo
echo "Example: eva-srun -p datamover run python --version"
echo
echo "The following defaults are used:"
echo "-t 00:30:00 --mem=4G"
echo
echo "The command will be run using srun, unless dryrun is used in which case it will echo the command."
echo
echo "For full srun documentation, see https://slurm.schedmd.com/srun.html"
exit 0
fi

if [[ $# -lt 2 ]]; then
echo "Usage: eva-srun [optional srun args] [run|dryrun] [code to run]"
echo "Example: eva-srun -p datamover run python --version"
exit 1
fi

# Defaults
WALL_TIME="00:30:00"
MEMORY="4G"

# Iterate through args until you reach "run" or "dryrun"
# Everything beforehand is an srun argument
OTHER_ARGS=()
while (($#)); do
if [[ "$1" == "run" ]] || [[ "$1" == "dryrun" ]]; then
break

# Overwrite defaults when needed
elif [[ "$1" == -t ]]; then
WALL_TIME="$2"
shift; shift
elif [[ "$1" == --time* ]]; then
WALL_TIME="${1#*=}"
shift

elif [[ "$1" == --mem* ]]; then # note mem doesn't have a short form
MEMORY="${1#*=}"
shift

# Otherwise save to OTHER_ARGS
else
OTHER_ARGS+=("$1")
shift
while (($#)); do
if [[ "$1" == -* ]] || [[ "$1" == "run" ]] || [[ "$1" == "dryrun" ]]; then
break
fi
OTHER_ARGS+=("$1")
shift
done
fi

done

# Everything afterwards is the command to submit
RUN_OR_DRYRUN="$1"
shift
COMMAND=$@

SRUN_CMD="srun -t ${WALL_TIME} --mem=${MEMORY} ${OTHER_ARGS[*]} --pty ${COMMAND}"

# Run the command if requested
if [[ $RUN_OR_DRYRUN == "run" ]]; then
${SRUN_CMD}
else
echo "Run using: ${SRUN_CMD}"
fi
Loading