From 5f4baba012a548417519a6bd3d5d068939794255 Mon Sep 17 00:00:00 2001
From: panyy <yf_pan.yuanyin@unisinsight.com>
Date: Mon, 6 Jul 2026 20:33:54 +0800
Subject: [PATCH] =?UTF-8?q?feat(=E6=80=9D=E7=BB=B4=E5=AF=BC=E5=9B=BE?=
 =?UTF-8?q?=E5=8A=A9=E6=89=8B)=EF=BC=9A=E8=AE=BE=E7=BD=AE=E5=8D=A0?=
 =?UTF-8?q?=E7=94=A8gpu=E9=99=90=E9=A2=9D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 docker/china/Dockerfile.source-patch |  1 +
 mineru/backend/vlm/utils.py          | 12 +++++++++++-
 scripts/start-mindmap.sh             |  2 ++
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/docker/china/Dockerfile.source-patch b/docker/china/Dockerfile.source-patch
index d74867e..a58177e 100644
--- a/docker/china/Dockerfile.source-patch
+++ b/docker/china/Dockerfile.source-patch
@@ -1,3 +1,4 @@
 FROM unis-mindmap-api:latest
 
 COPY mineru/backend/vlm/vlm_analyze.py /app/mineru/backend/vlm/vlm_analyze.py
+COPY mineru/backend/vlm/utils.py /app/mineru/backend/vlm/utils.py
diff --git a/mineru/backend/vlm/utils.py b/mineru/backend/vlm/utils.py
index c64defa..b6fdee2 100644
--- a/mineru/backend/vlm/utils.py
+++ b/mineru/backend/vlm/utils.py
@@ -71,6 +71,16 @@ def set_lmdeploy_backend(device_type: str) -> str:
 
 
 def set_default_gpu_memory_utilization() -> float:
+    override = os.getenv("MINERU_GPU_MEMORY_UTILIZATION")
+    if override:
+        try:
+            value = float(override)
+        except ValueError:
+            raise ValueError("MINERU_GPU_MEMORY_UTILIZATION must be a float between 0 and 1")
+        if not 0 < value < 1:
+            raise ValueError("MINERU_GPU_MEMORY_UTILIZATION must be a float between 0 and 1")
+        return value
+
     from vllm import __version__ as vllm_version
     device = get_device()
     gpu_memory = get_vram(device)
@@ -96,4 +106,4 @@ def set_default_batch_size() -> int:
     except Exception as e:
         logger.warning(f'Error determining VRAM: {e}, using default batch_ratio: 1')
         batch_size = 1
-    return batch_size
\ No newline at end of file
+    return batch_size
diff --git a/scripts/start-mindmap.sh b/scripts/start-mindmap.sh
index 51c5deb..2c37224 100644
--- a/scripts/start-mindmap.sh
+++ b/scripts/start-mindmap.sh
@@ -6,6 +6,7 @@ API_NAME=${API_NAME:-mineru-api}
 WEB_NAME=${WEB_NAME:-mindmap-web}
 WEB_PORT=${WEB_PORT:-9898}
 MINERU_API_URL=${MINERU_API_URL:-http://${API_NAME}:8000}
+MINERU_GPU_MEMORY_UTILIZATION=${MINERU_GPU_MEMORY_UTILIZATION:-0.5}
 
 sudo docker network create "$NETWORK" >/dev/null 2>&1 || true
 sudo docker rm -f "$API_NAME" "$WEB_NAME" >/dev/null 2>&1 || true
@@ -14,6 +15,7 @@ sudo docker run -d --name "$API_NAME" --network "$NETWORK" \
   --privileged --ipc=host --shm-size=32g \
   -e MINERU_MODEL_SOURCE=local \
   -e PYTHONPATH=/app:/usr/local/corex-4.4.0/lib64/python3/dist-packages \
+  -e MINERU_GPU_MEMORY_UTILIZATION="$MINERU_GPU_MEMORY_UTILIZATION" \
   unis-mindmap-api:latest \
   mineru-api --host 0.0.0.0 --port 8000