fix(trtllm): restore deps, default CMD, and reduce layer count after upstream base switch

container/templates/trtllm_runtime.Dockerfile

@@ -7,58 +7,92 @@
 ########## Runtime Image #########
 ##################################
 
+# Transport stage — runtime pulls /workspace_src/ in one bind-mount cp.
+FROM scratch AS workspace_files
+COPY --chmod=775 tests /workspace_src/tests
+COPY --chmod=775 examples /workspace_src/examples
+COPY --chmod=775 deploy /workspace_src/deploy
+COPY --chmod=775 dev /workspace_src/dev
+COPY --chmod=775 components/src/dynamo/common /workspace_src/components/src/dynamo/common
+COPY --chmod=775 components/src/dynamo/frontend /workspace_src/components/src/dynamo/frontend
+COPY --chmod=775 components/src/dynamo/trtllm /workspace_src/components/src/dynamo/trtllm
+COPY --chmod=775 components/src/dynamo/mocker /workspace_src/components/src/dynamo/mocker
+COPY --chmod=775 lib /workspace_src/lib
+COPY --chmod=664 ATTRIBUTION* LICENSE /workspace_src/
+
+# Transport stage for dynamo_base artifacts. uv/uvx go to /usr/bin (not /bin)
+# because upstream is usrmerged and cross-stage COPY chokes on the symlink.
+FROM scratch AS dynamo_base_export
+COPY --from=dynamo_base /usr/bin/nats-server /usr/bin/nats-server
+COPY --from=dynamo_base /usr/local/bin/etcd/ /usr/local/bin/etcd/
+COPY --from=dynamo_base /bin/uv /usr/bin/uv
+COPY --from=dynamo_base /bin/uvx /usr/bin/uvx
+
+{% if target == "runtime" %}
+# Renamed `runtime` → `runtime_full` so the final stage can re-FROM upstream
+# and overlay our changes as a single layer (cuts depth for downstream wrappers).
+FROM ${RUNTIME_IMAGE}:${RUNTIME_IMAGE_TAG} AS runtime_full
+{% else %}
 FROM ${RUNTIME_IMAGE}:${RUNTIME_IMAGE_TAG} AS runtime
+{% endif %}
 
 ARG ENABLE_KVBM
 ARG ENABLE_GPU_MEMORY_SERVICE
 ARG TARGETARCH
 
-WORKDIR /workspace
-
 # DYNAMO_HOME points at /workspace so bundled TRT-LLM scripts that reference
-# $DYNAMO_HOME/examples/... resolve. Examples and tests are copied below.
-ENV DYNAMO_HOME=/workspace
-ENV HOME=/home/dynamo
-ENV PATH=/usr/local/bin/etcd:${PATH}
-
-# Workaround for ai-dynamo/nixl#1668: nixl-cu13's bundled UCX 1.20.0 hangs in
+# $DYNAMO_HOME/examples/... resolve. LD_PRELOAD/NIXL_PLUGIN_DIR are a workaround
+# for ai-dynamo/nixl#1668: nixl-cu13's bundled UCX 1.20.0 hangs in
 # `uct_md_query_tl_resources` (md_resources realloc loop, >1 GiB) when two NIXL
