Dockerfile.tensorrt_llm 21 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

16
ARG BASE_IMAGE="nvcr.io/nvidia/pytorch"
17
ARG BASE_IMAGE_TAG="25.05-py3"
18
ARG RELEASE_BUILD
19
20
ARG RUNTIME_IMAGE="nvcr.io/nvidia/cuda"
ARG RUNTIME_IMAGE_TAG="12.9.0-runtime-ubuntu24.04"
21

22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
# Define general architecture ARGs for supporting both x86 and aarch64 builds.
#   ARCH: Used for package suffixes (e.g., amd64, arm64)
#   ARCH_ALT: Used for Rust targets, manylinux suffix (e.g., x86_64, aarch64)
#
# Default values are for x86/amd64:
#   --build-arg ARCH=amd64 --build-arg ARCH_ALT=x86_64
#
# For arm64/aarch64, build with:
#   --build-arg ARCH=arm64 --build-arg ARCH_ALT=aarch64
#
# NOTE: There isn't an easy way to define one of these values based on the other value
# without adding if statements everywhere, so just define both as ARGs for now.
ARG ARCH=amd64
ARG ARCH_ALT=x86_64

37
38
39
40
41
##################################
########## Build Image ###########
##################################

FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS build
42

43
44
45
# Redeclare ARCH and ARCH_ALT so they're available in this build stage
ARG ARCH
ARG ARCH_ALT
46
47

ARG NIXL_UCX_REF=v1.19.x
48
ARG NIXL_REF=3c47a48955e6f96bd5d4fb43a9d80bb64722f8e4
49

50
51
52
53
54
55
ENV NIXL_SRC_DIR=/opt/nixl
ENV NIXL_PREFIX=/opt/nvidia/nvda_nixl
ENV NIXL_LIB_DIR=$NIXL_PREFIX/lib/${ARCH_ALT}-linux-gnu
ENV NIXL_PLUGIN_DIR=$NIXL_LIB_DIR/plugins
ENV LD_LIBRARY_PATH=$NIXL_LIB_DIR:$NIXL_PLUGIN_DIR:$LD_LIBRARY_PATH

56
57
58
USER root

# Install utilities
59
60
61
62
63
64
65
66
67
RUN apt update -y && \
    apt install -y \
    git \
    wget \
    curl \
    nvtop \
    tmux \
    vim \
    ## NIXL dependencies
68
    cmake \
69
    meson \
70
71
72
73
74
75
    ninja-build \
    pybind11-dev \
    # These headers are missing with the hpcx installer, required
    # by UCX to find RDMA devices
    libibverbs-dev rdma-core ibverbs-utils libibumad-dev \
    libnuma-dev librdmacm-dev ibverbs-providers
76
77

### UCX EFA Setup ###
78
79
80
81
RUN rm -rf /opt/hpcx/ucx && \
    rm -rf /usr/local/ucx && \
    echo "Building UCX with reference $NIXL_UCX_REF" && \
    cd /usr/local/src &&                            \
82
    git clone https://github.com/openucx/ucx.git && \
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
    cd ucx &&                                       \
    git checkout $NIXL_UCX_REF &&                   \
    ./autogen.sh && ./configure                     \
    --prefix=/usr/local/ucx                         \
    --enable-shared                                 \
    --disable-static                                \
    --disable-doxygen-doc                           \
    --enable-optimizations                          \
    --enable-cma                                    \
    --enable-devel-headers                          \
    --with-cuda=/usr/local/cuda                     \
    --with-verbs                                    \
    --with-efa                                      \
    --with-dm                                       \
    --with-gdrcopy=/usr/local                       \
    --enable-mt &&                                  \
    make -j &&                                      \
    make -j install-strip &&                        \
101
102
    ldconfig

103
104
105
106
ENV LD_LIBRARY_PATH=\
/usr/lib:/usr/local/ucx/lib:\
/usr/local/ucx/lib/ucx:\
$LD_LIBRARY_PATH
107
108
109
110
111
ENV CPATH=/usr/include:$CPATH
ENV PATH=/usr/bin:$PATH
ENV PKG_CONFIG_PATH=/usr/lib/pkgconfig:$PKG_CONFIG_PATH
SHELL ["/bin/bash", "-c"]

