使用UV在打包Docker镜像时自动管理CPU或GPU环境

  1. 在pyproject.toml中修改如下
#  project.optional-dependencies中定义好的依赖不需要在dependencies中重复声明
[project.optional-dependencies]
cpu = [
    "torch~=2.3.1",
    "onnxruntime~=1.21.1"
]
cu121 = [
    "torch>=2.3.1",
    "onnxruntime-gpu~=1.21.1"
]

[tool.uv]
conflicts = [
    [
        { extra = "cpu" },
        { extra = "cu121" },
    ],
]

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-cu121", extra = "cu121" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu121"
url = "https://download.pytorch.org/whl/cu121"
explicit = true
  1. 同步依赖
# 使用cpu环境
uv sync --extra cpu
# 使用gpu环境
uv sync --extra cu121

#  uv run 命令同样需要extra指定
posted @ 2025-04-22 15:42  墨雨听风  阅读(164)  评论(0)    收藏  举报