Upload folder using huggingface_hub

0bda409 verified 4 months ago

9.21 kB

#!/bin/bash
# run.sh - MACE 训练统一入口脚本
#
# 用法:
#   bash run.sh --config configs/DMC.yaml              # 直接运行训练
#   bash run.sh --config configs/DMC.yaml --submit      # 提交 SLURM 作业
#   bash run.sh --config configs/DMC.yaml --dry-run     # 仅打印命令，不执行
#
set -euo pipefail

# ============================================================
# 参数解析
# ============================================================
CONFIG=""
SUBMIT=false
DRY_RUN=false

while [[ $# -gt 0 ]]; do
    case "$1" in
        --config)
            CONFIG="$2"; shift 2 ;;
        --config=*)
            CONFIG="${1#*=}"; shift ;;
        --submit)
            SUBMIT=true; shift ;;
        --dry-run)
            DRY_RUN=true; shift ;;
        -h|--help)
            echo "用法: bash run.sh --config <config.yaml> [--submit] [--dry-run]"
            echo ""
            echo "选项:"
            echo "  --config <file>   YAML 配置文件路径 (必需)"
            echo "  --submit          生成 SLURM 脚本并提交作业"
            echo "  --dry-run         仅打印训练命令，不执行"
            exit 0 ;;
        *)
            echo "[ERROR] 未知参数: $1"
            exit 1 ;;
    esac
done

if [ -z "$CONFIG" ]; then
    echo "[ERROR] 请指定配置文件: bash run.sh --config configs/xxx.yaml"
    exit 1
fi

if [ ! -f "$CONFIG" ]; then
    echo "[ERROR] 配置文件不存在: $CONFIG"
    exit 1
fi

# ============================================================
# 路径设置
# ============================================================
DEMO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$DEMO_DIR/../.." && pwd)"
CONFIG_ABS="$(cd "$(dirname "$CONFIG")" && pwd)/$(basename "$CONFIG")"
PARSE_PY="$DEMO_DIR/_parse_config.py"

# 如果用户没有设置 ONESCIENCE_DATASETS_DIR，默认指向仓库根目录
export ONESCIENCE_DATASETS_DIR="${ONESCIENCE_DATASETS_DIR:-$REPO_ROOT}"

# ============================================================
# 解析配置
# ============================================================
EXP_NAME=$(python3 "$PARSE_PY" "$CONFIG_ABS" name)
TRAIN_CMD=$(python3 "$PARSE_PY" "$CONFIG_ABS" command)
ENV_EXPORTS=$(python3 "$PARSE_PY" "$CONFIG_ABS" env)
ENV_ARGS=$(python3 "$PARSE_PY" "$CONFIG_ABS" env-args)
DATA_FILES=$(python3 "$PARSE_PY" "$CONFIG_ABS" data-files)

# ============================================================
# Dry-run 模式
# ============================================================
if $DRY_RUN; then
    TIMESTAMP=$(date +%Y%m%d_%H%M%S)
    OUTPUT_DIR="$DEMO_DIR/outputs/${EXP_NAME}_${TIMESTAMP}"
    echo "========================================="
    echo "Dry-run: $EXP_NAME"
    echo "Config:  $CONFIG_ABS"
    echo "Output:  $OUTPUT_DIR (未创建)"
    echo "========================================="
    echo ""
    echo "# 环境变量:"
    echo "$ENV_EXPORTS"
    echo ""
    echo "# 训练命令:"
    echo "$TRAIN_CMD"
    echo ""
    echo "# env_setup.sh 参数: $ENV_ARGS"
    echo "# 数据文件:"
    echo "$DATA_FILES" | sed 's/^/#   /'
    exit 0
fi

# ============================================================
# 创建输出目录
# ============================================================
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
OUTPUT_DIR="$DEMO_DIR/outputs/${EXP_NAME}_${TIMESTAMP}"
mkdir -p "$OUTPUT_DIR"
cp "$CONFIG_ABS" "$OUTPUT_DIR/config.yaml"

