summarylogtreecommitdiffstats
path: root/PKGBUILD
blob: f532e8c52d465b2be32a99c2fb5b22c4d7c7732d (plain)
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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
# Maintainer: Smoolak <smoolak [at] gmail.com>

pkgname=python-torch-tensorrt
_pkgname=torch-tensorrt
pkgver=2.12.0
pkgrel=2
pkgdesc="Easily achieve the best inference performance for any PyTorch model on the NVIDIA platform"
arch=('x86_64')
url="https://github.com/pytorch/TensorRT"
license=('BSD-3-Clause')
depends=(
    'python'
    'python-pytorch-cuda'
    'python-numpy'
    'python-packaging'
    'python-typing_extensions'
    'python-dllist'
    'python-tensorrt'
    'python-rich'
    'python-nvidia-modelopt'
    'python-torchvision'
    'cuda'
    'tensorrt'
)
makedepends=(
    'python-build'
    'python-installer'
    'python-wheel'
    'python-setuptools'
    'python-cffi'
    'pybind11'
    'python-ninja'
)
optdepends=(
    'python-pydot: for engine visualization features'
    'python-graphviz: for graph utilities'
    'tensorrt-llm: for LLM inference support'
)
source=("${pkgname}-${pkgver}.tar.gz::https://github.com/pytorch/TensorRT/archive/refs/tags/v${pkgver}.tar.gz"
        "fix-glog-0.7.patch")
source_x86_64=("bazel-8.1.1::https://github.com/bazelbuild/bazel/releases/download/8.1.1/bazel-8.1.1-linux-x86_64")
sha256sums=('7276b1098638e293da27a978046c2ac6ee455e8656a7692061fc19167e475e64'
            '4a7413873885f73d7e3fa142308953225a6b3af4d0969707e75bfcadf046e9b8')
sha256sums_x86_64=('a2a095d7006ea70bdfdbe90a71f99f957fee1212d4614cfcfdbe3aadae681def')

prepare() {
    cd "TensorRT-${pkgver}"

    # Apply patch for glog 0.7.x compatibility
    patch -p1 < "${srcdir}/fix-glog-0.7.patch"

    # Fix pybind11 version requirement to work with system pybind11
    sed -i 's/pybind11==2.6.2/pybind11>=2.6.2/' pyproject.toml

    # Make downloaded Bazel executable
    chmod +x "${srcdir}/bazel-8.1.1"

    # Create a bazel symlink that setup.py can find
    ln -sf "${srcdir}/bazel-8.1.1" "${srcdir}/bazel"

    # TensorRT lib architecture path for x86_64
    _trt_arch="x86_64-linux-gnu"

    # Create a local directory structure that matches what Bazel's third_party/tensorrt/local/BUILD expects
    mkdir -p "tensorrt_local/lib/${_trt_arch}"
    mkdir -p "tensorrt_local/include/${_trt_arch}"

    # Link system TensorRT libraries
    ln -sf /usr/lib/libnvinfer*.so* "tensorrt_local/lib/${_trt_arch}/"
    ln -sf /usr/lib/libnvparsers.so* "tensorrt_local/lib/${_trt_arch}/"
    ln -sf /usr/lib/libnvonnxparser*.so* "tensorrt_local/lib/${_trt_arch}/"
    ln -sf /usr/lib/libnvinfer_plugin*.so* "tensorrt_local/lib/${_trt_arch}/"

    # Link system TensorRT headers
    ln -sf /usr/include/NvInfer*.h "tensorrt_local/include/${_trt_arch}/"
    ln -sf /usr/include/NvOnnx*.h "tensorrt_local/include/${_trt_arch}/"
    ln -sf /usr/include/NvCaffeParser.h "tensorrt_local/include/${_trt_arch}/" 2>/dev/null || true
    ln -sf /usr/include/NvUffParser.h "tensorrt_local/include/${_trt_arch}/" 2>/dev/null || true

    # Modify MODULE.bazel to use system TensorRT instead of downloading it.
    # Block positions shift between releases, so match by content (not line
    # numbers): comment out the active http_archive(name="tensorrt") and
    # uncomment the new_local_repository(name="tensorrt"), pointing it at the
    # symlink tree created above. tensorrt_rtx blocks are left untouched.
    python - "MODULE.bazel" "${srcdir}/TensorRT-${pkgver}/tensorrt_local" << 'PYEOF'
import re, sys
path, trt = sys.argv[1], sys.argv[2]
src = open(path).read().splitlines(keepends=True)
out, i = [], 0
while i < len(src):
    if src[i].strip() == 'http_archive(' and i+1 < len(src) \
       and re.match(r'\s*name = "tensorrt",\s*$', src[i+1]):
        while True:
            out.append('#' + src[i]); end = src[i].strip() == ')'; i += 1
            if end: break
        continue
    if src[i].lstrip('#').strip() == 'new_local_repository(' and i+1 < len(src) \
       and 'name = "tensorrt"' in src[i+1]:
        while True:
            u = re.sub(r'^#\s?', '', src[i]).replace('path = "/usr/"', f'path = "{trt}"')
            out.append(u); end = src[i].lstrip('#').strip() == ')'; i += 1
            if end: break
        continue
    out.append(src[i]); i += 1
open(path, 'w').writelines(out)
PYEOF

    # Also update CUDA path in MODULE.bazel to use /opt/cuda
    sed -i 's|path = "/usr/local/cuda-13.0/"|path = "/opt/cuda/"|' MODULE.bazel

    # Patch setup.py to use system TensorRT instead of querying Bazel
    # This replaces the Bazel query lambdas with simple lambdas that return "/usr"
    sed -i '/tensorrt_x86_64_external_dir = (/,/^    )$/c\    tensorrt_x86_64_external_dir = lambda: "/usr"' setup.py
    sed -i '/tensorrt_rtx_external_dir = (/,/^    )$/c\    tensorrt_rtx_external_dir = lambda: "/usr"' setup.py
    sed -i '/tensorrt_sbsa_external_dir = (/,/^    )$/c\    tensorrt_sbsa_external_dir = lambda: "/usr"' setup.py
    sed -i '/tensorrt_jetpack_external_dir = (/,/^    )$/c\    tensorrt_jetpack_external_dir = lambda: "/usr"' setup.py

    # glog 0.7 needs GLOG_USE_GLOG_EXPORT defined for every Bazel compile,
    # including the one setup.py runs during the wheel build (which doesn't get
    # our command-line --copt). Put it in .bazelrc so all invocations inherit it.
    printf '\nbuild --copt=-DGLOG_USE_GLOG_EXPORT --host_copt=-DGLOG_USE_GLOG_EXPORT\n' >> .bazelrc
}

