diff --git a/bash/images_load/preload_deep_swe_images.sh b/bash/images_load/preload_deep_swe_images.sh new file mode 100755 index 0000000..922a22f --- /dev/null +++ b/bash/images_load/preload_deep_swe_images.sh @@ -0,0 +1,108 @@ +#!/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 - </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 "============================================================"