brun

#!/bin/bash
(( DEBUG ))  && set -x
function nop()
{
    return -1
}

declare -i VERBOSE=${VERBOSE:-0}
typeset -r -A QUEUES=( ['v100']='v100-localqueue'
		       ['a100']='a100-localqueue'
		       ['h100']='h100-localqueue'
		       ['none']='dummy-localqueue')

typeset -r -i CONTEXT=${CONTEXT:-1}
typeset -r DEFAULT_GPU='v100'
typeset -r GPU=${GPU:-"${DEFAULT_GPU}"}
# name must be lower case and no funny characters
typeset -r NAME=${NAME:-"job"}
typeset -r JOB_ID=${JOBID:-$$}
typeset -r JOB=${JOB:-"${NAME}-${GPU}-${JOB_ID}"}
typeset -r JOBS_DIR_NAME="${JOBS_DIR_NAME:-jobs}"
typeset  -r JOBS_DIR="${JOBS_DIR:-$(pwd)/${JOBS_DIR_NAME}}"
typeset -r CONTAINER="${JOB}-container"
typeset -i -r MAX_SEC="${MAX_SEC:-$(( 60 * 15 ))}"
typeset -i -r JOB_DELETE="${JOB_DELETE:-1}"
typeset -i -r WAIT="${WAIT:-1}"
typeset -i -r JOB_WAIT_TMOUT="${JOB_WAIT_TMOUT:-$(( MAX_SEC * 2 ))}"
typeset OUT=/dev/null 
typeset -i GPU_NUMREQ="${GPU_NUMREQ:-1}"
typeset -i GPU_NUMLIM="${GPU_NUMREQ:-1}"
typeset RESOURCES_YAML=""
typeset COMMAND_YAML=""
#typeset IMAGE=${IMAGE:-"image-registry.openshift-image-registry.svc:5000/redhat-ods-applications/minimal-gpu:2025.1"}
typeset IMAGE="${IMAGE:-"image-registry.openshift-image-registry.svc:5000/redhat-ods-applications/csw-run-f25:latest"}"
typeset CONTEXT_DIR="${CONTEXT_DIR:-$(pwd)}"
typeset OUTPUT_DIR="${JOBS_DIR}/${JOB}"
typeset -r GETLIST="${OUTPUT_DIR}/getlist"
typeset -i rc=0 complete_pid=0 failed_pid=0
typeset POD='' POD_INFO=''

# default behavior relies on HOSTNAME set to pod name of where we are running (true on openshift RHOAI)
typeset DEVPOD_NAME=${DEVPOD_NAME:-$HOSTNAME}
typeset DEVCONTAINER=$(oc get pods/${DEVPOD_NAME} -o=jsonpath="{.spec.containers[0].name}")



[[ $1 == -h ]] && {
    cat <<EOF
   brun [-h] <command line>
     brun creates and submits a batch job to a GPU batch queue.  The arguments are treated as a bash command line
     that will execute as the batch job.  The behaviour of the job submission can be controlled via several
     environment variables.  Please, feel free to improve this documentation and fix 
     bugs ;-). By default the files and subdirectories of your current working directory form a context for 
     the job.  The context is copied to the container in which the batch job will execute.  Thus the commandline
     can reference files in the directory.  Additinoally, files created in the working directory of commandline
     in the container will be copied back so that you can inspect output of your job (eg. profiles, logs and outputs).
     These files and directories are placed in a directory make in job specific subdirectory of a directory called jobs.
     Eg.

1. The simple usage run a binary that exsits in the current directory on the default GPU type
$ br ./hello
Hello from CPU
Hello from GPU
Hello from GPU
Hello from GPU
Hello from GPU
Hello from GPU
Hello from GPU
Hello from GPU
Hello from GPU
Hello from GPU
Hello from GPU
RUNDIR: jobs/job-v100-9215

Note by default brun waits for the job to complete and displays after the standard output and error of the command line.
After that it display the directory where the outputs for the job where copied.

See repository README.md for more documentation and examples.

EOF
    exit 0
}

(( VERBOSE )) && OUT=/dev/stderr

[[ ! -v QUEUES[$GPU] ]] && {
    echo "ERROR: unsupported GPU $GPU : no queue found" > /dev/stderr
    exit -1
}
typeset QUEUE=${QUEUES[$GPU]}

if [[ "$GPU" == 'none' ]]; then
    RESOURCES_YAML=\
"          resources:
            requests:
              cpu: \"1\"
              memory: \"1Gi\"
            limits:
              cpu: \"1\"
              memory: \"1Gi\"
"
else
    RESOURCES_YAML=\
"          resources:
            requests:
              nvidia.com/gpu: ${GPU_NUMREQ}
            limits:
              nvidia.com/gpu: ${GPU_NUMLIM}
"
fi

cmd="$@"

