企业官网建设流程全解析

SkyPilot 大规模图片语义搜索实战：CLIP + ChromaDB 全链路构建图像向量数据库

【免费下载链接】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 仓库中的 Image Vector Database 应用示例 为主线，完整讲解如何在云端构建大规模图片语义搜索系统：从 OpenAI CLIP 生成图片向量，到用 ChromaDB 构建向量数据库，再到通过 Sky Serve 发布可供调用的搜索 API。读者完成后将掌握一套可直接运行的三阶段流水线（向量计算 → 建库 → 服务），并理解 SkyPilot 如何通过管理任务、存储挂载和自动扩缩容，把百万级图片的向量化与检索变成一条简单的命令行操作。

一、为什么需要向量数据库做图片搜索

1.1 传统搜索的局限与语义搜索的优势

随着图片数据量增长，传统基于关键词或元数据的搜索方式难以捕捉图片的真实语义。例如用户想找"一朵云的图片"，如果图片没有预设标签，文本检索就会失效。而向量数据库支持语义搜索：将图片和文本映射到同一个向量空间，通过计算余弦相似度直接检索概念上匹配的图片。

该示例 README 明确指出三个关键价值点：

• 可扩展性（Scalability）：现代应用需要处理数百万乃至数十亿张图片，传统数据库方案会变慢且难以管理；
• 灵活性（Flexibility）：把图片嵌入向量后，可灵活适配多种检索场景，从"找相似商品"到"找特定物体或风格"的图片；
• 性能（Performance）：向量数据库针对高维空间的最近邻查询做了专门优化，能在大规模数据集上实现实时或近实时检索。

1.2 SkyPilot 在本方案中的角色

SkyPilot 将上述大规模云上作业的运维复杂度抽象掉，通过managed jobs（管理任务）帮助用户高效、低成本地运行计算密集型任务，自动完成机器选择、调度与容错，无需手动管理任何云资源。整套方案在仓库中的代码结构如下：

examples/vector_database/ ├── README.md # 完整的分步指南 ├── batch_compute_vectors.py # 分区并批量提交向量计算任务（SDK 用法） ├── compute_vectors.yaml # 向量计算任务定义 ├── build_vectordb.yaml # 向量入库任务定义 ├── serve_vectordb.py # 服务发布（SDK 用法） ├── serve_vectordb.yaml # 服务任务定义 └── scripts/ ├── compute_vectors.py # CLIP 向量计算实现 ├── build_vectordb.py # ChromaDB 建库实现 └── serve_vectordb.py # FastAPI 搜索服务实现

二、Step 0：环境准备与配置

2.1 前提条件

开始前需要满足两项前提：

1. 安装 SkyPilot：确保已安装 SkyPilot，并运行sky check验证各云平台配置可用（若尚未配置，可参考仓库根目录的安装与云配置文档，如 README.md 中的安装说明）；
2. Hugging Face Token：示例使用 ImageNet-1k 数据集（ILSVRC/imagenet-1k），需要从 Hugging Face Hub 下载数据，因此必须配置 token。

2.2 配置 HF_TOKEN

将 token 写入~/.env文件：

HF_TOKEN=hf_xxxxx

或者直接设置环境变量HF_TOKEN。批量提交脚本 的解析逻辑印证了这两种方式：它先读取os.environ.get('HF_TOKEN')，若不存在则解析~/.env文件中以HF_TOKEN=开头的行。

三、Step 1：用 OpenAI CLIP 计算图片向量

3.1 为什么用 CLIP

图片必须被转换为向量表示（embedding）才能存入向量数据库。OpenAI 的 CLIP 模型学习了把图片和文本映射到同一向量空间的表示，这使得语义相似度计算成为可能——例如查询"一朵云的图片"（"a photo of a cloud"）可以匹配到概念上相关的图片。

在 scripts/compute_vectors.py 中，默认模型配置为ViT-bigG-14，预训练权重使用laion2b_s39b_b160k，数据集默认为ILSVRC/imagenet-1k。该脚本按流水线处理 ImageNet 图片：先初始化模型（setup_model），再迭代数据集（get_dataset_iterator），加载并预处理单张图片（do_data_loading），最后批量交给do_batch_processing完成向量计算，并通过tqdm输出进度。

