haproxy/dev/patchbot/scripts/submit-ai.sh

80 lines
2.5 KiB
Bash
Executable File

#!/bin/bash
# note: the program may re-execute itself: when it has more than one patch to
# process, it will call itself with one patch only in argument. When called
# with a single patch in argument, it will always start the analysis directly.
# The program uses several environment variables:
# - EXT file name extension for the response
# - MODEL path to the model file (GGUF format)
# - FORCE force to re-process existing patches
# - PROGRAM path to the script to be called
# - CACHE path to the prompt cache (optional)
# - CACHE_RO force cache to remain read-only
# - PROMPT_PFX path to the prompt prefix (before the patch)
# - PROMPT_SFX path to the prompt suffix (after the patch)
# - TOT_CPUS total number of usable CPUs (def: nproc or 1)
# - SLOT_CPUS if defined, it's an array of CPU sets for each running slot
# - CPU_SLOT passed by the first level to the second one to allow binding
# to a specific CPU set based on the slot number from 0 to N-1.
die() {
[ "$#" -eq 0 ] || echo "$*" >&2
exit 1
}
err() {
echo "$*" >&2
}
quit() {
[ "$#" -eq 0 ] || echo "$*"
exit 0
}
#### Main
# detect if running under -x, pass it down to sub-processes
#opt=; set -o | grep xtrace | grep -q on && opt=-x
USAGE="Usage: ${0##*/} [ -s slots ] patch..."
MYSELF="$0"
TOT_CPUS=${TOT_CPUS:-$(nproc)}
TOT_CPUS=${TOT_CPUS:-1}
SLOTS=1
while [ -n "$1" -a -z "${1##-*}" ]; do
case "$1" in
-s) SLOTS="$2" ; shift 2 ;;
-h|--help) quit "$USAGE" ;;
*) die "$USAGE" ;;
esac
done
[ -n "$EXT" ] || die "Missing extension name (EXT)"
[ -n "$MODEL" ] || die "Missing model name (MODEL)"
[ -n "$PROGRAM" ] || die "Missing program name (PROGRAM)"
[ -n "$PROMPT_PFX" ] || die "Missing prompt prefix (PROMPT_PFX)"
[ -n "$PROMPT_SFX" ] || die "Missing prompt suffix (PROMPT_SFX)"
PATCHES=( "$@" )
if [ ${#PATCHES[@]} = 0 ]; then
die "$USAGE"
elif [ ${#PATCHES[@]} = 1 ]; then
# really execute
taskset_cmd=""
if [ -n "$CPU_SLOT" ] && [ -n "${SLOT_CPUS[$CPU_SLOT]}" ]; then
taskset_cmd="taskset -c ${SLOT_CPUS[$CPU_SLOT]}"
fi
export CPU=$TOT_CPUS
${taskset_cmd} ${PROGRAM} "${PATCHES[0]}"
else
# divide CPUs by number of slots
export TOT_CPUS=$(( (TOT_CPUS + SLOTS - 1) / SLOTS ))
# reexecute ourselves in parallel with a single patch each
xargs -n 1 -P "${SLOTS}" --process-slot-var=CPU_SLOT "${MYSELF}" -s 1 <<< "${PATCHES[@]}"
fi