@@ -14,9 +14,9 @@ ARG TARGETARCH
1414# when updating past that version, either remove all TensorRT related packages,
1515# or investigate if there is a way to keep TensorRT support
1616ARG TENSORFLOW_VERSION=2.17.1
17- ARG TENSORRT_VERSION=8.6.1.post1
18- ARG TENSORRT_LIBS_VERSION=8.6.1
19- ARG TENSORRT_BINDINGS_VERSION=8.6.1
17+ ARG TENSORRT_VERSION=10.2.0
18+ ARG TENSORRT_LIBS_VERSION=10.2.0
19+ ARG TENSORRT_BINDINGS_VERSION=10.2.0
2020
2121# nvidia container toolkit
2222# https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/docker-specialized.html
@@ -30,22 +30,22 @@ ENV NVIDIA_REQUIRE_CUDA "cuda>=12.3"
3030# to `tensorrt-cu12`, `tensorrt-cu12-libs`, and `tensorrt-cu12-bindings` respectively
3131RUN python3 -m pip install --quiet --no-cache-dir --extra-index-url https://pypi.nvidia.com \
3232 tensorflow[and-cuda]==${TENSORFLOW_VERSION} \
33- tensorrt==${TENSORRT_VERSION} \
34- tensorrt-libs==${TENSORRT_LIBS_VERSION} \
35- tensorrt-bindings==${TENSORRT_BINDINGS_VERSION}
33+ tensorrt-cu12 ==${TENSORRT_VERSION} \
34+ tensorrt-cu12- libs==${TENSORRT_LIBS_VERSION} \
35+ tensorrt-cu12- bindings==${TENSORRT_BINDINGS_VERSION}
3636
3737# create symlinks for TensorRT libs
3838# - https://github.com/tensorflow/tensorflow/issues/61986#issuecomment-1880489731
3939# - We are creating symlinks for the following libs, as this is where TF looks for them:
4040# - libnvinfer.so.8.6.1 -> libnvinfer.so.8
4141# - libnvinfer_plugin.so.8.6.1 -> libnvinfer_plugin.so.8
42- ENV PYTHON_SITE_PACKAGES /opt/conda/lib/python3.11 /site-packages
42+ ENV PYTHON_SITE_PACKAGES /opt/conda/lib/python3.12 /site-packages
4343ENV TENSORRT_LIBS ${PYTHON_SITE_PACKAGES}/tensorrt_libs
4444RUN ln -s ${TENSORRT_LIBS}/libnvinfer.so.${TENSORRT_LIBS_VERSION%%.*} ${TENSORRT_LIBS}/libnvinfer.so.${TENSORRT_LIBS_VERSION} \
4545 && ln -s ${TENSORRT_LIBS}/libnvinfer_plugin.so.${TENSORRT_LIBS_VERSION%%.*} ${TENSORRT_LIBS}/libnvinfer_plugin.so.${TENSORRT_LIBS_VERSION}
4646
4747# envs - cudnn, tensorrt
48- ENV LD_LIBRARY_PATH ${LD_LIBRARY_PATH}/nvidia/cudnn/lib: ${TENSORRT_LIBS}
48+ ENV LD_LIBRARY_PATH= ${LD_LIBRARY_PATH:+${LD_LIBRARY_PATH}:} ${TENSORRT_LIBS}
4949
5050# install - requirements.txt
5151COPY --chown=${NB_USER}:${NB_GID} requirements.txt /tmp
0 commit comments