mirror of
https://github.com/mudler/LocalAI.git
synced 2026-10-29 09:38:51 -04:00
fix(qdrant-store): upsert points in chunks with a timeout per chunk
StoresSet sent the whole batch in one Upsert under a single request_timeout_ms, so a large batch could hit the timeout even though every point was fine. Send at most 500 points per request and apply the timeout to each request. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
2 files changed
+50
-2
No files matched your search
@@ -29,6 +29,9 @@ const (
|
||||
|
||||
// Results carry up to 50 MiB of base64-encoded values.
|
||||
maxRecvMsgBytes = 128 << 20
|
||||
|
||||
// upsertChunkSize bounds the points sent in one Upsert request.
|
||||
upsertChunkSize = 500
|
||||
)
|
||||
|
||||
var pointNamespace = uuid.NewSHA1(uuid.NameSpaceURL, []byte("https://localai.io/qdrant-store"))
|
||||
@@ -135,12 +138,12 @@ func (s *QdrantStore) StoresSet(opts *pb.StoresSetOptions) error {
|
||||
}
|
||||
}
|
||||
|
||||
err := s.upsert(points)
|
||||
err := upsertChunked(points, upsertChunkSize, s.upsert)
|
||||
if isNotFound(err) {
|
||||
if err := s.createCollection(len(opts.Keys[0].Floats)); err != nil {
|
||||
return fmt.Errorf("qdrant-store: Set: %w", err)
|
||||
}
|
||||
err = s.upsert(points)
|
||||
err = upsertChunked(points, upsertChunkSize, s.upsert)
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("qdrant-store: Set: %w", err)
|
||||
@@ -148,6 +151,17 @@ func (s *QdrantStore) StoresSet(opts *pb.StoresSetOptions) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// upsertChunked calls send with batches of at most size points; send applies
|
||||
// its own request timeout per call, so a large batch is not bounded by one request_timeout_ms.
|
||||
func upsertChunked(points []*qdrant.PointStruct, size int, send func([]*qdrant.PointStruct) error) error {
|
||||
for start := 0; start < len(points); start += size {
|
||||
if err := send(points[start:min(start+size, len(points))]); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *QdrantStore) upsert(points []*qdrant.PointStruct) error {
|
||||
ctx, cancel := s.ctx()
|
||||
defer cancel()
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"math"
|
||||
"time"
|
||||
|
||||
@@ -73,3 +74,36 @@ var _ = Describe("QdrantStore", func() {
|
||||
Expect(err).To(MatchError(ContainSubstring("topK")))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("upsertChunked", func() {
|
||||
points := func(n int) []*qdrant.PointStruct {
|
||||
return make([]*qdrant.PointStruct, n)
|
||||
}
|
||||
|
||||
It("splits a batch into chunks of at most the given size", func() {
|
||||
var sizes []int
|
||||
err := upsertChunked(points(1201), 500, func(c []*qdrant.PointStruct) error {
|
||||
sizes = append(sizes, len(c))
|
||||
return nil
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(sizes).To(Equal([]int{500, 500, 201}))
|
||||
})
|
||||
|
||||
It("sends a small batch in one call", func() {
|
||||
calls := 0
|
||||
Expect(upsertChunked(points(3), 500, func([]*qdrant.PointStruct) error { calls++; return nil })).To(Succeed())
|
||||
Expect(calls).To(Equal(1))
|
||||
})
|
||||
|
||||
It("stops at the first failing chunk", func() {
|
||||
boom := errors.New("boom")
|
||||
calls := 0
|
||||
err := upsertChunked(points(1200), 500, func([]*qdrant.PointStruct) error {
|
||||
calls++
|
||||
return boom
|
||||
})
|
||||
Expect(err).To(MatchError(boom))
|
||||
Expect(calls).To(Equal(1))
|
||||
})
|
||||
})
|
||||
Reference in new issue
Block a user