#!/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="$REPO_ROOT/dataset/pretrain/data/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701"

if [ ! -d "$DATA_DIR" ]; then
  echo "pretrain data is missing under repo; run scripts/sync_pretrain_data_into_repo.sh after the 200B build finishes" >&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"
