#!/usr/bin/env bash
# Drop-in PYTHON_BIN for scripts/train_2node_16gpu.slurm. One container per node;
# train_large.py launches eight workers using container sys.executable.
set +x  # Never trace runtime credentials, including if launched with bash -x.
set -euo pipefail
HERE="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
if [[ "${1:-}" == --help-container ]]; then
  echo 'IMAGE=/path/environment.sif PROJECT_DIR=/shared/checkout apptainer/python.sh [Python arguments]'
  echo 'Optional: DATA_DIR MODEL_DIR ASSET_DIR OUTPUT_DIR CONTAINER_CACHE_DIR WANDB_NETRC_FILE GPUS=all|none DRY_RUN=1.'
  echo 'Stage0 external assets: LORAT_ROOT LORAT_WEIGHT_PATH TORCH_HOME (existing host paths, automatically bound).'
  echo 'Credentials are forwarded through the environment only; dry-run redacts values.'
  exit 0
fi
ROOT="$(cd -- "${PROJECT_DIR:-${TRACK_ANYTHING_ROOT:-$HERE/..}}" && pwd)"
test -f "$ROOT/scripts/train_large.py" || { echo 'PROJECT_DIR must be the large-training checkout.' >&2; exit 2; }
test -f "$ROOT/apptainer/entrypoint.sh" || { echo 'PROJECT_DIR must include apptainer/entrypoint.sh.' >&2; exit 2; }
: "${IMAGE:?Set IMAGE to the built large-cu128-wandb-v1 SIF}"
IMAGE="$(realpath -e -- "$IMAGE")"
test -f "$IMAGE" || { echo 'IMAGE must be a regular SIF file.' >&2; exit 2; }
CACHE="$(realpath -m -- "${CONTAINER_CACHE_DIR:-$ROOT/.container-cache}")"
check_path() { [[ "$2" != *:* && "$2" != *,* && "$2" != *$'\n'* ]] || { echo "$1 cannot contain colon, comma or newline." >&2; exit 2; }; }
for name in ROOT CACHE; do check_path "$name" "${!name}"; done
args=(exec --cleanenv --no-eval --no-home --no-mount home,cwd,hostfs,bind-paths
      --pwd "$ROOT" --bind "$ROOT:$ROOT" --bind "$CACHE:/cache")
for name in DATA_DIR MODEL_DIR ASSET_DIR OUTPUT_DIR; do
  value="${!name:-}"
  if [[ -n "$value" ]]; then
    test -d "$value" || { echo "$name must already exist on each node." >&2; exit 2; }
    value="$(cd -- "$value" && pwd)"; check_path "$name" "$value"
    if [[ "$name" == OUTPUT_DIR ]]; then mode=rw; else mode=ro; fi
    args+=(--bind "$value:$value:$mode")
  fi
done
# --cleanenv drops host asset settings, and --no-mount disables implicit host
# binds. Resolve both the forwarded value and its bind from the SAME path;
# canonicalization also handles symlinks whose targets leave the checkout.
# Bind only the weight file, not its possibly private parent directory.
for name in LORAT_ROOT LORAT_WEIGHT_PATH TORCH_HOME; do
  value="${!name:-}"
  if [[ -z "$value" && "$name" == LORAT_ROOT && -d "$ROOT/third_party/LoRAT/trackit" ]]; then
    value="$ROOT/third_party/LoRAT"
  fi
  if [[ -n "$value" ]]; then
    if [[ "$name" == LORAT_WEIGHT_PATH ]]; then
      test -f "$value" || { echo "$name must be an existing file on each node." >&2; exit 2; }
    else
      test -d "$value" || { echo "$name must be an existing directory on each node." >&2; exit 2; }
    fi
    value="$(realpath -e -- "$value")"; check_path "$name" "$value"
    if [[ "$name" == LORAT_ROOT ]]; then
      test -d "$value/trackit" || { echo 'LORAT_ROOT must contain the LoRAT trackit package.' >&2; exit 2; }
    fi
    if [[ "$name" == TORCH_HOME ]]; then mode=rw; else mode=ro; fi
    args+=(--bind "$value:$value:$mode")
    printf -v "$name" '%s' "$value"
    export "$name"
  else
    # An empty optional setting means use the image/application default; it
    # must not overwrite TORCH_HOME=/cache/torch with an empty string.
    unset "$name"
  fi