# ============================================================
# SLURM 提交模式
# ============================================================
if $SUBMIT; then
    SLURM_VARS=$(python3 "$PARSE_PY" "$CONFIG_ABS" slurm)
    eval "$SLURM_VARS"

    # 检查是否是多节点
    SLURM_SCRIPT="$OUTPUT_DIR/submit.sh"

    # 生成 SLURM header
    sed -e "s|{{JOB_NAME}}|${JOB_NAME}|g" \
        -e "s|{{PARTITION}}|${PARTITION}|g" \
        -e "s|{{NODES}}|${NODES}|g" \
        -e "s|{{NTASKS_PER_NODE}}|${NTASKS_PER_NODE}|g" \
        -e "s|{{CPUS_PER_TASK}}|${CPUS_PER_TASK}|g" \
        -e "s|{{GPUS_PER_NODE}}|${GPUS_PER_NODE}|g" \
        -e "s|{{TIME}}|${TIME}|g" \
        "$DEMO_DIR/templates/slurm_header.template" > "$SLURM_SCRIPT"

    # 添加环境初始化
    cat >> "$SLURM_SCRIPT" << 'SETUP_BLOCK'

# 环境初始化
SETUP_BLOCK
    cat >> "$SLURM_SCRIPT" << 'ENV_BLOCK'
set +u
if [ -n "${MACE_ENV_SCRIPT:-}" ] && [ -f "${MACE_ENV_SCRIPT}" ]; then
    source "${MACE_ENV_SCRIPT}"
else
    echo "[WARN] MACE_ENV_SCRIPT 未设置或文件不存在，跳过环境初始化。请自行确保 conda/matchem 环境已激活。"
fi
set -u
ENV_BLOCK

    # 添加预检
    echo "" >> "$SLURM_SCRIPT"
    echo "# 预检" >> "$SLURM_SCRIPT"
    # 将数据文件列表转为参数
    DATA_ARGS=""
    while IFS= read -r line; do
        [ -z "$line" ] && continue
        DATA_ARGS="$DATA_ARGS \"$line\""
    done <<< "$DATA_FILES"
    echo "bash $DEMO_DIR/templates/preflight_check.sh $DATA_ARGS" >> "$SLURM_SCRIPT"

    # 添加环境变量和训练命令
    echo "" >> "$SLURM_SCRIPT"
    echo "# 工作目录" >> "$SLURM_SCRIPT"
    echo "cd $OUTPUT_DIR" >> "$SLURM_SCRIPT"
    echo "" >> "$SLURM_SCRIPT"
    echo "# 环境变量" >> "$SLURM_SCRIPT"
    echo "$ENV_EXPORTS" >> "$SLURM_SCRIPT"
    echo "" >> "$SLURM_SCRIPT"
    echo "# 将仓库根目录加入 PYTHONPATH，确保能 import 本地 model 包" >> "$SLURM_SCRIPT"
    echo "export PYTHONPATH=\"$REPO_ROOT:\${PYTHONPATH:-}\"" >> "$SLURM_SCRIPT"
    echo "" >> "$SLURM_SCRIPT"
    echo "# 屏蔽 e3nn FutureWarning 和 TorchScript UserWarning" >> "$SLURM_SCRIPT"
    echo 'export PYTHONWARNINGS="ignore::FutureWarning:e3nn.o3._wigner,ignore::UserWarning:torch.jit._check"' >> "$SLURM_SCRIPT"
    echo "" >> "$SLURM_SCRIPT"
    echo '# 屏蔽 PyTorch NCCL C++ INFO 日志' >> "$SLURM_SCRIPT"
    echo 'export TORCH_CPP_LOG_LEVEL=WARNING' >> "$SLURM_SCRIPT"
    echo 'export NCCL_DEBUG=ERROR' >> "$SLURM_SCRIPT"
    echo '' >> "$SLURM_SCRIPT"
    echo '# 屏蔽 glog INFO（ProcessGroupNCCL.cpp 初始化信息）' >> "$SLURM_SCRIPT"
    echo 'export GLOG_minloglevel=1' >> "$SLURM_SCRIPT"
    echo '' >> "$SLURM_SCRIPT"
    echo '# AMD DCU: 避免 RCCL "Missing HSA_FORCE_FINE_GRAIN_PCIE" 警告' >> "$SLURM_SCRIPT"
    echo 'export HSA_FORCE_FINE_GRAIN_PCIE=1' >> "$SLURM_SCRIPT"

    # 多节点特殊处理
    if [ "$NODES" -gt 1 ]; then
        cat >> "$SLURM_SCRIPT" << 'MULTI_NODE'

