Files
Velocity-OS/infrastructure/model-hydration/hydrate_models.sh

84 lines
3.3 KiB
Bash

#!/usr/bin/env bash
# ============================================================
# Velocity-OS — Model Hydration Script
# Uses s5cmd for high-throughput parallel S3 → NVMe sync.
# Run once at initial install; safe to re-run for updates.
#
# Models synced (all stored on NVMe, never in Docker images):
# - Wan 2.2 (ComfyUI video/image model)
# - Qwen-Image 2512 (ComfyUI poster/image model)
# - Qwen3.6 35B A3B (SGLang LLM — MIG slice 0)
#
# Requires: s5cmd, AWS credentials with S3 read access
# ============================================================
set -euo pipefail
# ── Configuration ────────────────────────────────────────────
S3_BUCKET="${S3_MODEL_BUCKET:-s3://desineuron-models}"
LOCAL_BASE="/opt/dlami/nvme/models"
S5CMD_CONCURRENCY="${S5CMD_CONCURRENCY:-32}" # Tune to NVMe write IOPS
# ── Ensure directories exist ─────────────────────────────────
mkdir -p \
"${LOCAL_BASE}/comfy/wan2.2" \
"${LOCAL_BASE}/comfy/qwen-image-2512" \
"${LOCAL_BASE}/llm/qwen3.6-35b-a3b"
# ── Check s5cmd installed ────────────────────────────────────
if ! command -v s5cmd &> /dev/null; then
echo "ERROR: s5cmd not found. Install from https://github.com/peak/s5cmd"
echo " curl -L https://github.com/peak/s5cmd/releases/latest/download/s5cmd_Linux_x86_64.tar.gz | tar xz -C /usr/local/bin"
exit 1
fi
echo "=== Velocity-OS Model Hydration ==="
echo "Source: ${S3_BUCKET}"
echo "Target: ${LOCAL_BASE}"
echo "s5cmd workers: ${S5CMD_CONCURRENCY}"
echo ""
# ── Wan 2.2 (ComfyUI — MIG slice 1) ──────────────────────────
echo "[1/3] Syncing Wan 2.2..."
s5cmd \
--numworkers "${S5CMD_CONCURRENCY}" \
--credentials-file /etc/velocity/aws-credentials \
sync \
"${S3_BUCKET}/wan2.2/*" \
"${LOCAL_BASE}/comfy/wan2.2/"
echo " ✓ Wan 2.2 synced."
# ── Qwen-Image 2512 (ComfyUI — MIG slice 1) ──────────────────
echo "[2/3] Syncing Qwen-Image 2512..."
s5cmd \
--numworkers "${S5CMD_CONCURRENCY}" \
--credentials-file /etc/velocity/aws-credentials \
sync \
"${S3_BUCKET}/qwen-image-2512/*" \
"${LOCAL_BASE}/comfy/qwen-image-2512/"
echo " ✓ Qwen-Image 2512 synced."
# ── Qwen3.6 35B A3B (SGLang — MIG slice 0) ───────────────────
echo "[3/3] Syncing Qwen3.6 35B A3B (LLM — ~70GB, be patient)..."
s5cmd \
--numworkers "${S5CMD_CONCURRENCY}" \
--credentials-file /etc/velocity/aws-credentials \
sync \
"${S3_BUCKET}/qwen3.6-35b-a3b/*" \
"${LOCAL_BASE}/llm/qwen3.6-35b-a3b/"
echo " ✓ Qwen3.6 35B synced."
# ── Verify checksums (optional — if .sha256 files exist in S3) ─
echo ""
echo "=== Verifying checksums ==="
for dir in "${LOCAL_BASE}/comfy/wan2.2" "${LOCAL_BASE}/comfy/qwen-image-2512" "${LOCAL_BASE}/llm/qwen3.6-35b-a3b"; do
if ls "${dir}"/*.sha256 2>/dev/null | head -1 | grep -q sha256; then
echo " Checking ${dir}..."
(cd "${dir}" && sha256sum -c ./*.sha256 --quiet) && echo "${dir} checksums OK"
fi
done
echo ""
echo "=== Model hydration complete ==="
echo "NVMe usage:"
du -sh "${LOCAL_BASE}"/*/* 2>/dev/null || true