From 5c5aee1366c6578ae48893f20fdb37a021873af6 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 18 Feb 2025 11:15:01 -0800 Subject: [PATCH] add trtllm disagg scripts --- examples/python_rs/llm/trtllm/README.md | 74 +++++++ .../python_rs/llm/trtllm/common/__init__.py | 0 .../python_rs/llm/trtllm/common/client.py | 81 +++++++ .../python_rs/llm/trtllm/common/parser.py | 86 +++++++ .../python_rs/llm/trtllm/common/protocol.py | 29 +++ .../python_rs/llm/trtllm/disagg/__init__.py | 0 .../llm/trtllm/disagg/disagg_config.yaml | 18 ++ .../python_rs/llm/trtllm/disagg/router.py | 122 ++++++++++ .../python_rs/llm/trtllm/disagg/worker.py | 209 ++++++++++++++++++ examples/python_rs/llm/trtllm/model.json | 24 ++ 10 files changed, 643 insertions(+) create mode 100644 examples/python_rs/llm/trtllm/README.md create mode 100644 examples/python_rs/llm/trtllm/common/__init__.py create mode 100644 examples/python_rs/llm/trtllm/common/client.py create mode 100644 examples/python_rs/llm/trtllm/common/parser.py create mode 100644 examples/python_rs/llm/trtllm/common/protocol.py create mode 100644 examples/python_rs/llm/trtllm/disagg/__init__.py create mode 100644 examples/python_rs/llm/trtllm/disagg/disagg_config.yaml create mode 100644 examples/python_rs/llm/trtllm/disagg/router.py create mode 100644 examples/python_rs/llm/trtllm/disagg/worker.py create mode 100644 examples/python_rs/llm/trtllm/model.json diff --git a/examples/python_rs/llm/trtllm/README.md b/examples/python_rs/llm/trtllm/README.md new file mode 100644 index 000000000..1d8e89502 --- /dev/null +++ b/examples/python_rs/llm/trtllm/README.md @@ -0,0 +1,74 @@ + + +# TensorRT-LLM Integration with Triton Distributed + +This example demonstrates how to use Triton Distributed to serve large language models with the tensorrt_llm engine, enabling efficient model serving with monolithic option. + +## Prerequisites + +Start required services (etcd and NATS): + + Option A: Using [Docker Compose](/runtime/rust/docker-compose.yml) (Recommended) + ```bash + docker-compose up -d + ``` + + Option B: Manual Setup + + - [NATS.io](https://docs.nats.io/running-a-nats-service/introduction/installation) server with [Jetstream](https://docs.nats.io/nats-concepts/jetstream) + - example: `nats-server -js --trace` + - [etcd](https://etcd.io) server + - follow instructions in [etcd installation](https://etcd.io/docs/v3.5/install/) to start an `etcd-server` locally + +*Note*: This example is work in progress. + +## Building the Environment [WIP] + +Use the following image with tritonserver, triton distributed and tensorrt_llm. + +```bash +nvcr.io/pfteb4cqjzrs/playground/tritondistributed:trtllm_disagg-shreyas +``` + +## Launching the Environment +``` +# Run image interactively +docker run -it --network host --shm-size=64G --ulimit memlock=-1 --ulimit stack=67108864 -e HF_HOME=/home/scratch.shreyasm_gpu/hf_cache --gpus=all -v .:/workspace nvcr.io/pfteb4cqjzrs/playground/tritondistributed:trtllm_disagg-shreyas +``` + +## Deployment Options + +### 1. Disaggregated Deployment + +NOTE: This will start both context and generation servers. +```bash +mpirun --allow-run-as-root -n 2 python3 -m disagg.worker --engine_args model.json & +``` + +Run router. This is still a WIP. It needs information about which endpoints are ctx/gen. +```bash +python3 -m disagg.router +``` + +Run client +```bash +python3 -m common.client \ + --prompt "Describe the capital of France" \ + --max-tokens 10 \ + --temperature 0.5 +``` \ No newline at end of file diff --git a/examples/python_rs/llm/trtllm/common/__init__.py b/examples/python_rs/llm/trtllm/common/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/examples/python_rs/llm/trtllm/common/client.py b/examples/python_rs/llm/trtllm/common/client.py new file mode 100644 index 000000000..d7b17d781 --- /dev/null +++ b/examples/python_rs/llm/trtllm/common/client.py @@ -0,0 +1,81 @@ +# 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. + + +import argparse +import asyncio + +import uvloop +from triton_distributed_rs import DistributedRuntime, triton_worker + +from .protocol import Request + + +@triton_worker() +async def worker( + runtime: DistributedRuntime, + prompt: str, + max_tokens: int, + temperature: float, + streaming: bool, +): + """ + Instantiate a `backend` client and call the `generate` endpoint + """ + # get endpoint + endpoint = ( + runtime.namespace("triton-init").component("router").endpoint("generate") + ) + + # create client + client = await endpoint.client() + + # list the endpoints + print(client.endpoint_ids()) + + # issue request + tasks = [] + for _ in range(1): + tasks.append( + client.generate( + Request( + prompt=prompt, + sampling_params={ + "temperature": temperature, + "max_tokens": max_tokens, + }, + streaming=streaming, + ).model_dump_json() + ) + ) + streams = await asyncio.gather(*tasks) + + # process response + for stream in streams: + async for resp in stream: + print(resp) + + +if __name__ == "__main__": + uvloop.install() + + parser = argparse.ArgumentParser() + parser.add_argument("--prompt", type=str, default="what is the capital of france?") + parser.add_argument("--max-tokens", type=int, default=10) + parser.add_argument("--temperature", type=float, default=0.5) + parser.add_argument("--streaming", type=bool, default=True) + args = parser.parse_args() + + asyncio.run(worker(args.prompt, args.max_tokens, args.temperature, args.streaming)) diff --git a/examples/python_rs/llm/trtllm/common/parser.py b/examples/python_rs/llm/trtllm/common/parser.py new file mode 100644 index 000000000..d732bea28 --- /dev/null +++ b/examples/python_rs/llm/trtllm/common/parser.py @@ -0,0 +1,86 @@ +# 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. + +import argparse +import json +import os +from typing import Any, Dict, Tuple + +# Define the expected keys for each config +# TODO: Add more keys as needed +PYTORCH_CONFIG_KEYS = { + "use_cuda_graph", + "cuda_graph_batch_sizes", + "cuda_graph_max_batch_size", + "cuda_graph_padding_enabled", + "enable_overlap_scheduler", + "kv_cache_dtype", + "torch_compile_enabled", + "torch_compile_fullgraph", + "torch_compile_inductor_enabled", +} + +LLM_ENGINE_KEYS = { + "model", + "tokenizer", + "tokenizer_model", + "skip_tokenizer_init", + "trust_remote_code", + "tensor_parallel_size", + "dtype", + "revision", + "tokenizer_revision", + "speculative_model", + "enable_chunked_prefill", + "disagg_config", +} + + +def _get_llm_args(args_dict): + pytorch_config_args = { + k: v for k, v in args_dict.items() if k in PYTORCH_CONFIG_KEYS and v is not None + } + llm_engine_args = { + k: v for k, v in args_dict.items() if k in LLM_ENGINE_KEYS and v is not None + } + if "model" not in llm_engine_args: + raise ValueError("Model name is required in the TRT-LLM engine config.") + + return (pytorch_config_args, llm_engine_args) + + +def _init_engine_args(engine_args_filepath): + """Initialize engine arguments from config file.""" + if not os.path.isfile(engine_args_filepath): + raise ValueError( + f"'{engine_args_filepath}' containing TRT-LLM engine args must be provided in when launching the worker" + ) + + try: + with open(engine_args_filepath) as file: + trtllm_engine_config = json.load(file) + except json.JSONDecodeError as e: + raise RuntimeError(f"Failed to parse engine config: {e}") + + return _get_llm_args(trtllm_engine_config) + + +def parse_tensorrt_llm_args() -> Tuple[Dict[str, Any], Dict[str, Any]]: + parser = argparse.ArgumentParser(description="A TensorRT-LLM Worker parser") + parser.add_argument( + "--engine_args", type=str, required=True, help="Path to the engine args file" + ) + args = parser.parse_args() + return _init_engine_args(args.engine_args) diff --git a/examples/python_rs/llm/trtllm/common/protocol.py b/examples/python_rs/llm/trtllm/common/protocol.py new file mode 100644 index 000000000..30eb6a438 --- /dev/null +++ b/examples/python_rs/llm/trtllm/common/protocol.py @@ -0,0 +1,29 @@ +# 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. + +from tensorrt_llm.llmapi import DisaggregatedParams +from pydantic import BaseModel + + +class Request(BaseModel): + prompt: str + sampling_params: dict + streaming: bool = True + disaggregated_params: DisaggregatedParams = {} + + +class Response(BaseModel): + text: str + disaggregated_params: DisaggregatedParams = {} \ No newline at end of file diff --git a/examples/python_rs/llm/trtllm/disagg/__init__.py b/examples/python_rs/llm/trtllm/disagg/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/examples/python_rs/llm/trtllm/disagg/disagg_config.yaml b/examples/python_rs/llm/trtllm/disagg/disagg_config.yaml new file mode 100644 index 000000000..f8a275b78 --- /dev/null +++ b/examples/python_rs/llm/trtllm/disagg/disagg_config.yaml @@ -0,0 +1,18 @@ +model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 +hostname: localhost +port: 8000 +backend: "pytorch" +context_servers: + num_instances: 1 + gpu_fraction: 0.25 + tp_size: 1 + pp_size: 1 + urls: + - "localhost:8001" +generation_servers: + num_instances: 1 + gpu_fraction: 0.25 + tp_size: 1 + pp_size: 1 + urls: + - "localhost:8002" diff --git a/examples/python_rs/llm/trtllm/disagg/router.py b/examples/python_rs/llm/trtllm/disagg/router.py new file mode 100644 index 000000000..9171372a0 --- /dev/null +++ b/examples/python_rs/llm/trtllm/disagg/router.py @@ -0,0 +1,122 @@ +# 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. + + +import argparse +import asyncio +import copy +import uvloop +from triton_distributed_rs import DistributedRuntime, triton_worker +from tensorrt_llm.llmapi import DisaggregatedParams +from tensorrt_llm.logger import logger +from tensorrt_llm.llmapi.disagg_utils import (CtxGenServerConfig, + parse_disagg_config_file) +from triton_distributed_rs import DistributedRuntime, triton_endpoint, triton_worker + +from common.protocol import Request, Response + +logger.set_level("info") + +def get_server_urls(server_configs): + ctx_server_urls = [] + gen_server_urls = [] + for cfg in server_configs: + if cfg.type == "ctx": + ctx_server_urls.append(f"http://{cfg.hostname}:{cfg.port}") + else: + gen_server_urls.append(f"http://{cfg.hostname}:{cfg.port}") + + return ctx_server_urls, gen_server_urls + + +class Router: + def __init__(self, ctx_client, gen_client): + self.ctx_server_idx = 0 + self.gen_server_idx = 0 + # TODO: Add support for multiple clients + # would need a different way of obtaining endpoints from a component + self.ctx_clients = [ctx_client] + self.gen_clients = [gen_client] + logger.info("INITIALIZED ROUTER") + + def get_next_server(self, servers, server_type): + """Round-robin selection of next available server""" + if not servers: + raise ValueError(f"No {server_type} servers available") + + if server_type == "ctx": + server = servers[self.ctx_server_idx] + self.ctx_server_idx = (self.ctx_server_idx + 1) % len(servers) + else: + server = servers[self.gen_server_idx] + self.gen_server_idx = (self.gen_server_idx + 1) % len(servers) + + return server + + @triton_endpoint(Request, Response) + async def generate(self, request): + gen_req = copy.deepcopy(request) + + # Pick a context server + ctx_client = self.get_next_server(self.ctx_clients, "ctx") + logger.info(f"Sending request to ctx server: {ctx_client}") + + # Send request to context server + request.sampling_params["max_tokens"] = 1 + request.disaggregated_params = DisaggregatedParams(request_type="context_only") + logger.info(f"Request to ctx server: {request}") + + async for ctx_resp in await ctx_client.generate(request.model_dump_json()): + gen_req.disaggregated_params = ctx_resp.disaggregated_params + gen_req.disaggregated_params.request_type = "generation_only" + break + + # Pick a generation server + gen_client = self.get_next_server(self.gen_clients, "gen") + logger.info(f"Sending request to gen server: {gen_client}") + logger.info(f"Request to gen server: {gen_req}") + + async for response in await gen_client.generate(gen_req.model_dump_json()): + yield response + +@triton_worker() +async def worker( + runtime: DistributedRuntime, +): + """ + Instantiate a `backend` component and serve the `generate` endpoint + A `Component` can serve multiple endpoints + """ + component = runtime.namespace("triton-init").component("router") + await component.create_service() + + client = await runtime.namespace("triton-init").component("tensorrt-llm").endpoint("generate").client() + # TODO: need information about which endpoints are ctx/gen + print("client endpoints: ", client.endpoint_ids()) + + endpoint = component.endpoint("generate") + await endpoint.serve_endpoint(Router(client).generate) + +if __name__ == "__main__": + uvloop.install() + + parser = argparse.ArgumentParser() + parser.add_argument("--disagg-config", type=str, + default="disagg/disagg_config.yaml") + args = parser.parse_args() + disagg_config = parse_disagg_config_file(args.disagg_config) + ctx_server_urls, gen_server_urls = get_server_urls(disagg_config.server_configs) + + asyncio.run(worker()) \ No newline at end of file diff --git a/examples/python_rs/llm/trtllm/disagg/worker.py b/examples/python_rs/llm/trtllm/disagg/worker.py new file mode 100644 index 000000000..1ce15b326 --- /dev/null +++ b/examples/python_rs/llm/trtllm/disagg/worker.py @@ -0,0 +1,209 @@ +# 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. + + +import asyncio +import os +import threading +from contextlib import asynccontextmanager +from typing import Any, Dict, Tuple + +import uvloop +from common.parser import parse_tensorrt_llm_args +from common.protocol import Request, Response +from tensorrt_llm import SamplingParams +from tensorrt_llm._torch import LLM +from tensorrt_llm._utils import set_mpi_comm +from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.llmapi import LLM, BuildConfig, KvCacheConfig, MpiCommSession +from tensorrt_llm.llmapi import DisaggregatedParams +from tensorrt_llm.llmapi.disagg_utils import (CtxGenServerConfig, + parse_disagg_config_file, + split_world_comm) +from triton_distributed_rs import DistributedRuntime, triton_endpoint, triton_worker + +from mpi4py.futures import MPICommExecutor + +logger.set_level("info") + +class TensorrtLLMEngine: + """ + Request handler for the generate endpoint + """ + + def __init__(self, engine_args: Tuple[Dict[str, Any], Dict[str, Any]], + disagg_config: CtxGenServerConfig, instance_idx: int, sub_comm): + self.pytorch_config_args, self.llm_engine_args = engine_args + self.disagg_config = disagg_config + self.instance_idx = instance_idx + self.mpi_session = MpiCommSession(sub_comm) + self._init_engine() + + def _init_engine(self): + logger.info("Initializing engine") + + os.environ['TRTLLM_USE_MPI_KVCACHE'] = "1" + # Run the engine in a separate thread running the AsyncIO event loop. + self._llm_engine = None + self._llm_engine_start_cv = threading.Condition() + self._llm_engine_shutdown_event = asyncio.Event() + self._event_thread = threading.Thread( + target=asyncio.run, args=(self._run_llm_engine(),) + ) + self._event_thread.start() + with self._llm_engine_start_cv: + while self._llm_engine is None: + self._llm_engine_start_cv.wait() + + # The 'threading.Thread()' will not raise the exception here should the engine + # failed to start, so the exception is passed back via the engine variable. + if isinstance(self._llm_engine, Exception): + e = self._llm_engine + logger.error(f"Failed to start engine: {e}") + if self._event_thread is not None: + self._event_thread.join() + self._event_thread = None + raise e + + async def _run_llm_engine(self): + # Counter to keep track of ongoing request counts. + self._ongoing_request_count = 0 + + @asynccontextmanager + async def async_llm_wrapper(): + # Create LLM in a thread to avoid blocking + loop = asyncio.get_running_loop() + try: + pytorch_config = PyTorchConfig(**self.pytorch_config_args) + llm = await loop.run_in_executor( + None, + lambda: LLM( + **self.llm_engine_args, + # TODO: no hardcode + tensor_parallel_size=1, + pipeline_parallel_size=1, + gpus_per_node=None, + trust_remote_code=True, + mpi_session=self.mpi_session, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.25), + pytorch_backend_config=pytorch_config, + backend="pytorch" + ), + ) + yield llm + finally: + if "llm" in locals(): + # Run shutdown in a thread to avoid blocking + await loop.run_in_executor(None, llm.shutdown) + + try: + async with async_llm_wrapper() as engine: + # Capture the engine event loop and make it visible to other threads. + self._event_loop = asyncio.get_running_loop() + + # Signal the engine is started and make it visible to other threads. + with self._llm_engine_start_cv: + self._llm_engine = engine + self._llm_engine_start_cv.notify_all() + + logger.info("Engine loaded and ready to serve...") + + # Wait for the engine shutdown signal. + await self._llm_engine_shutdown_event.wait() + + # Wait for the ongoing requests to complete. + while self._ongoing_request_count > 0: + logger.info( + "Awaiting remaining {} requests".format( + self._ongoing_request_count + ) + ) + await asyncio.sleep(1) + + # Cancel all tasks in the event loop. + for task in asyncio.all_tasks(loop=self._event_loop): + if task is not asyncio.current_task(): + task.cancel() + + except Exception as e: + # Signal and pass the exception back via the engine variable if the engine + # failed to start. If the engine has started, re-raise the exception. + with self._llm_engine_start_cv: + if self._llm_engine is None: + self._llm_engine = e + self._llm_engine_start_cv.notify_all() + return + raise e + + self._llm_engine = None + logger.info("Shutdown complete") + + @triton_endpoint(Request, Response) + async def generate(self, request): + self._ongoing_request_count += 1 + logger.info(f"Received request: {request}") + sampling_params = SamplingParams(**request.sampling_params) + + assert request.disaggregated_params is not None + + async for response in self._llm_engine.generate_async( + request.prompt, sampling_params, streaming=request.streaming, + disaggregated_params=request.disaggregated_params + ): + logger.info(f"Generated response: {response}") + yield Response(text=response.outputs[0].text, + disaggregated_params=response.outputs[0].disaggregated_params) + self._ongoing_request_count -= 1 + + +@triton_worker() +async def worker( + runtime: DistributedRuntime, + engine_args: Tuple[Dict[str, Any], Dict[str, Any]], + disagg_config: CtxGenServerConfig, + instance_idx: int, + sub_comm +): + """ + Instantiate a `backend` component and serve the `generate` endpoint + A `Component` can serve multiple endpoints + """ + component = runtime.namespace("triton-init").component("tensorrt-llm") + await component.create_service() + + endpoint = component.endpoint("generate") + await endpoint.serve_endpoint(TensorrtLLMEngine(engine_args, disagg_config, instance_idx, sub_comm).generate) + + +if __name__ == "__main__": + uvloop.install() + args, engine_args = parse_tensorrt_llm_args() + disagg_config = parse_disagg_config_file(engine_args[1].pop('disagg_config')) + + logger.info(f"disagg_config: {disagg_config}") + + is_leader, instance_idx, sub_comm = split_world_comm( + disagg_config.server_configs) + logger.info(f"is_leader: {is_leader}, instance_idx: {instance_idx}") + + if is_leader: + asyncio.run(worker(engine_args, component_type, disagg_config, instance_idx, sub_comm)) + + else: + set_mpi_comm(sub_comm) + with MPICommExecutor(sub_comm) as executor: + if not is_leader and executor is not None: + raise RuntimeError(f"rank{COMM_WORLD} should not have executor") diff --git a/examples/python_rs/llm/trtllm/model.json b/examples/python_rs/llm/trtllm/model.json new file mode 100644 index 000000000..3e1d08579 --- /dev/null +++ b/examples/python_rs/llm/trtllm/model.json @@ -0,0 +1,24 @@ +{ + "model":"TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "tokenizer": null, + "tokenizer_model": null, + "skip_tokenizer_init": null, + "trust_remote_code": null, + "tensor_parallel_size": null, + "dtype": null, + "revision": null, + "tokenizer_revision": null, + "speculative_model": null, + "enable_chunked_prefill": null, + "disagg_config": "disagg/disagg_config.yaml", + + "use_cuda_graph": null, + "cuda_graph_batch_sizes": null, + "cuda_graph_max_batch_size": null, + "cuda_graph_padding_enabled": null, + "enable_overlap_scheduler": null, + "kv_cache_dtype": null, + "torch_compile_enabled": null, + "torch_compile_fullgraph": null, + "torch_compile_inductor_enabled": null +} \ No newline at end of file