#!/bin/bash
## Example slurm prolog script that can be used to reject job allocation and drain nodes based on health status.
## It exits the node with the number of errors.
## Running healthagent prolog tests.

response=$(/usr/bin/health -p -c gpumemorycheck gpu_id=$SLURM_JOB_GPUS -c gpudiagnosticcheck gpu_id=$SLURM_JOB_GPUS)
sum=$(echo $response | jq '[.gpu | to_entries[] | select(.value.status == "Error") | .value.error_count // 1] | add')
if [[ $sum -ge 1 ]] ; then
    exit $sum
else
    exit 0
fi