diff --git a/babeldoc/docvision/rpc_doclayout.py b/babeldoc/docvision/rpc_doclayout.py index 8e068ab..05cee0c 100644 --- 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 index 19c23a9..b4a9ba4 100644 --- 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}"' -] \ No newline at end of file +]