109 lines
3.9 KiB
Bash
Executable File
109 lines
3.9 KiB
Bash
Executable File
#!/bin/bash
|
||
# ============================================================
|
||
# 预加载 DeepSWE 任务 Docker 镜像
|
||
#
|
||
# 用途：
|
||
# Pier 会按 task.toml 里的 docker_image 起容器。镜像在
|
||
# public.ecr.aws，首次现拉会很慢。本脚本先下载数据集、抽出
|
||
# 113 个任务镜像 + mars-base，再并行 docker pull。
|
||
#
|
||
# 用法：
|
||
# bash bash/images_load/preload_deep_swe_images.sh
|
||
#
|
||
# 环境变量：
|
||
# DEEP_SWE_CACHE 数据集根目录（含 tasks/manifest.json）
|
||
# PULL_CONCURRENCY 并发 pull 数，默认 4
|
||
# ============================================================
|
||
|
||
set -euo pipefail
|
||
|
||
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
||
CACHE_DIR="${DEEP_SWE_CACHE:-${HOME}/.cache/evalscope/deep_swe/snapshots/evalscope/deep-swe}"
|
||
DATASET_ID="evalscope/deep-swe"
|
||
CONCURRENCY="${PULL_CONCURRENCY:-4}"
|
||
LOG_DIR="${PROJECT_ROOT}/logs"
|
||
LOG_FILE="${LOG_DIR}/pull_deep_swe_images.log"
|
||
IMAGE_FILE="/tmp/deep_swe_images.txt"
|
||
MARS_BASE="public.ecr.aws/x8v8d7g8/mars-base:latest"
|
||
|
||
mkdir -p "${LOG_DIR}"
|
||
: > "${LOG_FILE}"
|
||
|
||
echo "============================================================"
|
||
echo "1. 准备 DeepSWE 数据集快照"
|
||
echo "============================================================"
|
||
if [[ ! -f "${CACHE_DIR}/tasks/manifest.json" ]]; then
|
||
echo "本地没有 ${CACHE_DIR}，从 ModelScope 下载 ${DATASET_ID}"
|
||
python3 - <<PY
|
||
from modelscope import dataset_snapshot_download
|
||
path = dataset_snapshot_download("${DATASET_ID}", cache_dir="${HOME}/.cache/evalscope/deep_swe")
|
||
print(path)
|
||
PY
|
||
if [[ ! -f "${CACHE_DIR}/tasks/manifest.json" ]]; then
|
||
# modelscope 可能把快照放到带 hash 的目录，找一份有 manifest 的
|
||
FOUND="$(find "${HOME}/.cache/evalscope/deep_swe" -path '*/tasks/manifest.json' -print -quit 2>/dev/null || true)"
|
||
if [[ -n "${FOUND}" ]]; then
|
||
CACHE_DIR="$(cd "$(dirname "${FOUND}")/.." && pwd)"
|
||
echo "使用快照: ${CACHE_DIR}"
|
||
else
|
||
echo "ERROR: 下载后仍找不到 tasks/manifest.json"
|
||
exit 1
|
||
fi
|
||
fi
|
||
else
|
||
echo "使用已有快照: ${CACHE_DIR}"
|
||
fi
|
||
|
||
echo ""
|
||
echo "============================================================"
|
||
echo "2. 提取任务镜像列表"
|
||
echo "============================================================"
|
||
{
|
||
echo "${MARS_BASE}"
|
||
grep -hR '^docker_image' "${CACHE_DIR}/tasks" --include='task.toml' 2>/dev/null \
|
||
| sed 's/docker_image = //; s/"//g' \
|
||
| sed 's/^[[:space:]]*//; s/[[:space:]]*$//'
|
||
} | awk 'NF && !seen[$0]++' > "${IMAGE_FILE}"
|
||
|
||
IMAGE_COUNT="$(wc -l < "${IMAGE_FILE}" | tr -d ' ')"
|
||
echo "找到 ${IMAGE_COUNT} 个唯一镜像，列表: ${IMAGE_FILE}"
|
||
|
||
echo ""
|
||
echo "============================================================"
|
||
echo "3. 并行 docker pull（并发 ${CONCURRENCY}）"
|
||
echo "============================================================"
|
||
echo "日志: ${LOG_FILE}"
|
||
|
||
PULL_FAILED=0
|
||
while IFS= read -r image; do
|
||
[[ -z "${image}" ]] && continue
|
||
(
|
||
echo "[$(date '+%Y-%m-%d %H:%M:%S')] Pulling ${image} ..." >> "${LOG_FILE}"
|
||
if docker image inspect "${image}" >/dev/null 2>&1; then
|
||
echo "[$(date '+%Y-%m-%d %H:%M:%S')] SKIP already present ${image}" >> "${LOG_FILE}"
|
||
exit 0
|
||
fi
|
||
if docker pull "${image}" >> "${LOG_FILE}" 2>&1; then
|
||
echo "[$(date '+%Y-%m-%d %H:%M:%S')] OK ${image}" >> "${LOG_FILE}"
|
||
else
|
||
echo "[$(date '+%Y-%m-%d %H:%M:%S')] FAILED ${image}" >> "${LOG_FILE}"
|
||
exit 1
|
||
fi
|
||
) &
|
||
if (( $(jobs -r -p | wc -l) >= CONCURRENCY )); then
|
||
wait -n || PULL_FAILED=1
|
||
fi
|
||
done < "${IMAGE_FILE}"
|
||
|
||
wait || PULL_FAILED=1
|
||
|
||
echo ""
|
||
echo "============================================================"
|
||
if [[ "${PULL_FAILED}" -eq 0 ]]; then
|
||
echo "DeepSWE 镜像拉取完成"
|
||
else
|
||
echo "部分镜像拉取失败，请查看: ${LOG_FILE}"
|
||
exit 1
|
||
fi
|
||
echo "============================================================"
|