TensorRT/.github/scripts/install-torch-tensorrt.sh at 5b3c6647bfed1e70d84e65700309d98bd8076c1e · pytorch/TensorRT · GitHub

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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
#set -exou pipefail
set -x

TORCH=$(grep "^torch>" ${PWD}/py/requirements.txt)
INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION}
PLATFORM=$(python -c "import sys; print(sys.platform)")

if [[ $(uname -m) == "aarch64" ]]; then
    # install cuda for aarch64
    source .github/scripts/install-cuda-aarch64.sh
    install_cuda_aarch64
fi

# Install all the dependencies required for Torch-TensorRT
python -m pip install --upgrade "pip>=25.1" "tomli>=1.1.0; python_version < '3.11'"
python -m pip install \
    --pre \
    --extra-index-url https://pypi.nvidia.com \
    --extra-index-url https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION} \
    --group test \
    --group test-ext \
    --group quantization
TORCHVISION=$(python - <<'PY'
try:
    import tomllib
except ModuleNotFoundError:
    import tomli as tomllib

with open("pyproject.toml", "rb") as f:
    deps = tomllib.load(f)["dependency-groups"]["test-ext"]

for dep in deps:
    if dep.startswith("torchvision"):
        print(dep)
        break
else:
    raise SystemExit("torchvision was not found in dependency group test-ext")
PY
)
# test dependencies might install a different version of torch or torchvision
# eg. timm will install the latest torchvision, however we want to use the torchvision from nightly
# reinstall torch torchvision to make sure we have the correct version
python -m pip uninstall -y torch torchvision
python -m pip install --force-reinstall --pre ${TORCHVISION} --index-url ${INDEX_URL} --extra-index-url https://pypi.org/simple
python -m pip install --force-reinstall --pre ${TORCH} --index-url ${INDEX_URL} --extra-index-url https://pypi.org/simple

# If CUDA 13 (cu13), prepend venv's NVIDIA CUDA 13 libs to LD_LIBRARY_PATH
if [[ "${CU_VERSION}" == cu13* ]]; then
    SITE_PACKAGES="$(python -c 'import sysconfig; print(sysconfig.get_path("platlib"))')"
    export LD_LIBRARY_PATH="${SITE_PACKAGES}/nvidia/cu13/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
fi

# Install Torch-TensorRT
if [[ ${PLATFORM} == win32 ]]; then
    # pin-check: no-nightly -- this glob also matches the Linux-only ExecuTorch runtime wheel, but
    # ExecuTorch publishes no win32 nightly and the [executorch] extra is Linux-only, so this
    # platform's plain torch-tensorrt install needs no nightly index.
    python -m pip install ${RUNNER_ARTIFACT_DIR}/torch_tensorrt*.whl
else
    # The nightly channel is needed because this glob also matches the ExecuTorch runtime wheel,
    # whose install_requires names an ExecuTorch dev build that is published only there.
    # Hardcoded rather than ${CHANNEL} like the lines above: a .dev wheel exists on no other
    # channel, so deriving it would break this install on exactly the test and release runs the
    # index was added for. It is an extra index, not a replacement, and torch is already
    # force-reinstalled from ${INDEX_URL} above, so the pinned torch is not at risk from it.
    # || exit 1 because line 1's `set -exou pipefail` is commented out and linux-test.yml
    # concatenates this file ahead of the user script, so a failure here would otherwise be
    # discarded and the job would die later with an unrelated-looking ImportError. Scoped to the
    # line this change is responsible for; re-enabling set -e for the whole file is a
    # pre-existing hazard worth a separate change.
    python -m pip install /opt/torch-tensorrt-builds/torch_tensorrt*.whl --use-deprecated=legacy-resolver \
        --extra-index-url "https://download.pytorch.org/whl/nightly/${CU_VERSION}" || exit 1
fi

echo -e "Running test script";