build() {
    cd "TensorRT-${pkgver}"

    # Set environment variables for CUDA and TensorRT
    export CUDA_HOME=/opt/cuda
    export CUDACXX=/opt/cuda/bin/nvcc
    export TENSORRT_DIR=/usr
    export LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH}

    # Use the specific Bazel version we downloaded
    export PATH="${srcdir}:${PATH}"
    export BAZEL_EXE="${srcdir}/bazel-8.1.1"
    export USE_HOST_BAZEL=1

    # Set compilation mode to optimized
    export COMPILATION_MODE=opt

    # Fix glog 0.7 compilation issue - ensure GLOG_USE_GLOG_EXPORT is defined.
    # Bazel ignores the shell CXXFLAGS, so the define must also be passed to its
    # C++ compiles via --copt/--host_copt or glog/logging.h fails to compile
    # (GLOG_EXPORT undefined -> "two or more data types").
    export CXXFLAGS="${CXXFLAGS} -DGLOG_USE_GLOG_EXPORT"

    # Keep Bazel's output base off the (quota'd / NFS) home cache by rooting it
    # under $srcdir; otherwise it defaults to ~/.cache/bazel and can fail with
    # "Disk quota exceeded" on networked homes.
    export HOME="${srcdir}"

    # First, build the C++ library with Bazel
    "${srcdir}/bazel-8.1.1" --output_user_root="${srcdir}/.bazelroot" \
        build //:libtorchtrt --compilation_mode=opt \
        --copt=-DGLOG_USE_GLOG_EXPORT --host_copt=-DGLOG_USE_GLOG_EXPORT

    # Extract libraries from the tarball built by Bazel
    tar -xzf bazel-bin/libtorchtrt.tar.gz

    # Copy the built libraries to the location expected by setup.py
    mkdir -p py/torch_tensorrt/lib
    cp -v torch_tensorrt/lib/*.so py/torch_tensorrt/lib/

    # Build Python wheel using modern build system
    python -m build --wheel --no-isolation
}

check() {
    cd "TensorRT-${pkgver}"

    # Basic import test (may fail if CUDA/TensorRT not available during build)
    local _pyver=$(python -c 'import sys; print(f"{sys.version_info[0]}{sys.version_info[1]}")')
    PYTHONPATH="${PWD}/build/lib.linux-${CARCH}-cpython-${_pyver}:${PYTHONPATH}" \
        python -c "import torch_tensorrt; print(torch_tensorrt.__version__)" || \
        echo "Warning: Import test failed - this is expected if CUDA/TensorRT are not available"
}

package() {
    cd "TensorRT-${pkgver}"

    # Install the package using modern installer
    python -m installer --destdir="$pkgdir" dist/*.whl

    # Install license
    install -Dm644 LICENSE "${pkgdir}/usr/share/licenses/${pkgname}/LICENSE"

    # Install documentation
    install -Dm644 README.md "${pkgdir}/usr/share/doc/${pkgname}/README.md"

    # Install C++ headers if built
    if [ -d "py/torch_tensorrt/include" ]; then
        install -dm755 "${pkgdir}/usr/include"
        cp -r py/torch_tensorrt/include "${pkgdir}/usr/include/torch_tensorrt"
    fi
}