-# agents init on the same host — blocks every TRT-LLM native-disagg multi-process
-# test. Force-load TRT-LLM's bundled libnixl 0.9.0 (uses system UCX, no bug).
-# LD_PRELOAD is the only lever: nixl-cu13's _bindings.so has DT_RPATH which beats
-# LD_LIBRARY_PATH. Drop this block when the upstream issue is fixed.
-ENV LD_PRELOAD=/usr/local/lib/python3.12/dist-packages/tensorrt_llm/libs/nixl/libnixl.so
-ENV NIXL_PLUGIN_DIR=/usr/local/lib/python3.12/dist-packages/tensorrt_llm/libs/nixl/plugins
-# Fail the build loudly if upstream moves these paths — otherwise LD_PRELOAD
-# silently logs "cannot be preloaded: ignored" and the hang returns at runtime.
-RUN test -f "${LD_PRELOAD}" && test -d "${NIXL_PLUGIN_DIR}"
-
-# Upstream's /etc/shinit_v2 prepends /usr/local/tensorrt/lib (and others) to
-# LD_LIBRARY_PATH, but only when a shell starts. K8s spawns python3 directly,
-# so the libs aren't found and `import tensorrt` fails. Register the paths
-# with ldconfig instead so they resolve regardless of launch method.
-RUN ARCH_ALT=$([ "${TARGETARCH}" = "amd64" ] && echo "x86_64" || echo "aarch64") && \
+# agents init on the same host. Force-load TRT-LLM's bundled libnixl 0.9.0
+# (uses system UCX, no bug). LD_PRELOAD is the only lever: nixl-cu13's
+# _bindings.so has DT_RPATH which beats LD_LIBRARY_PATH. Drop the two NIXL
+# vars when the upstream issue is fixed.
+ENV DYNAMO_HOME=/workspace \
+    HOME=/home/dynamo \
+    PATH=/usr/local/bin/etcd:${PATH} \
+    LD_PRELOAD=/opt/dynamo/libstdc++.so.6:/usr/local/lib/python3.12/dist-packages/tensorrt_llm/libs/nixl/libnixl.so \
+    NIXL_PLUGIN_DIR=/usr/local/lib/python3.12/dist-packages/tensorrt_llm/libs/nixl/plugins
+
+WORKDIR /workspace
+
+# Install packages missing from upstream, sanity-check libnixl, register
+# TRT-LLM lib paths with ldconfig (upstream's /etc/shinit_v2 only sets them
+# for shells, not K8s python3 launches), swap upstream's single-binary etcd
+# for dynamo_base's directory, and symlink system libstdc++ to a stable
+# path for LD_PRELOAD — keeps PyInstaller-bundled tools (specifically `jet`,
+# NVIDIA's internal PyInstaller-packaged CI runner) from shadowing it with an
+# older copy.
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt,sharing=locked \
+    apt-get update && \
+    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+        openssh-server \
+        librdmacm1 \
+        rdma-core && \
+    test -f /usr/local/lib/python3.12/dist-packages/tensorrt_llm/libs/nixl/libnixl.so && \
+    test -d "${NIXL_PLUGIN_DIR}" && \
+    ARCH_ALT=$([ "${TARGETARCH}" = "amd64" ] && echo "x86_64" || echo "aarch64") && \
     printf '%s\n' \
         "/usr/local/tensorrt/lib" \
         "/usr/local/cuda/lib64" \
         "/usr/local/ucx/lib" \
         "/opt/nvidia/nvda_nixl/lib/${ARCH_ALT}-linux-gnu" \
         "/opt/nvidia/nvda_nixl/lib64" \
         > /etc/ld.so.conf.d/00-dynamo-trtllm.conf && \
-    ldconfig
-
-# Upstream ships /usr/local/bin/etcd as a single binary; remove it so we can
-# install dynamo_base's etcd directory (etcd+etcdctl+etcdutl) at the same path.
-RUN rm -f /usr/local/bin/etcd
-COPY --from=dynamo_base /usr/bin/nats-server /usr/bin/nats-server
-COPY --from=dynamo_base /usr/local/bin/etcd/ /usr/local/bin/etcd/
-# Copy uv from dynamo_base so the venv symlink below resolves even if a future
-# upstream tensorrt-llm/release tag drops /usr/local/bin/uv from its image.
-COPY --from=dynamo_base /bin/uv /bin/uvx /bin/
-
-# Create dynamo user with group 0 for OpenShift compatibility.
-# Also clear upstream's /workspace baggage (README.md, tutorials/, docker-examples/,
-# license.txt) — pytest collection picks up broken tutorial test files otherwise.
+    ldconfig && \
+    rm -f /usr/local/bin/etcd && \
+    mkdir -p /opt/dynamo && \
+    LIBSTDCPP=/usr/lib/${ARCH_ALT}-linux-gnu/libstdc++.so.6 && \
+    test -f "$LIBSTDCPP" && ln -sf "$LIBSTDCPP" /opt/dynamo/libstdc++.so.6
+
+# One COPY pulls nats-server, etcd/, uv, uvx into their final paths.
+COPY --from=dynamo_base_export / /
+
+# dynamo user (group 0 for OpenShift), clear upstream /workspace baggage
+# (otherwise pytest collects broken tutorial test files), and create the
+# Dynamo venv on non-dev. --system-site-packages keeps upstream's solve
+# importable since system Python is PEP 668 externally-managed.
 RUN userdel -r ubuntu > /dev/null 2>&1 || true \
     && useradd -m -s /bin/bash -g 0 dynamo \
     && [ `id -u dynamo` -eq 1000 ] \
