#!/bin/bash -e

wd=$(pwd)
jobid=$(squeue --me | head -2 | tail -n1 | awk '{print $1}')


#
# Example assume allocation was created, e.g.:
# N=1 ; salloc -p standard-g  --threads-per-core 1 --exclusive -N $N --gpus $((N*8)) -t 4:00:00 --mem 0
#

module purge
module load CrayEnv
module load PrgEnv-cray/8.6.0
module load craype-x86-trento
module load craype-accel-amd-gfx90a

module use module use /appl/local/containers/test-modules
module load rocm/7.0.3.lua

set -x

# Install pytorch if it doesn't exist already.
if [ ! -d $wd/miniconda3/envs/pytorch ] ; then
    curl -LO https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
    bash ./Miniconda3-* -b -p $wd/miniconda3 -s
    rm -rf ./Miniconda3-*

    sed -i 's/defaults/conda-forge/g' $wd/miniconda3/.condarc

    source $wd/miniconda3/bin/activate base
    conda create -y -n pytorch python=3.12
    source $wd/miniconda3/bin/activate pytorch
    conda install -y pip
    pip install --no-cache-dir  --pre torch==2.10.0+rocm7.0 --index-url https://download.pytorch.org/whl/
else
    source $wd/miniconda3/bin/activate pytorch
fi

srun --jobid=$jobid -n1 --gpus 8 \
    python -c 'import torch; print("I have this many devices:", torch.cuda.device_count())'