112
113
114
115
116
117
118
### NIXL SETUP ###
# Clone nixl source
# TEMP: disable gds backend for arm64
RUN git clone "https://github.com/ai-dynamo/nixl.git" ${NIXL_SRC_DIR} && \
    cd ${NIXL_SRC_DIR} && \
    git checkout ${NIXL_REF} && \
    if [ "$ARCH" = "arm64" ]; then \
119
        nixl_build_args="-Ddisable_gds_backend=true"; \
120
    else \
121
122
123
124
125
126
127
        nixl_build_args=""; \
    fi && \
    mkdir build && \
    meson setup build/ --buildtype=release --prefix=$NIXL_PREFIX $nixl_build_args && \
    cd build/ && \
    ninja && \
    ninja install;
128

129
# nats
130
131
RUN wget --tries=3 --waitretry=5 https://github.com/nats-io/nats-server/releases/download/v2.10.28/nats-server-v2.10.28-${ARCH}.deb && \
    dpkg -i nats-server-v2.10.28-${ARCH}.deb && rm nats-server-v2.10.28-${ARCH}.deb
132

133
# etcd
134
ENV ETCD_VERSION="v3.5.21"
135
RUN wget https://github.com/etcd-io/etcd/releases/download/$ETCD_VERSION/etcd-$ETCD_VERSION-linux-${ARCH}.tar.gz -O /tmp/etcd.tar.gz && \
136
137
138
    mkdir -p /usr/local/bin/etcd && \
    tar -xvf /tmp/etcd.tar.gz -C /usr/local/bin/etcd --strip-components=1 && \
    rm /tmp/etcd.tar.gz
139
140
ENV PATH=/usr/local/bin/etcd/:$PATH

141
142
143
144
145
146
147
148
149
ARG HAS_TRTLLM_CONTEXT=0
ARG TENSORRTLLM_PIP_WHEEL="tensorrt-llm"
ARG TENSORRTLLM_INDEX_URL="https://pypi.python.org/simple"

COPY --from=trtllm_wheel . /trtllm_wheel/

# TODO: Currently, ABI compatibility issues with TRTLLM wheel and NGC PyTorch prevent us
# from using the TRTLLM wheel in a uv venv. Once the issues are resolved, we can
# use uv to install TensorRT-LLM wheel within the uv venv.
150
151
152
# Note: TensorRT needs to be uninstalled before installing the TRTLLM wheel
# because there might be mismatched versions of TensorRT between the NGC PyTorch
# and the TRTLLM wheel.
153
RUN [ -f /etc/pip/constraint.txt ] && : > /etc/pip/constraint.txt || true && \
154
    pip uninstall -y tensorrt && \
155
156
157
158
159
160
161
162
163
164
165
    if [ "$HAS_TRTLLM_CONTEXT" = "1" ]; then \
        # Install from local wheel directory in build context
        WHEEL_FILE=$(find /trtllm_wheel -name "*.whl" | head -n 1); \
        if [ -n "$WHEEL_FILE" ]; then \
            pip install "$WHEEL_FILE"; \
        else \
            echo "No wheel file found in /trtllm_wheel directory."; \
            exit 1; \
        fi; \
    else \
         # Install TensorRT-LLM wheel from the provided index URL, allow dependencies from PyPI
166
         pip install --extra-index-url "${TENSORRTLLM_INDEX_URL}" \
167
168
         "${TENSORRTLLM_PIP_WHEEL}" ; \
    fi
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183

# Install test dependencies
RUN --mount=type=bind,source=./container/deps/requirements.test.txt,target=/tmp/requirements.txt \
    pip install --requirement /tmp/requirements.txt

### MISC UTILITY SETUP ###

# Finish pyright install
RUN pyright --help > /dev/null 2>&1

