From 028e784a40b4de9c36a5454754656f290a7a0400 Mon Sep 17 00:00:00 2001 From: vishnukvmd Date: Wed, 8 May 2024 17:34:37 +0530 Subject: [PATCH] Retry requests to fetch embeddings thrice before throwing an error --- server/pkg/controller/embedding/controller.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/server/pkg/controller/embedding/controller.go b/server/pkg/controller/embedding/controller.go index d6e78209fa..342411ea31 100644 --- a/server/pkg/controller/embedding/controller.go +++ b/server/pkg/controller/embedding/controller.go @@ -350,6 +350,10 @@ func (c *Controller) getEmbeddingObjectsParallelV2(userID int64, dbEmbeddingRows } func (c *Controller) getEmbeddingObject(objectKey string, downloader *s3manager.Downloader) (ente.EmbeddingObject, error) { + return c.getEmbeddingObjectWithRetries(objectKey, downloader, 3) +} + +func (c *Controller) getEmbeddingObjectWithRetries(objectKey string, downloader *s3manager.Downloader, retryCount int) (ente.EmbeddingObject, error) { var obj ente.EmbeddingObject buff := &aws.WriteAtBuffer{} _, err := downloader.Download(buff, &s3.GetObjectInput{ @@ -358,6 +362,9 @@ func (c *Controller) getEmbeddingObject(objectKey string, downloader *s3manager. }) if err != nil { log.Error(err) + if retryCount > 0 { + return c.getEmbeddingObjectWithRetries(objectKey, downloader, retryCount-1) + } return obj, stacktrace.Propagate(err, "") } err = json.Unmarshal(buff.Bytes(), &obj)