# 多节点分布式设置
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29517
export WORLD_SIZE=$SLURM_NTASKS

echo "MASTER_ADDR: $MASTER_ADDR"
echo "WORLD_SIZE: $WORLD_SIZE"

# 使用 srun 启动分布式训练
MULTI_NODE
        # srun 包裹训练命令
        echo "srun --export=ALL bash -c '" >> "$SLURM_SCRIPT"
        echo "  export RANK=\$SLURM_PROCID" >> "$SLURM_SCRIPT"
        echo "  export LOCAL_RANK=\$SLURM_LOCALID" >> "$SLURM_SCRIPT"
        echo "  exec $TRAIN_CMD" >> "$SLURM_SCRIPT"
        echo "'" >> "$SLURM_SCRIPT"
    else
        echo "" >> "$SLURM_SCRIPT"
        echo "# 训练命令" >> "$SLURM_SCRIPT"
        echo "$TRAIN_CMD" >> "$SLURM_SCRIPT"
    fi

    echo "========================================="
    echo "SLURM 脚本已生成: $SLURM_SCRIPT"
    echo "配置快照已保存: $OUTPUT_DIR/config.yaml"
    echo "========================================="
    echo ""
    echo "提交作业..."
    sbatch "$SLURM_SCRIPT"
    exit 0
fi

# ============================================================
# 直接运行模式
# ============================================================
echo "========================================="
echo "实验: $EXP_NAME"
echo "配置: $CONFIG_ABS"
echo "输出: $OUTPUT_DIR"
echo "========================================="

# 加载用户环境初始化脚本（路径通过 MACE_ENV_SCRIPT 指定）
set +u
if [ -n "${MACE_ENV_SCRIPT:-}" ] && [ -f "${MACE_ENV_SCRIPT}" ]; then
    source "${MACE_ENV_SCRIPT}"
else
    echo "[WARN] MACE_ENV_SCRIPT 未设置或文件不存在，跳过环境初始化。请自行确保 conda/matchem 环境已激活。"
fi
set -u

# 预检
DATA_ARGS=""
while IFS= read -r line; do
    [ -z "$line" ] && continue
    DATA_ARGS="$DATA_ARGS \"$line\""
done <<< "$DATA_FILES"
eval "bash $DEMO_DIR/templates/preflight_check.sh $DATA_ARGS"

# 设置环境变量
eval "$ENV_EXPORTS"

# 将仓库根目录加入 PYTHONPATH，确保能 import 本地 model 包
export PYTHONPATH="$REPO_ROOT:${PYTHONPATH:-}"

# 屏蔽 e3nn FutureWarning 和 TorchScript UserWarning
export PYTHONWARNINGS="ignore::FutureWarning:e3nn.o3._wigner,ignore::UserWarning:torch.jit._check"

# 屏蔽 PyTorch NCCL C++ INFO 日志
export TORCH_CPP_LOG_LEVEL=WARNING
export NCCL_DEBUG=ERROR

# 屏蔽 glog INFO（ProcessGroupNCCL.cpp 初始化信息）
export GLOG_minloglevel=1

# AMD DCU: 避免 RCCL "Missing HSA_FORCE_FINE_GRAIN_PCIE" 警告
export HSA_FORCE_FINE_GRAIN_PCIE=1

# 切换到输出目录
cd "$OUTPUT_DIR"

# 执行训练
echo "========================================="
echo "开始训练..."
echo "========================================="
if [ -n "${TRAIN_CMD:-}" ]; then
    eval "$TRAIN_CMD"
else
    echo "[FATAL] TRAIN_CMD 为空，无法执行任务"
    exit 1
