enhance: retry on inconsistent requery (#1975) (#1976)

Signed-off-by: longjiquan <[email protected]>
milvus-io · Mar 18, 2024 · 2110394 · 2110394
1 parent 5332f09
commit 2110394
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py
@@ -734,7 +734,7 @@ def _execute_search(
                 return SearchFuture(None, None, e)
             raise e from e
 
-    @retry_on_rpc_failure()
+    @retry_on_rpc_failure(retry_on_inconsistent_requery=True)
     def search(
         self,
         collection_name: str,

diff --git a/pymilvus/decorators.py b/pymilvus/decorators.py
@@ -45,6 +45,7 @@ def retry_on_rpc_failure(
     initial_back_off: float = 0.01,
     max_back_off: float = 3,
     back_off_multiplier: int = 3,
+    retry_on_inconsistent_requery: bool = False,
 ):
     def wrapper(func: Any):
         @functools.wraps(func)
@@ -106,8 +107,14 @@ def timeout(start_time: Optional[float] = None) -> bool:
                         raise MilvusException(
                             code=e.code, message=f"{to_msg}, message={e.message}"
                         ) from e
-                    if _retry_on_rate_limit and (
-                        e.code == ErrorCode.RATE_LIMIT or e.compatible_code == common_pb2.RateLimit
+                    if (
+                        _retry_on_rate_limit
+                        and (
+                            e.code == ErrorCode.RATE_LIMIT
+                            or e.compatible_code == common_pb2.RateLimit
+                        )
+                        or retry_on_inconsistent_requery
+                        and e.code == 2200
                     ):
                         time.sleep(back_off)
                         back_off = min(back_off * back_off_multiplier, max_back_off)