基于 SkyPilot 与 OpenAI CLIP 构建大规模图像语义搜索的向量数据库实战
【免费下载链接】skypilotThe AI Compute Platform for frontier teams. SkyPilot turns fragmented AI compute into one AI supercomputer, so frontier AI teams build custom intelligence faster.项目地址: https://gitcode.com/GitHub_Trending/sk/skypilot
本文以 SkyPilot 仓库中的examples/vector_database为例,完整讲解如何把海量图像通过 CLIP 模型转换为向量嵌入(Embedding),使用 ChromaDB 构建可实时查询的向量数据库,并以 HTTP 服务的形式对外提供"以文搜图"的语义检索能力。读完本文,你将掌握一套可复用的三步工作流:分布式向量计算(SkyPilot Managed Jobs)→ 向量入库(ChromaDB)→ 服务化部署(Sky Launch / Sky Serve),并理解每一步背后的 YAML 配置与源码实现细节。
一、为什么大规模图像搜索需要向量数据库
随着图像数据量的增长,基于关键词或元数据的传统检索方式越来越难以捕获图像的完整语义。例如查询"云朵的照片",传统标签搜索只能命中显式打了"cloud"标签的图片,而向量数据库则能把图像和文本映射到同一个语义空间中,实现真正的语义匹配。
在本示例中,向量数据库带来的核心收益可以归纳为三点:
- 可扩展性:现代应用可能面对数百万甚至数十亿张图片,传统数据库方案在大数据量下会变慢且难以管理,而向量数据库专为高维最近邻(Nearest Neighbor)查询而设计。
- 灵活性:把图像存储为向量嵌入后,可以灵活适配不同检索场景,从"找相似商品"到"找包含特定物体或风格的图片"。
- 性能:向量数据库针对高维空间的最近邻查询做了深度优化,能够在千万级甚至亿级向量上实现实时或近实时的检索。
而 SkyPilot 的价值在于把"跑大规模计算任务"这件事的云上基础设施复杂度抽象掉:它负责找机器、挂载数据、调度任务,帮助用户以高效且经济的方式运行这类计算密集型任务。
整个示例的工作流清晰分为三步,这也是本文的主线:
| 步骤 | 目标 | 产物 |
|---|---|---|
| Step 1 | 用 OpenAI CLIP 将图像批量编码为向量 | Parquet 格式的嵌入文件 |
| Step 2 | 将嵌入批量写入 ChromaDB | 持久化的向量数据库 |
| Step 3 | 将数据库封装为 HTTP 服务对外提供语义检索 | 可查询的/searchAPI 端点 |
对应的完整代码与配置均位于仓库的 examples/vector_database 目录下。
二、Step 0:环境准备
2.1 安装 SkyPilot 并通过云端检查
首先确保 SkyPilot 已正确安装,并运行以下命令验证云端凭据与资源配置可用:
sky checksky check会探测当前环境可用的所有云提供商(AWS、GCP、Azure、Kubernetes 等)及其配额情况。该命令必须成功返回,后续所有sky jobs launch、sky launch、sky serve up操作才具备执行前提。
2.2 配置 Hugging Face Token
本示例的数据集来自 Hugging Face Hub,下载数据集需要 HF Token。有两种配置方式:
方式一:在~/.env文件中写入:
HF_TOKEN=hf_xxxxx方式二:直接设置环境变量:
export HF_TOKEN=hf_xxxxx值得注意的是,这个 Token 不仅用于本地下载,还会被安全地传递到云端任务。在 batch_compute_vectors.py 的源码中可以看到:脚本会优先从环境变量HF_TOKEN读取,若不存在则解析~/.env文件;两种来源都取不到时会直接抛出ValueError终止执行。而 compute_vectors.yaml 中通过secrets: HF_TOKEN: null声明该密钥,由task_copy.update_secrets({'HF_TOKEN': hf_token})在启动任务时注入,避免了把 Token 硬编码进 YAML 或明文写入日志。
三、Step 1:用 OpenAI CLIP 计算图像向量
3.1 一键启动分布式向量计算
CLIP 模型能够把图像与文本映射到同一个嵌入空间,从而让"一张云朵的照片"这样的文本查询与相关图像在向量空间中距离更近。本步骤的目标就是把整个图像数据集的每一张图片编码成向量。
直接运行仓库中的启动脚本即可:
python3 batch_compute_vectors.py脚本会自动寻找可用的云上机器并行计算。运行后你会在任务日志中看到类似输出:
(clip-batch-compute-vectors, pid=2523) 2025-01-27 23:57:27,387 - root - INFO - Saved partition 2 to /output/embeddings_90000_100000.parquet_part_2/data.parquet (clip-batch-compute-vectors, pid=2523) 2025-01-27 23:59:39,720 - root - INFO - Saved partition 3 to /output/embeddings_90000_100000.parquet_part_3/data.parquet (clip-batch-compute-vectors, pid=2523) 2025-01-28 00:01:56,707 - root - INFO - Saved partition 4 to /output/embeddings_90000_100000.parquet_part_4/data.parquet注意embeddings_90000_100000这类命名:每个分区负责数据集的一个索引区间,编码结果以 Parquet 格式落盘。
3.2 任务分区逻辑:把大任务拆成并行小任务
batch_compute_vectors.py 是整个 Step 1 的调度核心,它做了三件事:
- 按索引区间切分数据:
calculate_job_range()将全局[start_idx, end_idx)区间按任务数均分,并把余数均匀分配给前几个任务,保证负载均衡。脚本默认参数为--start-idx 0、--end-idx 1000000、--num-jobs 100,即默认把 100 万张图片切给 100 个并行任务。 - 动态注入每个任务的环境变量:通过
task.update_envs({'START_IDX': ..., 'END_IDX': ...})为每个任务写入各自的索引区间。 - 批量提交 Managed Jobs:循环调用
sky.jobs.launch(task_copy, name=f'vector-compute-{job_start}-{job_end}'),一次性向云上提交多个托管任务。
3.3 compute_vectors.yaml 配置解读
Step 1 的任务定义在 compute_vectors.yaml,这是理解整个流程的关键配置文件,值得逐段拆解:
name: clip-batch-compute-vectors workdir: . resources: accelerators: # 按价格排序(最便宜到最贵) T4: 1 L4: 1 A10G: 1 A10: 1 V100: 1 memory: 32+ any_of: - use_spot: true - use_spot: false num_nodes: 1accelerators列表:SkyPilot 会按声明顺序优先选用最便宜的可用 GPU(这里是 T4),只有最便宜的不可用时才依次尝试 L4、A10G 等。这正体现了 SkyPilot"找到性价比最优机器"的调度策略,与 serve_vectordb.yaml 中的声明方式完全一致。any_of: use_spot二元选择:允许任务在 Spot(抢占式实例,更便宜)与按需实例之间自由选择,让调度器以成本为导向决策。memory: 32+:要求机器至少 32GB 内存,为 CLIP 模型加载与批量编码预留足够空间。
文件挂载与任务运行部分同样关键:
file_mounts: /output: name: sky-demo-embedding # 必须与 build_vectordb.yaml 中的 source 一致 mode: MOUNT /images: name: sky-demo-image mode: MOUNT envs: START_IDX: '' END_IDX: '' secrets: HF_TOKEN: null # 启动时注入 run: | python scripts/compute_vectors.py \ --output-path "/output/embeddings_${START_IDX}_${END_IDX}.parquet" \ --start-idx ${START_IDX} \ --end-idx ${END_IDX} \ --batch-size 64 \ --checkpoint-size 1000file_mounts把云端对象存储桶挂载到任务节点:/images提供原始图像数据,/output承接编码结果。因为多个任务写同一个桶,SkyPilot 使用sky-demo-embedding与sky-demo-image两个命名存储桶(Sky Storage)来跨任务共享数据。mode: MOUNT表示以挂载方式访问存储,而非每次启动拷贝快照,避免海量数据的传输开销。- 断点续算:
--checkpoint-size 1000让计算脚本每处理 1000 张图片写一次检查点,Spot 实例被回收后可从中断处继续,配合 Managed Jobs 的重试机制实现弹性容错。
3.4 监控计算任务
计算任务提交后,可以通过 SkyPilot 的两条命令查看状态:
sky jobs queue # 查看所有托管任务的排队与运行状态 sky dashboard # 打开可视化 Dashboard,可看到任务分布在哪些区域从 batch_compute_vectors.py 的实现可以看到,这些任务是彼此独立的 Managed Jobs,SkyPilot 会在不同区域、不同可用区自动调度它们,充分利用 Spot 价格差降低总成本。
四、Step 2:把嵌入构建成向量数据库
4.1 为什么需要专用向量引擎
有了图像嵌入之后,还需要一个专门的高维向量检索引擎来支撑毫秒级的相似度查询。本示例选用ChromaDB作为向量数据库,负责把 Step 1 产出的嵌入批量写入并持久化,从而支持在百万级向量上的实时或近实时检索。
4.2 build_vectordb.yaml 配置解读
构建数据库的任务定义在 build_vectordb.yaml:
name: vectordb-build workdir: . file_mounts: /clip_embeddings: name: sky-demo-embedding # 与 compute_vectors.yaml 的 /output 是同一个桶 mode: MOUNT /vectordb: name: sky-vectordb # 与 serve_vectordb.yaml 的 /vectordb 是同一个桶 mode: MOUNT /images: name: sky-demo-image mode: MOUNT setup: | pip install chromadb pandas tqdm pyarrow run: | python scripts/build_vectordb.py \ --collection-name clip_embeddings \ --persist-dir /vectordb/chroma \ --embeddings-dir /clip_embeddings \ --batch-size 1000三个挂载桶的分工非常清晰,且相互之间存在严格的命名约定(YAML 注释也反复强调"必须与另一份 YAML 中的 source 一致"):
/clip_embeddings→sky-demo-embedding:读入 Step 1 的计算结果;/vectordb→sky-vectordb:写入构建好的 ChromaDB 持久化目录,供 Step 3 读取;/images→sky-demo-image:保留原始图片,供服务端返回检索结果的图像内容。
这条"一个桶写、另一个桶读"的传递链是三步流水线能在不同任务、不同集群间无缝衔接的关键。
4.3 源码实现细节
scripts/build_vectordb.py 的实现揭示了构建过程的核心逻辑:
- 扫描 Parquet 文件:
list_local_parquet_files()用glob递归搜索挂载桶下所有**/*.parquet文件,说明 Step 1 的多个分区输出会被一次性全部纳入。 - 多进程批处理:
process_parquet_file()通过ProcessPoolExecutor并行处理每个 Parquet 文件——用pandas.read_parquet读取,按--batch-size(默认 1000)分片,解包其中的 pickle 序列化数据,拆出idx、向量和 base64 编码的图像。 --batch-size的调优含义:该参数"需要能放进内存"(源码注释原话),批量越大单次写库效率越高,但对节点内存要求越高,需与 YAML 中的memory: 32+配合权衡。
运行构建任务:
sky jobs launch build_vectordb.yaml日志会显示逐个处理 Parquet 文件的过程:
(vectordb-build, pid=2457) INFO:__main__:Processing /clip_embeddings/embeddings_0_500.parquet_part_0/data.parquet Processing batches: 100%|██████████| 1/1 [00:00<00:00, 1.19it/s] Processing files: 100%|██████████| 12/12 [00:05<00:00, 2.04it/s]4.4 注意
ChromaDB 采用--collection-name clip_embeddings作为集合名,此名称在 Step 3 的建库与服务配置中必须保持一致(build_vectordb.yaml 与 serve_vectordb.yaml 均使用clip_embeddings),否则服务端将无法找到集合。
五、Step 3:把向量数据库服务化
构建完成后,需要把数据库封装成一个 API 服务,让本地客户端或其他应用(如图像搜索引擎、推荐系统)能够调用它执行语义检索。本示例提供两种部署方式:CLI与SDK。
5.1 方式一:CLI 部署
模式 A:常规集群部署
sky launch -c vecdb_serve serve_vectordb.yaml这会启动一个名为vecdb_serve的集群,在集群上常驻运行向量数据库服务。查询部署地址:
sky status --ip vecdb_serve模式 B:Sky Serve 服务化部署
sky serve up serve_vectordb.yaml -n vectordb这种方式把向量数据库部署为云上的托管服务,Sky Serve 会自动完成健康检查与弹性扩缩容,并通过公网端点对外提供服务。查询服务端点:
sky serve status vectordb --endpoint5.2 方式二:SDK 部署
serve_vectordb.py 提供了等价的 Python SDK 方式:
python3 serve_vectordb.py # 以集群方式启动 python3 serve_vectordb.py --serve # 以 Sky Serve 服务方式启动脚本内部通过sky.Task.from_yaml('serve_vectordb.yaml')加载任务定义,然后依据是否带--serve参数分别调用serve_sdk.up(task, service_name='vectordb-serve')或sky.launch(task, cluster_name='vectordb-serve'),最后用sky.stream_and_get(req_id)同步等待部署完成,并自动从状态中解析出endpoint打印出来:
if args.serve: serve_status = sky.get(serve_sdk.status(service_names=_SERVICE_NAME)) endpoint = serve_status[0]['endpoint'] else: cluster_status = sky.get(sky.endpoints(cluster=_SERVICE_NAME)) endpoint = cluster_status[int(port)] print(f'endpoint: {endpoint}')脚本还会自动从任务资源中解析出ports字段(只支持单个端口),无需手动硬编码端口号。
5.3 serve_vectordb.yaml:服务端配置详解
serve_vectordb.yaml 与前两个配置最明显的差异是多了service段:
resources: accelerators: T4: 1 L4: 1 A10G: 1 A10: 1 V100: 1 memory: 32+ ports: 8000 use_spot: true file_mounts: /vectordb: name: sky-vectordb mode: MOUNT /images: name: sky-demo-image mode: MOUNT service: replicas: 1 readiness_probe: path: /healthports: 8000:声明服务监听端口,Sky Serve 据此自动完成负载均衡器的端口映射。service.readiness_probe.path: /health:Sky Serve 通过轮询/health路径判断副本是否就绪,只有就绪的副本才会接收流量——这正是 README 中"自动健康检查"的落地配置。- 为什么服务端必须要有 GPU:YAML 注释明确说明"serve requires a GPU to compute the embeddings"——在线推理阶段,用户的文本查询需要实时用 CLIP 编码为向量,因此需要 GPU 而非仅做存储查询。
5.4 在线查询的源码链路
scripts/serve_vectordb.py 实现了完整的 FastAPI 服务,其查询链路可以拆解为四步:
- 文本编码:
encode_text()用 CLIP 的tokenizer处理查询文本,model.encode_text(text_tokens)得到特征向量,并在torch.no_grad()下做 L2 归一化(text_features /= text_features.norm(dim=-1, keepdim=True)),确保与建库阶段存储的向量处于同一度量空间。 - 向量检索:
query_collection()调用 ChromaDB 的collection.query(query_embeddings=..., n_results=..., include=['metadatas', 'distances', 'documents'])找出最近邻向量。 - 距离转相似度:由于 ChromaDB 默认返回 L2 距离,源码用
similarity = 1 - (distance / 2)将其换算为余弦相似度分值,返回给调用方一个直观的匹配分数。 - API 暴露:
@app.post('/search')端点接收SearchQuery(text查询词 + 可选n_results,默认返回 5 条结果),返回SearchResult列表,每项包含命中的image_path与similarity。
拿到上一步打印的 endpoint 后,即可用任意 HTTP 客户端发起语义检索,例如向http://<endpoint>/searchPOST 一条{"text": "a photo of a cloud"},服务会返回语义最接近的图像及其相似度。整个服务还可嵌入更大的应用(如图像搜索引擎、推荐系统)作为检索后端。
六、三步流水线背后的设计要点
6.1 存储桶贯穿始终:跨任务的数据传递契约
细读三份 YAML 可以发现,整个流水线没有走任何本地文件传输,全靠三个命名存储桶接力:
| 存储桶名 | 写入方 | 读取方 | 承载数据 |
|---|---|---|---|
sky-demo-image | 用户(数据集) | Step 1 / Step 2 / Step 3 | 原始图像 |
sky-demo-embedding | Step 1(/output) | Step 2(/clip_embeddings) | CLIP 嵌入(Parquet) |
sky-vectordb | Step 2(/vectordb) | Step 3(/vectordb) | ChromaDB 持久化目录 |
由于不同步骤运行在不同的集群与任务上,这种以云端对象存储为中介的设计让任务之间彻底解耦:任何一步都可以独立重跑、横向扩容,而不会破坏后续步骤的数据完整性。
6.2 成本优化贯穿始终
从 Step 1 到 Step 3,成本优化的手段层层叠加:
- 多 GPU 候选 + 按价格排序:T4 → L4 → A10G → A10 → V100,SkyPilot 自动选择当前可用且最便宜的加速器;
- Spot 优先:Step 1 通过
any_of允许 Spot/按需二选一,Step 3 直接use_spot: true,充分利用可抢占实例的低价; - 断点续算:
--checkpoint-size配合 Managed Jobs 自动重启,让 Spot 回收不再导致整段任务返工。
6.3 托管任务与托管服务的分工
- 批处理(Step 1/2)使用
sky jobs launch的Managed Jobs:适合有始有终、可并行切分、支持自动重试的离线任务; - 在线服务(Step 3)使用
sky serve up的Sky Serve:适合需要常驻、对外暴露端点、具备就绪探针与扩缩容能力的服务型负载。
两者在 compute_vectors.yaml、build_vectordb.yaml 与 serve_vectordb.yaml 三份配置中形成了清晰的分工模式,可作为同类"离线建库 + 在线检索"应用的通用范本。
七、总结
从一张张原始图片到可实时语义检索的在线 API,这个示例完整覆盖了向量检索应用的全生命周期:
python3 batch_compute_vectors.py:把百万级图像数据集切分成上百个并行任务,用 CLIP 编码为向量,落盘到共享存储桶;sky jobs launch build_vectordb.yaml:读取嵌入,以多进程批处理方式构建 ChromaDB 集合并持久化;sky launch/sky serve up/ SDK:把数据库包装成带健康检查与扩缩容能力的 HTTP 服务,通过/search端点实现"以文搜图"。
本示例的完整代码、三份 YAML 配置与全部 Python 脚本均可在 examples/vector_database 目录中复现运行,其中 batch_compute_vectors.py、scripts/build_vectordb.py 与 scripts/serve_vectordb.py 分别对应上述三个步骤的核心实现,可以作为构建自有图像语义检索系统的起点。
【免费下载链接】skypilotThe AI Compute Platform for frontier teams. SkyPilot turns fragmented AI compute into one AI supercomputer, so frontier AI teams build custom intelligence faster.项目地址: https://gitcode.com/GitHub_Trending/sk/skypilot
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考