#!/bin/bash

set -euo pipefail

CONTAINER_NAME=${CONTAINER_NAME:-layer_wise_benchmarks}
TRTLLM_ROOT=$(realpath -- "$(dirname -- "$0")"/../..)
CONTAINER_MOUNTS="$TRTLLM_ROOT:$TRTLLM_ROOT"

if [ -z "${SLURM_JOB_ID:-}" ]; then
    echo "Please set SLURM_JOB_ID"
    exit 1
fi

WORKDIR=$(realpath -- "$(pwd)")

set -x
srun --mpi=pmix \
    -N "$NODES" \
    --ntasks-per-node $((NP / NODES)) \
    --container-name "$CONTAINER_NAME" \
    --container-mounts "$CONTAINER_MOUNTS" \
    --container-workdir "$WORKDIR" \
    "$@"