# Enable Git operations in the /workspace directory
RUN printf "[safe]\n      directory=/workspace\n" > /root/.gitconfig

# Rust build/dev dependencies
RUN apt-get update && \
184
185
186
187
188
    apt-get install --no-install-recommends -y \
    gdb \
    protobuf-compiler \
    cmake \
    libssl-dev \
189
190
    pkg-config \
    libclang-dev
191
192
193
194

ENV RUSTUP_HOME=/usr/local/rustup \
    CARGO_HOME=/usr/local/cargo \
    PATH=/usr/local/cargo/bin:$PATH \
195
    RUST_VERSION=1.87.0
196

197
198
199
200
# Define Rust target based on ARCH_ALT ARG
ARG RUSTARCH=${ARCH_ALT}-unknown-linux-gnu

# Install Rust using RUSTARCH derived from ARCH_ALT
201
RUN wget --tries=3 --waitretry=5 "https://static.rust-lang.org/rustup/archive/1.28.1/${RUSTARCH}/rustup-init" && \
202
    # TODO: Add SHA check back based on RUSTARCH
203
    chmod +x rustup-init && \
204
    ./rustup-init -y --no-modify-path --profile default --default-toolchain $RUST_VERSION --default-host ${RUSTARCH} && \
205
206
    rm rustup-init && \
    chmod -R a+w $RUSTUP_HOME $CARGO_HOME
207

208
ARG CARGO_BUILD_JOBS
209
210
211
212
# Set CARGO_BUILD_JOBS to 16 if not provided
# This is to prevent cargo from building $(nproc) jobs in parallel,
# which might exceed the number of opened files limit.
ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16}
213

214
ENV CARGO_TARGET_DIR=/workspace/target
215
216
217

# Install uv, create virtualenv for general use, and build dynamo wheel
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
218
219
220
221

### VIRTUAL ENVIRONMENT SETUP ###
RUN mkdir /opt/dynamo && \
    uv venv /opt/dynamo/venv --python 3.12
222
223
224
ENV VIRTUAL_ENV=/opt/dynamo/venv

# Install NIXL Python module
225
226
227
228
# TODO: Move gds_path selection based on arch into NIXL build
# TEMP: disable gds backend for arm64
RUN if [ "$ARCH" = "arm64" ]; then \
        cd ${NIXL_SRC_DIR} && uv build . --out-dir /workspace/wheels/nixl \
229
        --config-settings=setup-args="-Ddisable_gds_backend=true"; \
