#!/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}"
CONFIG="${CONFIG:-$REPO_ROOT/training/nemo_megatron/pretrain_2b_moe_200b.yaml}"
DATA_DIR="/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701"

if [ ! -d "$DATA_DIR" ]; then
  echo "pretrain source data is missing; check dataset/pretrain/manifests/rebalanced_182b_20260701_summary.json" >&2
  exit 1
fi

mkdir -p "$REPO_ROOT/runs" "$REPO_ROOT/logs"

docker run --rm --gpus all --ipc=host --network=host \
  -v /mnt/beegfs:/mnt/beegfs \
  -w "$REPO_ROOT" \
  "$IMAGE" \
  bash -lc "echo TODO: launch NeMo/Megatron pretraining with $CONFIG; ls -lh $CONFIG"