背景介绍

本文将讨论 Milvus 如何调度查询任务。我们还将讨论实现 Milvus 调度的问题、解决方案和未来方向。

背景介绍

我们从《大规模向量搜索引擎中的数据管理》中了解到，向量相似性搜索是通过两个向量在高维空间中的距离来实现的。向量搜索的目标是找到与目标向量最接近的 K 个向量。

衡量向量距离的方法有很多，比如欧氏距离：

1-euclidean-distance.png

其中 x 和 y 是两个向量。n 是向量的维数。

为了在数据集中找到 K 个最近的向量，需要计算目标向量与要搜索的数据集中所有向量之间的欧氏距离。然后，根据距离对向量进行排序，以获得 K 个最近的向量。计算工作量与数据集的大小成正比。数据集越大，查询所需的计算工作量就越大。专门用于图形处理的 GPU 恰好拥有大量内核，可以提供所需的计算能力。因此，在 Milvus 的实施过程中也考虑到了多 GPU 支持。

基本概念

数据块（表文件）

为了提高对大规模数据搜索的支持，我们对 Milvus 的数据存储进行了优化。Milvus 将表中的数据按大小分割成多个数据块。在向量搜索过程中，Milvus 在每个数据块中搜索向量并合并结果。一次向量搜索操作包括 N 次独立的向量搜索操作（N 为数据块的数量）和 N-1 次结果合并操作。

任务队列（任务表）

每个资源都有一个任务数组，用于记录属于该资源的任务。每个任务都有不同的状态，包括开始、加载、载入、执行和已执行。计算设备中的加载器和执行器共享同一个任务队列。

查询调度

2-query-scheduling.png

1. 当 Milvus 服务器启动时，Milvus 会通过server_config.yaml 配置文件中的gpu_resource_config 参数启动相应的 GpuResource。DiskResource 和 CpuResource 仍无法在server_config.yaml 中编辑。GpuResource 是search_resources 和build_index_resources 的组合，在下面的示例中称为{gpu0, gpu1} ：

3 示例代码.png

3-example.png

2. Milvus 接收到一个请求。表元数据存储在外部数据库中，单主机数据库为 SQLite 或 MySQl，分布式数据库为 MySQL。收到搜索请求后，Milvus 会验证表是否存在，维度是否一致。然后，Milvus 读取表的 TableFile 列表。

4-milvus-reads-tablefile-list.png

3. Milvus 创建一个 SearchTask。由于每个 TableFile 的计算都是独立进行的，因此 Milvus 会为每个 TableFile 创建一个 SearchTask。作为任务调度的基本单位，SearchTask 包含目标向量、搜索参数和 TableFile 的文件名。

5-table-file-list-task-creator.png

4. Milvus 选择计算设备。搜索任务执行计算的设备取决于每个设备的预计完成时间。预计完成时间是指当前时间与预计计算完成时间之间的预计间隔。

例如，当一个 SearchTask 的数据块加载到 CPU 内存时，下一个 SearchTask 正在 CPU 计算任务队列中等待，而 GPU 计算任务队列处于空闲状态。CPU 的预计完成时间等于前一个搜索任务和当前搜索任务的预计时间成本之和。GPU 的预计完成时间等于数据块加载到 GPU 的时间与当前搜索任务的预计时间成本之和。资源中搜索任务的预计完成时间等于资源中所有搜索任务的平均执行时间。然后，Milvus 会选择估计完成时间最少的设备，并将 SearchTask 分配给该设备。

这里我们假设 GPU1 的估计完成时间较短。

6-GPU1-shorter-estimated-completion-time.png

5. Milvus 将 SearchTask 添加到 DiskResource 的任务队列中。

6. Milvus 将 SearchTask 移至 CpuResource 的任务队列。CpuResource 中的加载线程按顺序从任务队列中加载每个任务。CpuResource 将相应的数据块读入 CPU 内存。

7. Milvus 将 SearchTask 移至 GpuResource。GpuResource 中的加载线程将数据从 CPU 内存复制到 GPU 内存。GpuResource 将相应的数据块读入 GPU 内存。

8. Milvus 在 GpuResource 中执行 SearchTask。由于 SearchTask 的结果相对较小，因此结果会直接返回 CPU 内存。

7-scheduler.png

9. Milvus 会将 SearchTask 的结果合并为整个搜索结果。

8-milvus-merges-searchtast-result.png

所有 SearchTask 完成后，Milvus 会将整个搜索结果返回给客户端。

建立索引

建立索引的过程与搜索过程基本相同，但没有合并过程。我们将不再详细讨论。

性能优化

缓存

如前所述，数据块需要在计算前加载到相应的存储设备，如 CPU 内存或 GPU 内存。为了避免重复加载数据，Milvus 引入了 LRU（最近最少使用）缓存。当缓存满时，新数据块会推走旧数据块。你可以根据当前内存大小，通过配置文件自定义缓存大小。建议使用较大的缓存来存储搜索数据，以有效节省数据加载时间并提高搜索性能。

数据加载和计算重叠

缓存并不能满足我们提高搜索性能的需求。当内存不足或数据集过大时，需要重新加载数据。我们需要减少数据加载对搜索性能的影响。数据加载，无论是从磁盘到 CPU 内存，还是从 CPU 内存到 GPU 内存，都属于 IO 操作，几乎不需要处理器做任何计算工作。因此，我们考虑并行执行数据加载和计算，以更好地利用资源。

我们将数据块的计算分为 3 个阶段（从磁盘加载到 CPU 内存、CPU 计算、结果合并）或 4 个阶段（从磁盘加载到 CPU 内存、从 CPU 内存加载到 GPU 内存、GPU 计算和结果检索、结果合并）。以 3 阶段计算为例，我们可以启动 3 个线程，分别负责 3 个阶段，以实现指令流水线功能。由于结果集大多较小，结果合并不需要太多时间。在某些情况下，数据加载和计算的重叠可以将搜索时间缩短 1/2。

9-sequential-overlapping-load-milvus.png

问题和解决方案

不同的传输速度

此前，Milvus 在多 GPU 任务调度中使用的是循环罗宾（Round Robin）策略。这一策略在我们的 4 GPU 服务器上运行完美，搜索性能提高了 4 倍。但是，对于我们的 2 GPU 主机，性能却没有提高 2 倍。我们做了一些实验，发现一个 GPU 的数据复制速度为 11 GB/秒。而另一个 GPU 的数据复制速度为 3 GB/秒。参考主板文档后，我们确认主板通过 PCIe x16 与一个 GPU 连接，通过 PCIe x4 与另一个 GPU 连接。也就是说，这些 GPU 的复制速度不同。随后，我们增加了复制时间，以测量每个 SearchTask 的最佳设备。

未来工作

增加复杂性的硬件环境

在实际条件下，硬件环境可能会更加复杂。对于拥有多个 CPU、NUMA 架构内存、NVLink 和 NVSwitch 的硬件环境，跨 CPU/GPU 的通信会带来很多优化机会。

查询优化

在实验过程中，我们发现了一些提高性能的机会。例如，当服务器收到针对同一个表的多个查询时，在某些条件下可以合并查询。通过使用数据局部性，我们可以提高性能。我们将在未来的开发中实现这些优化。 现在，我们已经知道在单主机、多 GPU 的情况下如何调度和执行查询。我们将在接下来的文章中继续介绍 Milvus 的更多内部机制。