59 lines
2.7 KiB
Python
59 lines
2.7 KiB
Python
"""Скачать колёса PyTorch с CUDA для сборки образа без сети (в docker/wheels).
|
||
|
||
Нужно, когда у машины, где собирается образ, нет выхода в интернет (или он
|
||
сломан, как у WSL за VPN), а у соседней машины есть. Скачивать можно хоть на
|
||
Windows: колёса берутся под Linux x86_64 и Python 3.10 — ровно под образ.
|
||
|
||
Тонкость: `pip download --platform` не пересчитывает маркеры окружения и на
|
||
Windows молча пропускает зависимости «только для Linux» — все библиотеки CUDA
|
||
и triton. Поэтому их список берётся из метаданных самого колеса torch и
|
||
докачивается явно, с точными версиями.
|
||
|
||
python docker/fetch_wheels.py
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import zipfile
|
||
from pathlib import Path
|
||
|
||
TORCH = "torch==2.9.1"
|
||
INDEX = "https://download.pytorch.org/whl/cu128"
|
||
OUT = Path(__file__).resolve().parent / "wheels"
|
||
TARGET = ["--platform", "manylinux_2_28_x86_64", "--platform", "manylinux_2_27_x86_64",
|
||
"--platform", "manylinux_2_17_x86_64", "--platform", "manylinux2014_x86_64",
|
||
"--python-version", "3.10", "--implementation", "cp", "--abi", "cp310",
|
||
"--only-binary=:all:"]
|
||
|
||
|
||
def pip_download(*args: str) -> None:
|
||
cmd = [sys.executable, "-m", "pip", "download", "--index-url", INDEX,
|
||
"--extra-index-url", "https://pypi.org/simple", "-d", str(OUT), *args]
|
||
subprocess.run(cmd, check=True)
|
||
|
||
|
||
def main() -> None:
|
||
OUT.mkdir(exist_ok=True)
|
||
pip_download(TORCH, *TARGET)
|
||
wheel = next(OUT.glob("torch-" + TORCH.split("==")[1] + "*.whl"))
|
||
with zipfile.ZipFile(wheel) as z:
|
||
meta = next(n for n in z.namelist() if n.endswith(".dist-info/METADATA"))
|
||
lines = z.read(meta).decode("utf-8").splitlines()
|
||
linux = []
|
||
for line in lines:
|
||
m = re.match(r"Requires-Dist: ([A-Za-z0-9_.\-]+==[^;\s]+); platform_system == \"Linux\"$",
|
||
line)
|
||
if m:
|
||
linux.append(m.group(1))
|
||
print(f"зависимости только для Linux: {len(linux)}")
|
||
pip_download("--no-deps", *TARGET, *linux)
|
||
subprocess.run([sys.executable, "-m", "pip", "download", "pip", "--only-binary=:all:",
|
||
"--no-deps", "-d", str(OUT)], check=True)
|
||
size = sum(p.stat().st_size for p in OUT.glob("*.whl")) / 2**30
|
||
print(f"готово: {len(list(OUT.glob('*.whl')))} колёс, {size:.1f} ГБ в {OUT}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|