93 lines
2.6 KiB
Python
93 lines
2.6 KiB
Python
"""根据目标 CUDA 版本创建独立的 GPU uv 环境。"""
|
||
|
||
import argparse
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
PADDLE_INDEXES = {
|
||
"cu118": "https://www.paddlepaddle.org.cn/packages/stable/cu118/",
|
||
"cu126": "https://www.paddlepaddle.org.cn/packages/stable/cu126/",
|
||
}
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser(
|
||
description="安装 PaddleOCR-VL-1.6 GPU 子项目依赖",
|
||
)
|
||
parser.add_argument(
|
||
"--cuda",
|
||
choices=sorted(PADDLE_INDEXES),
|
||
required=True,
|
||
help="目标机器的 CUDA Wheel 类型；必须依据 PaddlePaddle 官方兼容表选择",
|
||
)
|
||
parser.add_argument(
|
||
"--dry-run",
|
||
action="store_true",
|
||
help="只显示命令，不创建环境",
|
||
)
|
||
parser.add_argument(
|
||
"--allow-no-gpu",
|
||
action="store_true",
|
||
help="允许在未检测到 nvidia-smi 时创建环境（仅用于准备/CI，不代表可运行）",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
uv = shutil.which("uv")
|
||
if not uv:
|
||
print("[ERROR] 未找到 uv，请先安装 uv。", file=sys.stderr)
|
||
return 1
|
||
|
||
nvidia_smi = shutil.which("nvidia-smi")
|
||
if not args.dry_run and not nvidia_smi and not args.allow_no_gpu:
|
||
print(
|
||
"[ERROR] 未检测到 nvidia-smi，拒绝在无 NVIDIA GPU 的机器安装 CUDA 依赖。\n"
|
||
"如仅准备环境，请显式添加 --allow-no-gpu。",
|
||
file=sys.stderr,
|
||
)
|
||
return 2
|
||
|
||
project_dir = Path(__file__).resolve().parent
|
||
index_url = PADDLE_INDEXES[args.cuda]
|
||
command = [
|
||
uv,
|
||
"sync",
|
||
"--project",
|
||
str(project_dir),
|
||
"--index",
|
||
index_url,
|
||
]
|
||
|
||
print(f"目标 CUDA Wheel: {args.cuda}")
|
||
print(f"PaddlePaddle 索引: {index_url}")
|
||
print("执行命令:")
|
||
print(" " + " ".join(command))
|
||
|
||
if args.dry_run:
|
||
return 0
|
||
|
||
ready_marker = project_dir / ".gpu-ready"
|
||
if ready_marker.exists():
|
||
ready_marker.unlink()
|
||
|
||
completed = subprocess.run(command, check=False)
|
||
if completed.returncode != 0:
|
||
print(
|
||
"[ERROR] 依赖安装失败。请检查 GPU、驱动、Python 和 PaddlePaddle Wheel 兼容性。",
|
||
file=sys.stderr,
|
||
)
|
||
return completed.returncode
|
||
|
||
ready_marker.write_text(
|
||
f"cuda={args.cuda}\nindex={index_url}\n",
|
||
encoding="utf-8",
|
||
)
|
||
print("\n[OK] GPU 子项目环境已创建。下一步从仓库根目录运行：")
|
||
print(" python ocr.py verify --device gpu")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|