diff --git a/tests/README.md b/tests/README.md index 4e03df09d8f2..70f980f543f7 100644 --- a/tests/README.md +++ b/tests/README.md @@ -185,7 +185,7 @@ l0_a10: linux_distribution_name: ubuntu* tests: # ------------- PyTorch tests --------------- - - disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_cuda_graph[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_overlap[TinyLlama-1.1B-Chat-v1.0] @@ -200,7 +200,7 @@ l0_a10: 2. Copy all items in `tests` field to a text file, for example, `a10_list.txt`. Don't forget to remove extra characters like comments and the dash marks. ``` -disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_cuda_graph[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_overlap[TinyLlama-1.1B-Chat-v1.0] @@ -225,7 +225,7 @@ To set a timeout for specific long-running test cases, follow these steps: 1. Locate the test case line in the corresponding test-db YAML file (e.g., `tests/integration/test_lists/test-db/l0_a10.yml`). 2. Append `TIMEOUT (...)` to the test case line, as shown below: ```yaml - - disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] TIMEOUT (30) + - disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] TIMEOUT (30) ``` - Ensure there is **at least one space** before and after the `TIMEOUT` keyword. - The time value inside the parentheses `()` must be a **number** representing the timeout in **minutes**. @@ -233,7 +233,7 @@ To set a timeout for specific long-running test cases, follow these steps: ### For Local Testing (TXT files): 1. If you are running the tests locally using a prepared `.txt` file (e.g., `a10_list.txt`), append the `TIMEOUT` setting to the test case line in the same way: ``` - disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] TIMEOUT (30) + disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] TIMEOUT (30) ``` ## 6. Set isolated execution for cases individually @@ -251,7 +251,7 @@ Add `ISOLATION` to the test case line with proper spacing: **For CI (test-db YAML files):** ```yaml -- disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] ISOLATION +- disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] ISOLATION ``` ## 7. Combining test markers diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index c1695b49a029..07fe6d8aa5a1 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -340,10 +340,10 @@ "disaggregated/test_disaggregated.py::test_disaggregated_kv_cache_time_output[TinyLlama-1.1B-Chat-v1.0]": 52.78952780482359, "disaggregated/test_disaggregated.py::test_disaggregated_load_balance[TinyLlama-1.1B-Chat-v1.0]": 73.48997121001594, "disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0]": 67.3897166326642, - "disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0]": 54.22262764698826, + "disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0]": 54.22262764698826, "disaggregated/test_disaggregated.py::test_disaggregated_overlap[TinyLlama-1.1B-Chat-v1.0]": 98.97588296607137, - "disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0]": 67.9668476767838, - "disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun_trt_backend[TinyLlama-1.1B-Chat-v1.0]": 82.28277984517626, + "disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0]": 67.9668476767838, + "disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_trt_backend[TinyLlama-1.1B-Chat-v1.0]": 82.28277984517626, "disaggregated/test_disaggregated.py::test_disaggregated_trtllm_sampler[TinyLlama-1.1B-Chat-v1.0]": 62.51559329708107, "disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_llama_context_capacity[False-False-DeepSeek-V3-Lite-fp8/fp8]": 238.76137515995651, "disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_deepseek[False-False-DeepSeek-V3-Lite-fp8/fp8]": 78.98068026197143, diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index d74f829f380f..35c51a090c9a 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -16,7 +16,7 @@ import pytest import requests import yaml -from defs.common import revise_disaggregated_server_config_urls_with_free_ports +from defs.common import get_free_port_in_ci as get_free_port from tensorrt_llm.executor.result import GenerationResultBase from tensorrt_llm.llmapi import CompletionOutput, RequestOutput, SamplingParams @@ -170,8 +170,38 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): _apply_perf_flags(ctx_server_config) _apply_perf_flags(gen_server_config) - disaggregated_server_config = revise_disaggregated_server_config_urls_with_free_ports( - disaggregated_server_config) + # Always assign free port dynamically for service discovery + serve_port = get_free_port() + disaggregated_server_config["port"] = serve_port + + # Use HTTP service discovery + cluster_uri = f"http://localhost:{serve_port}" + print(f"Using HTTP service discovery at {cluster_uri}") + + # Create service discovery config + disagg_cluster = { + "cluster_uri": cluster_uri, + "cluster_name": "test_cluster", + "heartbeat_interval_sec": 1, + "inactive_timeout_sec": 2, + } + + # Auto-deduce minimal_instances from num_instances + num_ctx_instances = disaggregated_server_config["context_servers"][ + "num_instances"] + num_gen_instances = disaggregated_server_config["generation_servers"][ + "num_instances"] + disagg_cluster["minimal_instances"] = { + "context_servers": num_ctx_instances, + "generation_servers": num_gen_instances + } + + # Inject disagg_cluster into server config (for minimal_instances and is_ready check) + disaggregated_server_config["disagg_cluster"] = disagg_cluster + + # Inject into worker configs + ctx_server_config = {**ctx_server_config, "disagg_cluster": disagg_cluster} + gen_server_config = {**gen_server_config, "disagg_cluster": disagg_cluster} with open(disaggregated_serving_config_path, "w") as f: yaml.dump(disaggregated_server_config, f) @@ -221,12 +251,9 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): ctx_total_gpus = ctx_tp * ctx_pp * ctx_cp gen_total_gpus = gen_tp * gen_pp * gen_cp - ctx_urls = disaggregated_server_config["context_servers"]["urls"] - gen_urls = disaggregated_server_config["generation_servers"]["urls"] - - serve_port = disaggregated_server_config["port"] - ctx_ports = [int(url.split(":")[1]) for url in ctx_urls] - gen_ports = [int(url.split(":")[1]) for url in gen_urls] + # Auto-assign ports for workers (port=0 means dynamic assignment) + ctx_ports = [0] * num_ctx_instances + gen_ports = [0] * num_gen_instances ctx_servers = [] current_gpu_offset = 0 @@ -256,8 +283,9 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): ctx_server_args = ctx_args + [ "--port", - str(port), "--config", ctx_server_config_path, - f"--tp_size={ctx_tp}", f"--pp_size={ctx_pp}", f"--cp_size={ctx_cp}" + str(port), "--config", ctx_server_config_path, "--server_role", + "context", f"--tp_size={ctx_tp}", f"--pp_size={ctx_pp}", + f"--cp_size={ctx_cp}" ] if "max_num_tokens" in ctx_server_config: ctx_server_args.append( @@ -285,8 +313,9 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): gen_server_args = gen_args + [ "--port", - str(port), "--config", gen_server_config_path, - f"--tp_size={gen_tp}", f"--pp_size={gen_pp}", f"--cp_size={gen_cp}" + str(port), "--config", gen_server_config_path, "--server_role", + "generation", f"--tp_size={gen_tp}", f"--pp_size={gen_pp}", + f"--cp_size={gen_cp}" ] if "max_num_tokens" in gen_server_config: gen_server_args.append( @@ -346,11 +375,15 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): f"process {process.pid} exited with code {process.returncode}" ) try: - print("Checking health endpoint") - response = requests.get(f"http://localhost:{serve_port}/health") + print("Checking cluster_info endpoint for worker registration") + response = requests.get( + f"http://localhost:{serve_port}/cluster_info") if response.status_code == 200: - server_is_ready = True - break + cluster_info = response.json() + if cluster_info.get("is_ready"): + print(f"Cluster ready: {cluster_info}") + server_is_ready = True + break except requests.exceptions.ConnectionError: continue if not server_is_ready: @@ -420,7 +453,7 @@ def generate_async(prompt: str, def _get_perf_metrics(): path = "/perf_metrics" - perf_url = f"http://localhost:8000{path}" + perf_url = f"http://localhost:{serve_port}{path}" try: print(f"Fetching perf metrics from {perf_url}") resp = requests.get(perf_url, timeout=10) @@ -523,20 +556,15 @@ def run_parallel_test(model_name: str, } } - ctx_urls = [f"localhost:{8001 + i * 2}" for i in range(ctx_instances)] - gen_urls = [f"localhost:{8002 + i * 2}" for i in range(gen_instances)] - + # No need to generate URLs - workers will register via service discovery disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { "num_instances": ctx_instances, - "urls": ctx_urls }, "generation_servers": { "num_instances": gen_instances, - "urls": gen_urls } } with launch_disaggregated_llm(disaggregated_server_config, @@ -582,15 +610,12 @@ def test_auto_dtype(self, ctx_disable_overlap_scheduler, gen_server_config["cache_transceiver_config"] = {"backend": "DEFAULT"} disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -629,15 +654,12 @@ def test_ngram(self): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -688,15 +710,12 @@ def test_eagle3(self, overlap_scheduler, eagle3_one_model): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -724,15 +743,12 @@ def test_guided_decoding(self, backend: str, mocker): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -780,15 +796,12 @@ def test_guided_decoding_with_eagle3(self, backend: str, } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -860,15 +873,12 @@ def test_auto_dtype(self, overlap_scheduler): gen_server_config["max_seq_len"] = 8192 disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -903,15 +913,12 @@ def test_nixl_backend(self): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -939,15 +946,12 @@ def test_auto_dtype(self, overlap_scheduler, mtp_nextn): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -1040,15 +1044,12 @@ def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config, } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -1093,15 +1094,12 @@ def test_guided_decoding(self, backend: str, mtp_nextn: int, mocker): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -1146,15 +1144,12 @@ def test_auto_dtype(self, block_reuse): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -1207,15 +1202,12 @@ def test_auto_dtype(self, block_reuse, mocker): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -1279,15 +1271,12 @@ def test_auto_dtype(self, overlap_scheduler): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -1322,15 +1311,12 @@ def test_nixl_backend(self): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -1365,15 +1351,12 @@ def test_auto_dtype(self, overlap_scheduler, enable_partial_reuse): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -1410,15 +1393,12 @@ def _test_chunked_prefill_helper(self, *, ctx_pp: int): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, @@ -1598,15 +1578,12 @@ def test_nvfp4(self): } disaggregated_server_config = { "hostname": "localhost", - "port": 8000, "backend": "pytorch", "context_servers": { - "num_instances": 1, - "urls": ["localhost:8001"] + "num_instances": 1 }, "generation_servers": { - "num_instances": 1, - "urls": ["localhost:8002"] + "num_instances": 1 } } with launch_disaggregated_llm(disaggregated_server_config, diff --git a/tests/integration/defs/disaggregated/disagg_test_utils.py b/tests/integration/defs/disaggregated/disagg_test_utils.py new file mode 100644 index 000000000000..d9691a0b1127 --- /dev/null +++ b/tests/integration/defs/disaggregated/disagg_test_utils.py @@ -0,0 +1,455 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2024 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. + +"""Shared utilities for disaggregated tests.""" + +import asyncio +import os +import shutil +import subprocess +import sys +import tempfile +import time +import traceback +import uuid +from functools import wraps + +import openai +import pytest +import requests +import yaml +from defs.common import get_free_port_in_ci as get_free_port + +from tensorrt_llm.logger import logger + +# Service discovery constants +HEARTBEAT_INTERVAL = 1 +INACTIVE_TIMEOUT = 2 +# Check cluster status with a larger interval than inactive timeout to avoid flaky tests +CHECK_STATUS_INTERVAL = 3 + + +class ProcessWrapper: + """Wrapper for subprocess with log file and port information.""" + + def __init__(self, process, log_file=None, log_path=None, port=0): + self.process = process + self.log_file = log_file + self.log_path = log_path + self.port = port + + +def periodic_check(timeout=300, interval=3): + """Decorator for periodic checking with timeout. + + Retries the decorated async function until it returns True or timeout is reached. + Sleeps for interval seconds between retries. + + Args: + timeout: Maximum time to wait in seconds + interval: Time to sleep between checks in seconds + + Raises: + TimeoutError: If timeout is reached without success + """ + + def decorator(func): + @wraps(func) + async def wrapper(*args, **kwargs): + start_time = time.time() + while time.time() - start_time < timeout: + try: + result = await func(*args, **kwargs) + if result: + return result + except Exception as e: + logger.debug(f"Check failed: {e}") + await asyncio.sleep(interval) + raise TimeoutError(f"Timeout after {timeout}s waiting for {func.__name__}") + + return wrapper + + return decorator + + +def _run_worker( + model_name, worker_config, role, port, work_dir, device=-1, save_log=False, env=None +): + """Run a worker process (context or generation). + + Args: + model_name: Path to the model + worker_config: Worker configuration dict + role: Role name (ctx/gen) + port: Port number + work_dir: Working directory for config files + device: CUDA device ID (-1 for default) + save_log: Whether to save logs to file + env: Environment variables for the subprocess + + Returns: + ProcessWrapper: Wrapped subprocess + """ + worker_config_path = os.path.join(work_dir, f"{role}_{port}_config.yaml") + with open(worker_config_path, "w+") as f: + yaml.dump(worker_config, f) + f.flush() + cmd = [ + "trtllm-serve", + "serve", + model_name, + "--host", + "localhost", + "--port", + str(port), + "--config", + worker_config_path, + "--server_role", + "context" if role.startswith("ctx") else "generation", + ] + if env is None: + env = os.environ.copy() + else: + env = env.copy() + log_file = None + log_path = None + if save_log: + log_path = os.path.join(work_dir, f"worker_{role}_{port}.log") + log_file = open(log_path, "w+") + stdout = log_file + stderr = log_file + else: + stdout = sys.stdout + stderr = sys.stderr + if device != -1: + env["CUDA_VISIBLE_DEVICES"] = str(device) + print(f"Running {role} on port {port}") + return ProcessWrapper( + subprocess.Popen(cmd, env=env, stdout=stdout, stderr=stderr), + log_file=log_file, + log_path=log_path, + port=port, + ) + + +def run_ctx_worker(model_name, ctx_worker_config, work_dir, port=0, device=0, env=None): + """Launch a context worker with service discovery. + + Use port=0 to let the worker choose a free port. + """ + return _run_worker(model_name, ctx_worker_config, "ctx", port, work_dir, device, env=env) + + +def run_gen_worker(model_name, gen_worker_config, work_dir, port=0, device=1, env=None): + """Launch a generation worker with service discovery. + + Use port=0 to let the worker choose a free port. + """ + return _run_worker(model_name, gen_worker_config, "gen", port, work_dir, device, env=env) + + +def run_disagg_server(disagg_cluster_config, work_dir, port=0, save_log=False, env=None, cwd=None): + """Launch the disaggregated server. + + Args: + disagg_cluster_config: Server configuration dict + work_dir: Working directory for config files + port: Port number + save_log: Whether to save logs to file + env: Environment variables for the subprocess + + Returns: + ProcessWrapper: Wrapped subprocess + """ + disagg_server_config_path = os.path.join(work_dir, "disagg_server_config.yaml") + disagg_cluster_config["port"] = port + with open(disagg_server_config_path, "w+") as f: + yaml.dump(disagg_cluster_config, f) + cmds = ["trtllm-serve", "disaggregated", "-c", disagg_server_config_path] + log_file = None + log_path = None + if save_log: + log_path = os.path.join(work_dir, "disagg_server.log") + log_file = open(log_path, "w+") + stdout = log_file + stderr = log_file + else: + stdout = sys.stdout + stderr = sys.stderr + p = subprocess.Popen(cmds, env=env, stdout=stdout, stderr=stderr, cwd=cwd) + return ProcessWrapper(p, log_file=log_file, log_path=log_path, port=port) + + +async def _wait_for_disagg_server_status(port, ready, min_ctx_workers=-1, min_gen_workers=-1): + """Check disagg server status via /cluster_info endpoint. + + Args: + port: Server port + ready: Whether to check is_ready flag + min_ctx_workers: Minimum context workers (-1 to skip check) + min_gen_workers: Minimum generation workers (-1 to skip check) + + Returns: + bool: True if all conditions are met + """ + try: + info_resp = requests.get(f"http://localhost:{port}/cluster_info", timeout=5) + if info_resp.status_code != 200: + return False + info = info_resp.json() + + if ready and not info.get("is_ready", False): + return False + + if min_ctx_workers != -1: + ctx_count = len(info.get("current_workers", {}).get("context_servers", [])) + if ctx_count < min_ctx_workers: + return False + + if min_gen_workers != -1: + gen_count = len(info.get("current_workers", {}).get("generation_servers", [])) + if gen_count < min_gen_workers: + return False + + return True + except Exception as e: + logger.debug(f"Failed to check server status: {e}") + return False + + +async def wait_for_disagg_server_ready(port, timeout=300): + """Wait for disagg server to be ready.""" + + @periodic_check(timeout=timeout, interval=3) + async def _check(): + return await _wait_for_disagg_server_status(port, True) + + return await _check() + + +@periodic_check(timeout=300, interval=3) +async def wait_for_disagg_server_status(port, min_ctx_workers=-1, min_gen_workers=-1): + """Wait for disagg server to have minimum number of workers.""" + return await _wait_for_disagg_server_status(port, False, min_ctx_workers, min_gen_workers) + + +@periodic_check(timeout=300, interval=3) +async def wait_for_worker_ready(port): + """Wait for worker to be ready via /health endpoint.""" + logger.info(f"Waiting for worker {port} to be ready") + try: + info_resp = requests.get(f"http://localhost:{port}/health", timeout=5) + return info_resp.status_code == 200 + except Exception: + return False + + +@periodic_check(timeout=300, interval=3) +async def wait_for_port_released(port): + """Wait for port to be released after killing a process. + + When we kill a server, the port is not released immediately. + If the port is not released, bind will fail with OSError: [Errno 98] Address already in use. + """ + import socket + + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("localhost", port)) + print(f"Port {port} is released") + return True + except OSError: + return False + + +def verify_cluster_info(ready, ctx_workers=-1, gen_workers=-1, port=0, expected_code=200): + """Verify cluster info from /cluster_info endpoint. + + Args: + ready: Expected is_ready status + ctx_workers: Expected number of context workers (-1 to skip check) + gen_workers: Expected number of generation workers (-1 to skip check) + port: Server port + expected_code: Expected HTTP status code + """ + assert port > 0, "port must be positive" + info_resp = requests.get(f"http://localhost:{port}/cluster_info") + assert info_resp.status_code == expected_code + info = info_resp.json() + logger.info(f"verify_cluster_info: {info}, ready={ready}, ctx={ctx_workers}, gen={gen_workers}") + assert info["is_ready"] == ready + if ctx_workers != -1: + assert len(info["current_workers"]["context_servers"]) == ctx_workers + if gen_workers != -1: + assert len(info["current_workers"]["generation_servers"]) == gen_workers + + +def tail(file_path, n): + """Read last n lines from a file. + + Args: + file_path: Path to file + n: Number of lines to read + + Returns: + str: Last n lines of the file + """ + try: + proc = subprocess.Popen(["tail", "-n", str(n), file_path], stdout=subprocess.PIPE) + return proc.stdout.read().decode("utf-8") + except Exception as e: + print(f"Failed to tail {file_path}: {e}") + print(f"Traceback: {traceback.format_exc()}") + return "" + + +def terminate(*args, show_log_lines=30): + """Terminate processes and show their logs. + + Args: + *args: ProcessWrapper instances to terminate + show_log_lines: Number of log lines to show for debugging + """ + for arg in args: + if arg and isinstance(arg, ProcessWrapper): + try: + # Print log tail for debugging + if arg.log_path and os.path.exists(arg.log_path): + print(f"-------------{arg.log_path}---------------") + try: + print(tail(arg.log_path, show_log_lines)) + except Exception as e: + print(f"Failed to read log: {e}") + except Exception as e: + print(f"Failed to tail {arg.log_path}: {e}") + + if arg.process: + print(f"Killing process {arg.process.pid}") + try: + arg.process.kill() + arg.process.wait(timeout=10) + arg.process = None + if arg.log_file: + arg.log_file.close() + arg.log_file = None + except Exception as e: + print(f"Failed to terminate process {arg.process.pid}: {e}") + else: + print(f"Process is None on port {arg.port}") + + +def request_completion(model_name, prompt, port): + """Make a completion request to the disagg server. + + Args: + model_name: Model name for the request + prompt: Prompt text + port: Server port + + Returns: + Completion response from OpenAI client + """ + client = openai.OpenAI(api_key="tensorrt_llm", base_url=f"http://localhost:{port}/v1") + return client.completions.create( + model=model_name, prompt=prompt, max_tokens=10, temperature=0.0 + ) + + +# ============================================================================ +# Pytest Fixtures +# ============================================================================ + + +@pytest.fixture +def disagg_port(): + """Get a free port for disaggregated server.""" + return get_free_port() + + +@pytest.fixture +def work_dir(): + """Create a temporary working directory.""" + d = tempfile.mkdtemp() + yield d + shutil.rmtree(d, ignore_errors=True) + + +@pytest.fixture +def router(request): + """Parameterized router fixture.""" + return request.param + + +@pytest.fixture +def service_discovery(request, disagg_port, work_dir): + """Setup service discovery (etcd or http). + + Args: + request.param: "etcd" or "http" + + Yields: + tuple: (process or None, uri string) + """ + if request.param == "etcd": + data_dir = f"{work_dir}/disagg_test-etcd-{uuid.uuid4()}" + etcd = subprocess.Popen(["etcd", "--data-dir", data_dir]) + yield etcd, "etcd://localhost:2379" + try: + etcd.kill() + etcd.wait(timeout=10) + shutil.rmtree(data_dir) + except Exception: + print(f"Failed to kill etcd: {traceback.format_exc()}") + else: + yield None, f"http://localhost:{disagg_port}" + + +@pytest.fixture +def disagg_cluster_config(service_discovery): + """Create cluster config for workers and proxy server.""" + _, uri = service_discovery + return { + "cluster_uri": uri, + "cluster_name": "test_cluster", + "heartbeat_interval_sec": HEARTBEAT_INTERVAL, + "inactive_timeout_sec": INACTIVE_TIMEOUT, + } + + +@pytest.fixture +def disagg_server_config(disagg_cluster_config, router, disagg_port): + """Create disaggregated server configuration.""" + return { + "hostname": "localhost", + "port": disagg_port, + "disagg_cluster": disagg_cluster_config, + "context_servers": {"router": {"type": router}}, + "generation_servers": {"router": {"type": router}}, + } + + +@pytest.fixture +def worker_config(disagg_cluster_config): + """Create worker configuration.""" + return { + "disagg_cluster": disagg_cluster_config, + "disable_overlap_scheduler": True, + "cache_transceiver_config": {"backend": "DEFAULT"}, + "kv_cache_config": { + "free_gpu_memory_fraction": 0.2, + "enable_partial_reuse": False, + }, + "cuda_graph_config": {}, + } diff --git a/tests/integration/defs/disaggregated/test_auto_scaling.py b/tests/integration/defs/disaggregated/test_auto_scaling.py index 4821ee847638..3a1873d8be34 100644 --- a/tests/integration/defs/disaggregated/test_auto_scaling.py +++ b/tests/integration/defs/disaggregated/test_auto_scaling.py @@ -1,25 +1,17 @@ import asyncio import os -import shutil -import subprocess -import tempfile -import traceback -import uuid -from functools import wraps -import openai import pytest import requests -import yaml -from defs.common import get_free_port_in_ci as get_free_port from defs.conftest import llm_models_root +from disagg_test_utils import (CHECK_STATUS_INTERVAL, request_completion, + run_ctx_worker, run_disagg_server, + run_gen_worker, terminate, verify_cluster_info, + wait_for_disagg_server_ready, + wait_for_disagg_server_status, + wait_for_port_released) -from tensorrt_llm.logger import logger - -HEARTBEAT_INTERVAL = 1 -INACTIVE_TIMEOUT = 2 -# check cluster status with a larger interval than inactive timeout to avoid flaky tests -CHECK_STATUS_INTERVAL = 3 +pytest_plugins = ["disagg_test_utils"] ROUTER_TYPES = ["round_robin", "load_balancing", "kv_cache_aware"] @@ -32,312 +24,6 @@ def model_name(): return model_path -@pytest.fixture -def disagg_port(): - return get_free_port() - - -@pytest.fixture -def work_dir(): - return tempfile.mkdtemp() - - -@pytest.fixture -def service_discovery(request, disagg_port, work_dir): - if request.param == "etcd": - data_dir = f"{work_dir}/disagg_test-etcd-{uuid.uuid4()}" - etcd = subprocess.Popen(["etcd", "--data-dir", data_dir]) - yield etcd, f"etcd://localhost:2379" - try: - etcd.kill() - etcd.wait(timeout=10) - shutil.rmtree(data_dir) - except Exception: - print(f"Failed to kill etcd: {traceback.format_exc()}") - else: - yield None, f"http://localhost:{disagg_port}" - - -@pytest.fixture -def disagg_cluster_config(service_discovery): - # same cluster config for workers and proxy server - _, uri = service_discovery - return { - "cluster_uri": uri, - "cluster_name": "test_cluster", - "heartbeat_interval_sec": HEARTBEAT_INTERVAL, - "inactive_timeout_sec": INACTIVE_TIMEOUT, - } - - -@pytest.fixture -def router(request): - return request.param - - -@pytest.fixture -def disagg_server_config(disagg_cluster_config, router, disagg_port): - return { - "hostname": "localhost", - "port": disagg_port, - "disagg_cluster": disagg_cluster_config, - "context_servers": { - "router": { - "type": router - } - }, - "generation_servers": { - "router": { - "type": router - } - }, - } - - -@pytest.fixture -def worker_config(disagg_cluster_config): - return { - "disagg_cluster": disagg_cluster_config, - "disable_overlap_scheduler": True, - "cache_transceiver_config": { - "backend": "DEFAULT" - }, - "kv_cache_config": { - "free_gpu_memory_fraction": 0.2, - "enable_partial_reuse": False, - }, - "cuda_graph_config": {}, - } - - -class ProcessWrapper: - - def __init__(self, process, log_file=None, log_path=None, port=0): - self.process = process - self.log_file = log_file - self.log_path = log_path - self.port = port - - -def _run_worker(model_name, - worker_config, - role, - port, - work_dir, - device=-1, - save_log=False): - worker_config_path = os.path.join(work_dir, f"{role}_{port}_config.yaml") - with open(worker_config_path, "w+") as f: - yaml.dump(worker_config, f) - f.flush() - cmd = [ - "trtllm-serve", - "serve", - model_name, - "--host", - "localhost", - "--port", - str(port), - "--config", - worker_config_path, - "--server_role", - "context" if role.startswith("ctx") else "generation", - ] - env = os.environ.copy() - log_file = None - log_path = None - stdout = None - stderr = None - if save_log: - log_path = os.path.join(work_dir, f"worker_{role}_{port}.log") - log_file = open(log_path, "w+") - stdout = log_file - stderr = log_file - - if device != -1: - env["CUDA_VISIBLE_DEVICES"] = str(device) - print(f"Running {role} on port {port}") - return ProcessWrapper(subprocess.Popen(cmd, - env=env, - stdout=stdout, - stderr=stderr), - log_file=log_file, - log_path=log_path, - port=port) - - -# Use 0 as the port and provide disagg_cluster_config to let the worker choose a free port -def run_ctx_worker(model_name, ctx_worker_config, work_dir, port=0, device=0): - return _run_worker(model_name, ctx_worker_config, "ctx", port, work_dir, - device) - - -def run_gen_worker(model_name, gen_worker_config, work_dir, port=0, device=1): - return _run_worker(model_name, gen_worker_config, "gen", port, work_dir, - device) - - -def run_disagg_server(disagg_cluster_config, work_dir, port=0, save_log=False): - disagg_server_config_path = os.path.join(work_dir, - "disagg_server_config.yaml") - disagg_cluster_config["port"] = port - with open(disagg_server_config_path, "w+") as f: - yaml.dump(disagg_cluster_config, f) - cmds = ["trtllm-serve", "disaggregated", "-c", disagg_server_config_path] - log_file = None - log_path = None - stdout = None - stderr = None - if save_log: - log_path = os.path.join(work_dir, "disagg_server.log") - log_file = open(log_path, "w+") - stdout = log_file - stderr = log_file - p = subprocess.Popen(cmds, stdout=stdout, stderr=stderr) - return ProcessWrapper(p, log_file=log_file, log_path=log_path, port=port) - - -# wait until decorated function returns true, otherwise sleep for interval seconds and try again -# if timeout seconds is reached, then raise TimeoutError -def periodic_check(timeout=300, interval=3): - - def decorator(func): - - @wraps(func) - async def wrapper(*args, **kwargs): - elapsed_time = 0 - while elapsed_time < timeout: - elapsed_time += interval - await asyncio.sleep(interval) - try: - if ret := await func(*args, **kwargs): - return ret - except Exception as e: - print( - f"Failed to check {func.__name__} after {elapsed_time} seconds: {e}" - ) - raise TimeoutError( - f"Timeout waiting for {func.__name__} to complete after {timeout} seconds" - ) - - return wrapper - - return decorator - - -async def _wait_for_disagg_server_status(port, - ready=True, - min_ctx_workers=-1, - min_gen_workers=-1): - info_resp = requests.get(f"http://localhost:{port}/cluster_info") - logger.info( - f"Waiting for disagg server {port} to be ready: {info_resp.json()}") - if info_resp.status_code == 200: - info = info_resp.json() - if ready: - return info["is_ready"] - else: - return len(info["current_workers"] - ["context_servers"]) >= min_ctx_workers and len( - info["current_workers"] - ["generation_servers"]) >= min_gen_workers - return False - - -@periodic_check(timeout=300, interval=3) -async def wait_for_disagg_server_ready(port): - return await _wait_for_disagg_server_status(port, True) - - -@periodic_check(timeout=300, interval=3) -async def wait_for_disagg_server_status(port, - min_ctx_workers=-1, - min_gen_workers=-1): - return await _wait_for_disagg_server_status(port, False, min_ctx_workers, - min_gen_workers) - - -@periodic_check(timeout=300, interval=3) -async def wait_for_worker_ready(port): - logger.info(f"Waiting for worker {port} to be ready") - info_resp = requests.get(f"http://localhost:{port}/health") - return info_resp.status_code == 200 - - -def verify_cluster_info(ready, - ctx_workers=-1, - gen_workers=-1, - port=0, - expected_code=200): - assert port > 0, "port must be positive" - info_resp = requests.get(f"http://localhost:{port}/cluster_info") - assert info_resp.status_code == expected_code - info = info_resp.json() - print("verify_cluster_info", info, ready, ctx_workers, gen_workers) - assert info["is_ready"] == ready - if ctx_workers != -1: - assert len(info["current_workers"]["context_servers"]) == ctx_workers - if gen_workers != -1: - assert len(info["current_workers"]["generation_servers"]) == gen_workers - - -def tail(f, n): - try: - proc = subprocess.Popen(['tail', '-n', str(n), f], - stdout=subprocess.PIPE) - return proc.stdout.read().decode('utf-8') - except Exception as e: - print(f"Failed to tail {f}: {e}") - print(f"Traceback: {traceback.format_exc()}") - return "" - - -def terminate(*args, show_log_lines=30, release_port=True): - for arg in args: - if arg and isinstance(arg, ProcessWrapper): - try: - # tail the log file for better debugging on CI - if arg.log_path and os.path.exists(arg.log_path): - print(f"-------------{arg.log_path}---------------") - print(tail(arg.log_path, show_log_lines)) - except Exception as e: - print(f"Failed to tail {arg.log_path}: {e}") - print(f"Traceback: {traceback.format_exc()}") - if arg.process: - print(f"Killing process {arg.process.pid}") - try: - arg.process.kill() - arg.process.wait(timeout=10) - arg.process = None - if arg.log_file: - arg.log_file.close() - arg.log_file = None - except Exception: - print(f"Failed to terminate process {arg.process.pid}") - else: - print(f"Process is None on port {arg.port}") - - -# When we kill a server, the port is not released immediately -# If the port is not released, the bind will fail with OSError: [Errno 98] Address already in use -@periodic_check(timeout=300, interval=3) -async def wait_for_port_released(port): - import socket - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("localhost", port)) - print(f"Port {port} is released") - return True - - -def request_completion(model_name, prompt, port): - client = openai.OpenAI(api_key="tensorrt_llm", - base_url=f"http://localhost:{port}/v1") - return client.completions.create(model=model_name, - prompt=prompt, - max_tokens=10, - temperature=0.0) - - @pytest.mark.skip_less_device(2) @pytest.mark.parametrize("router", ROUTER_TYPES, indirect=True) @pytest.mark.asyncio(loop_scope="module") @@ -448,7 +134,7 @@ async def test_worker_restart(model_name, disagg_server_config, worker_config, port=disagg_port) print(response) # kill gen1, the request should fail - terminate(gen_worker1, release_port=True) + terminate(gen_worker1) await asyncio.sleep(CHECK_STATUS_INTERVAL) verify_cluster_info(False, 1, 0, port=disagg_port) with pytest.raises(Exception): @@ -463,7 +149,7 @@ async def test_worker_restart(model_name, disagg_server_config, worker_config, worker_config, work_dir, port=0, - device=0) + device=1) await wait_for_disagg_server_status(disagg_port, 1, 1) await asyncio.sleep(CHECK_STATUS_INTERVAL) verify_cluster_info(True, 1, 1, port=disagg_port) @@ -474,7 +160,7 @@ async def test_worker_restart(model_name, disagg_server_config, worker_config, assert len(response.choices[0].text) >= 1 # kill ctx1, the request should fail - terminate(ctx_worker1, release_port=True) + terminate(ctx_worker1) await asyncio.sleep(CHECK_STATUS_INTERVAL) verify_cluster_info(False, 0, 1, port=disagg_port) with pytest.raises(Exception): @@ -485,7 +171,7 @@ async def test_worker_restart(model_name, disagg_server_config, worker_config, worker_config, work_dir, port=0, - device=1) + device=0) await wait_for_disagg_server_status(disagg_port, 1, 1) await asyncio.sleep(CHECK_STATUS_INTERVAL) verify_cluster_info(True, 1, 1, port=disagg_port) diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config.yaml new file mode 100644 index 000000000000..a29c2a5303f8 --- /dev/null +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config.yaml @@ -0,0 +1,19 @@ +hostname: localhost +model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 +free_gpu_memory_fraction: 0.25 +backend: pytorch +disable_overlap_scheduler: true +context_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + kv_cache_config: + free_gpu_memory_fraction: 0.2 + cache_transceiver_config: + backend: DEFAULT +generation_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: DEFAULT diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_aware_balance.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_aware_balance.yaml index d64bac8763b7..a9bf2587d23e 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_aware_balance.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_aware_balance.yaml @@ -1,11 +1,10 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.1 -disable_overlap_scheduler: True -enable_autotuner: False +disable_overlap_scheduler: true +enable_autotuner: false context_servers: num_instances: 2 router: @@ -16,15 +15,12 @@ context_servers: tensor_parallel_size: 1 pipeline_parallel_size: 1 kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: False + enable_block_reuse: true + enable_partial_reuse: false event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" - - "localhost:8002" generation_servers: num_instances: 2 router: @@ -37,10 +33,7 @@ generation_servers: cache_transceiver_config: backend: DEFAULT kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: False + enable_block_reuse: true + enable_partial_reuse: false event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.1 - urls: - - "localhost:8003" - - "localhost:8004" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_aware_balance_deepseek_v3.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_aware_balance_deepseek_v3.yaml index fe15f70085c7..615bf8b74d41 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_aware_balance_deepseek_v3.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_aware_balance_deepseek_v3.yaml @@ -1,10 +1,9 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/bf16 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True -enable_autotuner: False +disable_overlap_scheduler: true +enable_autotuner: false context_servers: num_instances: 2 router: @@ -12,15 +11,12 @@ context_servers: tensor_parallel_size: 1 pipeline_parallel_size: 1 kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.1 cache_transceiver_config: - backend: "DEFAULT" - urls: - - "localhost:8001" - - "localhost:8002" + backend: DEFAULT generation_servers: num_instances: 2 router: @@ -28,12 +24,9 @@ generation_servers: tensor_parallel_size: 1 pipeline_parallel_size: 1 kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.1 cache_transceiver_config: - backend: "DEFAULT" - urls: - - "localhost:8003" - - "localhost:8004" + backend: DEFAULT diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_reuse.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_reuse.yaml index 26444b1ab237..e7b371a6479e 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_reuse.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_reuse.yaml @@ -1,23 +1,20 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True -enable_autotuner: False +disable_overlap_scheduler: true +enable_autotuner: false context_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.15 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -25,11 +22,9 @@ generation_servers: router: type: kv_cache_aware kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.05 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_reuse_deepseek_v3.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_reuse_deepseek_v3.yaml index 06a4c154b46b..36b459dabc1e 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_reuse_deepseek_v3.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_cache_reuse_deepseek_v3.yaml @@ -1,23 +1,20 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/bf16 free_gpu_memory_fraction: 0.15 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True -enable_autotuner: False +disable_overlap_scheduler: true +enable_autotuner: false context_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -25,11 +22,9 @@ generation_servers: router: type: kv_cache_aware kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.05 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_cancel_stress_test.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_cancel_stress_test.yaml index 2795ca46bd30..d888c2a878ca 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_cancel_stress_test.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_cancel_stress_test.yaml @@ -1,44 +1,39 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/bf16 -backend: "pytorch" -enable_autotuner: False +backend: pytorch +enable_autotuner: false context_servers: - disable_overlap_scheduler: True + disable_overlap_scheduler: true num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 max_num_tokens: 16384 max_seq_len: 32768 - enable_chunked_prefill: True + enable_chunked_prefill: true kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true free_gpu_memory_fraction: 0.3 cache_transceiver_config: - backend: "DEFAULT" + backend: DEFAULT max_tokens_in_buffer: 32768 cuda_graph_config: - enable_padding: True + enable_padding: true max_batch_size: 1 - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 max_num_tokens: 2048 max_seq_len: 32768 - enable_chunked_prefill: True + enable_chunked_prefill: true kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true free_gpu_memory_fraction: 0.85 cache_transceiver_config: - backend: "DEFAULT" + backend: DEFAULT max_tokens_in_buffer: 32768 cuda_graph_config: - enable_padding: True + enable_padding: true max_batch_size: 64 - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_cancel_stress_test_large.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_cancel_stress_test_large.yaml index 5a538d1f7145..b4c6fabd8c57 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_cancel_stress_test_large.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_cancel_stress_test_large.yaml @@ -1,44 +1,39 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-0324-FP4 -backend: "pytorch" -enable_autotuner: False +backend: pytorch +enable_autotuner: false context_servers: - disable_overlap_scheduler: True + disable_overlap_scheduler: true num_instances: 1 tensor_parallel_size: 4 pipeline_parallel_size: 1 max_num_tokens: 12000 max_seq_len: 262144 - enable_chunked_prefill: True + enable_chunked_prefill: true kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true free_gpu_memory_fraction: 0.2 cache_transceiver_config: - backend: "DEFAULT" + backend: DEFAULT max_tokens_in_buffer: 262144 cuda_graph_config: - enable_padding: True + enable_padding: true max_batch_size: 1 - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 4 pipeline_parallel_size: 1 max_num_tokens: 2048 max_seq_len: 262144 - enable_chunked_prefill: True + enable_chunked_prefill: true kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true free_gpu_memory_fraction: 0.3 cache_transceiver_config: - backend: "DEFAULT" + backend: DEFAULT max_tokens_in_buffer: 262144 cuda_graph_config: - enable_padding: True + enable_padding: true max_batch_size: 11 - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional.yaml index 28816380fe46..26aaeac42d90 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional.yaml @@ -1,26 +1,23 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.15 conditional_disagg_config: max_local_prefill_length: 100 -disable_overlap_scheduler: True -enable_autotuner: False +disable_overlap_scheduler: true +enable_autotuner: false context_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.15 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -28,11 +25,9 @@ generation_servers: router: type: kv_cache_aware kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.15 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3.yaml index b7f342027240..7887fd2725fb 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3.yaml @@ -1,26 +1,23 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/bf16 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.15 conditional_disagg_config: max_local_prefill_length: 100 -disable_overlap_scheduler: True -enable_autotuner: False +disable_overlap_scheduler: true +enable_autotuner: false context_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.15 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -28,11 +25,9 @@ generation_servers: router: type: kv_cache_aware kv_cache_config: - enable_block_reuse: True - enable_partial_reuse: True + enable_block_reuse: true + enable_partial_reuse: true event_buffer_max_size: 1024 free_gpu_memory_fraction: 0.15 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp2_genpp2.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp2_genpp2.yaml index b7f03c0f9f5c..c04b34238c6b 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp2_genpp2.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp2_genpp2.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: @@ -13,13 +12,11 @@ context_servers: pipeline_parallel_size: 2 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - enable_block_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + enable_block_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -29,10 +26,8 @@ generation_servers: max_seq_len: 4096 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - enable_block_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + enable_block_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp2_gentp2.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp2_gentp2.yaml index 892b4e8b31f0..76e44e23a12d 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp2_gentp2.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp2_gentp2.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: @@ -13,13 +12,11 @@ context_servers: pipeline_parallel_size: 2 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - enable_block_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + enable_block_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 @@ -29,9 +26,7 @@ generation_servers: max_seq_len: 4096 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp4_genpp4.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp4_genpp4.yaml index 2c7a67e1cbfe..ffee6430abcc 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp4_genpp4.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp4_genpp4.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: @@ -13,13 +12,11 @@ context_servers: pipeline_parallel_size: 4 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - enable_block_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + enable_block_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -29,10 +26,8 @@ generation_servers: max_seq_len: 4096 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - enable_block_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + enable_block_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp4_gentp4.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp4_gentp4.yaml index a1e4ad50a9c2..c176aa863b61 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp4_gentp4.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxpp4_gentp4.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: @@ -13,12 +12,10 @@ context_servers: pipeline_parallel_size: 4 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 4 @@ -28,9 +25,7 @@ generation_servers: max_seq_len: 4096 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite.yaml index 83f9b3a3e877..ce4c9b3917bf 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite.yaml @@ -1,23 +1,18 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.1 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp.yaml index 57eb4ea00410..a7ecc70fedd1 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp.yaml @@ -1,10 +1,9 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.1 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True +disable_overlap_scheduler: true speculative_config: decoding_type: MTP num_nextn_predict_layers: 1 @@ -15,8 +14,6 @@ context_servers: enable_attention_dp: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -24,5 +21,3 @@ generation_servers: enable_attention_dp: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_attention_dp_overlap.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_attention_dp_overlap.yaml index 4343850c77f3..1c3ed4091a24 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_attention_dp_overlap.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_attention_dp_overlap.yaml @@ -1,8 +1,7 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.1 -backend: "pytorch" +backend: pytorch cuda_graph_config: null speculative_config: decoding_type: MTP @@ -12,18 +11,14 @@ context_servers: tensor_parallel_size: 1 pipeline_parallel_size: 1 enable_attention_dp: true - disable_overlap_scheduler: True + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 enable_attention_dp: true - disable_overlap_scheduler: False + disable_overlap_scheduler: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_ctxpp2_gentp2.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_ctxpp2_gentp2.yaml index 4a61497e94e9..f75e014e858b 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_ctxpp2_gentp2.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_ctxpp2_gentp2.yaml @@ -1,11 +1,9 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.1 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True - +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 1 @@ -16,8 +14,6 @@ context_servers: num_nextn_predict_layers: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 @@ -28,5 +24,3 @@ generation_servers: num_nextn_predict_layers: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_two_mtp.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_two_mtp.yaml index 837e5df8e335..154d65be9e7a 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_two_mtp.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_two_mtp.yaml @@ -1,10 +1,9 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.1 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True +disable_overlap_scheduler: true speculative_config: decoding_type: MTP num_nextn_predict_layers: 2 @@ -15,14 +14,10 @@ context_servers: enable_attention_dp: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 enable_attention_dp: false - urls: - - "localhost:8002" cache_transceiver_config: backend: DEFAULT diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_genpp2.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_genpp2.yaml index ce53fd4626bd..8d6821cd996c 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_genpp2.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_genpp2.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: @@ -13,12 +12,10 @@ context_servers: pipeline_parallel_size: 1 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -28,9 +25,7 @@ generation_servers: max_seq_len: 4096 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1.yaml index 1335d63adfe8..840ba25e021d 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1.yaml @@ -1,24 +1,18 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 free_gpu_memory_fraction: 0.25 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 2 tensor_parallel_size: 1 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" - - "localhost:8003" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1_trt_backend.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1_trt_backend.yaml index fa5dffa518b8..da9ed69ee458 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1_trt_backend.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1_trt_backend.yaml @@ -1,22 +1,16 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 free_gpu_memory_fraction: 0.25 -backend: "trt" +backend: trt context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 2 tensor_parallel_size: 1 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" - - "localhost:8003" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1cp2_deepseek_v3_lite_bf16_tllm_gen.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1cp2_deepseek_v3_lite_bf16_tllm_gen.yaml index f7e879bb4c79..8ed8f5e7eb08 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1cp2_deepseek_v3_lite_bf16_tllm_gen.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp1cp2_deepseek_v3_lite_bf16_tllm_gen.yaml @@ -1,36 +1,31 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/bf16 free_gpu_memory_fraction: 0.25 -backend: "pytorch" +backend: pytorch cuda_graph_config: null context_servers: num_instances: 1 - enable_chunked_prefill: False + enable_chunked_prefill: false kv_cache_config: - enable_block_reuse: False - enable_partial_reuse: False + enable_block_reuse: false + enable_partial_reuse: false tokens_per_block: 32 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: - backend: "UCX" - urls: - - "localhost:8001" + backend: UCX generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 context_parallel_size: 2 - enable_chunked_prefill: False + enable_chunked_prefill: false cp_config: cp_type: HELIX tokens_per_block: 32 kv_cache_config: - enable_block_reuse: False - enable_partial_reuse: False + enable_block_reuse: false + enable_partial_reuse: false tokens_per_block: 32 cache_transceiver_config: - backend: "UCX" - urls: - - "localhost:8002" + backend: UCX diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite.yaml index 6b22665e9f17..0d50737cc267 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite.yaml @@ -1,23 +1,18 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.25 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp.yaml index 80a1a3636a80..bfec04d70572 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp.yaml @@ -1,25 +1,20 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.25 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 - enable_attention_dp: True + enable_attention_dp: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 - enable_attention_dp: True + enable_attention_dp: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one.yaml index 9dfb092151a7..4aa309cf6cf1 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one.yaml @@ -1,10 +1,9 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.25 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 2 @@ -12,8 +11,6 @@ context_servers: enable_attention_dp: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 @@ -21,5 +18,3 @@ generation_servers: enable_attention_dp: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one_mtp.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one_mtp.yaml index 4b6bc571dab4..7c44075406f7 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one_mtp.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one_mtp.yaml @@ -1,10 +1,9 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.25 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True +disable_overlap_scheduler: true speculative_config: decoding_type: MTP num_nextn_predict_layers: 1 @@ -15,8 +14,6 @@ context_servers: enable_attention_dp: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 @@ -24,6 +21,3 @@ generation_servers: enable_attention_dp: false cache_transceiver_config: backend: DEFAULT - - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap.yaml index 26218586f492..af8f62e920e1 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap.yaml @@ -1,26 +1,21 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 - enable_attention_dp: True - disable_overlap_scheduler: True + enable_attention_dp: true + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 - enable_attention_dp: True - disable_overlap_scheduler: False + enable_attention_dp: true + disable_overlap_scheduler: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap_cuda_graph.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap_cuda_graph.yaml index 99034f8a1a3e..acc41bc2dc89 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap_cuda_graph.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap_cuda_graph.yaml @@ -1,27 +1,22 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.25 -backend: "pytorch" +backend: pytorch context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 enable_attention_dp: true - disable_overlap_scheduler: True + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 enable_attention_dp: true cuda_graph_config: - enable_padding: False - disable_overlap_scheduler: False + enable_padding: false + disable_overlap_scheduler: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_mpi.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_mpi.yaml index 4cfe18ebaf67..bfc7372adfcb 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_mpi.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_mpi.yaml @@ -1,22 +1,17 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.25 -backend: "pytorch" -disable_overlap_scheduler: True +backend: pytorch +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: - backend: "MPI" - urls: - - "localhost:8001" + backend: MPI generation_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: - backend: "MPI" - urls: - - "localhost:8002" + backend: MPI diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_nixl.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_nixl.yaml index 3b1aa8fc0e34..e1628021af57 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_nixl.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_nixl.yaml @@ -1,22 +1,17 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.25 -backend: "pytorch" -disable_overlap_scheduler: True +backend: pytorch +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: - backend: "NIXL" - urls: - - "localhost:8001" + backend: NIXL generation_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: - backend: "NIXL" - urls: - - "localhost:8002" + backend: NIXL diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_overlap_cuda_graph.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_overlap_cuda_graph.yaml index 4c601fbb868c..b9d3f29b40b0 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_overlap_cuda_graph.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_overlap_cuda_graph.yaml @@ -1,25 +1,20 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.25 -backend: "pytorch" +backend: pytorch context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 - disable_overlap_scheduler: True + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cuda_graph_config: - enable_padding: False - disable_overlap_scheduler: False + enable_padding: false + disable_overlap_scheduler: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_ucx.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_ucx.yaml index d3395938cae6..6ab1ba3e6174 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_ucx.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_ucx.yaml @@ -1,22 +1,17 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/fp8 free_gpu_memory_fraction: 0.25 -backend: "pytorch" -disable_overlap_scheduler: True +backend: pytorch +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: - backend: "UCX" - urls: - - "localhost:8001" + backend: UCX generation_servers: num_instances: 1 tensor_parallel_size: 2 pipeline_parallel_size: 1 cache_transceiver_config: - backend: "UCX" - urls: - - "localhost:8002" + backend: UCX diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_tllm.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_tllm.yaml index 0dc7550f0e1d..dc90d3bf6d36 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_tllm.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_tllm.yaml @@ -1,8 +1,6 @@ model: gpt_oss/gpt-oss-120b hostname: localhost -port: 8100 backend: pytorch - context_servers: num_instances: 1 tensor_parallel_size: 2 @@ -16,7 +14,7 @@ context_servers: enable_chunked_prefill: true kv_cache_config: enable_block_reuse: false - free_gpu_memory_fraction: 0.80 + free_gpu_memory_fraction: 0.8 dtype: fp8 disable_overlap_scheduler: true moe_config: @@ -26,9 +24,6 @@ context_servers: cache_transceiver_config: backend: DEFAULT max_tokens_in_buffer: 16384 - urls: - - "localhost:8101" - generation_servers: num_instances: 1 tensor_parallel_size: 2 @@ -42,17 +37,27 @@ generation_servers: enable_chunked_prefill: true kv_cache_config: enable_block_reuse: false - free_gpu_memory_fraction: 0.80 + free_gpu_memory_fraction: 0.8 dtype: fp8 disable_overlap_scheduler: true moe_config: backend: TRTLLM cuda_graph_config: enable_padding: true - batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 768, 1024] + batch_sizes: + - 1 + - 2 + - 4 + - 8 + - 16 + - 32 + - 64 + - 128 + - 256 + - 512 + - 768 + - 1024 print_iter_log: true cache_transceiver_config: backend: DEFAULT max_tokens_in_buffer: 16384 - urls: - - "localhost:8102" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2pp2_gentp2pp2.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2pp2_gentp2pp2.yaml index ce47009aaadc..d80795b727ac 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2pp2_gentp2pp2.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2pp2_gentp2pp2.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: @@ -13,13 +12,11 @@ context_servers: pipeline_parallel_size: 2 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - enable_block_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + enable_block_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 @@ -29,10 +26,8 @@ generation_servers: max_seq_len: 4096 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - enable_block_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + enable_block_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp4_gentp4_deepseek_r1_v2_fp4_tllm.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp4_gentp4_deepseek_r1_v2_fp4_tllm.yaml index 1d1535d1ae32..58053cc0013c 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp4_gentp4_deepseek_r1_v2_fp4_tllm.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp4_gentp4_deepseek_r1_v2_fp4_tllm.yaml @@ -1,8 +1,6 @@ model: DeepSeek-R1/DeepSeek-R1-0528-FP4-v2 hostname: localhost -port: 8100 backend: pytorch - context_servers: num_instances: 1 tensor_parallel_size: 4 @@ -16,7 +14,7 @@ context_servers: enable_chunked_prefill: true kv_cache_config: enable_block_reuse: false - free_gpu_memory_fraction: 0.80 + free_gpu_memory_fraction: 0.8 dtype: fp8 moe_config: backend: TRTLLM @@ -25,9 +23,6 @@ context_servers: cache_transceiver_config: backend: DEFAULT max_tokens_in_buffer: 16384 - urls: - - "localhost:8101" - generation_servers: num_instances: 1 tensor_parallel_size: 4 @@ -41,16 +36,26 @@ generation_servers: enable_chunked_prefill: true kv_cache_config: enable_block_reuse: false - free_gpu_memory_fraction: 0.80 + free_gpu_memory_fraction: 0.8 dtype: fp8 moe_config: backend: TRTLLM cuda_graph_config: enable_padding: true - batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 768, 1024] + batch_sizes: + - 1 + - 2 + - 4 + - 8 + - 16 + - 32 + - 64 + - 128 + - 256 + - 512 + - 768 + - 1024 print_iter_log: true cache_transceiver_config: backend: DEFAULT max_tokens_in_buffer: 16384 - urls: - - "localhost:8102" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_cuda_graph_padding.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_cuda_graph_padding.yaml index 56db3df76976..1f9e42d73237 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_cuda_graph_padding.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_cuda_graph_padding.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch context_servers: num_instances: 1 max_batch_size: 1 @@ -11,14 +10,14 @@ context_servers: pipeline_parallel_size: 1 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: false cuda_graph_config: - batch_sizes: [1,3000] - disable_overlap_scheduler: True + batch_sizes: + - 1 + - 3000 + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -28,12 +27,16 @@ generation_servers: max_seq_len: 4096 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: false cuda_graph_config: - enable_padding: True - batch_sizes: [1,4,8,16,24,32] - disable_overlap_scheduler: True + enable_padding: true + batch_sizes: + - 1 + - 4 + - 8 + - 16 + - 24 + - 32 + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_deepseek_v3_lite_empty_batch.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_deepseek_v3_lite_empty_batch.yaml index 7fce3cfe6e64..abf05da97d9d 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_deepseek_v3_lite_empty_batch.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_deepseek_v3_lite_empty_batch.yaml @@ -1,7 +1,6 @@ hostname: localhost -port: 8000 model: DeepSeek-V3-Lite/bf16 -backend: "pytorch" +backend: pytorch context_servers: num_instances: 1 max_batch_size: 10 @@ -19,8 +18,6 @@ context_servers: cache_transceiver_config: max_tokens_in_buffer: 8448 backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -44,5 +41,3 @@ generation_servers: backend: DEFAULT stream_interval: 1 num_postprocess_workers: 1 - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_diff_max_tokens.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_diff_max_tokens.yaml index 26d1f6b6c154..c07260248822 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_diff_max_tokens.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_diff_max_tokens.yaml @@ -1,23 +1,18 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 free_gpu_memory_fraction: 0.25 -backend: "pytorch" +backend: pytorch cuda_graph_config: null -disable_overlap_scheduler: True +disable_overlap_scheduler: true context_servers: num_instances: 1 max_num_tokens: 512 max_batch_size: 64 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 max_num_tokens: 256 max_batch_size: 32 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only.yaml index 92b138376440..9253f421cfcd 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only.yaml @@ -1,7 +1,6 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 -backend: "pytorch" +backend: pytorch cuda_graph_config: null context_servers: num_instances: 0 @@ -11,11 +10,8 @@ generation_servers: pipeline_parallel_size: 1 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_block_reuse: False - enable_partial_reuse: False + enable_block_reuse: false + enable_partial_reuse: false cache_transceiver_config: backend: DEFAULT - print_iter_log: True - urls: - - "localhost:8002" - - "localhost:8003" + print_iter_log: true diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_bs1.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_bs1.yaml index 19d1eca714fd..67494b24ff0b 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_bs1.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_bs1.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: @@ -14,12 +13,10 @@ context_servers: enable_attention_dp: true kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 2 @@ -30,8 +27,6 @@ generation_servers: max_seq_len: 4096 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_trt_backend.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_trt_backend.yaml index ad706f8bf1f4..24b888b665be 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_trt_backend.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_gen_only_trt_backend.yaml @@ -1,7 +1,6 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 -backend: "trt" +backend: trt context_servers: num_instances: 0 generation_servers: @@ -10,10 +9,7 @@ generation_servers: pipeline_parallel_size: 1 kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_block_reuse: False - enable_partial_reuse: False + enable_block_reuse: false + enable_partial_reuse: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" - - "localhost:8003" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_llama4_kv_cache_overflow.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_llama4_kv_cache_overflow.yaml index 3295116acab7..fa65a710981c 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_llama4_kv_cache_overflow.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_llama4_kv_cache_overflow.yaml @@ -1,8 +1,6 @@ model: llama4-models/nvidia/Llama-4-Maverick-17B-128E-Instruct-FP8 hostname: localhost -port: 8000 backend: pytorch - context_servers: num_instances: 1 tensor_parallel_size: 4 @@ -24,9 +22,6 @@ context_servers: backend: UCX # Intentionally small to reproduce buffer overflow bug max_tokens_in_buffer: 2048 - urls: - - "localhost:8001" - generation_servers: num_instances: 1 tensor_parallel_size: 4 @@ -48,5 +43,3 @@ generation_servers: backend: UCX # Intentionally small to reproduce buffer overflow bug max_tokens_in_buffer: 2048 - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_load_balance.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_load_balance.yaml index f0593d9ef606..8540c6f555f6 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_load_balance.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_load_balance.yaml @@ -1,14 +1,13 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.15 context_servers: num_instances: 2 router: type: load_balancing - use_tokens: True + use_tokens: true max_batch_size: 1 max_num_tokens: 3000 max_seq_len: 4096 @@ -16,18 +15,15 @@ context_servers: pipeline_parallel_size: 1 kv_cache_config: free_gpu_memory_fraction: 0.15 - enable_partial_reuse: False - disable_overlap_scheduler: True + enable_partial_reuse: false + disable_overlap_scheduler: true cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" - - "localhost:8002" generation_servers: num_instances: 2 router: type: load_balancing - use_tokens: False + use_tokens: false max_batch_size: 256 max_num_tokens: 4096 max_seq_len: 4096 @@ -35,10 +31,7 @@ generation_servers: pipeline_parallel_size: 1 kv_cache_config: free_gpu_memory_fraction: 0.15 - enable_partial_reuse: False - disable_overlap_scheduler: False + enable_partial_reuse: false + disable_overlap_scheduler: false cache_transceiver_config: - backend: "DEFAULT" - urls: - - "localhost:8003" - - "localhost:8004" + backend: DEFAULT diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_metrics.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_metrics.yaml index 6d566aa4f99b..48fc0a072f02 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_metrics.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_metrics.yaml @@ -1,5 +1,4 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 free_gpu_memory_fraction: 0.25 backend: "pytorch" @@ -14,8 +13,6 @@ context_servers: perf_metrics_max_requests: 1000 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -24,5 +21,3 @@ generation_servers: perf_metrics_max_requests: 1000 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_mixed.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_mixed.yaml index dcc40a6a8b38..cf7478ce8588 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_mixed.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_mixed.yaml @@ -1,5 +1,4 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 free_gpu_memory_fraction: 0.25 backend: "pytorch" @@ -11,15 +10,9 @@ context_servers: pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" - - "localhost:8002" generation_servers: num_instances: 2 tensor_parallel_size: 1 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ngram.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ngram.yaml index 4e3417c732a0..4d0e7f804368 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ngram.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ngram.yaml @@ -1,29 +1,24 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 free_gpu_memory_fraction: 0.1 backend: pytorch -disable_overlap_scheduler: True +disable_overlap_scheduler: true context_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 cache_transceiver_config: - backend: "DEFAULT" - urls: - - "localhost:8001" + backend: DEFAULT generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 cache_transceiver_config: - backend: "DEFAULT" - urls: - - "localhost:8002" + backend: DEFAULT speculative_config: decoding_type: NGram max_draft_len: 4 max_matching_ngram_size: 4 - is_keep_all: True - is_use_oldest: True - is_public_pool: True + is_keep_all: true + is_use_oldest: true + is_public_pool: true diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap.yaml index d51ffabf8a24..3a872fbbc95c 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: @@ -12,13 +11,11 @@ context_servers: tensor_parallel_size: 1 pipeline_parallel_size: 1 kv_cache_config: - enable_block_reuse: False + enable_block_reuse: false free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -27,10 +24,8 @@ generation_servers: max_num_tokens: 4096 max_seq_len: 4096 kv_cache_config: - enable_block_reuse: False + enable_block_reuse: false free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: false cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_trt_backend.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_trt_backend.yaml index 3eb275c87e04..38aba69fd822 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_trt_backend.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_trt_backend.yaml @@ -1,8 +1,7 @@ hostname: localhost -port: 8000 model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 free_gpu_memory_fraction: 0.25 -backend: "trt" +backend: trt context_servers: num_instances: 1 tensor_parallel_size: 1 @@ -11,13 +10,9 @@ context_servers: free_gpu_memory_fraction: 0.2 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8001" generation_servers: num_instances: 1 tensor_parallel_size: 1 pipeline_parallel_size: 1 cache_transceiver_config: backend: DEFAULT - urls: - - "localhost:8002" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_trtllm_sampler.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_trtllm_sampler.yaml index 287d1103a4fd..f972a655c860 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_trtllm_sampler.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_trtllm_sampler.yaml @@ -1,7 +1,6 @@ model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 hostname: localhost -port: 8000 -backend: "pytorch" +backend: pytorch cuda_graph_config: null free_gpu_memory_fraction: 0.2 context_servers: @@ -11,15 +10,13 @@ context_servers: max_seq_len: 4096 tensor_parallel_size: 1 pipeline_parallel_size: 1 - sampler_type: "TRTLLMSampler" + sampler_type: TRTLLMSampler kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: false cache_transceiver_config: - backend: "DEFAULT" - disable_overlap_scheduler: True - urls: - - "localhost:8001" + backend: DEFAULT + disable_overlap_scheduler: true generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -27,12 +24,10 @@ generation_servers: max_batch_size: 256 max_num_tokens: 4096 max_seq_len: 4096 - sampler_type: "TRTLLMSampler" + sampler_type: TRTLLMSampler kv_cache_config: free_gpu_memory_fraction: 0.2 - enable_partial_reuse: False + enable_partial_reuse: false cache_transceiver_config: - backend: "DEFAULT" - disable_overlap_scheduler: False - urls: - - "localhost:8002" + backend: DEFAULT + disable_overlap_scheduler: false diff --git a/tests/integration/defs/disaggregated/test_configs/etcd_config.yaml b/tests/integration/defs/disaggregated/test_configs/etcd_config.yaml new file mode 100644 index 000000000000..c2ea9e0a41a9 --- /dev/null +++ b/tests/integration/defs/disaggregated/test_configs/etcd_config.yaml @@ -0,0 +1,4 @@ +server_type: "etcd" +hostname: "localhost" +port: 2379 +health_check_timeout: 5.0 diff --git a/tests/integration/defs/disaggregated/test_configs/gen_extra-llm-api-config.yml b/tests/integration/defs/disaggregated/test_configs/gen_extra-llm-api-config.yml new file mode 100644 index 000000000000..bede04a9d0cb --- /dev/null +++ b/tests/integration/defs/disaggregated/test_configs/gen_extra-llm-api-config.yml @@ -0,0 +1,3 @@ +cache_transceiver_config: + backend: "DEFAULT" + max_tokens_in_buffer: 2048 diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 32e505a5e00a..ae6175b335b5 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -13,31 +13,27 @@ # See the License for the specific language governing permissions and # limitations under the License. -import contextlib +import asyncio import os import re +import shutil import subprocess import tempfile import time from collections import namedtuple from dataclasses import dataclass -from typing import Callable +from typing import Any import pytest - -try: - import ray -except ImportError: - import tensorrt_llm.ray_stub as ray - import yaml -from defs.common import (get_free_port_in_ci, parse_gsm8k_output, - revise_disagg_config_file_with_free_ports, - wait_for_server) +from defs.common import get_free_port_in_ci as get_free_port +from defs.common import parse_gsm8k_output, wait_for_server from defs.conftest import (get_sm_version, llm_models_root, skip_arm, skip_no_hopper, skip_pre_blackwell) -from defs.trt_test_alternative import (check_call, check_output, popen, - print_info) +from defs.trt_test_alternative import check_call, check_output, print_info +from disagg_test_utils import (ProcessWrapper, run_ctx_worker, + run_disagg_server, run_gen_worker, terminate, + wait_for_disagg_server_ready) from test_common.perf_metrics_utils import (get_timing_metrics, validate_timing_metrics) @@ -66,202 +62,195 @@ def cleanup_output_files(): pass -def get_disagg_server_url_from_cfg(config_file: str) -> tuple[str, int]: - with open(config_file, 'r') as file: - config = yaml.safe_load(file) - server_host = config.get('hostname', 'localhost') - server_port = config.get('port', 8000) - return server_host, server_port +def get_default_disagg_cluster_config(): + """Get default disaggregated cluster configuration.""" + return { + "cluster_name": "test_cluster", + "heartbeat_interval_sec": 1, + "inactive_timeout_sec": 2 + } + + +def build_worker_config(base_config: dict[str, Any], + server_type_config: dict[str, Any], + disagg_cluster: dict[str, Any]) -> dict[str, Any]: + """ + Build worker configuration by merging base config with server-type specific config. + + Args: + base_config: Full YAML config (top-level) + server_type_config: context_servers or generation_servers section + disagg_cluster: Service discovery config (injected by test) + + Returns: + dict: Worker configuration for trtllm-serve + """ + # Fields to exclude from worker configs (not worker execution settings) + EXCLUDE_FROM_WORKER = { + 'hostname', + 'port', + 'num_instances', + 'urls', + 'router', + 'model', + 'context_servers', + 'generation_servers', + 'conditional_disagg_config', + } + + # Start with top-level fields (exclude server-only) + worker_config = { + k: v + for k, v in base_config.items() if k not in EXCLUDE_FROM_WORKER + } + + # Merge server-type specific config (overrides top-level) + worker_config.update({ + k: v + for k, v in server_type_config.items() if k not in EXCLUDE_FROM_WORKER + }) + + # Convert top-level free_gpu_memory_fraction into kv_cache_config + if 'free_gpu_memory_fraction' in worker_config: + frac = worker_config.pop('free_gpu_memory_fraction') + if 'kv_cache_config' not in worker_config: + worker_config['kv_cache_config'] = {} + worker_config['kv_cache_config'].setdefault('free_gpu_memory_fraction', + frac) + + # Add service discovery config + worker_config['disagg_cluster'] = disagg_cluster + + return worker_config def get_test_config(test_desc, example_dir, test_root): - """Get test configuration based on test description.""" + """Get config file path for a test description.""" test_configs_root = f"{test_root}/test_configs" config_map = { "2_ranks_diff_max_tokens": - (2, f"{test_configs_root}/disagg_config_diff_max_tokens.yaml"), - "2_ranks": (2, f"{example_dir}/disagg_config.yaml"), + f"{test_configs_root}/disagg_config_diff_max_tokens.yaml", + "2_ranks": + f"{test_configs_root}/disagg_config.yaml", "2_ranks_trt_backend": - (2, f"{test_configs_root}/disagg_config_trt_backend.yaml"), - "gen_only": (2, f"{test_configs_root}/disagg_config_gen_only.yaml"), + f"{test_configs_root}/disagg_config_trt_backend.yaml", + "gen_only": + f"{test_configs_root}/disagg_config_gen_only.yaml", "gen_only_trt_backend": - (2, f"{test_configs_root}/disagg_config_gen_only_trt_backend.yaml"), + f"{test_configs_root}/disagg_config_gen_only_trt_backend.yaml", "gen_only_bs1": - (4, f"{test_configs_root}/disagg_config_gen_only_bs1.yaml"), - "4_ranks": (4, f"{test_configs_root}/disagg_config_ctxtp2_gentp1.yaml"), + f"{test_configs_root}/disagg_config_gen_only_bs1.yaml", + "4_ranks": + f"{test_configs_root}/disagg_config_ctxtp2_gentp1.yaml", "4_ranks_trt_backend": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp1_trt_backend.yaml"), + f"{test_configs_root}/disagg_config_ctxtp2_gentp1_trt_backend.yaml", "cuda_graph": - (2, f"{test_configs_root}/disagg_config_cuda_graph_padding.yaml"), - "mixed": (2, f"{test_configs_root}/disagg_config_mixed.yaml"), - "overlap": (2, f"{test_configs_root}/disagg_config_overlap.yaml"), + f"{test_configs_root}/disagg_config_cuda_graph_padding.yaml", + "mixed": + f"{test_configs_root}/disagg_config_mixed.yaml", + "overlap": + f"{test_configs_root}/disagg_config_overlap.yaml", "overlap_transceiver_runtime_python": - (2, - f"{test_configs_root}/disagg_config_overlap_transceiver_runtime_python.yaml" - ), - "tool_calls": (2, f"{test_configs_root}/disagg_config_overlap.yaml"), - "perf_metrics": (2, f"{test_configs_root}/disagg_config_metrics.yaml"), + f"{test_configs_root}/disagg_config_overlap_transceiver_runtime_python.yaml", + "tool_calls": + f"{test_configs_root}/disagg_config_overlap.yaml", + "perf_metrics": + f"{test_configs_root}/disagg_config_metrics.yaml", "trtllm_sampler": - (2, f"{test_configs_root}/disagg_config_trtllm_sampler.yaml"), + f"{test_configs_root}/disagg_config_trtllm_sampler.yaml", "load_balance": - (4, f"{test_configs_root}/disagg_config_load_balance.yaml"), + f"{test_configs_root}/disagg_config_load_balance.yaml", "cache_aware_balance": - (4, f"{test_configs_root}/disagg_config_cache_aware_balance.yaml"), - "conditional": (2, - f"{test_configs_root}/disagg_config_conditional.yaml"), - "ngram": (2, f"{test_configs_root}/disagg_config_ngram.yaml"), + f"{test_configs_root}/disagg_config_cache_aware_balance.yaml", + "conditional": + f"{test_configs_root}/disagg_config_conditional.yaml", + "ngram": + f"{test_configs_root}/disagg_config_ngram.yaml", "ctxpp2_genpp2": - (4, f"{test_configs_root}/disagg_config_ctxpp2_genpp2.yaml"), + f"{test_configs_root}/disagg_config_ctxpp2_genpp2.yaml", "ctxtp2_genpp2": - (4, f"{test_configs_root}/disagg_config_ctxtp2_genpp2.yaml"), + f"{test_configs_root}/disagg_config_ctxtp2_genpp2.yaml", "ctxpp2_gentp2": - (4, f"{test_configs_root}/disagg_config_ctxpp2_gentp2.yaml"), + f"{test_configs_root}/disagg_config_ctxpp2_gentp2.yaml", "ctxtp2pp2_gentp2pp2": - (8, f"{test_configs_root}/disagg_config_ctxtp2pp2_gentp2pp2.yaml"), + f"{test_configs_root}/disagg_config_ctxtp2pp2_gentp2pp2.yaml", "ctxpp4_genpp4": - (8, f"{test_configs_root}/disagg_config_ctxpp4_genpp4.yaml"), + f"{test_configs_root}/disagg_config_ctxpp4_genpp4.yaml", "ctxpp4_gentp4": - (8, f"{test_configs_root}/disagg_config_ctxpp4_gentp4.yaml"), + f"{test_configs_root}/disagg_config_ctxpp4_gentp4.yaml", "deepseek_v3_lite_fp8_mpi": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_mpi.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_mpi.yaml", "deepseek_v3_lite_fp8_ucx": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_ucx.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_ucx.yaml", "deepseek_v3_lite_fp8_nixl": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_nixl.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_nixl.yaml", "deepseek_v3_lite_fp8_transceiver_runtime_python": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_transceiver_runtime_python.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_transceiver_runtime_python.yaml", "deepseek_v3_lite_fp8_tp1": - (2, - f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite.yaml", "deepseek_v3_lite_fp8_tp1_mtp": - (2, - f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp.yaml" - ), - "deepseek_v3_lite_fp_8_overlap_dp": - (2, - f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_overlap_dp.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp.yaml", "deepseek_v3_lite_fp8_attention_dp": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp.yaml", "deepseek_v3_lite_fp8_attention_dp_gen_only": - (4, - f"{test_configs_root}/disagg_config_gentp2_deepseek_v3_lite_attention_dp_gen_only.yaml" - ), + f"{test_configs_root}/disagg_config_gentp2_deepseek_v3_lite_attention_dp_gen_only.yaml", "deepseek_v3_lite_fp_8_attention_dp_overlap": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap.yaml", "deepseek_v3_lite_fp8_attention_dp_overlap_cuda_graph": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap_cuda_graph.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_overlap_cuda_graph.yaml", "deepseek_v3_lite_fp8_overlap_cuda_graph": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_overlap_cuda_graph.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_overlap_cuda_graph.yaml", "deepseek_v3_lite_fp8_attention_dp_one": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one.yaml", "deepseek_v3_lite_fp8_attention_dp_one_mtp": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one_mtp.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_deepseek_v3_lite_attention_dp_one_mtp.yaml", "deepseek_v3_lite_fp8_tp1_attention_dp_overlap_one_mtp": - (2, - f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_attention_dp_overlap.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_attention_dp_overlap.yaml", "deepseek_v3_lite_bf16_cache_aware_balance": - (4, - f"{test_configs_root}/disagg_config_cache_aware_balance_deepseek_v3.yaml" - ), + f"{test_configs_root}/disagg_config_cache_aware_balance_deepseek_v3.yaml", "deepseek_v3_lite_bf16_conditional": - (2, f"{test_configs_root}/disagg_config_conditional_deepseek_v3.yaml"), + f"{test_configs_root}/disagg_config_conditional_deepseek_v3.yaml", "deepseek_v3_lite_fp8_tp1_two_mtp": - (2, - f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_two_mtp.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_two_mtp.yaml", "deepseek_v3_lite_fp8_ctxpp2_gentp2_one_mtp": - (4, - f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_ctxpp2_gentp2.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp1_gentp1_deepseek_v3_lite_one_mtp_ctxpp2_gentp2.yaml", "deepseek_v3_lite_bf16_empty_batch": - (3, - f"{test_configs_root}/disagg_config_deepseek_v3_lite_empty_batch.yaml" - ), + f"{test_configs_root}/disagg_config_deepseek_v3_lite_empty_batch.yaml", "llama4_kv_cache_overflow": - (8, f"{test_configs_root}/disagg_config_llama4_kv_cache_overflow.yaml"), + f"{test_configs_root}/disagg_config_llama4_kv_cache_overflow.yaml", "deepseek_v3_lite_bf16_tllm_gen_helix": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp1cp2_deepseek_v3_lite_bf16_tllm_gen.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp2_gentp1cp2_deepseek_v3_lite_bf16_tllm_gen.yaml", "deepseek_r1_v2_fp4_stress": - (8, - f"{test_configs_root}/disagg_config_ctxtp4_gentp4_deepseek_r1_v2_fp4_tllm.yaml" - ), + f"{test_configs_root}/disagg_config_ctxtp4_gentp4_deepseek_r1_v2_fp4_tllm.yaml", "gpt_oss_120b_stress": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_gptoss_tllm.yaml"), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_gptoss_tllm.yaml", "gpt_oss_120b_harmony": - (4, - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_gptoss_tllm.yaml"), + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_gptoss_tllm.yaml", "cancel_stress_test": - (2, f"{test_configs_root}/disagg_config_cancel_stress_test.yaml"), + f"{test_configs_root}/disagg_config_cancel_stress_test.yaml", "cancel_stress_test_large": - (8, f"{test_configs_root}/disagg_config_cancel_stress_test_large.yaml"), + f"{test_configs_root}/disagg_config_cancel_stress_test_large.yaml", } if test_desc not in config_map: raise ValueError(f"Invalid test description: {test_desc}, " f"valid descriptions are: {config_map.keys()}") - return (config_map[test_desc][0], - revise_disagg_config_file_with_free_ports(config_map[test_desc][1])) - + return config_map[test_desc] -def get_extra_llm_config(config, suffix, cwd): - extra_llm_config = { - 'orchestrator_type': 'ray', - } - for key, value in config.items(): - if key not in ['num_instances', 'urls']: - extra_llm_config[key] = value - temp_fd, extra_config_file = tempfile.mkstemp(suffix='_%s.yaml' % suffix, - dir=cwd) - with os.fdopen(temp_fd, 'w') as f: - yaml.dump(extra_llm_config, f) +def setup_model_symlink(llm_venv, model_root, dest_subpath): + """Create symlink for model in test working directory. - return extra_config_file - - -def generate_worker_commands(model_path, config, server_config, - extra_config_file, server_role): - worker_commands = [] - - assert model_path, "model path is required." - - for url in server_config['urls']: - host, port = url.split(':') - cmd = [ - 'trtllm-serve', model_path, '--host', host, '--port', port, - '--backend', config['backend'], '--config', extra_config_file, - '--server_role', server_role - ] - worker_commands.append(cmd) - return worker_commands + Args: + llm_venv: Virtual environment object with get_working_directory() + model_root: Source model directory path + dest_subpath: Destination subdirectory (relative to working dir) + """ + dst = f"{llm_venv.get_working_directory()}/{dest_subpath}" + if not os.path.islink(dst): + os.makedirs(os.path.dirname(dst), exist_ok=True) + os.symlink(model_root, dst, target_is_directory=True) ClientTestSet = namedtuple('ClientTestSet', [ @@ -342,8 +331,12 @@ def run_client_tests(example_dir, # Prepare poll processes worker_processes = [] if use_ray: - for proc_cm in workers_proc: - worker_processes.append(proc_cm.__enter__()) + for proc in workers_proc: + # Ray passes context managers, SD passes raw Popen objects + if hasattr(proc, '__enter__'): + worker_processes.append(proc.__enter__()) + else: + worker_processes.append(proc) else: worker_processes = [workers_proc] @@ -353,7 +346,7 @@ def run_client_tests(example_dir, if client_test_set.completion: check_call(client_cmd, env=env, poll_procs=poll_procs) - # Run streaming completion test + # Streaming client run if client_test_set.completion_streaming: streaming_client_cmd = client_cmd + [ '--streaming', '-o', 'output_streaming.json' @@ -436,183 +429,184 @@ def fetch_prometheus_metrics(server_url: str): return response.text +def setup_disagg_cluster( + config_file: str, + model_name: str | None = None, + env: dict[str, str] | None = None, + cwd: str | None = None, + server_start_timeout: int = 300, +) -> tuple[dict[str, Any], list[ProcessWrapper], list[ProcessWrapper], + ProcessWrapper, int, str]: + """Load config, launch workers + disagg server, wait for ready. + + Args: + config_file: Path to disaggregated server config YAML + model_name: Model path override (defaults to config's 'model' field) + env: Environment variables to pass to subprocess (workers and disagg server) + server_start_timeout: Timeout in seconds for server to become ready + + Returns: + tuple: (config, ctx_workers, gen_workers, disagg_server, server_port, work_dir) + """ + with open(config_file, 'r') as f: + config = yaml.safe_load(f) + + disagg_cluster = get_default_disagg_cluster_config() + server_host = config.get("hostname", "localhost") + server_port = get_free_port() + work_dir = tempfile.mkdtemp() + disagg_cluster["cluster_uri"] = f"http://{server_host}:{server_port}" + + # Auto-deduce minimal_instances from num_instances + ctx_servers = config.get("context_servers", {}) + gen_servers = config.get("generation_servers", {}) + num_ctx_instances = ctx_servers.get("num_instances", 1) + num_gen_instances = gen_servers.get("num_instances", 1) + disagg_cluster["minimal_instances"] = { + "context_servers": num_ctx_instances, + "generation_servers": num_gen_instances + } + + # Calculate GPUs per worker instance: tp * pp * cp + gpus_per_ctx = (ctx_servers.get("tensor_parallel_size", 1) * + ctx_servers.get("pipeline_parallel_size", 1) * + ctx_servers.get("context_parallel_size", 1)) + gpus_per_gen = (gen_servers.get("tensor_parallel_size", 1) * + gen_servers.get("pipeline_parallel_size", 1) * + gen_servers.get("context_parallel_size", 1)) + + # Build worker configs + ctx_worker_config = build_worker_config(config, ctx_servers, disagg_cluster) + gen_worker_config = build_worker_config(config, gen_servers, disagg_cluster) + + # Launch workers + model = model_name or config.get("model") + ctx_workers = [] + gen_workers = [] + disagg_server = None + next_device = 0 + + import torch + num_gpus = torch.cuda.device_count() + + try: + for i in range(num_ctx_instances): + device_ids = ",".join( + str(d) for d in dict.fromkeys((next_device + j) % num_gpus + for j in range(gpus_per_ctx))) + ctx_workers.append( + run_ctx_worker(model, + ctx_worker_config, + work_dir, + port=0, + device=device_ids, + env=env)) + next_device += gpus_per_ctx + + for i in range(num_gen_instances): + device_ids = ",".join( + str(d) for d in dict.fromkeys((next_device + j) % num_gpus + for j in range(gpus_per_gen))) + gen_workers.append( + run_gen_worker(model, + gen_worker_config, + work_dir, + port=0, + device=device_ids, + env=env)) + next_device += gpus_per_gen + + # Build minimal server config and launch + server_config = { + "hostname": + server_host, + "port": + server_port, + "disagg_cluster": + disagg_cluster, + "context_servers": { + "router": ctx_servers.get("router", {}) + }, + "generation_servers": { + "router": gen_servers.get("router", {}) + }, + "conditional_disagg_config": + config.get("conditional_disagg_config", None), + "perf_metrics_max_requests": + config.get("perf_metrics_max_requests", 0), + } + disagg_server = run_disagg_server(server_config, + work_dir, + server_port, + env=env, + cwd=cwd) + + asyncio.run( + wait_for_disagg_server_ready(server_port, + timeout=server_start_timeout)) + except Exception: + terminate(*ctx_workers, *gen_workers, disagg_server) + shutil.rmtree(work_dir, ignore_errors=True) + raise + + return config, ctx_workers, gen_workers, disagg_server, server_port, work_dir + + def run_disaggregated_test(example_dir, test_desc, num_iters=5, env=None, - cwd=None, prompt_file="prompts.json", - extra_endpoints_test: Callable[[str], None] = None, - model_path=None): - """Run disaggregated test with given configuration.""" - cleanup_output_files() - run_env = env.copy() - - # on some CI nodes , we set UCX_TLS to "^ib,gdr_copy" to avoid the issue that IB equipped but not available, and gdr_copy pin buffer failed. - # we set UCX_MM_ERROR_HANDLING to "y" to avoid the issue that NIXL cannot use IB or TCP for notify on some CI nodes, - # setting it to "y" will enable NIXL to use system memory for notify. + extra_endpoints_test=None, + model_path=None, + cwd=None): + """Run disaggregated test using service discovery instead of MPI.""" - run_env["UCX_TLS"] = "^ib,gdr_copy" - run_env["UCX_MM_ERROR_HANDLING"] = "y" - num_ranks, config_file = get_test_config(test_desc, example_dir, - os.path.dirname(__file__)) - - use_ray = mpi_disabled() - if not use_ray: - workers_cmd = [ - 'mpirun', '--allow-run-as-root', '--oversubscribe', '-n', - str(num_ranks), 'trtllm-serve', 'disaggregated_mpi_worker', '-c', - config_file - ] - else: + if mpi_disabled(): pytest.skip( "https://nvbugs/5584607 Ray orchestrator is not supported with NIXL(DEFAULT) cache transceiver backend." ) - with open(config_file, 'r') as f: - config = yaml.safe_load(f) - - if config['backend'] != "pytorch": - pytest.skip( - "Ray orchestrator is only supported with pytorch backend.") - - extra_config_files = [] - workers_cmds = [] - - # Generate ctx and gen server worker commands - ctx_extra_config_file = get_extra_llm_config(config['context_servers'], - "ctx", cwd) - extra_config_files.append(ctx_extra_config_file) - workers_cmds.extend( - generate_worker_commands(model_path, config, - config['context_servers'], - ctx_extra_config_file, 'context')) - - gen_extra_config_file = get_extra_llm_config( - config['generation_servers'], "gen", cwd) - extra_config_files.append(gen_extra_config_file) - workers_cmds.extend( - generate_worker_commands(model_path, config, - config['generation_servers'], - gen_extra_config_file, 'generation')) - - server_start_timeout = 1200 - server_cmd = [ - 'trtllm-serve', 'disaggregated', '--server_start_timeout', - str(server_start_timeout), '-c', config_file - ] - server_host, server_port = get_disagg_server_url_from_cfg(config_file) - server_url = f"http://{server_host}:{server_port}" - try: - if not use_ray: - with ( # Start workers - open('output_workers.log', 'w') as output_workers, - popen(workers_cmd, - stdout=output_workers, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd) as workers_proc, - # Start server - open('output_disagg.log', 'w') as output_disagg, - popen(server_cmd, - stdout=output_disagg, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd) as server_proc): - run_client_tests(example_dir, - config_file, - test_desc, - num_iters, - env, - server_start_timeout, - prompt_file, - extra_endpoints_test, - server_url, - workers_proc, - server_proc, - use_ray=False) + config_file = get_test_config(test_desc, example_dir, + os.path.dirname(__file__)) + config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ + setup_disagg_cluster(config_file, model_name=model_path, env=env, cwd=cwd) - else: - runtime_env = { - "env_vars": { - "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1" - } - } - ray.init(address="local", - include_dashboard=False, - ignore_reinit_error=True, - runtime_env=runtime_env) - gcs_addr = ray.get_runtime_context().gcs_address - ray_port = str(gcs_addr.split(":")[1]) - run_env.update({ - "RAY_ADDRESS": f"localhost:{ray_port}", - "TLLM_RAY_FORCE_LOCAL_CLUSTER": "0" - }) - workers_proc = [] - with contextlib.ExitStack() as stack: - workers_log = stack.enter_context( - open('output_workers.log', 'w')) - - for cmd in workers_cmds: - proc = stack.enter_context( - popen( - cmd, - stdout=workers_log, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd, - )) - workers_proc.append(proc) - - output_disagg = stack.enter_context( - open('output_disagg.log', 'w')) - server_proc = stack.enter_context( - popen(server_cmd, - stdout=output_disagg, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd)) - - if not wait_for_server(server_host, - server_port, - timeout_seconds=server_start_timeout): - raise RuntimeError( - f"Disaggregated server failed to start within {server_start_timeout} seconds" - ) - - run_client_tests(example_dir, - config_file, - test_desc, - num_iters, - env, - server_start_timeout, - prompt_file, - extra_endpoints_test, - server_url, - workers_proc, - server_proc, - use_ray=True) - except Exception: - # Print outputs on error - logger.error("-------- Workers output --------") - with open('output_workers.log', 'r') as f: - logger.error(f.read()) + server_host = config.get("hostname", "localhost") - logger.error("-------- Disagg server output --------") - with open('output_disagg.log', 'r') as f: - logger.error(f.read()) - raise + try: + server_url = f"http://{server_host}:{server_port}" + + # Create a temporary client config file with the correct server port + client_config = config.copy() + client_config["port"] = server_port + client_config["hostname"] = server_host + temp_fd, client_config_file = tempfile.mkstemp(suffix='.yaml', + dir=work_dir) + with os.fdopen(temp_fd, 'w') as f: + yaml.dump(client_config, f) + + # collect all worker processes for monitoring + all_worker_procs = [w.process for w in ctx_workers + ] + [w.process for w in gen_workers] + + # run client tests + run_client_tests( + example_dir, + client_config_file, + test_desc, + num_iters, + env, + 300, # timeout + prompt_file, + extra_endpoints_test, + server_url, + all_worker_procs, + disagg_server.process, + use_ray=True) finally: - if 'server_proc' in locals() and 'workers_proc' in locals(): - server_proc.terminate() - workers_proc.terminate() - server_proc.wait() - workers_proc.wait() - if use_ray: - ray.shutdown() - for extra_file in extra_config_files: - if os.path.exists(extra_file): - os.remove(extra_file) + terminate(*ctx_workers, *gen_workers, disagg_server) + shutil.rmtree(work_dir, ignore_errors=True) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], @@ -620,59 +614,45 @@ def run_disaggregated_test(example_dir, def test_disaggregated_diff_max_tokens(disaggregated_test_root, disaggregated_example_root, llm_venv, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "2_ranks_diff_max_tokens", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), - prompt_file="long_prompts.json") + prompt_file="long_prompts.json", + cwd=llm_venv.get_working_directory()) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], indirect=True) -def test_disaggregated_single_gpu_with_mpirun(disaggregated_test_root, - disaggregated_example_root, - llm_venv, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) +def test_disaggregated_single_gpu(disaggregated_test_root, + disaggregated_example_root, llm_venv, + llama_model_root): + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") + env = llm_venv._new_env.copy() + env["CUDA_VISIBLE_DEVICES"] = "0" run_disaggregated_test(disaggregated_example_root, "2_ranks", - env=llm_venv._new_env, + env=env, cwd=llm_venv.get_working_directory()) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], indirect=True) -def test_disaggregated_single_gpu_with_mpirun_trt_backend( - disaggregated_test_root, disaggregated_example_root, llm_venv, - llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) +def test_disaggregated_single_gpu_trt_backend(disaggregated_test_root, + disaggregated_example_root, + llm_venv, llama_model_root): + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") + env = llm_venv._new_env.copy() + env["CUDA_VISIBLE_DEVICES"] = "0" run_disaggregated_test(disaggregated_example_root, "2_ranks_trt_backend", - env=llm_venv._new_env, + env=env, cwd=llm_venv.get_working_directory()) @@ -681,14 +661,8 @@ def test_disaggregated_single_gpu_with_mpirun_trt_backend( def test_disaggregated_benchmark_gen_only(disaggregated_test_root, disaggregated_example_root, llm_venv, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") env = llm_venv._new_env.copy() env['TRTLLM_DISAGG_BENCHMARK_GEN_ONLY'] = '1' @@ -703,14 +677,8 @@ def test_disaggregated_benchmark_gen_only(disaggregated_test_root, def test_disaggregated_benchmark_gen_only_trt_backend( disaggregated_test_root, disaggregated_example_root, llm_venv, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") env = llm_venv._new_env.copy() env['TRTLLM_DISAGG_BENCHMARK_GEN_ONLY'] = '1' @@ -726,37 +694,25 @@ def test_disaggregated_benchmark_gen_only_trt_backend( def test_disaggregated_genbs1(disaggregated_test_root, disaggregated_example_root, llm_venv, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") env = llm_venv._new_env.copy() env['TRTLLM_DISAGG_BENCHMARK_GEN_ONLY'] = '1' run_disaggregated_test(disaggregated_example_root, "gen_only_bs1", - env=llm_venv._new_env, + env=env, cwd=llm_venv.get_working_directory()) @pytest.mark.skip_less_device(2) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], indirect=True) -def test_disaggregated_multi_gpu_with_mpirun(disaggregated_test_root, - disaggregated_example_root, - llm_venv, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) +def test_disaggregated_multi_gpu(disaggregated_test_root, + disaggregated_example_root, llm_venv, + llama_model_root): + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "4_ranks", @@ -767,17 +723,11 @@ def test_disaggregated_multi_gpu_with_mpirun(disaggregated_test_root, @pytest.mark.skip_less_device(2) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], indirect=True) -def test_disaggregated_multi_gpu_with_mpirun_trt_backend( - disaggregated_test_root, disaggregated_example_root, llm_venv, - llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) +def test_disaggregated_multi_gpu_trt_backend(disaggregated_test_root, + disaggregated_example_root, + llm_venv, llama_model_root): + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "4_ranks_trt_backend", @@ -789,14 +739,8 @@ def test_disaggregated_multi_gpu_with_mpirun_trt_backend( indirect=True) def test_disaggregated_cuda_graph(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "cuda_graph", @@ -808,14 +752,8 @@ def test_disaggregated_cuda_graph(disaggregated_test_root, llm_venv, indirect=True) def test_disaggregated_mixed(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "mixed", @@ -827,14 +765,8 @@ def test_disaggregated_mixed(disaggregated_test_root, llm_venv, indirect=True) def test_disaggregated_overlap(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "overlap", @@ -847,14 +779,8 @@ def test_disaggregated_overlap(disaggregated_test_root, llm_venv, def test_disaggregated_overlap_transceiver_runtime_python( disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "overlap_transceiver_runtime_python", @@ -867,14 +793,8 @@ def test_disaggregated_overlap_transceiver_runtime_python( def test_disaggregated_perf_metrics(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") def extra_endpoints_test(server_url: str): item = get_timing_metrics(server_url) @@ -884,8 +804,8 @@ def extra_endpoints_test(server_url: str): run_disaggregated_test(disaggregated_example_root, "perf_metrics", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), - extra_endpoints_test=extra_endpoints_test) + extra_endpoints_test=extra_endpoints_test, + cwd=llm_venv.get_working_directory()) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], @@ -894,14 +814,8 @@ def test_disaggregated_chat_completion_tool_calls(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "tool_calls", @@ -916,14 +830,8 @@ def test_disaggregated_chat_completion_tool_calls(disaggregated_test_root, def test_disaggregated_kv_cache_time_output(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") output_path = os.path.join(llm_venv.get_working_directory(), "cache_time") run_disaggregated_test(disaggregated_example_root, @@ -933,7 +841,7 @@ def test_disaggregated_kv_cache_time_output(disaggregated_test_root, llm_venv, cwd=llm_venv.get_working_directory()) assert os.path.isdir(output_path) send_file = os.path.join(output_path, "rank_0_send.csv") - recv_file = os.path.join(output_path, "rank_1_recv.csv") + recv_file = os.path.join(output_path, "rank_0_recv.csv") assert os.path.exists(send_file) assert os.path.exists(recv_file) with open(send_file, "r") as f: @@ -964,14 +872,8 @@ def test_disaggregated_kv_cache_time_output(disaggregated_test_root, llm_venv, def test_disaggregated_trtllm_sampler(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "trtllm_sampler", @@ -984,14 +886,8 @@ def test_disaggregated_trtllm_sampler(disaggregated_test_root, llm_venv, def test_disaggregated_load_balance(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "load_balance", @@ -1004,14 +900,8 @@ def test_disaggregated_load_balance(disaggregated_test_root, llm_venv, def test_disaggregated_cache_aware_balance(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "cache_aware_balance", @@ -1024,14 +914,8 @@ def test_disaggregated_cache_aware_balance(disaggregated_test_root, llm_venv, def test_disaggregated_conditional(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "conditional", @@ -1043,14 +927,8 @@ def test_disaggregated_conditional(disaggregated_test_root, llm_venv, indirect=True) def test_disaggregated_ngram(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "ngram", env=llm_venv._new_env, @@ -1063,19 +941,13 @@ def test_disaggregated_ngram(disaggregated_test_root, llm_venv, def test_disaggregated_ctxpp2_genpp2(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "ctxpp2_genpp2", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), - model_path=llama_model_root) + model_path=llama_model_root, + cwd=llm_venv.get_working_directory()) @pytest.mark.skip_less_device(4) @@ -1084,19 +956,13 @@ def test_disaggregated_ctxpp2_genpp2(disaggregated_test_root, llm_venv, def test_disaggregated_ctxtp2_genpp2(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "ctxtp2_genpp2", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), - model_path=llama_model_root) + model_path=llama_model_root, + cwd=llm_venv.get_working_directory()) @pytest.mark.skip_less_device(4) @@ -1105,19 +971,13 @@ def test_disaggregated_ctxtp2_genpp2(disaggregated_test_root, llm_venv, def test_disaggregated_ctxpp2_gentp2(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "ctxpp2_gentp2", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), - model_path=llama_model_root) + model_path=llama_model_root, + cwd=llm_venv.get_working_directory()) @pytest.mark.skip_less_device(8) @@ -1126,14 +986,8 @@ def test_disaggregated_ctxpp2_gentp2(disaggregated_test_root, llm_venv, def test_disaggregated_ctxtp2pp2_gentp2pp2(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "ctxtp2pp2_gentp2pp2", env=llm_venv._new_env, @@ -1146,14 +1000,8 @@ def test_disaggregated_ctxtp2pp2_gentp2pp2(disaggregated_test_root, llm_venv, def test_disaggregated_ctxpp4_genpp4(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "ctxpp4_genpp4", env=llm_venv._new_env, @@ -1167,42 +1015,34 @@ def test_disaggregated_ctxpp4_genpp4(disaggregated_test_root, llm_venv, def test_disaggregated_ctxpp4_gentp4(disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): - src_dst_dict = { - llama_model_root: - f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") run_disaggregated_test(disaggregated_example_root, "ctxpp4_gentp4", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), - model_path=llama_model_root) + model_path=llama_model_root, + cwd=llm_venv.get_working_directory()) @skip_no_hopper @pytest.mark.skip_less_device(4) +@pytest.mark.skip( + reason="MPI cache transceiver requires shared MPI process group, " + "incompatible with service discovery which launches separate subprocesses") @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-fp8'], indirect=True) def test_disaggregated_deepseek_v3_lite_fp8_mpi(disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") env = llm_venv._new_env.copy() env["TRTLLM_USE_MPI_KVCACHE"] = "1" run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_mpi", env=env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1212,18 +1052,13 @@ def test_disaggregated_deepseek_v3_lite_fp8_mpi(disaggregated_test_root, def test_disaggregated_deepseek_v3_lite_fp8_tp1_single_gpu( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_tp1", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1233,18 +1068,13 @@ def test_disaggregated_deepseek_v3_lite_fp8_tp1_single_gpu( def test_disaggregated_deepseek_v3_lite_fp8_tp1_single_gpu_mtp( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_tp1_mtp", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1256,20 +1086,14 @@ def test_disaggregated_deepseek_v3_lite_fp8_ctxpp2_gentp2_one_mtp( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): #add one mtp layer, pp rank0 will have 15 layer, pp rank 1 will have 16 layers. - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_ctxpp2_gentp2_one_mtp", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), - model_path=deepseek_v3_model_root) + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory()) @skip_no_hopper @@ -1282,22 +1106,16 @@ def test_disaggregated_deepseek_v3_lite_fp8_ucx(disaggregated_test_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") env = llm_venv._new_env.copy() env["TRTLLM_USE_UCX_KVCACHE"] = "1" env["UCX_TLS"] = "^ib,gdr_copy" run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_ucx", env=env, - cwd=llm_venv.get_working_directory(), - model_path=deepseek_v3_model_root) + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory()) @skip_no_hopper @@ -1309,14 +1127,8 @@ def test_disaggregated_deepseek_v3_lite_fp8_nixl(disaggregated_test_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") env = llm_venv._new_env.copy() env["TRTLLM_USE_NIXL_KVCACHE"] = "1" env["UCX_TLS"] = "^ib,gdr_copy" @@ -1324,8 +1136,8 @@ def test_disaggregated_deepseek_v3_lite_fp8_nixl(disaggregated_test_root, run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_nixl", env=env, - cwd=llm_venv.get_working_directory(), - model_path=deepseek_v3_model_root) + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory()) @skip_no_hopper @@ -1335,22 +1147,15 @@ def test_disaggregated_deepseek_v3_lite_fp8_nixl(disaggregated_test_root, def test_disaggregated_deepseek_v3_lite_fp8_transceiver_runtime_python( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") env = llm_venv._new_env.copy() env["UCX_TLS"] = "^ib,gdr_copy" run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_transceiver_runtime_python", env=env, - cwd=llm_venv.get_working_directory(), - model_path=deepseek_v3_model_root) + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory()) @skip_no_hopper @@ -1360,14 +1165,8 @@ def test_disaggregated_deepseek_v3_lite_fp8_transceiver_runtime_python( def test_disaggregated_deepseek_v3_lite_fp8_ucx_tp1_single_gpu( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") env = llm_venv._new_env.copy() env["TRTLLM_USE_UCX_KVCACHE"] = "1" env["UCX_TLS"] = "^ib,gdr_copy" @@ -1375,6 +1174,7 @@ def test_disaggregated_deepseek_v3_lite_fp8_ucx_tp1_single_gpu( run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_tp1", env=env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1385,19 +1185,13 @@ def test_disaggregated_deepseek_v3_lite_fp8_ucx_tp1_single_gpu( def test_disaggregated_deepseek_v3_lite_fp8_attention_dp( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_attention_dp", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1408,21 +1202,15 @@ def test_disaggregated_deepseek_v3_lite_fp8_attention_dp( def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_gen_only( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") env = llm_venv._new_env.copy() env['TRTLLM_DISAGG_BENCHMARK_GEN_ONLY'] = '1' run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_attention_dp_gen_only", env=env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1433,18 +1221,13 @@ def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_gen_only( def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap( disaggregated_test_root, llm_venv, disaggregated_example_root, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp_8_attention_dp_overlap", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1455,20 +1238,14 @@ def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap( def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap_cuda_graph( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test( disaggregated_example_root, "deepseek_v3_lite_fp8_attention_dp_overlap_cuda_graph", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1479,19 +1256,13 @@ def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap_cuda_graph( def test_disaggregated_deepseek_v3_lite_fp8_overlap_cuda_graph( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_overlap_cuda_graph", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1502,19 +1273,13 @@ def test_disaggregated_deepseek_v3_lite_fp8_overlap_cuda_graph( def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_one( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_attention_dp_one", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1525,19 +1290,13 @@ def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_one( def test_disaggregated_deepseek_v3_lite_fp8_attention_dp_one_mtp( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_attention_dp_one_mtp", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1549,22 +1308,15 @@ def test_disaggregated_deepseek_v3_lite_fp8_tp1_attention_dp_overlap_one_mtp( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test( disaggregated_example_root, "deepseek_v3_lite_fp8_tp1_attention_dp_overlap_one_mtp", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), - model_path=deepseek_v3_model_root) + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory()) @skip_no_hopper @@ -1573,18 +1325,13 @@ def test_disaggregated_deepseek_v3_lite_fp8_tp1_attention_dp_overlap_one_mtp( def test_disaggregated_deepseek_v3_lite_bf16_cache_aware_balance( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/bf16", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/bf16") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_bf16_cache_aware_balance", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1594,18 +1341,13 @@ def test_disaggregated_deepseek_v3_lite_bf16_cache_aware_balance( def test_disaggregated_deepseek_v3_lite_bf16_conditional( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/bf16", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/bf16") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_bf16_conditional", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1615,19 +1357,13 @@ def test_disaggregated_deepseek_v3_lite_bf16_conditional( def test_disaggregated_deepseek_v3_lite_fp8_tp1_two_mtp( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/fp8", - } - - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/fp8") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_fp8_tp1_two_mtp", env=llm_venv._new_env, + model_path=deepseek_v3_model_root, cwd=llm_venv.get_working_directory()) @@ -1668,126 +1404,84 @@ def run_disaggregated_benchmark(example_dir, benchmark_model_root, shared_gpt_path, env=None, - cwd=None, - num_ranks=2, random_input_len=16, random_output_len=64, num_prompts=100, max_concurrency=32, - skip_warmup=False): + skip_warmup=False, + model_path=None, + cwd=None): """Run disaggregated test with given configuration.""" - run_env = env.copy() + run_env = env.copy() if env else os.environ.copy() run_env["UCX_TLS"] = "^ib,gdr_copy" run_env["UCX_MM_ERROR_HANDLING"] = "y" - workers_cmd = [ - 'mpirun', '--allow-run-as-root', '--oversubscribe', '-n', - str(num_ranks), 'trtllm-serve', 'disaggregated_mpi_worker', '-c', - config_file - ] - server_start_timeout = 1200 - server_cmd = [ - 'trtllm-serve', 'disaggregated', '--server_start_timeout', - str(server_start_timeout), '-c', config_file - ] - server_host, server_port = get_disagg_server_url_from_cfg(config_file) + config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ + setup_disagg_cluster(config_file, model_name=model_path, env=run_env, cwd=cwd) + + server_host = config.get("hostname", "localhost") + try: - with ( # Start workers - open('output_workers.log', 'w') as output_workers, - popen(workers_cmd, - stdout=output_workers, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd) as workers_proc, - # Start server - open('output_disagg.log', 'w') as output_disagg, - popen(server_cmd, - stdout=output_disagg, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd) as server_proc): - # Ensure the sever has started - client_dir = f"{example_dir}/clients" - client_cmd = [ - 'python3', f'{client_dir}/disagg_client.py', '-c', config_file, - '-p', f'{client_dir}/prompts.json', '--ignore-eos', - '--server-start-timeout', - str(server_start_timeout) - ] - # Warm up - check_call(client_cmd, - env=env, - poll_procs=[workers_proc, server_proc]) - # Start Benchmark - benchmark_script = os.path.join(benchmark_root, - "benchmark_serving.py") - benchmark_cmd = [ - 'python3', - benchmark_script, - '--model', - benchmark_model_root, - '--tokenizer', - benchmark_model_root, - '--dataset-name', - 'random', - '--dataset-path', - shared_gpt_path, - '--random-input-len', - str(random_input_len), - '--random-output-len', - str(random_output_len), - '--random-prefix-len', - '0', - '--num-prompts', - str(num_prompts), - '--max-concurrency', - str(max_concurrency), - '--host', - server_host, - '--port', - str(server_port), - '--ignore-eos', - '--no-test-input', - '--percentile-metrics', - 'e2el,ttft', - ] - # warm up - if not skip_warmup: - check_call(benchmark_cmd, env=env) - output = check_output(benchmark_cmd, env=env) - e2el_pattern = r"Median E2EL \(ms\):\s*(\d+\.?\d*)" - ttft_pattern = r"Median TTFT \(ms\):\s*(\d+\.?\d*)" - e2el_match = re.search(e2el_pattern, output) - ttft_match = re.search(ttft_pattern, output) - if e2el_match and ttft_match: - median_e2el = float(e2el_match.group(1)) - median_ttft = float(ttft_match.group(1)) - return median_e2el, median_ttft - else: - raise ValueError("No benchmark result found") + # Start Benchmark + benchmark_script = os.path.join(benchmark_root, "benchmark_serving.py") + benchmark_cmd = [ + 'python3', + benchmark_script, + '--model', + benchmark_model_root, + '--tokenizer', + benchmark_model_root, + '--dataset-name', + 'random', + '--dataset-path', + shared_gpt_path, + '--random-input-len', + str(random_input_len), + '--random-output-len', + str(random_output_len), + '--random-prefix-len', + '0', + '--num-prompts', + str(num_prompts), + '--max-concurrency', + str(max_concurrency), + '--host', + server_host, + '--port', + str(server_port), + '--ignore-eos', + '--no-test-input', + '--percentile-metrics', + 'e2el,ttft', + ] + # warm up + if not skip_warmup: + check_call(benchmark_cmd, env=env) + output = check_output(benchmark_cmd, env=env) + e2el_pattern = r"Median E2EL \(ms\):\s*(\d+\.?\d*)" + ttft_pattern = r"Median TTFT \(ms\):\s*(\d+\.?\d*)" + e2el_match = re.search(e2el_pattern, output) + ttft_match = re.search(ttft_pattern, output) + if e2el_match and ttft_match: + median_e2el = float(e2el_match.group(1)) + median_ttft = float(ttft_match.group(1)) + return median_e2el, median_ttft + else: + raise ValueError("No benchmark result found") except Exception: - # Print outputs on error - logger.error("-------- Workers output --------") - with open('output_workers.log', 'r') as f: - logger.error(f.read()) - - logger.error("-------- Disagg server output --------") - with open('output_disagg.log', 'r') as f: - logger.error(f.read()) + logger.error("Benchmark test failed") raise finally: - server_proc.terminate() - workers_proc.terminate() - server_proc.wait() - workers_proc.wait() + terminate(*ctx_workers, *gen_workers, disagg_server) + shutil.rmtree(work_dir, ignore_errors=True) def get_config_for_benchmark(model_root, backend): serve_config = { "model": model_root, "hostname": "localhost", - "port": get_free_port_in_ci(), + "port": get_free_port(), "backend": "pytorch", "context_servers": { "num_instances": 1, @@ -1801,7 +1495,7 @@ def get_config_for_benchmark(model_root, backend): "backend": backend, "max_tokens_in_buffer": 512, }, - "urls": [f"localhost:{get_free_port_in_ci()}"] + "urls": [f"localhost:{get_free_port()}"] }, "generation_servers": { "num_instances": 1, @@ -1814,7 +1508,7 @@ def get_config_for_benchmark(model_root, backend): "backend": backend, "max_tokens_in_buffer": 512, }, - "urls": [f"localhost:{get_free_port_in_ci()}"] + "urls": [f"localhost:{get_free_port()}"] } } return serve_config @@ -1822,7 +1516,6 @@ def get_config_for_benchmark(model_root, backend): def run_disaggregated_aiperf(config_file, model_path, - num_ranks, server_start_timeout=1200, input_tokens=128, output_tokens=100, @@ -1841,7 +1534,6 @@ def run_disaggregated_aiperf(config_file, Args: config_file: Path to disaggregated server config YAML model_path: Path to model for tokenizer - num_ranks: Number of MPI ranks for workers server_start_timeout: Timeout in seconds for server startup input_tokens: Mean synthetic input tokens output_tokens: Mean output tokens to generate @@ -1856,107 +1548,87 @@ def run_disaggregated_aiperf(config_file, env: Environment variables dict cwd: Working directory """ + cleanup_output_files() run_env = env.copy() run_env["UCX_TLS"] = "^ib,gdr_copy" run_env["UCX_MM_ERROR_HANDLING"] = "y" - workers_cmd = [ - 'mpirun', '--allow-run-as-root', '--oversubscribe', '-n', - str(num_ranks), 'trtllm-serve', 'disaggregated_mpi_worker', '-c', - config_file - ] - - server_cmd = [ - 'trtllm-serve', 'disaggregated', '--server_start_timeout', - str(server_start_timeout), '-c', config_file - ] + config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ + setup_disagg_cluster(config_file, model_name=model_path, env=run_env, cwd=cwd, + server_start_timeout=server_start_timeout) + server_host = config.get("hostname", "localhost") artifact_dir = os.path.join(cwd or ".", "benchmark-results") - server_host, server_port = get_disagg_server_url_from_cfg(config_file) try: - with (open('output_workers.log', 'w') as output_workers, - popen(workers_cmd, - stdout=output_workers, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd) as workers_proc, open('output_disagg.log', 'w') as - output_disagg, - popen(server_cmd, - stdout=output_disagg, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd) as server_proc): - - # Wait for server to be ready - if not wait_for_server(server_host, - server_port, - timeout_seconds=server_start_timeout): - raise RuntimeError( - f"Disaggregated server did not become ready within {server_start_timeout} seconds" - ) - - # Build base command (using aiperf instead of genai-perf) - aiperf_cmd = [ - 'aiperf', 'profile', '--model', model_path, '--tokenizer', - model_path, '--endpoint-type', endpoint_type - ] - - # Add endpoint path based on type - if endpoint_type == 'chat': - aiperf_cmd.extend(['--endpoint', '/v1/chat/completions']) - - # Add streaming flag if enabled - if streaming: - aiperf_cmd.append('--streaming') - - # Add common parameters - aiperf_cmd.extend([ - '--url', f'{server_host}:{server_port}', - '--synthetic-input-tokens-mean', - str(input_tokens), '--synthetic-input-tokens-stddev', '0', - '--output-tokens-mean', - str(output_tokens), '--output-tokens-stddev', '0', - '--extra-inputs', f'max_tokens:{output_tokens}', - '--extra-inputs', f'min_tokens:{output_tokens}', - '--extra-inputs', 'ignore_eos:true', '--concurrency', - str(concurrency), '--warmup-request-count', - str(warmup_request_count) - ]) - - # Use request-count or num-dataset-entries - if request_count is not None: - aiperf_cmd.extend(['--request-count', str(request_count)]) - else: - # Default: use num-dataset-entries for compatibility - aiperf_cmd.extend(['--num-dataset-entries', '64']) - - aiperf_cmd.extend([ - '--random-seed', - str(random_seed), '--artifact-dir', artifact_dir - ]) + # Wait for server to be ready + if not wait_for_server( + server_host, server_port, timeout_seconds=server_start_timeout): + raise RuntimeError( + f"Disaggregated server did not become ready within {server_start_timeout} seconds" + ) + + # Build base command (using aiperf instead of genai-perf) + aiperf_cmd = [ + 'aiperf', 'profile', '--model', model_path, '--tokenizer', + model_path, '--endpoint-type', endpoint_type + ] - # Run aiperf - check_call(aiperf_cmd, - env=env, - poll_procs=[workers_proc, server_proc]) - - if accuracy_test: - accuracy_test_result, accuracy_value = run_accuracy_test( - model_path=model_path, - server_url=f"http://{server_host}:{server_port}", - concurrency=concurrency, - max_retries=3, - timeout=1200, - max_gen_toks=256, - max_length=4096) - - # only raise error if accuracy test passed and accuracy value is less than threshold - if accuracy_test_result and (accuracy_value < threshold): - raise AssertionError( - f"Accuracy test failed: accuracy value {accuracy_value} is less than test threshold {threshold}" - ) + # Add endpoint path based on type + if endpoint_type == 'chat': + aiperf_cmd.extend(['--endpoint', '/v1/chat/completions']) + + # Add streaming flag if enabled + if streaming: + aiperf_cmd.append('--streaming') + + # Add common parameters + aiperf_cmd.extend([ + '--url', f'{server_host}:{server_port}', + '--synthetic-input-tokens-mean', + str(input_tokens), '--synthetic-input-tokens-stddev', '0', + '--output-tokens-mean', + str(output_tokens), '--output-tokens-stddev', '0', '--extra-inputs', + f'max_tokens:{output_tokens}', '--extra-inputs', + f'min_tokens:{output_tokens}', '--extra-inputs', 'ignore_eos:true', + '--concurrency', + str(concurrency), '--warmup-request-count', + str(warmup_request_count) + ]) + + # Use request-count or num-dataset-entries + if request_count is not None: + aiperf_cmd.extend(['--request-count', str(request_count)]) + else: + # Default: use num-dataset-entries for compatibility + aiperf_cmd.extend(['--num-dataset-entries', '64']) + + aiperf_cmd.extend( + ['--random-seed', + str(random_seed), '--artifact-dir', artifact_dir]) + + # Run aiperf + all_worker_procs = [w.process for w in ctx_workers + gen_workers] + check_call(aiperf_cmd, + env=env, + poll_procs=all_worker_procs + [disagg_server.process]) + + if accuracy_test: + accuracy_test_result, accuracy_value = run_accuracy_test( + model_path=model_path, + server_url=f"http://{server_host}:{server_port}", + concurrency=concurrency, + max_retries=3, + timeout=1200, + max_gen_toks=256, + max_length=4096) + + # only raise error if accuracy test passed and accuracy value is less than threshold + if accuracy_test_result and (accuracy_value < threshold): + raise AssertionError( + f"Accuracy test failed: accuracy value {accuracy_value} is less than test threshold {threshold}" + ) except Exception: # Print outputs on error @@ -1981,10 +1653,8 @@ def run_disaggregated_aiperf(config_file, pass raise finally: - server_proc.terminate() - workers_proc.terminate() - server_proc.wait() - workers_proc.wait() + terminate(*ctx_workers, *gen_workers, disagg_server) + shutil.rmtree(work_dir, ignore_errors=True) def run_accuracy_test(model_path: str, server_url: str, concurrency: int, @@ -2115,6 +1785,7 @@ def test_disaggregated_benchmark_on_diff_backends( benchmark_model_root, shared_gpt_path, env=env, + model_path=benchmark_model_root, cwd=llm_venv.get_working_directory()) ucx_e2el, ucx_ttft = run_disaggregated_benchmark( disaggregated_example_root, @@ -2123,6 +1794,7 @@ def test_disaggregated_benchmark_on_diff_backends( benchmark_model_root, shared_gpt_path, env=env, + model_path=benchmark_model_root, cwd=llm_venv.get_working_directory()) print(f"Nixl E2EL: {nixl_e2el} ms, UCX E2EL: {ucx_e2el} ms") print(f"Nixl TTFT: {nixl_ttft} ms, UCX TTFT: {ucx_ttft} ms") @@ -2137,19 +1809,11 @@ def test_disaggregated_deepseek_v3_lite_bf16_empty_batch( disaggregated_example_root, llm_venv, benchmark_model_root, benchmark_root, shared_gpt_path): - src_dst_dict = { - benchmark_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/bf16", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, benchmark_model_root, "DeepSeek-V3-Lite/bf16") test_desc = "deepseek_v3_lite_bf16_empty_batch" - num_ranks, config_file = get_test_config(test_desc, - disaggregated_example_root, - os.path.dirname(__file__)) + config_file = get_test_config(test_desc, disaggregated_example_root, + os.path.dirname(__file__)) env = llm_venv._new_env.copy() e2el, ttft = run_disaggregated_benchmark( @@ -2159,13 +1823,13 @@ def test_disaggregated_deepseek_v3_lite_bf16_empty_batch( benchmark_model_root, shared_gpt_path, env=env, - cwd=llm_venv.get_working_directory(), - num_ranks=num_ranks, num_prompts=10, max_concurrency=10, random_input_len=384, random_output_len=1536, - skip_warmup=True) + skip_warmup=True, + model_path=benchmark_model_root, + cwd=llm_venv.get_working_directory()) print(f"E2EL: {e2el} ms, TTFT: {ttft} ms") assert e2el > 0 and ttft > 0 @@ -2187,21 +1851,14 @@ def test_llama4_long_context_kv_cache_overflow(disaggregated_test_root, llama4_model_root = os.path.join(models_root, model_path) # Create symlink to match config file path - src_dst_dict = { - llama4_model_root: f"{llm_venv.get_working_directory()}/{model_path}", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, llama4_model_root, model_path) - num_ranks, config_file = get_test_config("llama4_kv_cache_overflow", - disaggregated_example_root, - os.path.dirname(__file__)) + config_file = get_test_config("llama4_kv_cache_overflow", + disaggregated_example_root, + os.path.dirname(__file__)) run_disaggregated_aiperf(config_file=config_file, model_path=llama4_model_root, - num_ranks=num_ranks, server_start_timeout=1200, input_tokens=128000, output_tokens=100, @@ -2215,20 +1872,15 @@ def test_llama4_long_context_kv_cache_overflow(disaggregated_test_root, def test_disaggregated_deepseek_v3_lite_bf16_tllm_gen_helix( disaggregated_test_root, disaggregated_example_root, llm_venv, deepseek_v3_model_root): - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/bf16", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/bf16") run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_bf16_tllm_gen_helix", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), - prompt_file="long_prompts.json") + prompt_file="long_prompts.json", + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory()) @skip_pre_blackwell @@ -2238,17 +1890,12 @@ def test_disaggregated_gpt_oss_120b_harmony(disaggregated_test_root, disaggregated_example_root, llm_venv, model_path): model_dir = f"{llm_models_root()}/{model_path}" - src_dst_dict = { - model_dir: f"{llm_venv.get_working_directory()}/{model_path}", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, model_dir, model_path) run_disaggregated_test(disaggregated_example_root, "gpt_oss_120b_harmony", env=llm_venv._new_env, + model_path=model_dir, cwd=llm_venv.get_working_directory()) @@ -2279,21 +1926,13 @@ def test_disaggregated_stress_test(disaggregated_test_root, model_path = test_config.model_path test_desc = test_config.test_desc model_dir = f"{llm_models_root()}/{model_path}" - src_dst_dict = { - model_dir: f"{llm_venv.get_working_directory()}/{model_path}", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, model_dir, model_path) - num_ranks, config_file = get_test_config(test_desc, - disaggregated_example_root, - os.path.dirname(__file__)) + config_file = get_test_config(test_desc, disaggregated_example_root, + os.path.dirname(__file__)) run_disaggregated_aiperf(config_file=config_file, model_path=model_dir, - num_ranks=num_ranks, server_start_timeout=7200, input_tokens=input_tokens, output_tokens=output_tokens, @@ -2378,85 +2017,66 @@ async def run_bursts(): def run_disaggregated_cancel_test(example_dir, test_desc, env=None, - cwd=None, num_bursts=64, - requests_per_burst=64): + requests_per_burst=64, + server_start_timeout=1200, + model_path=None, + cwd=None): """Run disaggregated test with request cancellation stress test.""" cleanup_output_files() run_env = env.copy() run_env["UCX_TLS"] = "^ib,gdr_copy" - num_ranks, config_file = get_test_config(test_desc, example_dir, - os.path.dirname(__file__)) + config_file = get_test_config(test_desc, example_dir, + os.path.dirname(__file__)) + config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ + setup_disagg_cluster(config_file, model_name=model_path, env=run_env, cwd=cwd, + server_start_timeout=server_start_timeout) - workers_cmd = [ - 'mpirun', '--allow-run-as-root', '--oversubscribe', '-n', - str(num_ranks), 'trtllm-serve', 'disaggregated_mpi_worker', '-c', - config_file - ] - - server_start_timeout = 1200 - server_cmd = [ - 'trtllm-serve', 'disaggregated', '--server_start_timeout', - str(server_start_timeout), '-c', config_file - ] - server_host, server_port = get_disagg_server_url_from_cfg(config_file) + server_host = config.get("hostname", "localhost") server_url = f"http://{server_host}:{server_port}" try: - with (open('output_workers.log', 'w') as output_workers, - popen(workers_cmd, - stdout=output_workers, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd) as workers_proc, open('output_disagg.log', 'w') as - output_disagg, - popen(server_cmd, - stdout=output_disagg, - stderr=subprocess.STDOUT, - env=run_env, - cwd=cwd) as server_proc): - - # Wait for server to be ready - if not wait_for_server(server_host, - server_port, - timeout_seconds=server_start_timeout): - raise RuntimeError( - f"Disaggregated server did not become ready within {server_start_timeout} seconds" - ) - - # Run the cancel stress test - run_cancel_stress_test(server_url, - num_bursts=num_bursts, - requests_per_burst=requests_per_burst) - - # Verify server is still healthy after stress test by sending a normal request - client_dir = f"{example_dir}/clients" - client_cmd = [ - 'python3', f'{client_dir}/disagg_client.py', '-c', config_file, - '-p', f'{client_dir}/prompts.json', '--ignore-eos', - '--server-start-timeout', - str(server_start_timeout) - ] - check_call(client_cmd, - env=env, - poll_procs=[workers_proc, server_proc]) + # Wait for server to be ready + if not wait_for_server( + server_host, server_port, timeout_seconds=server_start_timeout): + raise RuntimeError( + f"Disaggregated server did not become ready within {server_start_timeout} seconds" + ) + + # Run the cancel stress test + run_cancel_stress_test(server_url, + num_bursts=num_bursts, + requests_per_burst=requests_per_burst) + + # Create a temporary client config with the correct dynamic port + client_config = config.copy() + client_config["port"] = server_port + client_config["hostname"] = server_host + temp_fd, client_config_file = tempfile.mkstemp(suffix='.yaml', + dir=work_dir) + with os.fdopen(temp_fd, 'w') as f: + yaml.dump(client_config, f) + + # Verify server is still healthy after stress test by sending a normal request + client_dir = f"{example_dir}/clients" + client_cmd = [ + 'python3', f'{client_dir}/disagg_client.py', '-c', + client_config_file, '-p', f'{client_dir}/prompts.json', + '--ignore-eos', '--server-start-timeout', + str(server_start_timeout) + ] + all_worker_procs = [w.process for w in ctx_workers + gen_workers] + check_call(client_cmd, + env=env, + poll_procs=all_worker_procs + [disagg_server.process]) except Exception: - logger.error("-------- Workers output --------") - with open('output_workers.log', 'r') as f: - logger.error(f.read()) - - logger.error("-------- Disagg server output --------") - with open('output_disagg.log', 'r') as f: - logger.error(f.read()) + logger.error("Cancel test failed") raise finally: - if 'server_proc' in locals() and 'workers_proc' in locals(): - server_proc.terminate() - workers_proc.terminate() - server_proc.wait() - workers_proc.wait() + terminate(*ctx_workers, *gen_workers, disagg_server) + shutil.rmtree(work_dir, ignore_errors=True) @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-bf16'], @@ -2471,21 +2091,16 @@ def test_disaggregated_cancel_large_context_requests(disaggregated_test_root, This test sends bursts of requests with large contexts and cancels them during prefill to stress test resource cleanup. """ - src_dst_dict = { - deepseek_v3_model_root: - f"{llm_venv.get_working_directory()}/DeepSeek-V3-Lite/bf16", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/bf16") run_disaggregated_cancel_test(disaggregated_example_root, "cancel_stress_test", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), num_bursts=5, - requests_per_burst=32) + requests_per_burst=32, + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory()) @pytest.mark.skip_less_device(8) @@ -2500,17 +2115,12 @@ def test_disaggregated_cancel_large_context_requests_long( during prefill to stress test resource cleanup. """ model_dir = f"{llm_models_root()}/{model_path}" - src_dst_dict = { - model_dir: f"{llm_venv.get_working_directory()}/{model_path}", - } - for src, dst in src_dst_dict.items(): - if not os.path.islink(dst): - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) + setup_model_symlink(llm_venv, model_dir, model_path) run_disaggregated_cancel_test(disaggregated_example_root, "cancel_stress_test_large", env=llm_venv._new_env, - cwd=llm_venv.get_working_directory(), num_bursts=1000, - requests_per_burst=32) + requests_per_burst=32, + model_path=model_dir, + cwd=llm_venv.get_working_directory()) diff --git a/tests/integration/defs/disaggregated/test_workers.py b/tests/integration/defs/disaggregated/test_workers.py index b1b537ec35d8..a1c4d3bf62fd 100644 --- a/tests/integration/defs/disaggregated/test_workers.py +++ b/tests/integration/defs/disaggregated/test_workers.py @@ -3,15 +3,18 @@ import copy import json import os -import subprocess -from typing import Generator, List, Optional, Tuple +import tempfile +from typing import List import aiohttp import pytest import yaml -from defs.common import revise_disagg_config_file_with_free_ports +from defs.common import get_free_port_in_ci as get_free_port from defs.conftest import skip_no_hopper -from defs.trt_test_alternative import popen +from disagg_test_utils import (HEARTBEAT_INTERVAL, INACTIVE_TIMEOUT, + run_ctx_worker, run_disagg_server, + run_gen_worker, terminate, + wait_for_disagg_server_ready) from transformers import AutoTokenizer from tensorrt_llm import logger @@ -23,45 +26,49 @@ block_key_hasher) -def get_ctx_gen_server_urls_from_cfg(config_file: str): - with open(config_file, 'r') as file: - config = yaml.safe_load(file) - ctx_servers = [] - gen_servers = [] - for server in config["context_servers"]["urls"]: - ctx_servers.append("http://" + server) - for server in config["generation_servers"]["urls"]: - gen_servers.append("http://" + server) - return ctx_servers, gen_servers - - -def run_disaggregated_workers( - config_file: str, - stdout=None, - env: Optional[dict] = None, - cwd: Optional[str] = None, - num_ranks: Optional[int] = None -) -> Tuple[Generator[subprocess.Popen, None, None], List[str], List[str]]: - - config_file = revise_disagg_config_file_with_free_ports(config_file) - ctx_servers, gen_servers = get_ctx_gen_server_urls_from_cfg(config_file) - - # TODO: auto detect num_ranks - assert num_ranks is not None - - # Start workers - workers_cmd = [ - 'mpirun', '--allow-run-as-root', '--oversubscribe', '-n', - str(num_ranks), 'trtllm-serve', 'disaggregated_mpi_worker', '-c', - config_file - ] - logger.info(f"Running workers with command: {' '.join(workers_cmd)}") - workers_proc = popen(workers_cmd, - stdout=stdout, - stderr=subprocess.STDOUT, - env=env, - cwd=cwd) - return workers_proc, ctx_servers, gen_servers +def build_worker_config(base_config, server_type_config, disagg_cluster): + """Build worker configuration by merging base config with server-type specific config. + + Args: + base_config: Full YAML config (top-level) + server_type_config: context_servers or generation_servers section + disagg_cluster: Service discovery config + + Returns: + dict: Worker configuration for trtllm-serve + """ + EXCLUDE_FROM_WORKER = { + 'hostname', + 'port', + 'num_instances', + 'urls', + 'router', + 'model', + 'context_servers', + 'generation_servers', + 'conditional_disagg_config', + } + + worker_config = { + k: v + for k, v in base_config.items() if k not in EXCLUDE_FROM_WORKER + } + + worker_config.update({ + k: v + for k, v in server_type_config.items() if k not in EXCLUDE_FROM_WORKER + }) + + if 'free_gpu_memory_fraction' in worker_config: + frac = worker_config.pop('free_gpu_memory_fraction') + if 'kv_cache_config' not in worker_config: + worker_config['kv_cache_config'] = {} + worker_config['kv_cache_config'].setdefault('free_gpu_memory_fraction', + frac) + + worker_config['disagg_cluster'] = disagg_cluster + + return worker_config DEFAULT_TIMEOUT_SERVER_START = 900 @@ -512,27 +519,100 @@ def load_default_prompts(disaggregated_example_root: str): @contextlib.contextmanager -def background_workers(llm_venv, config_file: str, num_ranks: int = None): +def background_workers(llm_venv, config_file: str): cwd = llm_venv.get_working_directory() os.chdir(cwd) - with open(os.path.join(cwd, 'output_workers.log'), 'w+') as log_file: - workers_proc, ctx_servers, gen_servers = run_disaggregated_workers( - config_file=config_file, - stdout=log_file, - env=llm_venv._new_env, - cwd=cwd, - num_ranks=num_ranks) - try: - with workers_proc as proc: - yield ctx_servers, gen_servers - except Exception: - log_file.seek(0) - logger.error("-------- Worker output --------") - logger.error(log_file.read()) - raise - finally: - proc.terminate() - proc.wait() + env = llm_venv._new_env + + with open(config_file, 'r') as f: + config = yaml.safe_load(f) + + model = config.get("model") + ctx_server_cfg = config.get("context_servers", {}) + gen_server_cfg = config.get("generation_servers", {}) + num_ctx = ctx_server_cfg.get("num_instances", 1) + num_gen = gen_server_cfg.get("num_instances", 1) + + disagg_port = get_free_port() + work_dir = tempfile.mkdtemp() + disagg_cluster = { + "cluster_uri": f"http://localhost:{disagg_port}", + "cluster_name": "test_cluster", + "heartbeat_interval_sec": HEARTBEAT_INTERVAL, + "inactive_timeout_sec": INACTIVE_TIMEOUT, + "minimal_instances": { + "context_servers": num_ctx, + "generation_servers": num_gen, + }, + } + + ctx_worker_config = build_worker_config(config, ctx_server_cfg, + disagg_cluster) + gen_worker_config = build_worker_config(config, gen_server_cfg, + disagg_cluster) + + gpus_per_ctx = (ctx_server_cfg.get("tensor_parallel_size", 1) * + ctx_server_cfg.get("pipeline_parallel_size", 1)) + gpus_per_gen = (gen_server_cfg.get("tensor_parallel_size", 1) * + gen_server_cfg.get("pipeline_parallel_size", 1)) + + ctx_workers = [] + gen_workers = [] + ctx_urls = [] + gen_urls = [] + next_device = 0 + + import torch + num_gpus = torch.cuda.device_count() + + for i in range(num_ctx): + port = get_free_port() + ctx_urls.append(f"http://localhost:{port}") + ctx_workers.append( + run_ctx_worker(model, + ctx_worker_config, + work_dir, + port=port, + device=next_device % num_gpus, + env=env)) + next_device += gpus_per_ctx + + for i in range(num_gen): + port = get_free_port() + gen_urls.append(f"http://localhost:{port}") + gen_workers.append( + run_gen_worker(model, + gen_worker_config, + work_dir, + port=port, + device=next_device % num_gpus, + env=env)) + next_device += gpus_per_gen + + server_config = { + "hostname": "localhost", + "port": disagg_port, + "disagg_cluster": disagg_cluster, + "context_servers": { + "router": ctx_server_cfg.get("router", {}) + }, + "generation_servers": { + "router": gen_server_cfg.get("router", {}) + }, + } + disagg_server = run_disagg_server(server_config, + work_dir, + disagg_port, + env=env) + + try: + asyncio.run(wait_for_disagg_server_ready(disagg_port)) + yield ctx_urls, gen_urls + except Exception: + logger.error("-------- Service discovery workers error --------") + raise + finally: + terminate(*ctx_workers, *gen_workers, disagg_server) @pytest.mark.skip(reason="https://nvbugs/5372970") @@ -545,8 +625,8 @@ def test_workers_conditional_disaggregation(disaggregated_test_root, 'test_configs/disagg_config_cache_reuse.yaml') prepare_llama_model(llama_model_root, llm_venv) - with background_workers(llm_venv, config_file, - 2) as (ctx_servers, gen_servers): + with background_workers(llm_venv, + config_file) as (ctx_servers, gen_servers): tester = ConditionalWorkerTester(ctx_servers, gen_servers) prompts = load_default_prompts(disaggregated_example_root) asyncio.run(tester.test_multi_round_request(prompts)) @@ -569,8 +649,8 @@ def test_workers_conditional_disaggregation_deepseek_v3_lite_bf16( os.makedirs(os.path.dirname(dst), exist_ok=True) os.symlink(src, dst, target_is_directory=True) - with background_workers(llm_venv, config_file, - 2) as (ctx_servers, gen_servers): + with background_workers(llm_venv, + config_file) as (ctx_servers, gen_servers): tester = ConditionalWorkerTester(ctx_servers, gen_servers) prompts = load_default_prompts(disaggregated_example_root) asyncio.run(tester.test_multi_round_request(prompts)) @@ -585,8 +665,8 @@ def test_workers_kv_cache_events(disaggregated_test_root, 'test_configs/disagg_config_cache_reuse.yaml') prepare_llama_model(llama_model_root, llm_venv) - with background_workers(llm_venv, config_file, - 2) as (ctx_servers, gen_servers): + with background_workers(llm_venv, + config_file) as (ctx_servers, gen_servers): tester = KvCacheEventWorkerTester(ctx_servers, gen_servers) prompts = load_default_prompts(disaggregated_example_root) asyncio.run(tester.test_multi_round_request(prompts, 6)) @@ -602,8 +682,8 @@ def test_workers_kv_cache_aware_router(disaggregated_test_root, 'test_configs/disagg_config_cache_aware_balance.yaml') prepare_llama_model(llama_model_root, llm_venv) - with background_workers(llm_venv, config_file, - 4) as (ctx_servers, gen_servers): + with background_workers(llm_venv, + config_file) as (ctx_servers, gen_servers): tester = KvCacheAwareRouterTester(ctx_servers, gen_servers) prompts = load_default_prompts(disaggregated_example_root) asyncio.run(tester.test_multi_round_request(prompts, 16, 4)) @@ -627,8 +707,8 @@ def test_workers_kv_cache_aware_router_deepseek_v3_lite_bf16( os.makedirs(os.path.dirname(dst), exist_ok=True) os.symlink(src, dst, target_is_directory=True) - with background_workers(llm_venv, config_file, - 4) as (ctx_servers, gen_servers): + with background_workers(llm_venv, + config_file) as (ctx_servers, gen_servers): tester = KvCacheAwareRouterTester(ctx_servers, gen_servers, model_name="DeepSeek-V3-Lite/bf16", @@ -646,7 +726,7 @@ def test_workers_kv_cache_aware_router_eviction(disaggregated_test_root, 'test_configs/disagg_config_cache_reuse.yaml') prepare_llama_model(llama_model_root, llm_venv) - with background_workers(llm_venv, config_file, - 2) as (ctx_servers, gen_servers): + with background_workers(llm_venv, + config_file) as (ctx_servers, gen_servers): tester = KvCacheAwareRouterTester(ctx_servers, gen_servers) asyncio.run(tester.test_eviction()) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 322da003a0e9..934103b179a4 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -457,9 +457,9 @@ examples/serve/test_serve_negative.py::test_missing_content_type_header examples/serve/test_serve_negative.py::test_extremely_large_batch # e2e disaggregated serving test -disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] -disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] -disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun_trt_backend[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_trt_backend[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_cuda_graph[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_mpi[DeepSeek-V3-Lite-fp8] disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_ucx[DeepSeek-V3-Lite-fp8] diff --git a/tests/integration/test_lists/qa/llm_function_rtx6k.txt b/tests/integration/test_lists/qa/llm_function_rtx6k.txt index fdf4e0658744..759b43bf057f 100644 --- a/tests/integration/test_lists/qa/llm_function_rtx6k.txt +++ b/tests/integration/test_lists/qa/llm_function_rtx6k.txt @@ -269,8 +269,8 @@ test_e2e.py::test_eagle3_output_repetition_4gpus[llama4-models/nvidia/Llama-4-Ma test_e2e.py::test_eagle3_output_repetition_4gpus[Qwen3/saved_models_Qwen3-235B-A22B_nvfp4_hf-Qwen3/qwen3-235B-eagle3] -disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] -disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_cuda_graph[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_load_balance[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_cache_aware_balance[TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 01b2f0748d82..8615e9e29a18 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -42,8 +42,8 @@ l0_a10: - unittest/disaggregated/test_extractor_v2.py - unittest/disaggregated/test_peer.py - unittest/disaggregated/region/test_block.py - - disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] - - disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun_trt_backend[TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_trt_backend[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_cuda_graph[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_overlap[TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index b134fc11fc81..021f130988ac 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -113,8 +113,8 @@ l0_dgx_h100: - test_e2e.py::test_trtllm_bench_llmapi_launch[pytorch_backend-llama-v3-llama3-8b] # ------------- Disaggregated serving tests --------------- - unittest/disaggregated/test_py_cache_transceiver_mp.py - - disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] - - disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu_with_mpirun_trt_backend[TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu_trt_backend[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_gentp2[TinyLlama-1.1B-Chat-v1.0]