[[ -z "$cmd" ]] && {
    cat > /dev/stderr <<EOF
USAGE: ${0##*/} <command line>
  Runs the specified command on a GPU node
  Supported GPU types are:  ${!QUEUES[@]}
  Default GPU type is $DEFAULT_GPU
  To overide default set GPU enviroment variable eg.
     $ GPU=a100 ${0##*/} nvidia-smi
EOF
    exit 0
}

(( VERBOSE )) && {
    echo "GPU:${GPU} JOB:${JOB} CONTAINER:${CONTAINER} QUEUE:$QUEUE" \
	 "MAX_SEC:${MAX_SEC} WAIT:${WAIT} JOB_WAIT:${JOB_WAIT} JOB_DELETE:${JOB_DELETE}" \
	 "CONTEXT_DIR:${CONTEXT_DIR} OUTPUT_DIR=${OUTPUT_DIR}" > /dev/stderr
    echo "CMD: $cmd" > /dev/stderr
}

if (( CONTEXT )); then
    [[ ! -d $CONTEXT_DIR ]] && {
	echo "ERROR: CONTEXT_DIR: $CONTEXT_DIR is not a directory"
	exit -1
    }
    [[ -a ${OUTPUT_DIR} ]] && {
	echo "ERROR: $OUTPUT_DIR directory already exits"
	exit -1
    }

    ! mkdir -p ${OUTPUT_DIR}  && {
	echo "ERROR: Failed to make output dir"
	exit -1
    }
    
    # create list of context to send files subdirectories for
    # of $CONTEXT_DIR excluding the ${JOBS_DIR} if a direct
    # child of $CONTEXT_DIR
    (
	jdir="${JOBS_DIR#${CONTEXT_DIR}/}"
	[[ "$jdir" != "$JOBS_DIR" ]] && {
	    jdir="./${jdir}"
	}
	cd ${CONTEXT_DIR}
	find . -mindepth 1 -maxdepth 1 -not -path ${jdir} -prune  > ${GETLIST}
    )
    (( ! VERBOSE )) && {
	RSYNC_VERBOSE="-q"
    }
    COMMAND_YAML=$(cat <<EOF
          command: ["/bin/sh", 
                    "-c", 
                    "export RSYNC_RSH='oc rsh -c ${DEVCONTAINER}';
                     mkdir ${JOB} &&
                     rsync ${RSYNC_VERBOSE} --archive --no-owner --no-group --omit-dir-times --numeric-ids ${DEVPOD_NAME}:${GETLIST} ${JOB}/getlist >/dev/null 2>&1 &&
                     rsync ${RSYNC_VERBOSE}  -r --archive --no-owner --no-group --omit-dir-times --numeric-ids --files-from=${JOB}/getlist ${DEVPOD_NAME}:${CONTEXT_DIR}/ ${JOB}/ &&
                     find ${JOB} -mindepth 1 -maxdepth 1 > ${JOB}/gotlist &&
                     cd ${JOB} && ${cmd} |& tee $JOB.log; cd ..;
                     rsync ${RSYNC_VERBOSE} --archive --no-owner --no-group --omit-dir-times --no-relative --numeric-ids --exclude-from=${JOB}/gotlist ${JOB} ${DEVPOD_NAME}:${JOBS_DIR}"]
EOF
		   )
else
    COMMAND_YAML=$(cat <<EOF 
          command: ["/bin/sh", "-c", "${cmd}"]
EOF
		)
fi

YAML=$(cat <<EOF
apiVersion: batch/v1
kind: Job
metadata:
  name: ${JOB}
  labels:
    kueue.x-k8s.io/queue-name: ${QUEUE}
    test_name: kueue_test
spec:
  parallelism: 1
  completions: 1
  backoffLimit: 0
  activeDeadlineSeconds: ${MAX_SEC}
  template:
    spec:
      restartPolicy: Never
      maximumExecutionTimeSeconds: ${MAX_SEC}
      containers:
        - name: ${CONTAINER}
          image: ${IMAGE}
${COMMAND_YAML}
${RESOURCES_YAML}
EOF
)

(( VERBOSE )) && echo "$YAML"

oc apply -f - >$OUT 2>&1 <<< "$YAML"

(( VERBOSE )) && {
    echo "Job Created: $JOB" > $OUT
    oc describe job $JOB > $OUT 2>&1
}


if (( WAIT )); then
    trap nop SIGINT

    oc wait --for=condition=complete --timeout=${JOB_WAIT_TMOUT}s jobs/$JOB >$OUT 2>&1 & complete_pid=$!   
    oc wait --for=condition=failed --timeout=${JOB_WAIT_TMOUT}s jobs/$JOB >$OUT 2>&1 && exit -1 & failed_pid=$!

    echo "Waiting $JOB for condition: complete or failed" > $OUT
    wait -n $complete_pid $failed_pid
    rc=$?
    
    if ((  $rc == 0 )); then
	echo "$JOB: Completed" >$OUT
    else
	POD=$(oc get pods -l job-name=$JOB -o name)
	POD_INFO="$(oc describe $POD 2>&1)"
	echo "*** ERROR: $JOB FAILED!!!! *****"
	if [[ -n $POD ]]; then
	    echo "$POD:"
	    grep -E 'State:|Reason:|Exit Code:|Started:|Finished:' <<<"$POD_INFO"
	else
	    echo "ERROR: Execution exceeded $MAX_SEC ???"
	fi
    fi
    
    oc logs jobs/$JOB
    
    (( JOB_DELETE )) && { 
	oc delete jobs/$JOB >$OUT 2>&1
    }
    
    (( CONTEXT )) && {
	[[ -a ${OUTPUT_DIR}/getlist ]] && rm ${OUTPUT_DIR}/getlist
	echo RUNDIR: ${OUTPUT_DIR#$(pwd)/} > /dev/stderr
    }
else
    echo "WARNING: Not waiting for $JOB to finish you must manually delete with 'gel $JOB'" > /dev/stderr
    if (( CONTEXT )); then 
	echo $JOB ${OUTPUT_DIR#$(pwd)/}
    else
	echo $JOB
    fi
fi