feat: 媒体机直传 + 向量回填
- media: spider 下载图片后按内容 sha1 生成 key，经后端预签名直传 S4，上报只带 key （spider 不持有 AK/SK；单张失败自动回退为上报 URL） - embed: 新增 embed 子命令，从后端领取缺向量的图片→下载→本地 CLIP 编码→回传 （媒体拆分后后端拿不到图片字节，向量只能靠回填补齐） - serve: 服务化常驻模式（领取 crawl 任务 → 抓取 → 回报） - config: 新增 media.upload/concurrency、vision.endpoint/concurrency
This commit is contained in:
toom1996
255
internal/embed/worker.go
Normal file
255
internal/embed/worker.go
Normal file
@ -0,0 +1,255 @@
|
||||
// Package embed 是「向量回填」的执行端。
|
||||
//
|
||||
// 为什么需要它：媒体拆分后 spider 直接传图到对象存储，后端拿到的是 key 而不是图片字节，
|
||||
// 于是入库时算不出向量（image_embeddings 缺行）。这个 worker 补上这一段：
|
||||
//
|
||||
// 向后端要「缺向量的图片」 → 从存储下载 → 送本地 CLIP 编码 → 把向量回传给后端
|
||||
//
|
||||
// 它跑在**有 GPU 的机器**上，所以放在 spider 仓库里（和爬虫同一份部署物），
|
||||
// 但逻辑上只在用 GPU，与抓取完全独立：不抓 Vogue、不写业务库。
|
||||
package embed
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"my-spiders/internal/ingest"
|
||||
)
|
||||
|
||||
// Worker 向量回填执行端。
|
||||
type Worker struct {
|
||||
Client *ingest.Client // 取任务 + 回传（HMAC）
|
||||
Vision string // 视觉服务地址，如 http://127.0.0.1:8900
|
||||
Concurrency int // 下载并发
|
||||
Batch int // 单次送编码的张数
|
||||
HTTP *http.Client
|
||||
}
|
||||
|
||||
// Stats 一轮处理的统计。
|
||||
type Stats struct {
|
||||
Fetched int // 后端派发的张数
|
||||
Saved int // 后端确认写入的张数
|
||||
Failed int // 失败张数（已计入失败计数）
|
||||
Remaining int // 后端统计的剩余量（-1 = 未知）
|
||||
}
|
||||
|
||||
// New 创建 worker。并发/批量传 0 时取默认值（4 / 8）。
|
||||
func New(client *ingest.Client, vision string, concurrency, batch int) *Worker {
|
||||
if concurrency <= 0 {
|
||||
concurrency = 4
|
||||
}
|
||||
if batch <= 0 {
|
||||
batch = 8
|
||||
}
|
||||
return &Worker{
|
||||
Client: client,
|
||||
Vision: vision,
|
||||
Concurrency: concurrency,
|
||||
Batch: batch,
|
||||
// 图片可能几百 KB~几 MB，60s 足够；向量 POST 走同一个 client。
|
||||
HTTP: &http.Client{Timeout: 60 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
// RunOnce 处理一轮：拉取 → 下载 → 编码 → 回传。
|
||||
//
|
||||
// 不会因为单张失败而中断：下载失败/编码跳过的图进 failures，由后端累加计数，
|
||||
// 连续失败到上限后自然退出队列（坏图不会无限占位）。
|
||||
func (w *Worker) RunOnce(ctx context.Context, kind string, limit int) (Stats, error) {
|
||||
var st Stats
|
||||
pending, remaining, err := w.Client.FetchPendingEmbeddings(ctx, kind, limit)
|
||||
if err != nil {
|
||||
return st, err
|
||||
}
|
||||
st.Fetched = len(pending)
|
||||
st.Remaining = remaining
|
||||
if len(pending) == 0 {
|
||||
log.Printf("[embed] %s 无待回填图片（取时剩余 %d）", kind, remaining)
|
||||
return st, nil
|
||||
}
|
||||
|
||||
// ① 并发下载：后端给的是可直接访问的 url，worker 不关心存储实现。
|
||||
data := make([][]byte, len(pending))
|
||||
errs := make([]error, len(pending))
|
||||
w.parallel(len(pending), func(i int) {
|
||||
body, derr := w.download(ctx, pending[i].URL)
|
||||
if derr != nil {
|
||||
errs[i] = derr
|
||||
return
|
||||
}
|
||||
data[i] = body
|
||||
})
|
||||
|
||||
// ② 编码：只把下载成功的送进去，按 Batch 分批（视觉服务对单次张数有上限）。
|
||||
type job struct {
|
||||
pos int
|
||||
img []byte
|
||||
}
|
||||
jobs := make([]job, 0, len(pending))
|
||||
for i := range pending {
|
||||
if errs[i] == nil {
|
||||
jobs = append(jobs, job{pos: i, img: data[i]})
|
||||
}
|
||||
}
|
||||
|
||||
items := make([]ingest.EmbeddingVector, 0, len(jobs))
|
||||
failures := make([]ingest.EmbeddingFailure, 0, len(pending)-len(jobs))
|
||||
|
||||
for start := 0; start < len(jobs); start += w.Batch {
|
||||
end := start + w.Batch
|
||||
if end > len(jobs) {
|
||||
end = len(jobs)
|
||||
}
|
||||
batch := jobs[start:end]
|
||||
imgs := make([][]byte, len(batch))
|
||||
for i := range batch {
|
||||
imgs[i] = batch[i].img
|
||||
}
|
||||
vecs, idxs, skipped, eerr := w.embed(ctx, imgs)
|
||||
if eerr != nil {
|
||||
for _, b := range batch {
|
||||
p := pending[b.pos]
|
||||
failures = append(failures, ingest.EmbeddingFailure{Kind: kind, ImageID: p.ImageID, Error: eerr.Error()})
|
||||
}
|
||||
continue
|
||||
}
|
||||
// indexes 缺省时按位置对应；存在时严格按它映射（视觉服务可能跳过坏图）。
|
||||
for j := range vecs {
|
||||
posInBatch := j
|
||||
if len(idxs) == len(vecs) {
|
||||
posInBatch = idxs[j]
|
||||
}
|
||||
if posInBatch < 0 || posInBatch >= len(batch) {
|
||||
continue
|
||||
}
|
||||
p := pending[batch[posInBatch].pos]
|
||||
items = append(items, ingest.EmbeddingVector{Kind: kind, ImageID: p.ImageID, Vector: vecs[j]})
|
||||
}
|
||||
for _, s := range skipped {
|
||||
if s < 0 || s >= len(batch) {
|
||||
continue
|
||||
}
|
||||
p := pending[batch[s].pos]
|
||||
failures = append(failures, ingest.EmbeddingFailure{Kind: kind, ImageID: p.ImageID, Error: "视觉服务跳过（图片无法解码）"})
|
||||
}
|
||||
}
|
||||
|
||||
// ③ 下载失败的也一并上报，让后端计数
|
||||
for i := range pending {
|
||||
if errs[i] != nil {
|
||||
failures = append(failures, ingest.EmbeddingFailure{Kind: kind, ImageID: pending[i].ImageID, Error: errs[i].Error()})
|
||||
}
|
||||
}
|
||||
|
||||
if len(items) == 0 && len(failures) == 0 {
|
||||
return st, nil
|
||||
}
|
||||
saved, failed, serr := w.Client.SubmitEmbeddings(ctx, items, failures)
|
||||
if serr != nil {
|
||||
return st, serr
|
||||
}
|
||||
st.Saved = saved
|
||||
st.Failed = failed
|
||||
log.Printf("[embed] %s 本轮：取到 %d 张 → 写入 %d，失败 %d（剩余 %d）",
|
||||
kind, st.Fetched, st.Saved, st.Failed, st.Remaining)
|
||||
return st, nil
|
||||
}
|
||||
|
||||
// embed 调本地视觉服务编码一批图片。
|
||||
// 返回的 idxs 是「vectors 对应入参下标」；skipped 是无法解码的下标。
|
||||
func (w *Worker) embed(ctx context.Context, imgs [][]byte) ([][]float32, []int, []int, error) {
|
||||
payload := struct {
|
||||
Images []string `json:"images"`
|
||||
}{Images: make([]string, len(imgs))}
|
||||
for i, b := range imgs {
|
||||
payload.Images[i] = base64.StdEncoding.EncodeToString(b)
|
||||
}
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, w.Vision+"/embed", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := w.HTTP.Do(req)
|
||||
if err != nil {
|
||||
return nil, nil, nil, fmt.Errorf("视觉服务不可用: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, nil, nil, fmt.Errorf("视觉服务返回 %d: %s", resp.StatusCode, truncate(string(raw), 120))
|
||||
}
|
||||
var out struct {
|
||||
Dim int `json:"dim"`
|
||||
Vectors [][]float32 `json:"vectors"`
|
||||
Indexes []int `json:"indexes"`
|
||||
Skipped []int `json:"skipped"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return nil, nil, nil, fmt.Errorf("视觉服务响应解析失败: %w", err)
|
||||
}
|
||||
return out.Vectors, out.Indexes, out.Skipped, nil
|
||||
}
|
||||
|
||||
func (w *Worker) download(ctx context.Context, url string) ([]byte, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := w.HTTP.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("下载失败 status=%d", resp.StatusCode)
|
||||
}
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(body) == 0 {
|
||||
return nil, fmt.Errorf("下载到空文件")
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
// parallel 以 Concurrency 路并发执行 fn(0..n-1)。
|
||||
func (w *Worker) parallel(n int, fn func(i int)) {
|
||||
if n == 0 {
|
||||
return
|
||||
}
|
||||
if w.Concurrency <= 0 {
|
||||
w.Concurrency = 1
|
||||
}
|
||||
sem := make(chan struct{}, w.Concurrency)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
sem <- struct{}{}
|
||||
defer func() { <-sem }()
|
||||
fn(i)
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n] + "..."
|
||||
}
|
||||
Reference in New Issue
Block a user