done
if [[ -n "${WANDB_NETRC_FILE:-}" ]]; then
  test -f "$WANDB_NETRC_FILE" || { echo 'WANDB_NETRC_FILE must be an existing private credentials file.' >&2; exit 2; }
  value="$(realpath -e -- "$WANDB_NETRC_FILE")"; check_path WANDB_NETRC_FILE "$value"
  args+=(--bind "$value:/cache/home/.netrc:ro")
fi
unset APPTAINER_HOME SINGULARITY_HOME APPTAINER_BIND APPTAINER_BINDPATH SINGULARITY_BIND SINGULARITY_BINDPATH
# Do not inherit arbitrary APPTAINERENV_ or SINGULARITYENV_ variables, especially
# an old PYTHONPATH, HOME, or an accidental conflicting CUDA_VISIBLE_DEVICES.
while IFS= read -r variable; do
  case "$variable" in APPTAINERENV_*|SINGULARITYENV_*) unset "$variable" ;; esac
done < <(compgen -e)
export APPTAINERENV_TRACK_ANYTHING_ROOT="$ROOT"
export APPTAINERENV_WANDB_MODE="${WANDB_MODE:-online}"
case "$APPTAINERENV_WANDB_MODE" in online|offline|disabled) ;; *) echo 'Invalid WANDB_MODE.' >&2; exit 2 ;; esac
case "${GPUS:-all}" in
  none) export APPTAINERENV_CUDA_VISIBLE_DEVICES=-1 ;;
  all)
    args+=(--nv)
    if [[ "${CUDA_VISIBLE_DEVICES+x}" ]]; then export APPTAINERENV_CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES"; fi
    ;;
  *) echo 'GPUS must be all (allocated set) or none; use the scheduler to allocate GPUs.' >&2; exit 2 ;;
esac
forward=(SLURM_NODEID SLURM_JOB_ID SLURM_JOB_NUM_NODES SLURM_NNODES SLURM_LOCALID
         MASTER_ADDR MASTER_PORT NODE_RANK TRACKING_SMOKE_RUN_ID
         NCCL_DEBUG NCCL_DEBUG_SUBSYS NCCL_SOCKET_IFNAME NCCL_SOCKET_FAMILY
         NCCL_IB_HCA NCCL_IB_GID_INDEX NCCL_IB_TC NCCL_IB_TIMEOUT NCCL_IB_DISABLE
         NCCL_NET_GDR_LEVEL NCCL_CUMEM_ENABLE NCCL_P2P_DISABLE
         TORCH_DISTRIBUTED_DEBUG TORCH_NCCL_ASYNC_ERROR_HANDLING
         LORAT_ROOT LORAT_WEIGHT_PATH TORCH_HOME
         WANDB_API_KEY WANDB_BASE_URL WANDB_ENTITY
         HTTP_PROXY HTTPS_PROXY ALL_PROXY NO_PROXY http_proxy https_proxy all_proxy no_proxy)
for variable in "${forward[@]}"; do
  if [[ "${!variable+x}" ]]; then
    export "APPTAINERENV_${variable}=${!variable}"
  fi
done
# Keep the environment-only SIF reusable: launch the checked-out entrypoint so
# runtime fixes do not depend on rebuilding an already published image.
args+=("$IMAGE" /bin/bash "$ROOT/apptainer/entrypoint.sh" /opt/venv/bin/python "$@")
if [[ "${DRY_RUN:-0}" == 1 ]]; then
  printf 'Container command (environment values omitted): '
  printf '%q ' "${APPTAINER:-apptainer}" "${args[@]}"; printf '\n'
  printf 'Forwarded environment names only:'
  for variable in "${forward[@]}"; do
    if [[ "${!variable+x}" ]]; then printf ' %s' "$variable"; fi
  done
  printf '\nNo process launched.\n'
  exit 0
fi
command -v "${APPTAINER:-apptainer}" >/dev/null || { echo 'Load/install Apptainer on each node first.' >&2; exit 127; }
umask 077
mkdir -p -- "$CACHE/home"
exec "${APPTAINER:-apptainer}" "${args[@]}"