3.2 一键提交：批量计算入口

python3 batch_compute_vectors.py

该脚本会自动为图片数据集分区，并为每个分区提交一个 SkyPilot managed job，实现数据并行。核心机制（见 batch_compute_vectors.py）：

• 通过calculate_job_range()把全局索引区间[start_idx, end_idx)均匀切分到多个 job，余数分配给前面的 job；
• 加载任务模板sky.Task.from_yaml('compute_vectors.yaml')；
• 对每个分区调用task.update_envs({'START_IDX': ..., 'END_IDX': ...})注入该 job 的数据范围；
• 通过task.update_secrets({'HF_TOKEN': hf_token})安全传递令牌；
• 最后调用sky.jobs.launch(task_copy, name=f'vector-compute-{job_start}-{job_end}')逐个提交。

脚本默认参数为：--start-idx 0、--end-idx 1000000（不包含）、--num-jobs 100。运行后终端会输出每个分区任务的日志，例如：

(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

也可通过sky jobs queue和sky dashboard查看任务状态与跨区域分布情况。

3.3 计算任务的任务定义：compute_vectors.yaml

任务定义文件 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: 1 file_mounts: /output: name: sky-demo-embedding # 必须与 build_vectordb.yaml 中的 source 一致 mode: MOUNT /images: name: sky-demo-image # 必须与 build_vectordb.yaml 中的 source 一致 mode: MOUNT envs: # 这些环境变量是必需的，但在启动时传入 START_IDX: '' END_IDX: '' secrets: HF_TOKEN: null # 在 CLI 中用 --secret HF_TOKEN 传入 setup: | pip install numpy==1.26.4 pip install torch==2.5.1 torchvision==0.20.1 ftfy regex tqdm pip install datasets webdataset requests Pillow open_clip_torch pip install fastapi uvicorn aiohttp pandas pyarrow tenacity 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 1000 echo "Processing complete. Results saved in node-specific files under /output/"

关键设计点：

• 多加速器回退：列出 T4/L4/A10G/A10/V100 五种 GPU，SkyPilot 会按价格从便宜到贵依次尝试，优先选用当前云上最便宜的可用加速器；
• Spot 回退：any_of块声明了"优先抢占式实例，失败则回退按需实例"的策略，有助于降低成本；
• 存储即状态：/output挂载名为sky-demo-embedding的 Sky Storage，/images挂载sky-demo-image。向量结果直接写入云存储，天然成为跨任务、跨集群共享的中间产物；
• 环境变量注入：START_IDX/END_IDX由batch_compute_vectors.py在提交时逐个 job 覆盖写入，实现数据分片；
• 模型与依赖固化：setup阶段精确锁定 numpy/torch/torchvision 版本，保证各分区任务环境一致；
• 运行参数：--batch-size 64控制每次推理的图片批量大小，--checkpoint-size 1000控制每处理 1000 张即落盘一次，避免任务中断造成大量重复计算。

四、Step 2：用 ChromaDB 构建向量数据库

4.1 建库任务定义：build_vectordb.yaml

得到图片向量后，需要一个专门的引擎在海量向量上进行快速相似度检索，本示例选用ChromaDB存储与查询向量。建库任务定义见 build_vectordb.yaml：

name: vectordb-build workdir: . file_mounts: /clip_embeddings: name: sky-demo-embedding # 必须与 compute_vectors.yaml 中的 source 一致 mode: MOUNT /vectordb: name: sky-vectordb # 必须与 serve_vectordb.yaml 中的 source 一致 mode: MOUNT /images: name: sky-demo-image # 必须与 compute_vectors.yaml 中的 source 一致 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

执行建库命令：

sky jobs launch build_vectordb.yaml

该任务会批量读取上一步生成的 CLIP 向量（parquet 分片），逐批写入 ChromaDB，日志形如：

(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]

4.2 建库实现要点（scripts/build_vectordb.py）

从 scripts/build_vectordb.py 源码可以看到几个值得注意的实现细节：

• 临时目录策略：代码明确注释"挂载的存储桶不支持追加操作，因此先在 tmpdir 中构建，再拷贝到最终位置"，即with tempfile.TemporaryDirectory() as temp_dir中初始化chromadb.PersistentClient(path=temp_dir)，建库完成后整体拷贝到/vectordb挂载点。这规避了对象存储随机写的限制；
• 集合去重创建：先尝试create_collection，若抛ValueError（已存在）则改用get_collection，实现幂等重跑；
• 并行处理：每个 parquet 文件由process_parquet_file独立处理，通过ProcessPoolExecutor并行执行，worker 数为max(1, cpu_count() - 1)（留一个 CPU 给调度）；文件内再按--batch-size（默认 1000）分批，解包 pickled 数据（images_base64+embeddings），以str(idx)作为 ChromaDB 的文档 ID；
• 命令行参数：--collection-name（默认clip_embeddings）、--persist-dir（默认/vectordb/chroma）、--embeddings-dir（默认/clip_embeddings）、--prefix（挂载桶内搜索 parquet 的前缀路径，默认空字符串即全桶扫描）。

五、Step 3：对外服务向量数据库

建库完成后，需要暴露一个 API 端点供其他应用或本地客户端调用语义搜索，并可验证数据库是否正常工作。SkyPilot 提供 CLI 与 SDK 两种方式。

5.1 Option 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 --endpoint

5.2 Option 2：SDK 方式

通过 Python SDK 也能完成同样的部署：

python3 serve_vectordb.py

运行后脚本会打印服务端点地址，可直接用于查询。若想以 Sky Serve 托管服务方式发布：

python3 serve_vectordb.py --serve

5.3 服务任务定义：serve_vectordb.yaml

serve_vectordb.yaml 定义服务所需资源与探针：

name: vectordb-serve workdir: . resources: accelerators: # 按价格排序（最便宜到最贵） # SkyPilot 会尝试使用最便宜的可用加速器 # 服务需要 GPU 来计算嵌入向量 T4: 1 L4: 1 A10G: 1 A10: 1 V100: 1 memory: 32+ ports: 8000 use_spot: true file_mounts: /vectordb: name: sky-vectordb # 必须与 build_vectordb.yaml 中的 source 一致 mode: MOUNT /images: name: sky-demo-image # 必须与 build_vectordb.yaml 中的 source 一致 mode: MOUNT setup: | pip install numpy==1.26.4 pip install torch==2.5.1 torchvision==0.20.1 ftfy regex tqdm pip install open_clip_torch chromadb pandas pip install fastapi uvicorn pydantic run: | python scripts/serve_vectordb.py \ --collection-name clip_embeddings \ --persist-dir /vectordb/chroma \ --images-dir /images \ --host 0.0.0.0 \ --port 8000 service: replicas: 1 readiness_probe: path: /health

要点说明：

• 服务端必须使用 GPU 重新加载 CLIP 模型以实时编码查询文本，因此同样配置了多加速器回退与use_spot: true；
• ports: 8000声明对外服务端口，run命令绑定0.0.0.0:8000；
• service.readiness_probe.path: /health让 Sky Serve 通过/health端点判断实例是否就绪，未就绪的副本不会被路由流量；
• /vectordb（sky-vectordb）与/images（sky-demo-image）分别挂载建库产物与原始图片，实现"建库集群写、服务集群读"的存储解耦。

5.4 服务实现要点（scripts/serve_vectordb.py）

scripts/serve_vectordb.py 基于 FastAPI 实现，核心端点与逻辑如下：

• GET /health：返回{'status': 'healthy', 'collection_size': collection.count()}，作为 Sky Serve 的就绪探针，同时暴露当前集合规模；
• POST /search：接收 JSON 请求体（text查询文本、可选n_results，默认 5），流程为：
  1. 用 CLIP 将文本编码为向量：open_clip.create_model_and_transforms('ViT-bigG-14', pretrained='laion2b_s39b_b160k')加载模型，encode_text()在torch.no_grad()下编码并对特征做 L2 归一化；
  2. 调用collection.query(query_embeddings=..., n_results=..., include=['metadatas', 'distances', 'documents'])检索最近邻；
  3. 将 ChromaDB 的距离转换为余弦相似度：similarity = 1 - distance / 2；
  4. 返回结构为[{image_path, similarity}]的SearchResult列表；
• GET /image/{subpath:path}：从挂载的/images目录读取并返回对应图片（FileResponse，媒体类型image/jpeg），供前端直接展示检索结果；
• GET /：内置一个带样式的搜索页面（HTML），提供搜索输入框与结果网格展示，方便快速验证检索效果。

六、数据流转与存储约定：三段任务的"连接线"

三个任务通过Sky Storage 命名约定紧密衔接，这是整个方案最需要理解的全局设计。下表汇总了各任务挂载的存储桶及其"必须一致"的关系：

各 YAML 文件中的注释反复强调"this needs to be the same as the source in ..."，原因在于：Sky Storage 是跨集群、跨任务共享的持久化数据层，只有名称完全一致，后一阶段才能读到前一阶段的产物。这一约定让"计算、建库、服务"三个阶段可以在不同时刻、不同集群、甚至不同云区域独立运行，彼此仅通过存储解耦。

七、进阶：将方案推广到自己的数据集

将本示例迁移到自有数据，主要改动集中在 scripts/compute_vectors.py 的几个参数与数据源：

1. 更换数据集：--dataset-name参数（默认ILSVRC/imagenet-1k）替换为 Hugging Face 上的任意图片数据集，并相应调整get_dataset_iterator的解析逻辑；
2. 调整计算预算：batch_compute_vectors.py的--num-jobs（默认 100）、--start-idx/--end-idx（默认[0, 1000000)）决定分区粒度；分区越细，并行度越高，但每个 job 的调度开销也越大；
3. 更换向量模型：--model-name（默认ViT-bigG-14）与--pretrained（默认laion2b_s39b_b160k）可替换为open_clip支持的其他 CLIP 变体，注意tokenizer与模型必须配套；
4. 更换向量库：如需替换 ChromaDB（如 Milvus、Qdrant），只需改写 scripts/build_vectordb.py 的写入端与 scripts/serve_vectordb.py 的查询端，SkyPilot 侧的任务编排、存储挂载与弹性调度完全不需要改动；
5. 成本控制：compute_vectors.yaml与serve_vectordb.yaml中的accelerators列表与any_of/use_spot策略可直接复用——SkyPilot 会尝试从列表中最便宜的 GPU 开始分配，抢不到再依次回退。

八、总结

本示例演示了一条完整的大规模图片语义搜索流水线，其技术要点可归纳为：

1. 数据并行：通过 batch_compute_vectors.py 的calculate_job_range把百万级图片切成 100 个独立区间，每个区间一个 managed job，互不阻塞；
2. 中间产物存储化：向量结果、建库产物与原始图片全部落在 Sky Storage（sky-demo-embedding、sky-vectordb、sky-demo-image），三个阶段只靠存储名称衔接，天然支持跨集群、跨区域运行；
3. 弹性且省成本：多 GPU 回退 + Spot/按需回退的组合，让昂贵的大规模向量计算任务自动跑在"当下最便宜可用的机器"上；
4. 生产级服务：Sky Serve 的readiness_probe（/health）与自动扩缩容能力，让向量检索服务一键具备对外可用的高可用形态。

读者可以按README.md的 Step 0 → Step 3 顺序自行复现：先配置 HF_TOKEN 与 SkyPilot 环境，执行python3 batch_compute_vectors.py计算向量，再sky jobs launch build_vectordb.yaml建库，最后用sky serve up serve_vectordb.yaml -n vectordb发布服务，并用POST /search或内置 Web 页面验证"一朵云的图片"这类语义查询的检索效果。相关实现细节可继续研读 compute_vectors.yaml、build_vectordb.yaml、serve_vectordb.yaml 及scripts/目录下的三个 Python 文件。

【免费下载链接】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