#!/bin/bash

#SBATCH --job-name={{ job_name }}
#SBATCH --nodes={{ nodes }}
{% if gpus_per_node > 0 -%}
#SBATCH --gpus-per-node={{ gpus_per_node }}
{% endif -%}
#SBATCH --ntasks-per-node={{ ntasks_per_node }}
{% if nodelist -%}
#SBATCH --nodelist={{ nodelist }}
{% endif -%}
{% if partition -%}
#SBATCH --partition={{ partition }}
{% endif -%}
#SBATCH --output={{ log_dir }}/%x_%j.log
#SBATCH --error={{ log_dir }}/%x_%j.log
#SBATCH --wait-all-nodes=1

set -euxo pipefail

# Environment setup
{% if environment_setup -%}
{{ environment_setup }}
{% endif -%}

# Working directory
{% if work_dir -%}
cd {{ work_dir }}
{% endif -%}

# Execute command
{% if srun_args or launch_prefix -%}
srun {{ ([srun_args, launch_prefix, command] | select | join(' ')) }}
{% else -%}
srun {{ command }}
{% endif -%}
