#!/bin/bash

# SPDX-FileCopyrightText: 2026 Samudra Authors
#
# SPDX-License-Identifier: Apache-2.0

#SBATCH --job-name=%j-samudra-perlmutter-om4-speed-inlineinf-torch
#SBATCH --output=train/logs/slurm-%j.out
#SBATCH --error=train/logs/slurm-%j.err
#SBATCH --constraint=gpu&hbm80g
#SBATCH --account=m4874
#SBATCH --qos=debug
#SBATCH --mem=110G
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=2
#SBATCH --cpus-per-task=32
#SBATCH --gpus-per-task=1
#SBATCH --gpu-bind=none     # This ensures all GPUs are visible to all tasks
#SBATCH --time=00:20:00

set -e

# Load conda environment
source .venv/bin/activate

# Set master node information
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29400

# Debug info
echo "MASTER_ADDR: $MASTER_ADDR"
echo "MASTER_PORT: $MASTER_PORT"
echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST"
echo "SLURM_NTASKS_PER_NODE: $SLURM_NTASKS_PER_NODE"
echo "SLURM_NNODES: $SLURM_NNODES"

# Run the training script
srun --cpu-bind=cores \
     python src/ocean_emulators/train.py \
     configs/slurm_perlmutter_train_om4.yaml \
     --experiment.name $SLURM_JOB_NAME