#!/usr/bin/env bash
# SUM Parts - verify PYTORCH_CUDA_ALLOC_CONF is accepted before committing a run
#
# torch 2.0.1 aborts at CUDA init on an unknown key, and the traceback points at
# model.to(device) rather than at the env var, so an invalid setting looks like
# a model problem. Check it in two seconds instead.
set -uo pipefail

source "$HOME/miniconda3/etc/profile.d/conda.sh"
conda activate sumparts

CONF="${1:-garbage_collection_threshold:0.7,max_split_size_mb:128}"

echo "torch  : $(python -c 'import torch; print(torch.__version__)')"
echo "testing: PYTORCH_CUDA_ALLOC_CONF=$CONF"

PYTORCH_CUDA_ALLOC_CONF="$CONF" python - <<'PY'
import os
import torch

print("env    :", os.environ.get("PYTORCH_CUDA_ALLOC_CONF"))
x = torch.zeros(1024, 1024, device="cuda")
y = (x + 1).sum().item()
print("cuda   : OK", torch.cuda.get_device_name(0), "| smoke sum =", y)
print("alloc  : %.1f MB" % (torch.cuda.memory_allocated() / 1024**2))
PY
rc=$?

if [ $rc -eq 0 ]; then
    echo "ACCEPTED"
else
    echo "REJECTED (rc=$rc) -- do not launch with this setting"
fi
exit $rc