fi

	#!/bin/bash
	# run.sh - MACE 训练统一入口脚本
	#
	# 用法:
	# bash run.sh --config configs/DMC.yaml # 直接运行训练
	# bash run.sh --config configs/DMC.yaml --submit # 提交 SLURM 作业
	# bash run.sh --config configs/DMC.yaml --dry-run # 仅打印命令，不执行
	#
	set -euo pipefail

	# ============================================================
	# 参数解析
	# ============================================================
	CONFIG=""
	SUBMIT=false
	DRY_RUN=false

	while [[ $# -gt 0 ]]; do
	case "$1" in
	--config)
	CONFIG="$2"; shift 2 ;;
	--config=*)
	CONFIG="${1#*=}"; shift ;;
	--submit)
	SUBMIT=true; shift ;;
	--dry-run)
	DRY_RUN=true; shift ;;
	-h\|--help)
	echo "用法: bash run.sh --config <config.yaml> [--submit] [--dry-run]"
	echo ""
	echo "选项:"
	echo " --config <file> YAML 配置文件路径 (必需)"
	echo " --submit 生成 SLURM 脚本并提交作业"
	echo " --dry-run 仅打印训练命令，不执行"
	exit 0 ;;
	*)
	echo "[ERROR] 未知参数: $1"
	exit 1 ;;
	esac
	done

	if [ -z "$CONFIG" ]; then
	echo "[ERROR] 请指定配置文件: bash run.sh --config configs/xxx.yaml"
	exit 1
	fi

	if [ ! -f "$CONFIG" ]; then
	echo "[ERROR] 配置文件不存在: $CONFIG"
	exit 1
	fi

	# ============================================================
	# 路径设置
	# ============================================================
	DEMO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
	REPO_ROOT="$(cd "$DEMO_DIR/../.." && pwd)"
	CONFIG_ABS="$(cd "$(dirname "$CONFIG")" && pwd)/$(basename "$CONFIG")"
	PARSE_PY="$DEMO_DIR/_parse_config.py"

	# 如果用户没有设置 ONESCIENCE_DATASETS_DIR，默认指向仓库根目录
	export ONESCIENCE_DATASETS_DIR="${ONESCIENCE_DATASETS_DIR:-$REPO_ROOT}"

	# ============================================================
	# 解析配置
	# ============================================================
	EXP_NAME=$(python3 "$PARSE_PY" "$CONFIG_ABS" name)
	TRAIN_CMD=$(python3 "$PARSE_PY" "$CONFIG_ABS" command)
	ENV_EXPORTS=$(python3 "$PARSE_PY" "$CONFIG_ABS" env)
	ENV_ARGS=$(python3 "$PARSE_PY" "$CONFIG_ABS" env-args)
	DATA_FILES=$(python3 "$PARSE_PY" "$CONFIG_ABS" data-files)

	# ============================================================
	# Dry-run 模式
	# ============================================================
	if $DRY_RUN; then
	TIMESTAMP=$(date +%Y%m%d_%H%M%S)
	OUTPUT_DIR="$DEMO_DIR/outputs/${EXP_NAME}_${TIMESTAMP}"
	echo "========================================="
	echo "Dry-run: $EXP_NAME"
	echo "Config: $CONFIG_ABS"
	echo "Output: $OUTPUT_DIR (未创建)"
	echo "========================================="
	echo ""
	echo "# 环境变量:"
	echo "$ENV_EXPORTS"
	echo ""
	echo "# 训练命令:"
	echo "$TRAIN_CMD"
	echo ""
	echo "# env_setup.sh 参数: $ENV_ARGS"
	echo "# 数据文件:"
	echo "$DATA_FILES" \| sed 's/^/# /'
	exit 0
	fi

	# ============================================================
	# 创建输出目录
	# ============================================================
	TIMESTAMP=$(date +%Y%m%d_%H%M%S)
	OUTPUT_DIR="$DEMO_DIR/outputs/${EXP_NAME}_${TIMESTAMP}"
	mkdir -p "$OUTPUT_DIR"
	cp "$CONFIG_ABS" "$OUTPUT_DIR/config.yaml"

	# ============================================================
	# SLURM 提交模式
	# ============================================================
	if $SUBMIT; then
	SLURM_VARS=$(python3 "$PARSE_PY" "$CONFIG_ABS" slurm)
	eval "$SLURM_VARS"

	# 检查是否是多节点
	SLURM_SCRIPT="$OUTPUT_DIR/submit.sh"

	# 生成 SLURM header
	sed -e "s\|{{JOB_NAME}}\|${JOB_NAME}\|g" \
	-e "s\|{{PARTITION}}\|${PARTITION}\|g" \
	-e "s\|{{NODES}}\|${NODES}\|g" \
	-e "s\|{{NTASKS_PER_NODE}}\|${NTASKS_PER_NODE}\|g" \
	-e "s\|{{CPUS_PER_TASK}}\|${CPUS_PER_TASK}\|g" \
	-e "s\|{{GPUS_PER_NODE}}\|${GPUS_PER_NODE}\|g" \
	-e "s\|{{TIME}}\|${TIME}\|g" \
	"$DEMO_DIR/templates/slurm_header.template" > "$SLURM_SCRIPT"

	# 添加环境初始化
	cat >> "$SLURM_SCRIPT" << 'SETUP_BLOCK'

	# 环境初始化
	SETUP_BLOCK
	cat >> "$SLURM_SCRIPT" << 'ENV_BLOCK'
	set +u
	if [ -n "${MACE_ENV_SCRIPT:-}" ] && [ -f "${MACE_ENV_SCRIPT}" ]; then
	source "${MACE_ENV_SCRIPT}"
	else
	echo "[WARN] MACE_ENV_SCRIPT 未设置或文件不存在，跳过环境初始化。请自行确保 conda/matchem 环境已激活。"
	fi
	set -u
	ENV_BLOCK

	# 添加预检
	echo "" >> "$SLURM_SCRIPT"
	echo "# 预检" >> "$SLURM_SCRIPT"
	# 将数据文件列表转为参数
	DATA_ARGS=""
	while IFS= read -r line; do
	[ -z "$line" ] && continue
	DATA_ARGS="$DATA_ARGS \"$line\""
	done <<< "$DATA_FILES"
	echo "bash $DEMO_DIR/templates/preflight_check.sh $DATA_ARGS" >> "$SLURM_SCRIPT"

	# 添加环境变量和训练命令
	echo "" >> "$SLURM_SCRIPT"
	echo "# 工作目录" >> "$SLURM_SCRIPT"
	echo "cd $OUTPUT_DIR" >> "$SLURM_SCRIPT"
	echo "" >> "$SLURM_SCRIPT"
	echo "# 环境变量" >> "$SLURM_SCRIPT"
	echo "$ENV_EXPORTS" >> "$SLURM_SCRIPT"
	echo "" >> "$SLURM_SCRIPT"
	echo "# 将仓库根目录加入 PYTHONPATH，确保能 import 本地 model 包" >> "$SLURM_SCRIPT"
	echo "export PYTHONPATH=\"$REPO_ROOT:\${PYTHONPATH:-}\"" >> "$SLURM_SCRIPT"
	echo "" >> "$SLURM_SCRIPT"
	echo "# 屏蔽 e3nn FutureWarning 和 TorchScript UserWarning" >> "$SLURM_SCRIPT"
	echo 'export PYTHONWARNINGS="ignore::FutureWarning:e3nn.o3._wigner,ignore::UserWarning:torch.jit._check"' >> "$SLURM_SCRIPT"
	echo "" >> "$SLURM_SCRIPT"
	echo '# 屏蔽 PyTorch NCCL C++ INFO 日志' >> "$SLURM_SCRIPT"
	echo 'export TORCH_CPP_LOG_LEVEL=WARNING' >> "$SLURM_SCRIPT"
	echo 'export NCCL_DEBUG=ERROR' >> "$SLURM_SCRIPT"
	echo '' >> "$SLURM_SCRIPT"
	echo '# 屏蔽 glog INFO（ProcessGroupNCCL.cpp 初始化信息）' >> "$SLURM_SCRIPT"
	echo 'export GLOG_minloglevel=1' >> "$SLURM_SCRIPT"
	echo '' >> "$SLURM_SCRIPT"
	echo '# AMD DCU: 避免 RCCL "Missing HSA_FORCE_FINE_GRAIN_PCIE" 警告' >> "$SLURM_SCRIPT"
	echo 'export HSA_FORCE_FINE_GRAIN_PCIE=1' >> "$SLURM_SCRIPT"

	# 多节点特殊处理
	if [ "$NODES" -gt 1 ]; then
	cat >> "$SLURM_SCRIPT" << 'MULTI_NODE'

	# 多节点分布式设置
	export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" \| head -n 1)
	export MASTER_PORT=29517
	export WORLD_SIZE=$SLURM_NTASKS

	echo "MASTER_ADDR: $MASTER_ADDR"
	echo "WORLD_SIZE: $WORLD_SIZE"

	# 使用 srun 启动分布式训练
	MULTI_NODE
	# srun 包裹训练命令
	echo "srun --export=ALL bash -c '" >> "$SLURM_SCRIPT"
	echo " export RANK=\$SLURM_PROCID" >> "$SLURM_SCRIPT"
	echo " export LOCAL_RANK=\$SLURM_LOCALID" >> "$SLURM_SCRIPT"
	echo " exec $TRAIN_CMD" >> "$SLURM_SCRIPT"
	echo "'" >> "$SLURM_SCRIPT"
	else
	echo "" >> "$SLURM_SCRIPT"
	echo "# 训练命令" >> "$SLURM_SCRIPT"
	echo "$TRAIN_CMD" >> "$SLURM_SCRIPT"
	fi

	echo "========================================="
	echo "SLURM 脚本已生成: $SLURM_SCRIPT"
	echo "配置快照已保存: $OUTPUT_DIR/config.yaml"
	echo "========================================="
	echo ""
	echo "提交作业..."
	sbatch "$SLURM_SCRIPT"
	exit 0
	fi

	# ============================================================
	# 直接运行模式
	# ============================================================
	echo "========================================="
	echo "实验: $EXP_NAME"
	echo "配置: $CONFIG_ABS"
	echo "输出: $OUTPUT_DIR"
	echo "========================================="

	# 加载用户环境初始化脚本（路径通过 MACE_ENV_SCRIPT 指定）
	set +u
	if [ -n "${MACE_ENV_SCRIPT:-}" ] && [ -f "${MACE_ENV_SCRIPT}" ]; then
	source "${MACE_ENV_SCRIPT}"
	else
	echo "[WARN] MACE_ENV_SCRIPT 未设置或文件不存在，跳过环境初始化。请自行确保 conda/matchem 环境已激活。"
	fi
	set -u

	# 预检
	DATA_ARGS=""
	while IFS= read -r line; do
	[ -z "$line" ] && continue
	DATA_ARGS="$DATA_ARGS \"$line\""
	done <<< "$DATA_FILES"
	eval "bash $DEMO_DIR/templates/preflight_check.sh $DATA_ARGS"

	# 设置环境变量
	eval "$ENV_EXPORTS"

	# 将仓库根目录加入 PYTHONPATH，确保能 import 本地 model 包
	export PYTHONPATH="$REPO_ROOT:${PYTHONPATH:-}"

	# 屏蔽 e3nn FutureWarning 和 TorchScript UserWarning
	export PYTHONWARNINGS="ignore::FutureWarning:e3nn.o3._wigner,ignore::UserWarning:torch.jit._check"

	# 屏蔽 PyTorch NCCL C++ INFO 日志
	export TORCH_CPP_LOG_LEVEL=WARNING
	export NCCL_DEBUG=ERROR

	# 屏蔽 glog INFO（ProcessGroupNCCL.cpp 初始化信息）
	export GLOG_minloglevel=1

	# AMD DCU: 避免 RCCL "Missing HSA_FORCE_FINE_GRAIN_PCIE" 警告
	export HSA_FORCE_FINE_GRAIN_PCIE=1

	# 切换到输出目录
	cd "$OUTPUT_DIR"

	# 执行训练
	echo "========================================="
	echo "开始训练..."
	echo "========================================="
	if [ -n "${TRAIN_CMD:-}" ]; then
	eval "$TRAIN_CMD"
	else
	echo "[FATAL] TRAIN_CMD 为空，无法执行任务"
	exit 1
	fi