#!/usr/bin/env bash
set -euo pipefail

REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}"
DOCKER_PULL="${DOCKER_PULL:-0}"
USE_GPUS="${USE_GPUS:-0}"

cd "$REPO_ROOT"
python3 tools/validate_model_architecture.py

if [ "$DOCKER_PULL" = "1" ]; then
  docker pull "$IMAGE"
fi

GPU_ARGS=()
if [ "$USE_GPUS" = "1" ]; then
  GPU_ARGS=(--gpus all)
fi

docker run --rm "${GPU_ARGS[@]}" --ipc=host --network=host \
  -v /mnt/beegfs:/mnt/beegfs \
  -w "$REPO_ROOT" \
  "$IMAGE" \
  bash -lc 'set -euo pipefail
python - <<PY
import importlib
mods = ["torch", "nemo", "megatron.core"]
for name in mods:
    mod = importlib.import_module(name)
    version = getattr(mod, "__version__", "")
    print("IMPORT_OK", name, version)
PY
python tools/validate_model_architecture.py
'