diff --git a/.dockerignore b/.dockerignore index bc93f0a1ed80..aaccd81610d2 100644 --- a/.dockerignore +++ b/.dockerignore @@ -17,6 +17,7 @@ **/*.plan **/*.onnx **/*.plan +**/*.etcd **/.cache/* **/*onnx* # Engine must be allowed because code contains dynamo_engine.py @@ -37,4 +38,5 @@ **/*backup*/ .dockerignore **/target/* -**/*safetensors \ No newline at end of file +**/*safetensors +TensorRT-LLM \ No newline at end of file diff --git a/container/Dockerfile.tensorrt_llm b/container/Dockerfile.tensorrt_llm index b0eb3a134d48..436ee32daa05 100644 --- a/container/Dockerfile.tensorrt_llm +++ b/container/Dockerfile.tensorrt_llm @@ -13,8 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -ARG BASE_IMAGE="gitlab-master.nvidia.com:5005/dl/dgx/tritonserver/tensorrt-llm/amd64" -ARG BASE_IMAGE_TAG="krish-fix-trtllm-build.23766174" +ARG BASE_IMAGE="tensorrt_llm/release" +ARG BASE_IMAGE_TAG="latest" FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS dev diff --git a/container/build.sh b/container/build.sh index e1c0edccb302..11c549e7cd2f 100755 --- a/container/build.sh +++ b/container/build.sh @@ -51,10 +51,8 @@ DOCKERFILE=${SOURCE_DIR}/Dockerfile BUILD_CONTEXT=$(dirname "$(readlink -f "$SOURCE_DIR")") # Base Images -TENSORRTLLM_BASE_VERSION=25.01 -# FIXME: Need a public image for public consumption -TENSORRTLLM_BASE_IMAGE="gitlab-master.nvidia.com:5005/dl/dgx/tritonserver/tensorrt-llm/amd64" -TENSORRTLLM_BASE_IMAGE_TAG=krish-fix-trtllm-build.23766174 +TENSORRTLLM_BASE_IMAGE=tensorrt_llm/release +TENSORRTLLM_BASE_IMAGE_TAG=latest TENSORRTLLM_PIP_WHEEL_PATH="" VLLM_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base" diff --git a/container/build_trtllm_base_image.sh b/container/build_trtllm_base_image.sh new file mode 100755 index 000000000000..fc1e5f0c8f5f --- /dev/null +++ b/container/build_trtllm_base_image.sh @@ -0,0 +1,44 @@ +#!/bin/bash -e +# SPDX-FileCopyrightText: Copyright (c) 2024-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. + +# Build the TRT-LLM base image. + +# This script builds the TRT-LLM base image for Dynamo with TensorRT-LLM. +TRTLLM_COMMIT=9b931c0f6 + +while getopts "c:" opt; do + case ${opt} in + c) TRTLLM_COMMIT=$OPTARG ;; + *) echo "Invalid option" ;; + esac +done + +# Clone the TensorRT-LLM repository. +if [ ! -d "TensorRT-LLM" ]; then + git clone https://github.com/NVIDIA/TensorRT-LLM.git +fi + +cd TensorRT-LLM + +# Checkout the specified commit. +git checkout $TRTLLM_COMMIT + +# Update the submodules. +git submodule update --init --recursive +git lfs pull + +# Build the TRT-LLM base image. +make -C docker release_build diff --git a/examples/python_rs/llm/trtllm/README.md b/examples/python_rs/llm/trtllm/README.md index 7388721a3de0..b6a31189b7f6 100644 --- a/examples/python_rs/llm/trtllm/README.md +++ b/examples/python_rs/llm/trtllm/README.md @@ -39,48 +39,59 @@ Start required services (etcd and NATS): ## Building the Environment -TODO: Remove the internal references below. +### Build the Dynamo container with TensorRT-LLM +#### Step 1: Build TensorRT-LLM base container image -### Build the Dynamo container with latest TRT-LLM +Because of the known issue of C++11 ABI compatibility, we rebuild the TensorRT-LLM from source within NGC pytorch container. +See [here](https://nvidia.github.io/TensorRT-LLM/installation/linux.html) for more informantion. -#### Step 1:Build TRT-LLM wheel using latest tensorrt_llm main +Use the helper script to build a TensorRT-LLM container base image. The script uses a specific commit id from TensorRT-LLM main branch. +```bash +./container/build_trtllm_base_image.sh ``` -git clone https://github.com/NVIDIA/TensorRT-LLM.git -cd TensorRT-LLM - -# Start a dev docker container. Dont forget to mount your home directory to /home in the docker run command. -make -C docker jenkins_run LOCAL_USER=1 DOCKER_RUN_ARGS="-v /user/home:/home" -# Build wheel for the GPU architecture you are currently using ("native"). -# We use -f to run fast build which should speed up the build process. But it might not work for all GPUs and for full functionality you should disable it. -python3 scripts/build_wheel.py --clean --trt_root /usr/local/tensorrt -a native -i -p -ccache +For more information see [here](https://nvidia.github.io/TensorRT-LLM/installation/build-from-source-linux.html#option-1-build-tensorrt-llm-in-one-step) for more details on building from source. +If you already have a TensorRT-LLM container image, you can skip this step. -# Copy wheel to your local directory -cp build/tensorrt_llm-*.whl /home -``` +#### Step 2: Build the Dynamo container -####Step 2: Copy the TRT-LLM wheel to dynamo repository. ```bash -cp /home/tensorrt_llm-*.whl //dynamo/trtllm_wheel/ +# Build image +./container/build.sh --framework TENSORRTLLM ``` -####Step 3: Build the container +This build script internally points to the base container image built with step 1. If you skipped previous step because you already have the container image available, you can run the build script with that image as a base. + ```bash -# Build image -./container/build.sh --framework TENSORRTLLM --tensorrtllm-pip-wheel-path trtllm_wheel +# Build dynamo image with other TRTLLM base image. +./container/build.sh --framework TENSORRTLLM --base-image --base-image-tag ``` -We need to copy the TRT-LLM wheel to repository and point the build script to the path within -the repository so that it can be picked by the docker build context. - ## Launching the Environment -``` +```bash # Run image interactively from with the Dynamo root directory. ./container/run.sh --framework TENSORRTLLM -it ``` +## Quick Start using dynamo run + +```bash +# Run the server +dynamo run out=pystr:./monolith/dynamo_engine.py -- --engine_args llm_api_config.yaml +``` + +The above command should load the model specified in `llm_api_config.yaml` and start accepting +text input from the client. For more details on the `dynamo run` command, please refer to the +[dynamo run](/launch/README.md#python-bring-your-own-engine) documentation. + +Currently only monolithic deployment option is supported by `dynamo run` for TensorRT-LLM. +Adding support for disaggregated deployment is under development. This does *not* require +any other pre-requisites mentioned in the [Prerequisites](#prerequisites) section. + + + ## Deployment Options Note: NATS and ETCD servers should be running and accessible from the container as described in the [Prerequisites](#prerequisites) section. diff --git a/examples/python_rs/llm/trtllm/common/protocol.py b/examples/python_rs/llm/trtllm/common/protocol.py index 4244a5c4e6e5..c993e12d59d1 100644 --- a/examples/python_rs/llm/trtllm/common/protocol.py +++ b/examples/python_rs/llm/trtllm/common/protocol.py @@ -23,12 +23,23 @@ from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, ChatCompletionStreamResponse, + CompletionRequest, CompletionResponseStreamChoice, DisaggregatedParams, UsageInfo, ) +# The max_tokens is being deprecated in favor of max_completion_tokens. +# However, TRTLLM protocol might still refer it as max_tokens. +class AdaptedCompletionRequest(CompletionRequest): + max_completion_tokens: Optional[int] = None + + +class AdaptedChatCompletionRequest(ChatCompletionRequest): + max_completion_tokens: Optional[int] = None + + class Tokens(BaseModel): tokens: list[int] diff --git a/examples/python_rs/llm/trtllm/disaggregated/server.py b/examples/python_rs/llm/trtllm/disaggregated/server.py index 65f9983cda22..0277ec95f188 100644 --- a/examples/python_rs/llm/trtllm/disaggregated/server.py +++ b/examples/python_rs/llm/trtllm/disaggregated/server.py @@ -160,7 +160,7 @@ async def generate_chat(self, request): gen_req = copy.deepcopy(request) - request.max_completion_tokens = 1 + request.max_tokens = 1 ctx_resp = await self._get_ctx_resp(request, self.ctx_chat_client) ctx_resp_obj = DisaggChatCompletionStreamResponse.model_validate_json(ctx_resp) diff --git a/examples/python_rs/llm/trtllm/monolith/dynamo_engine.py b/examples/python_rs/llm/trtllm/monolith/dynamo_engine.py index ef4759727c5f..273ae6482dad 100644 --- a/examples/python_rs/llm/trtllm/monolith/dynamo_engine.py +++ b/examples/python_rs/llm/trtllm/monolith/dynamo_engine.py @@ -23,10 +23,7 @@ from pathlib import Path from tensorrt_llm.logger import logger -from tensorrt_llm.serve.openai_protocol import ( - ChatCompletionRequest, - ChatCompletionStreamResponse, -) +from tensorrt_llm.serve.openai_protocol import ChatCompletionStreamResponse from dynamo.runtime import dynamo_endpoint @@ -41,6 +38,7 @@ ) from common.generators import chat_generator # noqa: E402 from common.parser import parse_dynamo_run_args # noqa: E402 +from common.protocol import AdaptedChatCompletionRequest # noqa: E402 logger.set_level("info") @@ -67,8 +65,10 @@ def init_global_engine(args, engine_config): engine = DynamoTRTLLMEngine(trt_llm_engine_config) -@dynamo_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) +@dynamo_endpoint(AdaptedChatCompletionRequest, ChatCompletionStreamResponse) async def generate(request): + if request.max_completion_tokens is not None: + request.max_tokens = request.max_completion_tokens async for response in chat_generator(engine, request): yield response diff --git a/examples/python_rs/llm/trtllm/monolith/worker.py b/examples/python_rs/llm/trtllm/monolith/worker.py index e6d83c36b235..0103a3b929d6 100644 --- a/examples/python_rs/llm/trtllm/monolith/worker.py +++ b/examples/python_rs/llm/trtllm/monolith/worker.py @@ -19,11 +19,10 @@ from common.base_engine import BaseTensorrtLLMEngine, TensorrtLLMEngineConfig from common.generators import chat_generator, completion_generator from common.parser import LLMAPIConfig, parse_tensorrt_llm_args +from common.protocol import AdaptedChatCompletionRequest, AdaptedCompletionRequest from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ( - ChatCompletionRequest, ChatCompletionStreamResponse, - CompletionRequest, CompletionStreamResponse, ) @@ -41,14 +40,18 @@ class TensorrtLLMEngine(BaseTensorrtLLMEngine): def __init__(self, trt_llm_engine_config: TensorrtLLMEngineConfig): super().__init__(trt_llm_engine_config) - @dynamo_endpoint(ChatCompletionRequest, ChatCompletionStreamResponse) + @dynamo_endpoint(AdaptedChatCompletionRequest, ChatCompletionStreamResponse) async def generate_chat(self, request): + if request.max_completion_tokens is not None: + request.max_tokens = request.max_completion_tokens async for response in chat_generator(self, request): yield response self._start_threads() - @dynamo_endpoint(CompletionRequest, CompletionStreamResponse) + @dynamo_endpoint(AdaptedCompletionRequest, CompletionStreamResponse) async def generate_completion(self, request): + if request.max_completion_tokens is not None: + request.max_tokens = request.max_completion_tokens async for response in completion_generator(self, request): yield response self._start_threads()