From 8c2dec63ebbe01bd233e083135717f4e90c6bc1c Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 16 Jul 2026 09:19:22 -0700 Subject: [PATCH 1/9] [https://nvbugs/6465993][test] unwaive Nemotron 3 Super disaggregated tests Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index cef79620dc6f..4ceb7518d7c9 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -17,9 +17,6 @@ accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symm accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[GSM8K-tp2pp2] SKIP (https://nvbugs/6428069) accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[MMLU-tp1pp2] SKIP (https://nvbugs/6427411) accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[MMLU-tp2pp2] SKIP (https://nvbugs/6428069) -accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=False] SKIP (https://nvbugs/6465993) -accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=False] SKIP (https://nvbugs/6465993) -accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_ctx_dp2_gen_tp4 SKIP (https://nvbugs/6465993) accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=False] SKIP (https://nvbugs/6427411) accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=True] SKIP (https://nvbugs/6402054) accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664) From 307cc47f247afcea91e4bc09702e7f35cf2bf578 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:26:20 -0700 Subject: [PATCH 2/9] [https://nvbugs/6465993][test] disable session reuse for diagnosis Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/integration/defs/pytest.ini | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integration/defs/pytest.ini b/tests/integration/defs/pytest.ini index ee5e97c00f74..0c53ba744e7c 100644 --- a/tests/integration/defs/pytest.ini +++ b/tests/integration/defs/pytest.ini @@ -9,6 +9,8 @@ junit_family=legacy addopts = --ignore-glob="*perf/test_perf.py" --ignore-glob="*perf/disagg/*" --ignore-glob="*test_list_validation.py" --ignore-glob="*llm-test-workspace*" --durations=0 -W ignore::DeprecationWarning --unused-fixtures -p test_common.session_reuse_hooks pythonpath = ../../../examples/auto_deploy ../../ +env = + TRTLLM_TEST_REUSE_SESSION=0 norecursedirs = ./triton/perf ./perf/disagg unused_fixtures_ignore_path = cpp From 88956d597fc812c8a49532050a838de360bffdeb Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 16 Jul 2026 22:32:50 -0700 Subject: [PATCH 3/9] [NVBUG 6465993][test] add temporary disagg timeout diagnostics Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../accuracy/test_disaggregated_serving.py | 109 +++++++++++++++++- tests/integration/defs/pytest.ini | 2 +- 2 files changed, 105 insertions(+), 6 deletions(-) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index cd63c052375c..48770a2a3c38 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1,3 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. import concurrent import contextlib import functools @@ -157,7 +171,22 @@ def launch_disaggregated_llm( gen_extra_env: Optional[Dict[str, str]] = None, request_timeout_s: float = DEFAULT_REQUEST_TIMEOUT_S, request_max_retries: Optional[int] = None, + diagnostic_label: Optional[str] = None, ): + diagnostic_started_at = time.monotonic() + + def _diagnostic(phase: str, **details: Any) -> None: + if diagnostic_label is None: + return + details_text = " ".join(f"{key}={value!r}" + for key, value in sorted(details.items())) + print( + f"[DISAGG-DIAG] label={diagnostic_label!r} " + f"elapsed_s={time.monotonic() - diagnostic_started_at:.3f} " + f"phase={phase!r} {details_text}", + flush=True) + + _diagnostic("configuration_start", model=model_name) temp_dir = tempfile.TemporaryDirectory() disaggregated_serving_config_path = os.path.join( temp_dir.name, "disaggregated_serving_config.yaml") @@ -272,6 +301,8 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): base_env = os.environ.copy() if extra_env: base_env.update(extra_env) + if diagnostic_label is not None: + base_env["PYTHONUNBUFFERED"] = "1" kv_cache_perf_dir = os.path.join(temp_dir.name, "kv_cache_perf") @@ -350,6 +381,12 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): log_files = [] try: for i, (env, args) in enumerate(server_configs): + _diagnostic( + "process_launch_start", + role=server_name, + index=i, + cuda_visible_devices=env.get("CUDA_VISIBLE_DEVICES"), + command=" ".join(args)) if enable_redirect_log: f = open(f"output_{server_name}_{i}.log", "w+") proc = popen(args, env=env, stdout=f, stderr=f) @@ -362,6 +399,9 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): opened_processes = [ stack.enter_context(proc) for proc in processes ] + _diagnostic("process_group_started", + role=server_name, + pids=[process.pid for process in opened_processes]) yield opened_processes for f in log_files: f.close() @@ -376,6 +416,7 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): disaggregated_serving_config_path, "--server_start_timeout", str(server_waiting_timeout), "-r", "360000" ] + _diagnostic("process_groups_start") with ( MyThreadPoolExecutor(max_workers=max_workers) as thread_pool, temp_dir, @@ -383,13 +424,23 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): multi_popen(gen_servers, "gen") as gen_processes, multi_popen([(base_env, server_cmd)], "disagg") as server_processes, ): + _diagnostic( + "all_processes_started", + context_pids=[process.pid for process in ctx_processes], + generation_pids=[process.pid for process in gen_processes], + disaggregated_pids=[process.pid for process in server_processes]) start_time = time.time() server_is_ready = False + registration_poll = 0 while time.time() - start_time < server_waiting_timeout: time.sleep(5) + registration_poll += 1 for process in itertools.chain(ctx_processes, gen_processes, server_processes): if process.poll() is not None: + _diagnostic("process_exited_before_registration", + pid=process.pid, + returncode=process.returncode) raise Exception( f"process {process.pid} exited with code {process.returncode}" ) @@ -399,13 +450,30 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): f"http://localhost:{serve_port}/cluster_info") if response.status_code == 200: cluster_info = response.json() + if registration_poll == 1 or registration_poll % 6 == 0: + _diagnostic("service_registration_poll", + poll=registration_poll, + cluster_info=cluster_info) if cluster_info.get("is_ready"): print(f"Cluster ready: {cluster_info}") + _diagnostic("service_registration_complete", + poll=registration_poll, + cluster_info=cluster_info) server_is_ready = True break - except requests.exceptions.ConnectionError: + elif registration_poll == 1 or registration_poll % 6 == 0: + _diagnostic("service_registration_http_status", + poll=registration_poll, + status_code=response.status_code) + except requests.exceptions.ConnectionError as error: + if registration_poll == 1 or registration_poll % 6 == 0: + _diagnostic("service_registration_unreachable", + poll=registration_poll, + error=type(error).__name__) continue if not server_is_ready: + _diagnostic("service_registration_timeout", + timeout_s=server_waiting_timeout) pytest.fail( f"Server is not ready after {server_waiting_timeout} seconds. Please check the logs for more details." ) @@ -418,9 +486,15 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): if request_max_retries is not None: client_kwargs["max_retries"] = request_max_retries client = openai.OpenAI(**client_kwargs) + _diagnostic("client_created", base_url=client_kwargs["base_url"]) + request_counter = itertools.count(1) def send_request(prompt: str, sampling_params: SamplingParams, - streaming: bool): + streaming: bool, request_index: int): + if request_index == 1: + _diagnostic("first_inference_request_start", + prompt_length=len(prompt), + streaming=streaming) kwargs = {} if sampling_params is not None: extra_body = {} @@ -460,6 +534,9 @@ def send_request(prompt: str, sampling_params: SamplingParams, prompt=prompt, stream=streaming, **kwargs) + if request_index == 1: + _diagnostic("first_inference_request_complete", + choice_count=len(response.choices)) result = Result(id=0, sampling_params=sampling_params, outputs=[ @@ -474,8 +551,13 @@ def send_request(prompt: str, sampling_params: SamplingParams, def generate_async(prompt: str, sampling_params: Optional[SamplingParams] = None, streaming: bool = False): + request_index = next(request_counter) + if request_index == 1: + _diagnostic("first_inference_request_submit", + prompt_length=len(prompt), + streaming=streaming) future = thread_pool.submit(send_request, prompt, sampling_params, - streaming) + streaming, request_index) thread_pool.futures.append(future) return future @@ -509,10 +591,18 @@ def _show_kvcache_time(kv_cache_perf_dir, max_lines=100): for line in f.readlines()[-max_lines:]: print(line.strip()) + _diagnostic("tokenizer_load_start") tokenizer = load_hf_tokenizer(model_name) + _diagnostic("tokenizer_load_complete") + accuracy_body_completed = False try: + _diagnostic("accuracy_body_start") yield DuckLLM(args, tokenizer, generate_async) + accuracy_body_completed = True + _diagnostic("accuracy_body_complete") finally: + _diagnostic("shutdown_start", + accuracy_body_completed=accuracy_body_completed) if enable_perf: _show_kvcache_time(kv_cache_perf_dir) _get_perf_metrics() @@ -540,6 +630,9 @@ def _show_kvcache_time(kv_cache_perf_dir, max_lines=100): pass # already exited between timeout and kill except OSError: pass # process already gone + _diagnostic( + "shutdown_complete", + returncodes=[process.poll() for process in all_processes]) def run_parallel_test(model_name: str, @@ -2409,8 +2502,14 @@ def test_auto_dtype(self, use_py_transceiver, block_reuse, mtp_nextn): if block_reuse: ctx_cfg["kv_cache_config"]["enable_block_reuse"] = True gen_cfg["kv_cache_config"]["enable_block_reuse"] = True - with launch_disaggregated_llm(disagg_cfg, ctx_cfg, gen_cfg, - self.MODEL_PATH) as llm: + diagnostic_label = None + if not use_py_transceiver and block_reuse and mtp_nextn == 3: + diagnostic_label = "nemotron3-super-mtp3-block-reuse" + with launch_disaggregated_llm(disagg_cfg, + ctx_cfg, + gen_cfg, + self.MODEL_PATH, + diagnostic_label=diagnostic_label) as llm: run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) @pytest.mark.skip_less_device(8) diff --git a/tests/integration/defs/pytest.ini b/tests/integration/defs/pytest.ini index 0c53ba744e7c..4d4e07aee089 100644 --- a/tests/integration/defs/pytest.ini +++ b/tests/integration/defs/pytest.ini @@ -6,7 +6,7 @@ threadleak = True # outlive the test they start under. threadleak_exclude = asyncio_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+ junit_family=legacy -addopts = --ignore-glob="*perf/test_perf.py" --ignore-glob="*perf/disagg/*" --ignore-glob="*test_list_validation.py" --ignore-glob="*llm-test-workspace*" --durations=0 -W ignore::DeprecationWarning --unused-fixtures -p test_common.session_reuse_hooks +addopts = --ignore-glob="*perf/test_perf.py" --ignore-glob="*perf/disagg/*" --ignore-glob="*test_list_validation.py" --ignore-glob="*llm-test-workspace*" --durations=0 -W ignore::DeprecationWarning --unused-fixtures -p test_common.session_reuse_hooks -s pythonpath = ../../../examples/auto_deploy ../../ env = From ae69749d17bc39c525a4c7a0200619a30064e13f Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 17 Jul 2026 08:54:05 -0700 Subject: [PATCH 4/9] [NVBUG 6465993][test] trace fixture setup for disagg timeout Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/integration/defs/conftest.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index bba88cf07c01..4db99802a92f 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -2343,6 +2343,31 @@ def pytest_runtest_makereport(item, call): report.url = "" +_NEMOTRON_DISAGG_DIAGNOSTIC_TEST = ( + "accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::" + "test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=False]") + + +@pytest.hookimpl(tryfirst=True, hookwrapper=True) +def pytest_fixture_setup(fixturedef, request): + item = getattr(request, "_pyfuncitem", None) + nodeid = getattr(item, "nodeid", "") + if not nodeid.endswith(_NEMOTRON_DISAGG_DIAGNOSTIC_TEST): + yield + return + + print( + f"[DISAGG-DIAG] phase='fixture_setup_start' " + f"fixture={fixturedef.argname!r} scope={fixturedef.scope!r}", + flush=True) + outcome = yield + status = "error" if outcome.excinfo else "complete" + print( + f"[DISAGG-DIAG] phase='fixture_setup_finish' " + f"fixture={fixturedef.argname!r} status={status!r}", + flush=True) + + @pytest.fixture(scope="session") def all_pytest_items(): """ From 2c060e635cb734c00897ab70be9dfe8f4ba47f3a Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 17 Jul 2026 11:11:44 -0700 Subject: [PATCH 5/9] [NVBUG 6465993][test] expose disagg fixture diagnostics Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/integration/defs/conftest.py | 38 +++++++++++++++++++++++------- 1 file changed, 29 insertions(+), 9 deletions(-) diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 4db99802a92f..099413498460 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -2346,6 +2346,16 @@ def pytest_runtest_makereport(item, call): _NEMOTRON_DISAGG_DIAGNOSTIC_TEST = ( "accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::" "test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=False]") +_NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE = "_nemotron_disagg_diagnostic_fd" + + +def _emit_nemotron_disagg_diagnostic(item, message: str) -> None: + diagnostic_fd = getattr(item, _NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE, + None) + if diagnostic_fd is None: + print(f"[DISAGG-DIAG] {message}", flush=True) + return + os.write(diagnostic_fd, f"[DISAGG-DIAG] {message}\n".encode()) @pytest.hookimpl(tryfirst=True, hookwrapper=True) @@ -2356,16 +2366,14 @@ def pytest_fixture_setup(fixturedef, request): yield return - print( - f"[DISAGG-DIAG] phase='fixture_setup_start' " - f"fixture={fixturedef.argname!r} scope={fixturedef.scope!r}", - flush=True) + _emit_nemotron_disagg_diagnostic( + item, f"phase='fixture_setup_start' fixture={fixturedef.argname!r} " + f"scope={fixturedef.scope!r}") outcome = yield status = "error" if outcome.excinfo else "complete" - print( - f"[DISAGG-DIAG] phase='fixture_setup_finish' " - f"fixture={fixturedef.argname!r} status={status!r}", - flush=True) + _emit_nemotron_disagg_diagnostic( + item, f"phase='fixture_setup_finish' fixture={fixturedef.argname!r} " + f"status={status!r}") @pytest.fixture(scope="session") @@ -2599,7 +2607,19 @@ def collect_status(item: pytest.Item): @pytest.hookimpl(wrapper=True) def pytest_runtest_protocol(item, nextitem): - ret = yield + diagnostic = item.nodeid.endswith(_NEMOTRON_DISAGG_DIAGNOSTIC_TEST) + if diagnostic: + diagnostic_fd = os.dup(1) + setattr(item, _NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE, diagnostic_fd) + _emit_nemotron_disagg_diagnostic(item, "phase='runtest_protocol_start'") + try: + ret = yield + finally: + if diagnostic: + _emit_nemotron_disagg_diagnostic(item, + "phase='runtest_protocol_finish'") + os.close(diagnostic_fd) + delattr(item, _NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE) collect_status(item) return ret From 51ff2ca06e94a3371189d5a7ab99512aea7cac7b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 17 Jul 2026 17:22:30 -0700 Subject: [PATCH 6/9] [NVBUG 6465993][test] stream disagg diagnostic output Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/integration/defs/conftest.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 099413498460..1263a668a776 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -2609,6 +2609,11 @@ def collect_status(item: pytest.Item): def pytest_runtest_protocol(item, nextitem): diagnostic = item.nodeid.endswith(_NEMOTRON_DISAGG_DIAGNOSTIC_TEST) if diagnostic: + upload_log_plugin = item.config.pluginmanager.getplugin( + "upload_log_plugin") + if upload_log_plugin is not None: + original_echo_to_stdout = upload_log_plugin.echo_to_stdout + upload_log_plugin.echo_to_stdout = True diagnostic_fd = os.dup(1) setattr(item, _NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE, diagnostic_fd) _emit_nemotron_disagg_diagnostic(item, "phase='runtest_protocol_start'") @@ -2620,6 +2625,8 @@ def pytest_runtest_protocol(item, nextitem): "phase='runtest_protocol_finish'") os.close(diagnostic_fd) delattr(item, _NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE) + if upload_log_plugin is not None: + upload_log_plugin.echo_to_stdout = original_echo_to_stdout collect_status(item) return ret From 80ef1530698856672a2a15d49e91481fbd7dca07 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 17 Jul 2026 19:46:15 -0700 Subject: [PATCH 7/9] [NVBUG 6465993][test] disable worker monitor for diagnosis Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/executor/worker.py | 16 +++++++++++++--- tests/integration/defs/pytest.ini | 1 + 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 5958d296e1c4..fcb3563f8021 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -311,8 +311,15 @@ def notify_proxy_threads_to_quit(): # error to the error_queue in the main thread. mpi_comm().barrier() - worker_process_identities = mpi_comm().allgather( - capture_worker_process_identity(mpi_rank())) + disable_worker_process_monitor = os.getenv( + "TRTLLM_TEST_DISABLE_WORKER_PROCESS_MONITOR", "0") == "1" + if disable_worker_process_monitor: + worker_process_identities = None + logger.warning("Worker process monitoring disabled by " + "TRTLLM_TEST_DISABLE_WORKER_PROCESS_MONITOR") + else: + worker_process_identities = mpi_comm().allgather( + capture_worker_process_identity(mpi_rank())) logger_debug(f"Worker {mpi_rank()} ready to setup backend...\n", "green") try: @@ -353,7 +360,10 @@ def notify_proxy_threads_to_quit(): worker.set_result_queue(result_queue) # Send ready signal with confirmation - ready_msg = (ready_signal, None, worker_process_identities) + if worker_process_identities is None: + ready_msg = (ready_signal, None) + else: + ready_msg = (ready_signal, None, worker_process_identities) if not worker_init_status_queue.notify_with_retry(ready_msg): logger.warning( "Failed to deliver ready signal to proxy, continuing anyway" diff --git a/tests/integration/defs/pytest.ini b/tests/integration/defs/pytest.ini index 4d4e07aee089..87fc31800ff2 100644 --- a/tests/integration/defs/pytest.ini +++ b/tests/integration/defs/pytest.ini @@ -11,6 +11,7 @@ pythonpath = ../../../examples/auto_deploy ../../ env = TRTLLM_TEST_REUSE_SESSION=0 + TRTLLM_TEST_DISABLE_WORKER_PROCESS_MONITOR=1 norecursedirs = ./triton/perf ./perf/disagg unused_fixtures_ignore_path = cpp From 0a00237298b5c430944cfb4020ea28ded567656d Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:17:59 -0700 Subject: [PATCH 8/9] [NVBUG 6465993][fix] use attention dtype for KV transfer buffers Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheFormatter.cpp | 7 +- .../batch_manager/cacheTransBuffer.cpp | 147 ++++++++++++------ .../batch_manager/cacheTransBuffer.h | 8 +- .../batch_manager/kvCacheManagerTest.cpp | 75 +++++++-- 4 files changed, 172 insertions(+), 65 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 9dc44531d409..3930c462f34b 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -729,12 +729,7 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess { NVTX3_SCOPED_RANGE(formatInputRecvBuffer); - // TODO(disagg-multi-dtype): pool 0's dtype is treated as canonical for the wire - // transport here. Pools with differing dtypes are rejected up-front in - // CacheTransBufferManager's constructor (see cacheTransBuffer.cpp). When - // per-pool dtype dispatch lands, this single dataType variable must be replaced - // with a per-pool lookup keyed by the source pool of each block. - auto dataType = mCacheManager->getPrimaryPool(0)->getDataType(); + auto const dataType = mCacheTransBufferManager->getDataType(); bool layerWise = common::getEnvDisaggLayerwise() && numKvPools == 1; if (layerWise) { diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index 772c9555f0f3..59d88cbdad42 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -27,6 +27,72 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager { +namespace +{ + +bool isCachePool(BlockManager const& blockManager, SizeType32 poolIdx) +{ + auto const& pool = blockManager.getPool(poolIdx); + return !pool.containsBlockScales && !pool.containsIndexerKCache; +} + +bool isAttentionCachePool(BlockManager const& blockManager, SizeType32 poolIdx) +{ + return isCachePool(blockManager, poolIdx) + && !LinearAttentionMetadata::hasLinearCache(blockManager.getPoolWindowSize(poolIdx)); +} + +nvinfer1::DataType getTransferDataType(KVCacheManager::BaseKVCacheManager* cacheManager, bool transferIndexerKCache) +{ + TLLM_CHECK(cacheManager); + if (transferIndexerKCache) + { + auto const indexerKCachePool = cacheManager->getIndexerKCachePool(); + TLLM_CHECK(indexerKCachePool); + return indexerKCachePool->getDataType(); + } + + auto const& blockManager = cacheManager->getBlockManager(); + std::optional cacheDataType; + std::optional attentionDataType; + SizeType32 firstPoolIdx = -1; + // Recurrent-state pools have a separate transfer manager and formatter. Only + // attention pools determine the KV transfer-buffer dtype. + for (SizeType32 poolIdx = 0; poolIdx < blockManager.getNumPools(); ++poolIdx) + { + if (!isCachePool(blockManager, poolIdx)) + { + continue; + } + + auto const poolDataType = blockManager.getPrimaryPool(poolIdx)->getDataType(); + if (!cacheDataType.has_value()) + { + cacheDataType = poolDataType; + } + if (!isAttentionCachePool(blockManager, poolIdx)) + { + continue; + } + if (!attentionDataType.has_value()) + { + attentionDataType = poolDataType; + firstPoolIdx = poolIdx; + continue; + } + + TLLM_CHECK_WITH_INFO(poolDataType == attentionDataType.value(), + "Disaggregated KV cache transfer does not yet support attention pools with differing dtypes " + "(pool %d dtype=%d, pool %d dtype=%d). TODO(disagg-multi-dtype): per-pool dtype dispatch in formatter.", + firstPoolIdx, static_cast(attentionDataType.value()), poolIdx, static_cast(poolDataType)); + } + + TLLM_CHECK_WITH_INFO(cacheDataType.has_value(), "Disaggregated KV cache transfer requires a cache pool"); + return attentionDataType.value_or(cacheDataType.value()); +} + +} // namespace + // ============================================================================ // FabricMemory Implementation // ============================================================================ @@ -194,39 +260,38 @@ bool FabricMemory::supportFabricMemory() size_t CacheTransBufferManager::computeTransferBufferSize( KVCacheManager::BaseKVCacheManager* cacheManager, std::optional maxNumTokens, bool transferIndexerKCache) { - nvinfer1::DataType dataType; - if (transferIndexerKCache) - { - dataType = cacheManager->getIndexerKCachePool()->getDataType(); - } - else + auto const dataType = getTransferDataType(cacheManager, transferIndexerKCache); + + auto const& blockManager = cacheManager->getBlockManager(); + auto const tokensPerBlock = blockManager.getTokensPerBlock(); + bool hasAttentionCachePool = false; + for (SizeType32 poolIdx = 0; poolIdx < blockManager.getNumPools(); ++poolIdx) { - dataType = cacheManager->getPrimaryPool(0)->getDataType(); + hasAttentionCachePool |= isAttentionCachePool(blockManager, poolIdx); } - - auto tokensPerBlock = cacheManager->getBlockManager().getTokensPerBlock(); size_t bufferSizeFromMaxNumToken = 0; if (maxNumTokens.has_value()) { TLLM_CHECK(maxNumTokens.value() % tokensPerBlock == 0); - auto dataSize = common::getDTypeSize(dataType); - SizeType32 kvCacheByteSizePerTokenPerLayer = 0; + auto const dataSize = common::getDTypeSize(dataType); + SizeType32 indexerCacheByteSizePerTokenPerLayer = 0; if (transferIndexerKCache) { - kvCacheByteSizePerTokenPerLayer + indexerCacheByteSizePerTokenPerLayer = cacheManager->getIndexerKCachePool()->getDimension<-1>() * dataSize / tokensPerBlock; } - else - { - auto primaryPool = cacheManager->getPrimaryPool(0); - kvCacheByteSizePerTokenPerLayer - = primaryPool->getDimension<-1>() * primaryPool->getDimension<2>() * dataSize / tokensPerBlock; - } - for (auto layerId = 0; layerId < cacheManager->getBlockManager().getNumLayers(); layerId++) + for (auto layerId = 0; layerId < blockManager.getNumLayers(); layerId++) { - auto poolIdx = cacheManager->getBlockManager().getLayerPoolIdx(layerId); - auto windowSize = static_cast(cacheManager->getBlockManager().getPoolWindowSize(poolIdx)); + auto const poolIdx = blockManager.getLayerPoolIdx(layerId); + auto const encodedWindowSize = blockManager.getPoolWindowSize(poolIdx); + if (!transferIndexerKCache && hasAttentionCachePool + && LinearAttentionMetadata::hasLinearCache(encodedWindowSize)) + { + continue; + } + + auto const windowSize = static_cast(encodedWindowSize); auto alignedWindowSize = (windowSize + tokensPerBlock - 1) / tokensPerBlock * tokensPerBlock; auto validTokenNum = (alignedWindowSize < maxNumTokens.value() ? alignedWindowSize : maxNumTokens.value()); if (common::getEnvKVCacheTransferAllBlocksForWindow()) @@ -235,7 +300,17 @@ size_t CacheTransBufferManager::computeTransferBufferSize( } validTokenNum += tokensPerBlock; // add one more block - bufferSizeFromMaxNumToken += validTokenNum * kvCacheByteSizePerTokenPerLayer; + if (transferIndexerKCache) + { + bufferSizeFromMaxNumToken += validTokenNum * indexerCacheByteSizePerTokenPerLayer; + } + else + { + auto const primaryPool = blockManager.getPrimaryPool(poolIdx); + auto const kvCacheByteSizePerTokenPerLayer + = primaryPool->getDimension<-1>() * primaryPool->getDimension<2>() * dataSize / tokensPerBlock; + bufferSizeFromMaxNumToken += validTokenNum * kvCacheByteSizePerTokenPerLayer; + } } } @@ -245,36 +320,12 @@ size_t CacheTransBufferManager::computeTransferBufferSize( CacheTransBufferManager::CacheTransBufferManager( KVCacheManager::BaseKVCacheManager* cacheManager, std::optional maxNumTokens, bool transferIndexerKCache) : BaseTransBufferManager(computeTransferBufferSize(cacheManager, maxNumTokens, transferIndexerKCache), - transferIndexerKCache ? cacheManager->getIndexerKCachePool()->getDataType() - : cacheManager->getPrimaryPool(0)->getDataType(), - maxNumTokens) + getTransferDataType(cacheManager, transferIndexerKCache), maxNumTokens) , mCacheManager{cacheManager} , mTransferIndexerKCache{transferIndexerKCache} { // TODO: FP4 dataSize TLLM_CHECK(mCacheManager); - // TODO(disagg-multi-dtype): Per-pool dtype dispatch in formatter / transfer buffer - // not yet implemented. Disagg currently picks pool 0's dtype as the canonical - // transport type (above), so any KV pool with a different dtype would be silently - // miscoerced on the wire. Fail loudly until per-pool dispatch lands. We restrict - // the comparison to KV pools (getNumPools(false, false)) since block-scale and - // indexer-K pools legitimately have their own dtypes and travel through their own - // code paths. - if (!transferIndexerKCache) - { - auto const numKvPools = mCacheManager->getBlockManager().getNumPools( - /*includeBlockScalePools=*/false, /*includeIndexerKCachePools=*/false); - auto const dtype0 = mCacheManager->getPrimaryPool(0)->getDataType(); - for (SizeType32 i = 1; i < numKvPools; ++i) - { - auto const dtypeI = mCacheManager->getPrimaryPool(i)->getDataType(); - TLLM_CHECK_WITH_INFO(dtypeI == dtype0, - "Disaggregated KV cache transfer does not yet support pools with differing dtypes " - "(pool 0 dtype=%d, pool %d dtype=%d). TODO(disagg-multi-dtype): per-pool dtype " - "dispatch in formatter.", - static_cast(dtype0), i, static_cast(dtypeI)); - } - } TLLM_LOG_INFO("CacheTransBufferManager created for KV cache"); } diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h index 1635c11bc673..e12c38404514 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -74,6 +74,12 @@ class CacheTransBufferManager : public BaseTransBufferManager return mCacheManager; } + /// @brief Get the data type used by KV cache transfer buffers. + [[nodiscard]] nvinfer1::DataType getDataType() const noexcept + { + return mDataType; + } + [[nodiscard]] BufferKind getBufferKind() const override { return mTransferIndexerKCache ? BufferKind::kKV_INDEXER : BufferKind::kKV; diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 34cf8bc82def..a5d2d4bd60fa 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -10605,12 +10605,58 @@ TEST_F(KVCacheManagerTest, VswaMixedHeadDimReuseSmoke) } } +TEST_F(KVCacheManagerTest, HybridDisaggUsesAttentionPoolDtype) +{ + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 16; + auto constexpr tokensPerBlock = 4; + auto constexpr blocksInPrimaryPool = 4; + auto constexpr blocksInSecondaryPool = 0; + auto constexpr maxNumSequences = 2; + auto constexpr maxBeamWidth = 1; + auto constexpr maxAttentionWindow = 16; + auto constexpr recurrentStatesBytes = 64; + SizeType32 constexpr recurrentStatesWindow = LinearAttentionMetadata::LinearCacheType::kRecurrentStates; + + LinearAttentionMetadata const linearAttentionMetadata{ + .linearLayerIndices = {0}, + .cacheType = recurrentStatesWindow, + .allRecurrentStatesBytes = recurrentStatesBytes, + }; + auto const blocksPerWindow = BlocksPerWindow{ + {recurrentStatesWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}, + {maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}, + }; + auto const poolConfigurations = std::vector{ + {recurrentStatesWindow, sizePerHead, nvinfer1::DataType::kHALF}, + {maxAttentionWindow, sizePerHead, nvinfer1::DataType::kFP8}, + }; + auto const stream = std::make_shared(); + + auto kvCacheManager = std::make_unique(std::vector{0, numKvHeads}, sizePerHead, + tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, + std::vector{recurrentStatesWindow, maxAttentionWindow}, nvinfer1::DataType::kFP8, + /*sinkTokenLength=*/0, stream, maxAttentionWindow, /*chunkSize=*/0, /*enableBlockReuse=*/false, + CacheType::kSELF, std::nullopt, nullptr, /*enablePartialReuse=*/false, /*copyOnPartialReuse=*/true, nullptr, + /*enableIndexerKCache=*/false, /*indexerKCacheQuantBlockSize=*/128, /*indexerKCacheIndexHeadDim=*/0, + /*indexerKCacheUseFp4=*/false, linearAttentionMetadata, poolConfigurations); + kvCacheManager->allocatePools(/*useUvm=*/false); + + CacheTransBufferManager cacheTransBufferManager(kvCacheManager.get(), /*maxNumTokens=*/tokensPerBlock); + EXPECT_EQ(cacheTransBufferManager.getDataType(), nvinfer1::DataType::kFP8); + + auto const bufferId = cacheTransBufferManager.assignBufferIndexForSend(); + ASSERT_TRUE(bufferId.has_value()); + EXPECT_EQ(cacheTransBufferManager.getSendBuffer(bufferId)->getDataType(), nvinfer1::DataType::kFP8); + cacheTransBufferManager.freeBufferIndexForSend(bufferId); +} + // A6: VSWA + disagg dtype mismatch must fire the A4 guard. // -// The constructor of CacheTransBufferManager picks pool 0's dtype as canonical for -// the wire transport. When a KVCacheManager hosts pools with differing dtypes -// (mixed-precision per-window), that silent coercion would corrupt the wire format. -// The guard added in cacheTransBuffer.cpp must throw at construction time. +// CacheTransBufferManager uses a single dtype for the wire transport. When a +// KVCacheManager hosts attention pools with differing dtypes (mixed-precision +// per-window), that silent coercion would corrupt the wire format. The guard in +// cacheTransBuffer.cpp must throw at construction time. // // This test only exercises the helper / construction path that runs the guard; it // does not stand up a full disaggregated transfer (out of scope at unit-test @@ -10649,14 +10695,23 @@ TEST_F(KVCacheManagerTest, VswaDisaggDtypeMismatchTriggersGuard) kvCacheManager->allocatePools(/*useUvm=*/false); // Sanity: the manager really does host KV pools with two different dtypes. - auto const numKvPools = kvCacheManager->getBlockManager().getNumPools( - /*includeBlockScalePools=*/false, /*includeIndexerKCachePools=*/false); - ASSERT_GE(numKvPools, 2); - auto const dtype0 = kvCacheManager->getPrimaryPool(0)->getDataType(); + auto const& blockManager = kvCacheManager->getBlockManager(); + ASSERT_GE(blockManager.getNumPools(/*includeBlockScalePools=*/false, /*includeIndexerKCachePools=*/false), 2); + std::optional dtype0; bool foundMismatch = false; - for (SizeType32 i = 1; i < numKvPools; ++i) + for (SizeType32 poolIdx = 0; poolIdx < blockManager.getNumPools(); ++poolIdx) { - if (kvCacheManager->getPrimaryPool(i)->getDataType() != dtype0) + auto const& pool = blockManager.getPool(poolIdx); + if (pool.containsBlockScales || pool.containsIndexerKCache) + { + continue; + } + auto const dataType = blockManager.getPrimaryPool(poolIdx)->getDataType(); + if (!dtype0.has_value()) + { + dtype0 = dataType; + } + else if (dataType != dtype0.value()) { foundMismatch = true; break; From f0f316bdc0f4d9bace6b90513df3850c316ffd36 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sat, 18 Jul 2026 01:57:49 -0700 Subject: [PATCH 9/9] [NVBUG 6465993][test] remove temporary diagnostics Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/executor/worker.py | 16 +-- .../accuracy/test_disaggregated_serving.py | 109 +----------------- tests/integration/defs/conftest.py | 54 +-------- tests/integration/defs/pytest.ini | 5 +- 4 files changed, 10 insertions(+), 174 deletions(-) diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index fcb3563f8021..5958d296e1c4 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -311,15 +311,8 @@ def notify_proxy_threads_to_quit(): # error to the error_queue in the main thread. mpi_comm().barrier() - disable_worker_process_monitor = os.getenv( - "TRTLLM_TEST_DISABLE_WORKER_PROCESS_MONITOR", "0") == "1" - if disable_worker_process_monitor: - worker_process_identities = None - logger.warning("Worker process monitoring disabled by " - "TRTLLM_TEST_DISABLE_WORKER_PROCESS_MONITOR") - else: - worker_process_identities = mpi_comm().allgather( - capture_worker_process_identity(mpi_rank())) + worker_process_identities = mpi_comm().allgather( + capture_worker_process_identity(mpi_rank())) logger_debug(f"Worker {mpi_rank()} ready to setup backend...\n", "green") try: @@ -360,10 +353,7 @@ def notify_proxy_threads_to_quit(): worker.set_result_queue(result_queue) # Send ready signal with confirmation - if worker_process_identities is None: - ready_msg = (ready_signal, None) - else: - ready_msg = (ready_signal, None, worker_process_identities) + ready_msg = (ready_signal, None, worker_process_identities) if not worker_init_status_queue.notify_with_retry(ready_msg): logger.warning( "Failed to deliver ready signal to proxy, continuing anyway" diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 48770a2a3c38..cd63c052375c 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1,17 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. import concurrent import contextlib import functools @@ -171,22 +157,7 @@ def launch_disaggregated_llm( gen_extra_env: Optional[Dict[str, str]] = None, request_timeout_s: float = DEFAULT_REQUEST_TIMEOUT_S, request_max_retries: Optional[int] = None, - diagnostic_label: Optional[str] = None, ): - diagnostic_started_at = time.monotonic() - - def _diagnostic(phase: str, **details: Any) -> None: - if diagnostic_label is None: - return - details_text = " ".join(f"{key}={value!r}" - for key, value in sorted(details.items())) - print( - f"[DISAGG-DIAG] label={diagnostic_label!r} " - f"elapsed_s={time.monotonic() - diagnostic_started_at:.3f} " - f"phase={phase!r} {details_text}", - flush=True) - - _diagnostic("configuration_start", model=model_name) temp_dir = tempfile.TemporaryDirectory() disaggregated_serving_config_path = os.path.join( temp_dir.name, "disaggregated_serving_config.yaml") @@ -301,8 +272,6 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): base_env = os.environ.copy() if extra_env: base_env.update(extra_env) - if diagnostic_label is not None: - base_env["PYTHONUNBUFFERED"] = "1" kv_cache_perf_dir = os.path.join(temp_dir.name, "kv_cache_perf") @@ -381,12 +350,6 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): log_files = [] try: for i, (env, args) in enumerate(server_configs): - _diagnostic( - "process_launch_start", - role=server_name, - index=i, - cuda_visible_devices=env.get("CUDA_VISIBLE_DEVICES"), - command=" ".join(args)) if enable_redirect_log: f = open(f"output_{server_name}_{i}.log", "w+") proc = popen(args, env=env, stdout=f, stderr=f) @@ -399,9 +362,6 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): opened_processes = [ stack.enter_context(proc) for proc in processes ] - _diagnostic("process_group_started", - role=server_name, - pids=[process.pid for process in opened_processes]) yield opened_processes for f in log_files: f.close() @@ -416,7 +376,6 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): disaggregated_serving_config_path, "--server_start_timeout", str(server_waiting_timeout), "-r", "360000" ] - _diagnostic("process_groups_start") with ( MyThreadPoolExecutor(max_workers=max_workers) as thread_pool, temp_dir, @@ -424,23 +383,13 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): multi_popen(gen_servers, "gen") as gen_processes, multi_popen([(base_env, server_cmd)], "disagg") as server_processes, ): - _diagnostic( - "all_processes_started", - context_pids=[process.pid for process in ctx_processes], - generation_pids=[process.pid for process in gen_processes], - disaggregated_pids=[process.pid for process in server_processes]) start_time = time.time() server_is_ready = False - registration_poll = 0 while time.time() - start_time < server_waiting_timeout: time.sleep(5) - registration_poll += 1 for process in itertools.chain(ctx_processes, gen_processes, server_processes): if process.poll() is not None: - _diagnostic("process_exited_before_registration", - pid=process.pid, - returncode=process.returncode) raise Exception( f"process {process.pid} exited with code {process.returncode}" ) @@ -450,30 +399,13 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): f"http://localhost:{serve_port}/cluster_info") if response.status_code == 200: cluster_info = response.json() - if registration_poll == 1 or registration_poll % 6 == 0: - _diagnostic("service_registration_poll", - poll=registration_poll, - cluster_info=cluster_info) if cluster_info.get("is_ready"): print(f"Cluster ready: {cluster_info}") - _diagnostic("service_registration_complete", - poll=registration_poll, - cluster_info=cluster_info) server_is_ready = True break - elif registration_poll == 1 or registration_poll % 6 == 0: - _diagnostic("service_registration_http_status", - poll=registration_poll, - status_code=response.status_code) - except requests.exceptions.ConnectionError as error: - if registration_poll == 1 or registration_poll % 6 == 0: - _diagnostic("service_registration_unreachable", - poll=registration_poll, - error=type(error).__name__) + except requests.exceptions.ConnectionError: continue if not server_is_ready: - _diagnostic("service_registration_timeout", - timeout_s=server_waiting_timeout) pytest.fail( f"Server is not ready after {server_waiting_timeout} seconds. Please check the logs for more details." ) @@ -486,15 +418,9 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): if request_max_retries is not None: client_kwargs["max_retries"] = request_max_retries client = openai.OpenAI(**client_kwargs) - _diagnostic("client_created", base_url=client_kwargs["base_url"]) - request_counter = itertools.count(1) def send_request(prompt: str, sampling_params: SamplingParams, - streaming: bool, request_index: int): - if request_index == 1: - _diagnostic("first_inference_request_start", - prompt_length=len(prompt), - streaming=streaming) + streaming: bool): kwargs = {} if sampling_params is not None: extra_body = {} @@ -534,9 +460,6 @@ def send_request(prompt: str, sampling_params: SamplingParams, prompt=prompt, stream=streaming, **kwargs) - if request_index == 1: - _diagnostic("first_inference_request_complete", - choice_count=len(response.choices)) result = Result(id=0, sampling_params=sampling_params, outputs=[ @@ -551,13 +474,8 @@ def send_request(prompt: str, sampling_params: SamplingParams, def generate_async(prompt: str, sampling_params: Optional[SamplingParams] = None, streaming: bool = False): - request_index = next(request_counter) - if request_index == 1: - _diagnostic("first_inference_request_submit", - prompt_length=len(prompt), - streaming=streaming) future = thread_pool.submit(send_request, prompt, sampling_params, - streaming, request_index) + streaming) thread_pool.futures.append(future) return future @@ -591,18 +509,10 @@ def _show_kvcache_time(kv_cache_perf_dir, max_lines=100): for line in f.readlines()[-max_lines:]: print(line.strip()) - _diagnostic("tokenizer_load_start") tokenizer = load_hf_tokenizer(model_name) - _diagnostic("tokenizer_load_complete") - accuracy_body_completed = False try: - _diagnostic("accuracy_body_start") yield DuckLLM(args, tokenizer, generate_async) - accuracy_body_completed = True - _diagnostic("accuracy_body_complete") finally: - _diagnostic("shutdown_start", - accuracy_body_completed=accuracy_body_completed) if enable_perf: _show_kvcache_time(kv_cache_perf_dir) _get_perf_metrics() @@ -630,9 +540,6 @@ def _show_kvcache_time(kv_cache_perf_dir, max_lines=100): pass # already exited between timeout and kill except OSError: pass # process already gone - _diagnostic( - "shutdown_complete", - returncodes=[process.poll() for process in all_processes]) def run_parallel_test(model_name: str, @@ -2502,14 +2409,8 @@ def test_auto_dtype(self, use_py_transceiver, block_reuse, mtp_nextn): if block_reuse: ctx_cfg["kv_cache_config"]["enable_block_reuse"] = True gen_cfg["kv_cache_config"]["enable_block_reuse"] = True - diagnostic_label = None - if not use_py_transceiver and block_reuse and mtp_nextn == 3: - diagnostic_label = "nemotron3-super-mtp3-block-reuse" - with launch_disaggregated_llm(disagg_cfg, - ctx_cfg, - gen_cfg, - self.MODEL_PATH, - diagnostic_label=diagnostic_label) as llm: + with launch_disaggregated_llm(disagg_cfg, ctx_cfg, gen_cfg, + self.MODEL_PATH) as llm: run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) @pytest.mark.skip_less_device(8) diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 1263a668a776..bba88cf07c01 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -2343,39 +2343,6 @@ def pytest_runtest_makereport(item, call): report.url = "" -_NEMOTRON_DISAGG_DIAGNOSTIC_TEST = ( - "accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::" - "test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=False]") -_NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE = "_nemotron_disagg_diagnostic_fd" - - -def _emit_nemotron_disagg_diagnostic(item, message: str) -> None: - diagnostic_fd = getattr(item, _NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE, - None) - if diagnostic_fd is None: - print(f"[DISAGG-DIAG] {message}", flush=True) - return - os.write(diagnostic_fd, f"[DISAGG-DIAG] {message}\n".encode()) - - -@pytest.hookimpl(tryfirst=True, hookwrapper=True) -def pytest_fixture_setup(fixturedef, request): - item = getattr(request, "_pyfuncitem", None) - nodeid = getattr(item, "nodeid", "") - if not nodeid.endswith(_NEMOTRON_DISAGG_DIAGNOSTIC_TEST): - yield - return - - _emit_nemotron_disagg_diagnostic( - item, f"phase='fixture_setup_start' fixture={fixturedef.argname!r} " - f"scope={fixturedef.scope!r}") - outcome = yield - status = "error" if outcome.excinfo else "complete" - _emit_nemotron_disagg_diagnostic( - item, f"phase='fixture_setup_finish' fixture={fixturedef.argname!r} " - f"status={status!r}") - - @pytest.fixture(scope="session") def all_pytest_items(): """ @@ -2607,26 +2574,7 @@ def collect_status(item: pytest.Item): @pytest.hookimpl(wrapper=True) def pytest_runtest_protocol(item, nextitem): - diagnostic = item.nodeid.endswith(_NEMOTRON_DISAGG_DIAGNOSTIC_TEST) - if diagnostic: - upload_log_plugin = item.config.pluginmanager.getplugin( - "upload_log_plugin") - if upload_log_plugin is not None: - original_echo_to_stdout = upload_log_plugin.echo_to_stdout - upload_log_plugin.echo_to_stdout = True - diagnostic_fd = os.dup(1) - setattr(item, _NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE, diagnostic_fd) - _emit_nemotron_disagg_diagnostic(item, "phase='runtest_protocol_start'") - try: - ret = yield - finally: - if diagnostic: - _emit_nemotron_disagg_diagnostic(item, - "phase='runtest_protocol_finish'") - os.close(diagnostic_fd) - delattr(item, _NEMOTRON_DISAGG_DIAGNOSTIC_FD_ATTRIBUTE) - if upload_log_plugin is not None: - upload_log_plugin.echo_to_stdout = original_echo_to_stdout + ret = yield collect_status(item) return ret diff --git a/tests/integration/defs/pytest.ini b/tests/integration/defs/pytest.ini index 87fc31800ff2..ee5e97c00f74 100644 --- a/tests/integration/defs/pytest.ini +++ b/tests/integration/defs/pytest.ini @@ -6,12 +6,9 @@ threadleak = True # outlive the test they start under. threadleak_exclude = asyncio_\d+|Thread-\d+ \(_manager_spawn\)|session-reuse-\w+ junit_family=legacy -addopts = --ignore-glob="*perf/test_perf.py" --ignore-glob="*perf/disagg/*" --ignore-glob="*test_list_validation.py" --ignore-glob="*llm-test-workspace*" --durations=0 -W ignore::DeprecationWarning --unused-fixtures -p test_common.session_reuse_hooks -s +addopts = --ignore-glob="*perf/test_perf.py" --ignore-glob="*perf/disagg/*" --ignore-glob="*test_list_validation.py" --ignore-glob="*llm-test-workspace*" --durations=0 -W ignore::DeprecationWarning --unused-fixtures -p test_common.session_reuse_hooks pythonpath = ../../../examples/auto_deploy ../../ -env = - TRTLLM_TEST_REUSE_SESSION=0 - TRTLLM_TEST_DISABLE_WORKER_PROCESS_MONITOR=1 norecursedirs = ./triton/perf ./perf/disagg unused_fixtures_ignore_path = cpp