@@ -67,20 +101,22 @@ RUN userdel -r ubuntu > /dev/null 2>&1 || true \
     && rm -rf /workspace && mkdir /workspace \
     && chown dynamo:0 /home/dynamo /home/dynamo/.cache /opt/dynamo /workspace \
     && mkdir -p /etc/profile.d \
-    && echo 'umask 002' > /etc/profile.d/00-umask.sh
-
-COPY --chmod=664 --chown=dynamo:0 ATTRIBUTION* LICENSE /workspace/
-COPY --chmod=775 --chown=dynamo:0 --from=wheel_builder /opt/dynamo/dist/*.whl /opt/dynamo/wheelhouse/
+    && echo 'umask 002' > /etc/profile.d/00-umask.sh{% if target not in ("dev", "local-dev") %} \
+    && python3 -m venv --system-site-packages /opt/dynamo/venv \
+    && ln -sf /usr/bin/uv /opt/dynamo/venv/bin/uv{% endif %}
 
 {% if target not in ("dev", "local-dev") %}
-# Upstream tensorrt-llm/release marks system Python as PEP 668 externally-managed.
-# Install Dynamo wheels into a venv with --system-site-packages so upstream's
-# solve stays importable while our wheels live in their own namespace.
-RUN python3 -m venv --system-site-packages /opt/dynamo/venv \
-    && ln -sf /usr/bin/uv /opt/dynamo/venv/bin/uv
 ENV VIRTUAL_ENV=/opt/dynamo/venv \
     PATH=/opt/dynamo/venv/bin:${PATH}
+{% endif %}
 
+# Place wheels in /opt/dynamo/wheelhouse unconditionally — dev/local-dev images
+# install from source and skip the pip install RUN below, but they still need
+# the wheels on disk because tests/dependencies/test_kvbm_imports.py greps
+# this path and runs in dev-derived test images.
+COPY --chmod=775 --chown=dynamo:0 --from=wheel_builder /opt/dynamo/dist/*.whl /opt/dynamo/wheelhouse/
+
+{% if target not in ("dev", "local-dev") %}
 RUN --mount=type=cache,target=/root/.cache/uv,sharing=locked \
     --mount=type=bind,source=./container/deps/requirements.trtllm.txt,target=/tmp/requirements.trtllm.txt \
     export UV_CACHE_DIR=/root/.cache/uv && \
@@ -107,33 +143,54 @@ RUN --mount=type=cache,target=/root/.cache/uv,sharing=locked \
     fi
 {% endif %}
 
-USER dynamo
-
-# Copy the workspace surface needed by trtllm pre-merge tests.
-# Keep optional framework trees out of /workspace so the upstream runtime does
-# not look like a fully-expanded generic image.
-COPY --chmod=775 --chown=dynamo:0 tests /workspace/tests
-COPY --chmod=775 --chown=dynamo:0 examples /workspace/examples
-COPY --chmod=775 --chown=dynamo:0 deploy /workspace/deploy
-COPY --chmod=775 --chown=dynamo:0 dev /workspace/dev
-COPY --chmod=775 --chown=dynamo:0 components/src/dynamo/common /workspace/components/src/dynamo/common
-COPY --chmod=775 --chown=dynamo:0 components/src/dynamo/frontend /workspace/components/src/dynamo/frontend
-COPY --chmod=775 --chown=dynamo:0 components/src/dynamo/trtllm /workspace/components/src/dynamo/trtllm
-COPY --chmod=775 --chown=dynamo:0 components/src/dynamo/mocker /workspace/components/src/dynamo/mocker
-COPY --chmod=775 --chown=dynamo:0 lib /workspace/lib
-
-RUN --mount=type=bind,source=./container/launch_message/runtime.txt,target=/opt/dynamo/launch_message.txt \
-    sed '/^#\s/d' /opt/dynamo/launch_message.txt > /opt/dynamo/.launch_screen
-
-USER root
-RUN chmod 755 /opt/dynamo/.launch_screen && \
+# Pull /workspace_src (incl. ATTRIBUTION/LICENSE) from the transport stage and
+# wire up the launch screen in a single RUN — saves the standalone workspace COPY layer.
+RUN --mount=type=bind,from=workspace_files,source=/workspace_src,target=/tmp/workspace_src \
+    --mount=type=bind,source=./container/launch_message/runtime.txt,target=/opt/dynamo/launch_message.txt \
+    cp -a /tmp/workspace_src/. /workspace/ && \
+    chown -R dynamo:0 /workspace && \
+    sed '/^#\s/d' /opt/dynamo/launch_message.txt > /opt/dynamo/.launch_screen && \
+    chmod 755 /opt/dynamo/.launch_screen && \
     echo 'cat /opt/dynamo/.launch_screen' >> /etc/bash.bashrc
 
 USER dynamo
 
+# Kept at the bottom — SHA changes per build; layers above stay cached.
 ARG DYNAMO_COMMIT_SHA
 ENV DYNAMO_COMMIT_SHA=${DYNAMO_COMMIT_SHA}
 
 # Reset upstream TRT-LLM image's entrypoint so derived runtimes behave like
 # other Dynamo images and can execute arbitrary commands directly.
 ENTRYPOINT []
+CMD ["/bin/bash"]
+
+{% if target == "runtime" %}
+# Rebase on upstream so this stage inherits upstream's image config
+# (ENV/WORKDIR/USER/CMD) and then overlay runtime_full's filesystem as a
+# single layer. Only Dynamo-specific env needs redeclaring below.
+FROM ${RUNTIME_IMAGE}:${RUNTIME_IMAGE_TAG} AS runtime
+# Whiteout paths runtime_full removed — COPY can't represent deletions, so
+# without this, upstream's /workspace, /home/ubuntu, and single-file
+# /usr/local/bin/etcd would leak alongside our content.
+RUN rm -rf /workspace /home/ubuntu /usr/local/bin/etcd
+COPY --from=runtime_full / /
+
+# Mirrors runtime_full's ENV — must stay in sync. Re-declaration is required
+# because `FROM ${RUNTIME_IMAGE}` here does not inherit runtime_full's config.
+ENV DYNAMO_HOME=/workspace \
+    HOME=/home/dynamo \
+    VIRTUAL_ENV=/opt/dynamo/venv \
+    PATH=/opt/dynamo/venv/bin:/usr/local/bin/etcd:${PATH} \
+    LD_PRELOAD=/opt/dynamo/libstdc++.so.6:/usr/local/lib/python3.12/dist-packages/tensorrt_llm/libs/nixl/libnixl.so \
+    NIXL_PLUGIN_DIR=/usr/local/lib/python3.12/dist-packages/tensorrt_llm/libs/nixl/plugins
+
+WORKDIR /workspace
+
+ARG DYNAMO_COMMIT_SHA
+ENV DYNAMO_COMMIT_SHA=${DYNAMO_COMMIT_SHA}
+
+USER dynamo
+
+ENTRYPOINT []
+CMD ["/bin/bash"]
+{% endif %}