feat(思维导图助手)：设置占用gpu限额

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

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
+    return batch_size

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