0%

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

  1. 在pyproject.toml中修改如下
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
#  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. 同步依赖
1
2
3
4
5
6
# 使用cpu环境
uv sync --extra cpu
# 使用gpu环境
uv sync --extra cu121

# uv run 命令同样需要extra指定