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

DATASET_ID="${DATASET_ID:-eigentom/laoyao_2b_moe_pretrain_parquet_20260702}"
BASE="${BASE:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train}"
IMAGE="${IMAGE:-modelscope-registry.cn-hangzhou.cr.aliyuncs.com/modelscope-repo/modelscope:ubuntu22.04-cuda13.0.3-py312-torch2.11.0-vllm0.23.0-modelscope1.37.1-swift4.3.1}"
ENDPOINT="${ENDPOINT:-https://www.modelscope.cn}"
MAX_WORKERS="${MAX_WORKERS:-8}"

export http_proxy="${http_proxy:-http://10.20.34.2:3128}"
export https_proxy="${https_proxy:-http://10.20.34.2:3128}"
export HTTP_PROXY="${HTTP_PROXY:-$http_proxy}"
export HTTPS_PROXY="${HTTPS_PROXY:-$https_proxy}"

if [[ -z "${MODELSCOPE_API_TOKEN:-}" ]]; then
  echo "ERROR: set MODELSCOPE_API_TOKEN." >&2
  exit 2
fi

MAIN_DIR="$BASE/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701"
LOGIC_DIR="$BASE/logic_topup_proof_pile_17b_v1_20260701"

test -d "$MAIN_DIR"
test -d "$LOGIC_DIR"

docker run --rm --network=host \
  -e http_proxy="$http_proxy" \
  -e https_proxy="$https_proxy" \
  -e HTTP_PROXY="$HTTP_PROXY" \
  -e HTTPS_PROXY="$HTTPS_PROXY" \
  -e MODELSCOPE_API_TOKEN="$MODELSCOPE_API_TOKEN" \
  -v /mnt/beegfs:/mnt/beegfs \
  "$IMAGE" \
  bash -lc '
    set -euo pipefail
    modelscope upload --repo-type dataset \
      --token "$MODELSCOPE_API_TOKEN" \
      --endpoint "'"$ENDPOINT"'" \
      --max-workers "'"$MAX_WORKERS"'" \
      --commit-message "upload main pretokenization parquet data" \
      "'"$DATASET_ID"'" \
      "'"$MAIN_DIR"'" \
      train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701

    modelscope upload --repo-type dataset \
      --token "$MODELSCOPE_API_TOKEN" \
      --endpoint "'"$ENDPOINT"'" \
      --max-workers "'"$MAX_WORKERS"'" \
      --commit-message "upload logic topup parquet data" \
      "'"$DATASET_ID"'" \
      "'"$LOGIC_DIR"'" \
      train/logic_topup_proof_pile_17b_v1_20260701

    echo UPLOAD_DONE
  '