Skip to content

Commit

Permalink
Support Milvus 2.4 GPU indexes
Browse files Browse the repository at this point in the history
Closes #74
  • Loading branch information
roji committed Mar 29, 2024
1 parent a6e62e0 commit ede11f6
Show file tree
Hide file tree
Showing 5 changed files with 88 additions and 1 deletion.
2 changes: 1 addition & 1 deletion Milvus.Client.Tests/CollectionTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ public async Task Compact()
});

long compactionId = await collection.CompactAsync();
if ((await Client.GetVersionAsync()).StartsWith("v2.4.", StringComparison.Ordinal))
if (await Client.GetParsedMilvusVersion() >= new Version(2, 4))
{
// Milvus 2.4 returns -1 here as the compaction ID
return;
Expand Down
19 changes: 19 additions & 0 deletions Milvus.Client.Tests/IndexTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,25 @@ await Collection.CreateIndexAsync(
await Collection.WaitForIndexBuildAsync("float_vector");
}

[Theory]
[InlineData(IndexType.GpuCagra, """{ "nlist": "8" }""")]
[InlineData(IndexType.GpuIvfFlat, """{ "nlist": "8" }""")]
[InlineData(IndexType.GpuIvfPq, """{ "nlist": "8", "m": "4" }""")]
[InlineData(IndexType.GpuBruteForce, """{ "nlist": "8" }""")]
public async Task Index_types_float_gpu(IndexType indexType, string extraParamsString)
{
if (await Client.GetParsedMilvusVersion() < new Version(2, 4))
{
// GPU indexes were introduced in Milvus 2.4
return;
}

await Collection.CreateIndexAsync(
"float_vector", indexType, SimilarityMetricType.L2,
extraParams: JsonSerializer.Deserialize<Dictionary<string, string>>(extraParamsString));
await Collection.WaitForIndexBuildAsync("float_vector");
}

[Theory]
[InlineData(IndexType.BinFlat, """{ "n_trees": "10" }""")]
[InlineData(IndexType.BinIvfFlat, """{ "n_trees": "8", "nlist": "8" }""")]
Expand Down
22 changes: 22 additions & 0 deletions Milvus.Client.Tests/Utils.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
namespace Milvus.Client.Tests;

public static class Utils
{
public static async Task<Version> GetParsedMilvusVersion(this MilvusClient client)
{
string version = await client.GetVersionAsync();

if (version.StartsWith("v", StringComparison.Ordinal))
{
version = version[1..];
}

int dash = version.IndexOf('-');
if (dash != -1)
{
version = version[..dash];
}

return Version.Parse(version);
}
}
40 changes: 40 additions & 0 deletions Milvus.Client/IndexType.cs
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,46 @@ public enum IndexType
/// </summary>
DiskANN,

/// <summary>
/// A graph-based index optimized for GPUs, GPU_CAGRA performs well on inference GPUs. It's best suited for
/// situations with a small number of queries, where training GPUs with lower memory frequency may not yield optimal
/// results.
/// </summary>
/// <remarks>
/// <see href="https://milvus.io/docs/gpu_index.md" />
/// </remarks>
GpuCagra,

/// <summary>
/// This quantization-based index organizes vector data into clusters and employs product quantization for efficient
/// search. It is ideal for scenarios requiring fast queries and can manage limited memory resources while balancing
/// accuracy and speed..
/// </summary>
/// <remarks>
/// <see href="https://milvus.io/docs/gpu_index.md" />
/// </remarks>
GpuIvfFlat,

/// <summary>
/// This quantization-based index organizes vector data into clusters and employs product quantization for efficient
/// search. It is ideal for scenarios requiring fast queries and can manage limited memory resources while balancing
/// accuracy and speed..
/// </summary>
/// <remarks>
/// <see href="https://milvus.io/docs/gpu_index.md" />
/// </remarks>
GpuIvfPq,

/// <summary>
/// This index is tailored for cases where extremely high recall is crucial, guaranteeing a recall of 1 by comparing
/// each query with all vectors in the dataset. It only requires the metric type (metric_type) and top-k (limit) as
/// index building and search parameters.
/// </summary>
/// <remarks>
/// <see href="https://milvus.io/docs/gpu_index.md" />
/// </remarks>
GpuBruteForce,

/// <summary>
/// ANNOY (Approximate Nearest Neighbors Oh Yeah) is an index that uses a hyperplane to divide a high-dimensional
/// space into multiple subspaces, and then stores them in a tree structure.
Expand Down
6 changes: 6 additions & 0 deletions Milvus.Client/MilvusCollection.Index.cs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,12 @@ static string GetGrpcIndexType(IndexType indexType)
IndexType.IvfSq8 => "IVF_SQ8",
IndexType.Hnsw => "HNSW",
IndexType.DiskANN => "DISKANN",

IndexType.GpuCagra => "GPU_CAGRA",
IndexType.GpuIvfFlat => "GPU_IVF_FLAT",
IndexType.GpuIvfPq => "GPU_IVF_PQ",
IndexType.GpuBruteForce => "GPU_BRUTE_FORCE",

IndexType.RhnswFlat => "RHNSW_FLAT",
IndexType.RhnswPq => "RHNSW_PQ",
IndexType.RhnswSq => "RHNSW_SQ",
Expand Down

0 comments on commit ede11f6

Please sign in to comment.