feat(docvision): add retry mechanism to predict_layout

funstory-ai · Feb 27, 2025 · 8e7b333 · 8e7b333
1 parent be748e0
commit 8e7b333
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 46 deletions.
diff --git a/babeldoc/docvision/rpc_doclayout.py b/babeldoc/docvision/rpc_doclayout.py
@@ -5,6 +5,10 @@
 import httpx
 import msgpack
 import numpy as np
+from tenacity import retry
+from tenacity import retry_if_exception_type
+from tenacity import stop_after_attempt
+from tenacity import wait_exponential
 
 from babeldoc.docvision.doclayout import DocLayoutModel
 from babeldoc.docvision.doclayout import YoloBox
@@ -34,6 +38,17 @@ def encode_image(image) -> bytes:
     return encoded
 
 
+@retry(
+    stop=stop_after_attempt(3),  # 最多重试3次
+    wait=wait_exponential(
+        multiplier=1, min=1, max=10
+    ),  # 指数退避策略，初始1秒，最大10秒
+    retry=retry_if_exception_type((httpx.HTTPError, Exception)),  # 针对哪些异常重试
+    before_sleep=lambda retry_state: logger.warning(
+        f"Request failed, retrying in {retry_state.next_action.sleep} seconds... "
+        f"(Attempt {retry_state.attempt_number}/3)"
+    ),
+)
 def predict_layout(
     image,
     host: str = "http://localhost:8000",
@@ -50,53 +65,49 @@ def predict_layout(
     Returns:
         List of predictions containing bounding boxes and classes
     """
-    try:
-        # Prepare request data
-        if not isinstance(image, list):
-            image = [image]
-        image_data = [encode_image(image) for image in image]
-        data = {
-            "image": image_data,
-            "imgsz": imgsz,
-        }
-
-        # Pack data using msgpack
-        packed_data = msgpack.packb(data, use_bin_type=True)
-        logger.debug(f"Packed data size: {len(packed_data)} bytes")
-
-        # Send request
-        logger.debug(f"Sending request to {host}/inference")
-        response = httpx.post(
-            f"{host}/inference",
-            data=packed_data,
-            headers={
-                "Content-Type": "application/msgpack",
-                "Accept": "application/msgpack",
-            },
-            timeout=300,
-            follow_redirects=True,
+    # Prepare request data
+    if not isinstance(image, list):
+        image = [image]
+    image_data = [encode_image(image) for image in image]
+    data = {
+        "image": image_data,
+        "imgsz": imgsz,
+    }
+
+    # Pack data using msgpack
+    packed_data = msgpack.packb(data, use_bin_type=True)
+    logger.debug(f"Packed data size: {len(packed_data)} bytes")
+
+    # Send request
+    logger.debug(f"Sending request to {host}/inference")
+    response = httpx.post(
+        f"{host}/inference",
+        data=packed_data,
+        headers={
+            "Content-Type": "application/msgpack",
+            "Accept": "application/msgpack",
+        },
+        timeout=300,
+        follow_redirects=True,
+    )
+
+    logger.debug(f"Response status: {response.status_code}")
+    logger.debug(f"Response headers: {response.headers}")
+
+    if response.status_code == 200:
+        try:
+            result = msgpack.unpackb(response.content, raw=False)
+            return result
+        except Exception as e:
+            logger.exception(f"Failed to unpack response: {e!s}")
+            raise
+    else:
+        logger.error(f"Request failed with status {response.status_code}")
+        logger.error(f"Response content: {response.content}")
+        raise Exception(
+            f"Request failed with status {response.status_code}: {response.text}",
         )
 
-        logger.debug(f"Response status: {response.status_code}")
-        logger.debug(f"Response headers: {response.headers}")
-
-        if response.status_code == 200:
-            try:
-                result = msgpack.unpackb(response.content, raw=False)
-                return result
-            except Exception as e:
-                logger.exception(f"Failed to unpack response: {e!s}")
-                raise
-        else:
-            logger.error(f"Request failed with status {response.status_code}")
-            logger.error(f"Response content: {response.content}")
-            raise Exception(
-                f"Request failed with status {response.status_code}: {response.text}",
-            )
-    except Exception as e:
-        logger.exception(f"Unexpected error: {e!s}")
-        raise
-
 
 class RpcDocLayoutModel(DocLayoutModel):
     """DocLayoutModel implementation that uses RPC service."""

diff --git a/pyproject.toml b/pyproject.toml
@@ -30,6 +30,7 @@ dependencies = [
     "xsdata[cli,lxml,soap]>=24.12",
     "msgpack>=1.1.0",
     "pydantic>=2.10.6",
+    "tenacity>=9.0.0",
 ]
 
 [project.urls]
@@ -147,4 +148,4 @@ version_pattern = "MAJOR.MINOR.PATCH[.PYTAGNUM]"
 ]
 "babeldoc/const.py" = [
     '__version__ = "{version}"'
-]
+]