enhance: retry on inconsistent requery (#1975)

* enhance: retry on inconsistent requery Signed-off-by: longjiquan <[email protected]> * fix Signed-off-by: longjiquan <[email protected]> * fix lint Signed-off-by: longjiquan <[email protected]> * revert Signed-off-by: longjiquan <[email protected]> * fix Signed-off-by: longjiquan <[email protected]> * fix Signed-off-by: longjiquan <[email protected]> --------- Signed-off-by: longjiquan <[email protected]>
milvus-io · Mar 18, 2024 · 35b24d2 · 35b24d2
1 parent 3e77421
commit 35b24d2
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py
@@ -753,7 +753,7 @@ def _execute_hybrid_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,
@@ -793,7 +793,7 @@ def search(
         )
         return self._execute_search(request, timeout, round_decimal=round_decimal, **kwargs)
 
-    @retry_on_rpc_failure()
+    @retry_on_rpc_failure(retry_on_inconsistent_requery=True)
     def hybrid_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)