RavenDB-23471 Adjusted GenerateEmbeddings

Lwiel · Jan 9, 2025 · 109b66c · 109b66c
1 parent e5f58e6
commit 109b66c
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 9 deletions.
diff --git a/src/Raven.Server/Config/Categories/IndexingConfiguration.cs b/src/Raven.Server/Config/Categories/IndexingConfiguration.cs
@@ -620,6 +620,13 @@ public virtual bool RunInMemory
         [ConfigurationEntry("Indexing.Corax.VectorSearch.OrderByScoreAutomatically", ConfigurationEntryScope.ServerWideOrPerDatabaseOrPerIndex)]
         public bool CoraxVectorSearchOrderByScoreAutomatically { get; set; }
 
+        [Description("MaxPercentageOfThreadsForEmbeddings")]
+        [DefaultValue(25)]
+        [IndexUpdateType(IndexUpdateType.None)]
+        [ConfigurationEntry("Indexing.Corax.VectorSearch.MaxPercentageOfThreadsForEmbeddings", ConfigurationEntryScope.ServerWideOnly)]
+        public int MaxPercentageOfThreadsForEmbeddings { get; set; }
+
+
         protected override void ValidateProperty(PropertyInfo property)
         {
             base.ValidateProperty(property);

diff --git a/src/Raven.Server/Documents/Indexes/VectorSearch/GenerateEmbeddings.cs b/src/Raven.Server/Documents/Indexes/VectorSearch/GenerateEmbeddings.cs
@@ -15,6 +15,7 @@
 using Raven.Server.Config;
 using Sparrow;
 using Sparrow.Server;
+using InvalidOperationException = System.InvalidOperationException;
 using VectorValue = Corax.Utils.VectorValue;
 
 namespace Raven.Server.Documents.Indexes.VectorSearch;
@@ -34,15 +35,26 @@ static GenerateEmbeddings()
     {
         BertOnnxTextEmbeddingGenerationServiceCtor = typeof(BertOnnxTextEmbeddingGenerationService).GetConstructor(BindingFlags.NonPublic | BindingFlags.Instance, [typeof(InferenceSession), typeof(BertTokenizer), typeof(int), typeof(BertOnnxOptions)]);
         if (BertOnnxTextEmbeddingGenerationServiceCtor == null)
-            throw new InvalidOperationException("TODO");
+            throw new InvalidOperationException($"Could not find constructor for {typeof(BertOnnxTextEmbeddingGenerationService)}.");
     }
 
     public static void Configure(RavenConfiguration configuration)
     {
         if (Embedder.IsValueCreated)
-            throw new InvalidOperationException("TODO");
+            throw new InvalidOperationException("Embedder has already been initialized.");
+
+        var numberOfCoresToUse = CalculateNumberOfCoresForOnnx(configuration);
 
-        OnnxSessionOptions = new SessionOptions();
+        OnnxSessionOptions = new SessionOptions() { IntraOpNumThreads = numberOfCoresToUse };
+    }
+
+    private static int CalculateNumberOfCoresForOnnx(RavenConfiguration configuration)
+    {
+        var cores = (int)MathF.Floor(Environment.ProcessorCount * ((float)configuration.Indexing.MaxPercentageOfThreadsForEmbeddings / 100));
+
+        var numberOfCoresToUse = int.Max(1, cores);
+
+        return numberOfCoresToUse;
     }
 
     [ThreadStatic]
@@ -143,6 +155,8 @@ private static (IDisposable MemoryScope, Memory<byte> Memory, int UsedBytes) Cre
             var embedding = embeddings[0];
 
             var memoryScope = allocator.Allocate(dimensions, out System.Memory<byte> memory);
+
+            MemoryMarshal.AsBytes(embedding.Span).CopyTo(memory.Span);
 
             return (memoryScope, memory, dimensions);
         }

diff --git a/src/Raven.Server/Program.cs b/src/Raven.Server/Program.cs
@@ -14,6 +14,7 @@
 using Raven.Server.Config;
 using Raven.Server.Config.Settings;
 using Raven.Server.Documents.Indexes.Static.NuGet;
+using Raven.Server.Documents.Indexes.VectorSearch;
 using Raven.Server.EventListener;
 using Raven.Server.Logging;
 using Raven.Server.ServerWide;
@@ -126,6 +127,7 @@ public static unsafe int Main(string[] args)
                 configuration.AddCommandLine(configurationArgs);
 
             configuration.Initialize();
+            GenerateEmbeddings.Configure(configuration);
 
             GlobalFlushingBehavior.NumberOfConcurrentSyncsPerPhysicalDrive = configuration.Storage.NumberOfConcurrentSyncsPerPhysicalDrive;
 

diff --git a/test/SlowTests/Issues/RavenDB-23473.cs b/test/SlowTests/Issues/RavenDB-23473.cs
@@ -4,11 +4,13 @@
 using FastTests;
 using Raven.Client.Documents;
 using Raven.Client.Documents.Indexes;
+using Raven.Client.Documents.Indexes.Vector;
 using Raven.Client.Documents.Operations;
 using Raven.Client.Documents.Subscriptions;
-using SmartComponents.LocalEmbeddings;
+using Raven.Server.Documents.Indexes.VectorSearch;
 using Sparrow;
 using Sparrow.Server;
+using Sparrow.Threading;
 using Tests.Infrastructure;
 using Xunit;
 using Xunit.Abstractions;
@@ -44,11 +46,10 @@ public void CanIndexVectorWhenPreviousElementsAreNullWithoutExplicitVectorFieldC
         var errors = Indexes.WaitForIndexingErrors(store, errorsShouldExists: false);
         Assert.Null(errors);
     }
-
+    /*
     [RavenFact(RavenTestCategory.Vector | RavenTestCategory.Indexes)]
     public async Task CanUpdateNoExplicitlyConfiguredVectorFieldViaSubscriptionWithLoadDocument()
     {
-        using (LocalEmbedder localEmbedder = new())
         using (var store = GetDocumentStore(Options.ForSearchEngine(RavenSearchEngineMode.Corax)))
         {
             await store.ExecuteIndexAsync(new VectorIndex());
@@ -60,6 +61,7 @@ public async Task CanUpdateNoExplicitlyConfiguredVectorFieldViaSubscriptionWithL
             var t = worker.Run(async x =>
             {
                 using var session = x.OpenAsyncSession();
+                using var allocator = new ByteStringContext(SharedMultipleUseFlag.None);
                 foreach (var item in x.Items)
                 {
                     var q = item.Result;
@@ -69,9 +71,8 @@ public async Task CanUpdateNoExplicitlyConfiguredVectorFieldViaSubscriptionWithL
                     var localEmbedding = await session.LoadAsync<Embedding>(embeddingId);
                     if (localEmbedding == null)
                     {
-                        // ReSharper disable once AccessToDisposedClosure
-                        var embedding = localEmbedder.Embed(q.Body);
-                        vector = embedding.Values.ToArray().ToList();
+                        var embedding = GenerateEmbeddings.FromText(allocator, new VectorOptions(), q.Body);
+                        vector = embedding.GetEmbedding();
                         localEmbedding = new Embedding { Vector = vector };
                         await session.StoreAsync(localEmbedding, embeddingId);
                     }
@@ -119,6 +120,7 @@ public async Task CanUpdateNoExplicitlyConfiguredVectorFieldViaSubscriptionWithL
             }            
         }
     }
+    */
 
     private class VectorIndex : AbstractIndexCreationTask<Question>
     {