230
231
232
233
234
235
    else \
        cd ${NIXL_SRC_DIR} && uv build . --out-dir /workspace/wheels/nixl; \
    fi && \
    # Install the wheel
    # TODO: Move NIXL wheel install to the wheel_builder stage
    uv pip install /workspace/wheels/nixl/*.whl
236
237
238
239
240

###################################
####### WHEEL BUILD STAGE #########
###################################

241
242
243
244
# Redeclare ARCH_ALT ARG so it's available for interpolation in the FROM instruction
ARG ARCH_ALT

FROM quay.io/pypa/manylinux_2_28_${ARCH_ALT} AS wheel_builder
245
246
247
248
249
250
251
ARG RELEASE_BUILD
ARG CARGO_BUILD_JOBS
# Set CARGO_BUILD_JOBS to 16 if not provided
# This is to prevent cargo from building $(nproc) jobs in parallel,
# which might exceed the number of opened files limit.
ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16}

252
253
254
# Keep in sync with the base image.
ENV NIXL_PREFIX=/opt/nvidia/nvda_nixl

255
256
WORKDIR /workspace

257
RUN yum update -y \
258
    && yum install -y llvm-toolset python3.12-devel \
259
    && yum install -y protobuf-compiler \
260
261
262
263
264
    && yum clean all \
    && rm -rf /var/cache/yum

ENV RUSTUP_HOME=/usr/local/rustup \
    CARGO_HOME=/usr/local/cargo \
265
266
    CARGO_TARGET_DIR=/workspace/target \
    VIRTUAL_ENV=/opt/dynamo/venv
267
268
269

COPY --from=build $RUSTUP_HOME $RUSTUP_HOME
COPY --from=build $CARGO_HOME $CARGO_HOME
270
COPY --from=build $NIXL_PREFIX $NIXL_PREFIX
271
272
273
COPY --from=build /workspace /workspace
COPY --from=build $VIRTUAL_ENV $VIRTUAL_ENV
ENV PATH=$CARGO_HOME/bin:$VIRTUAL_ENV/bin:$PATH
274

275
276
277
278
279
280
281
282
283
284
285
286
287
# Copy configuration files
COPY pyproject.toml /workspace/
COPY README.md /workspace/
COPY LICENSE /workspace/
COPY Cargo.toml /workspace/
COPY Cargo.lock /workspace/
COPY rust-toolchain.toml /workspace/

# Copy source code
COPY lib/ /workspace/lib/
COPY components /workspace/components
COPY launch /workspace/launch

288
289
290
291
292
RUN cargo build \
	--release \
	--locked \
	--features dynamo-llm/block-manager \
	--workspace
293

294
# Build dynamo wheels
295
RUN uv build --wheel --out-dir /workspace/dist && \
296
    cd /workspace/lib/bindings/python && \
297
298
299
300
    uv build --wheel --out-dir /workspace/dist --python 3.12 && \
    if [ "$RELEASE_BUILD" = "true" ]; then \
        uv build --wheel --out-dir /workspace/dist --python 3.11 && \
        uv build --wheel --out-dir /workspace/dist --python 3.10; \
301
    fi
302
303
304
305
306
307
308
309

########################################
########## Development Image ###########
########################################

FROM build AS dev

WORKDIR /workspace
310
311

COPY --from=wheel_builder /workspace /workspace
312
313
314
# Copy Cargo cache to avoid re-downloading dependencies
COPY --from=wheel_builder $CARGO_HOME $CARGO_HOME

315
316
317
# Copy rest of the code
COPY . /workspace

318
319
320
# Package the bindings
RUN mkdir -p /opt/dynamo/bindings/wheels && \
    mkdir /opt/dynamo/bindings/lib && \
321
    cp dist/ai_dynamo*cp312*.whl /opt/dynamo/bindings/wheels/. && \
322
    cp target/release/metrics /usr/local/bin
323

324
# Install wheels
325
326
327
328
RUN . /opt/dynamo/venv/bin/activate && \
    uv pip install /workspace/dist/ai_dynamo_runtime*cp312*.whl && \
    uv pip install /workspace/dist/ai_dynamo*any.whl

329
# Install dynamo.runtime and dynamo.llm wheels globally in container for tests
330
# TODO: In future, we may use a virtualenv for everything and remove this.
331
332
RUN pip install dist/ai_dynamo_runtime*cp312*.whl  && \
    pip install dist/ai_dynamo*any.whl
333

334
ENV DYNAMO_HOME=/workspace
335

336
# Use UCX for TRTLLM KV Cache Transfer
337
ARG TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL
338
ENV TRTLLM_USE_UCX_KVCACHE=1
339

340
341
342
# Create a script that sets the environment variables and source it.
# If TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL is set to 1, unset TRTLLM_USE_UCX_KVCACHE
# and set TRTLLM_USE_NIXL_KVCACHE to 1.
343
344
RUN echo '#!/bin/bash' > /usr/local/bin/set_trtllm_env.sh && \
    if [ "$TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL" = "1" ]; then \
345
        echo 'unset TRTLLM_USE_UCX_KVCACHE' >> /usr/local/bin/set_trtllm_env.sh; \
346
347
        echo 'export TRTLLM_USE_NIXL_KVCACHE=1' >> /usr/local/bin/set_trtllm_env.sh; \
    else \
348
        echo '# TRTLLM_USE_UCX_KVCACHE should already be set to 1' >> /usr/local/bin/set_trtllm_env.sh; \
349
350
351
352
353
    fi && \
    chmod +x /usr/local/bin/set_trtllm_env.sh

# Source the script in bashrc
RUN echo 'source /usr/local/bin/set_trtllm_env.sh' >> /root/.bashrc
354
355
356
357
358
359

# Copy launch banner
RUN --mount=type=bind,source=./container/launch_message.txt,target=/workspace/launch_message.txt \
    sed '/^#\s/d' /workspace/launch_message.txt > ~/.launch_screen && \
    echo "cat ~/.launch_screen" >> ~/.bashrc

360
361
# FIXME: May want a modification with dynamo banner on entry
ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"]
362
CMD []
363
364
365
366
367
368
369

####################################
########## Runtime Image ###########
####################################

FROM ${RUNTIME_IMAGE}:${RUNTIME_IMAGE_TAG} AS runtime

370
371
ARG ARCH_ALT

372
373
374
375
376
377
378
379
380
381
382
WORKDIR /workspace
ENV DYNAMO_HOME=/workspace
ENV VIRTUAL_ENV=/opt/dynamo/venv
ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"

# Install apt dependencies
# openssh-client, openssh-server are needed for OpenMPI
RUN apt-get update && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
        build-essential \
        python3-dev \
383
384
        # jq and curl for polling various endpoints and health checks
        jq \
385
386
387
388
389
390
        curl \
        # For debugging
        vim \
        # Libraries required by UCX to find RDMA devices
        libibverbs1 rdma-core ibverbs-utils libibumad3 \
        libnuma1 librdmacm1 ibverbs-providers \
391
392
393
394
395
396
397
398
399
400
401
402
        openssh-client \
        openssh-server && \
    rm -rf /var/lib/apt/lists/*

# Copy all bindings (wheels, lib, include) from dev image
COPY --from=dev /opt/dynamo/bindings /opt/dynamo/bindings

# Copy nats and etcd from build image
COPY --from=build /usr/bin/nats-server /usr/bin/nats-server
COPY --from=build /usr/local/bin/etcd/ /usr/local/bin/etcd/

# Copy UCX from build image as plugin for NIXL
403
# Copy NIXL source from wheel_builder image
404
COPY --from=build /usr/local/ucx /usr/local/ucx
405
COPY --from=wheel_builder $NIXL_PREFIX $NIXL_PREFIX
406
407
408
409
410
411
412
413
414
415
416
# Copy OpenMPI from build image
COPY --from=build /opt/hpcx/ompi /opt/hpcx/ompi
# Copy NUMA library from build image
COPY --from=build /usr/lib/${ARCH_ALT}-linux-gnu/libnuma.so* /usr/lib/${ARCH_ALT}-linux-gnu/

# Setup the python environment
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
RUN uv venv $VIRTUAL_ENV --python 3.12 && \
    echo "source $VIRTUAL_ENV/bin/activate" >> ~/.bashrc

# Common dependencies
417
# TODO: Remove extra install and use pyproject.toml to define all dependencies
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
RUN --mount=type=bind,source=./container/deps/requirements.txt,target=/tmp/requirements.txt \
    uv pip install --requirement /tmp/requirements.txt

# Install test dependencies
# TODO: Remove this once we have a functional CI image built on top of the runtime image
RUN --mount=type=bind,source=./container/deps/requirements.test.txt,target=/tmp/requirements.txt \
    uv pip install --requirement /tmp/requirements.txt

# Copy CUDA toolkit components needed for nvcc, cudafe, cicc etc.
COPY --from=build /usr/local/cuda/bin/nvcc /usr/local/cuda/bin/nvcc
COPY --from=build /usr/local/cuda/bin/cudafe++ /usr/local/cuda/bin/cudafe++
COPY --from=build /usr/local/cuda/bin/ptxas /usr/local/cuda/bin/ptxas
COPY --from=build /usr/local/cuda/bin/fatbinary /usr/local/cuda/bin/fatbinary
COPY --from=build /usr/local/cuda/include/ /usr/local/cuda/include/
COPY --from=build /usr/local/cuda/lib64/libcudart.so* /usr/local/cuda/lib64/
COPY --from=build /usr/local/cuda/nvvm /usr/local/cuda/nvvm

# Copy pytorch installation from NGC PyTorch
ARG TORCH_VER=2.8.0a0+5228986c39.nv25.5
ARG TORCHVISION_VER=0.22.0a0
ARG SETUPTOOLS_VER=78.1.1
ARG PYTORCH_TRITON_VER=3.3.0+git96316ce52.nvinternal
ARG JINJA2_VER=3.1.6
ARG NETWORKX_VER=3.4.2
ARG SYMPY_VER=1.14.0
ARG PACKAGING_VER=23.2
ARG FLASH_ATTN_VER=2.7.3
ARG MPMATH_VER=1.3.0
COPY --from=build /usr/local/lib/lib* /usr/local/lib/
COPY --from=build /usr/local/lib/python3.12/dist-packages/torch /usr/local/lib/python3.12/dist-packages/torch
COPY --from=build /usr/local/lib/python3.12/dist-packages/torch-${TORCH_VER}.dist-info /usr/local/lib/python3.12/dist-packages/torch-${TORCH_VER}.dist-info
COPY --from=build /usr/local/lib/python3.12/dist-packages/torchgen /usr/local/lib/python3.12/dist-packages/torchgen
COPY --from=build /usr/local/lib/python3.12/dist-packages/torchvision /usr/local/lib/python3.12/dist-packages/torchvision
COPY --from=build /usr/local/lib/python3.12/dist-packages/torchvision-${TORCHVISION_VER}.dist-info /usr/local/lib/python3.12/dist-packages/torchvision-${TORCHVISION_VER}.dist-info
COPY --from=build /usr/local/lib/python3.12/dist-packages/torchvision.libs /usr/local/lib/python3.12/dist-packages/torchvision.libs
COPY --from=build /usr/local/lib/python3.12/dist-packages/setuptools /usr/local/lib/python3.12/dist-packages/setuptools
COPY --from=build /usr/local/lib/python3.12/dist-packages/setuptools-${SETUPTOOLS_VER}.dist-info /usr/local/lib/python3.12/dist-packages/setuptools-${SETUPTOOLS_VER}.dist-info
COPY --from=build /usr/local/lib/python3.12/dist-packages/functorch /usr/local/lib/python3.12/dist-packages/functorch
COPY --from=build /usr/local/lib/python3.12/dist-packages/triton /usr/local/lib/python3.12/dist-packages/triton
COPY --from=build /usr/local/lib/python3.12/dist-packages/pytorch_triton-${PYTORCH_TRITON_VER}.dist-info /usr/local/lib/python3.12/dist-packages/pytorch_triton-${PYTORCH_TRITON_VER}.dist-info
COPY --from=build /usr/local/lib/python3.12/dist-packages/jinja2 /usr/local/lib/python3.12/dist-packages/jinja2
COPY --from=build /usr/local/lib/python3.12/dist-packages/jinja2-${JINJA2_VER}.dist-info /usr/local/lib/python3.12/dist-packages/jinja2-${JINJA2_VER}.dist-info
COPY --from=build /usr/local/lib/python3.12/dist-packages/networkx /usr/local/lib/python3.12/dist-packages/networkx
COPY --from=build /usr/local/lib/python3.12/dist-packages/networkx-${NETWORKX_VER}.dist-info /usr/local/lib/python3.12/dist-packages/networkx-${NETWORKX_VER}.dist-info
COPY --from=build /usr/local/lib/python3.12/dist-packages/sympy /usr/local/lib/python3.12/dist-packages/sympy
COPY --from=build /usr/local/lib/python3.12/dist-packages/sympy-${SYMPY_VER}.dist-info /usr/local/lib/python3.12/dist-packages/sympy-${SYMPY_VER}.dist-info
COPY --from=build /usr/local/lib/python3.12/dist-packages/packaging /usr/local/lib/python3.12/dist-packages/packaging
COPY --from=build /usr/local/lib/python3.12/dist-packages/packaging-${PACKAGING_VER}.dist-info /usr/local/lib/python3.12/dist-packages/packaging-${PACKAGING_VER}.dist-info
COPY --from=build /usr/local/lib/python3.12/dist-packages/flash_attn /usr/local/lib/python3.12/dist-packages/flash_attn
COPY --from=build /usr/local/lib/python3.12/dist-packages/flash_attn-${FLASH_ATTN_VER}.dist-info /usr/local/lib/python3.12/dist-packages/flash_attn-${FLASH_ATTN_VER}.dist-info
COPY --from=build /usr/local/lib/python3.12/dist-packages/flash_attn_2_cuda.cpython-312-*-linux-gnu.so /usr/local/lib/python3.12/dist-packages/

# Setup environment variables
471
472
473
474
475
476
477
478
479
480
481
482
ARG ARCH_ALT
ENV NIXL_PREFIX=/opt/nvidia/nvda_nixl
ENV NIXL_LIB_DIR=$NIXL_PREFIX/lib/${ARCH_ALT}-linux-gnu
ENV NIXL_PLUGIN_DIR=$NIXL_LIB_DIR/plugins

ENV LD_LIBRARY_PATH=\
$NIXL_LIB_DIR:\
$NIXL_PLUGIN_DIR:\
/usr/local/ucx/lib:\
/usr/local/ucx/lib/ucx:\
/opt/hpcx/ompi/lib:\
$LD_LIBRARY_PATH
483
484
485
486
487
488
489
490
491
ENV PATH=/opt/hpcx/ompi/bin:/usr/local/bin/etcd/:/usr/local/cuda/nvvm/bin:$PATH
ENV OPAL_PREFIX=/opt/hpcx/ompi

# Install TensorRT-LLM (same as in build stage)
ARG HAS_TRTLLM_CONTEXT=0
ARG TENSORRTLLM_PIP_WHEEL="tensorrt-llm"
ARG TENSORRTLLM_INDEX_URL="https://pypi.python.org/simple"

# Copy Dynamo wheels into wheelhouse
492
# Copy metrics binary from wheel_builder image, not part of ai-dynamo wheel
493
494
COPY --from=dev /workspace/wheels/nixl/*.whl wheelhouse/
COPY --from=wheel_builder /workspace/dist/*.whl wheelhouse/
495
COPY --from=dev /workspace/target/release/metrics /usr/local/bin/metrics
496

497
498
499
500
# NOTE: If a package (tensorrt_llm) exists on both --index-url and --extra-index-url,
# uv will prioritize the --extra-index-url, unless --index-strategy unsafe-best-match
# is also specified. So set the configurable index as a --extra-index-url for prioritization.
RUN uv pip install --extra-index-url "${TENSORRTLLM_INDEX_URL}" \
501
    "${TENSORRTLLM_PIP_WHEEL}" && \
502
    uv pip install ai-dynamo nixl --find-links wheelhouse
503
504
505
506
507
508

# Setup TRTLLM environment variables, same as in dev image
ENV TRTLLM_USE_UCX_KVCACHE=1
COPY --from=dev /usr/local/bin/set_trtllm_env.sh /usr/local/bin/set_trtllm_env.sh
RUN echo 'source /usr/local/bin/set_trtllm_env.sh' >> /root/.bashrc

509
# Copy benchmarks, backends and tests for CI
510
511
512
# TODO: Remove this once we have a functional CI image built on top of the runtime image
COPY tests /workspace/tests
COPY benchmarks /workspace/benchmarks
513
COPY components/backends/trtllm /workspace/components/backends/trtllm
514
515
RUN uv pip install /workspace/benchmarks

516
517
518
# Copy files for legal compliance
COPY ATTRIBUTION* LICENSE /workspace/

519
520
521
522
523
524
525
# Copy launch banner
RUN --mount=type=bind,source=./container/launch_message.txt,target=/workspace/launch_message.txt \
    sed '/^#\s/d' /workspace/launch_message.txt > ~/.launch_screen && \
    echo "cat ~/.launch_screen" >> ~/.bashrc

ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"]
CMD []