From dc9eec7b05c9c36584ee6b62c5ce092be0d13063 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Wed, 3 Dec 2025 10:51:00 -0800 Subject: [PATCH 1/5] Fix flaky kv connector tests Signed-off-by: jthomson04 --- .../defs/llmapi/test_llm_api_connector.py | 38 ++++++++++++------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index cd85d4218508..caebd104af33 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -63,6 +63,15 @@ def enforce_single_worker(monkeypatch): yield +def generate_and_sleep(model, *args, **kwargs): + # When using the overlap scheduler, we generate an extra token. We want to be able to track calls to the connector made with this extra token. + # This is problematic since the final calls to the connector API occur after the final response has been returned. + # To get around this, we sleep for a short period of time to account for this additional time. + outputs = model.generate(*args, **kwargs) + time.sleep(1) + return outputs + + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_overlap_scheduler", [True, False]) def test_connector_simple(enforce_single_worker, model_with_connector, @@ -81,7 +90,7 @@ def test_connector_simple(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - model.generate(["Hello, world"], sampling_params) + generate_and_sleep(model, ["Hello, world"], sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 @@ -154,9 +163,10 @@ def test_connector_async_onboard(enforce_single_worker, model_with_connector, worker.get_finished.side_effect = lambda finished_gen, load_async: ( finished_gen, load_async) - model.generate([ - "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua." - ], SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True)) + outputs = generate_and_sleep( + model, [ + "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua." + ], SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True)) # Once for the initial poll, then once for each token. One extra token when using the overlap scheduler. assert worker.get_finished.call_count == NUM_TOKENS + 1 + int( @@ -187,7 +197,7 @@ def test_connector_async_save(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) - model.generate(["Hello, world"], sampling_params) + generate_and_sleep(model, ["Hello, world"], sampling_params) assert scheduler.request_finished.call_count == 1 @@ -227,7 +237,7 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) - model.generate([0] * NUM_INPUT_TOKENS, sampling_params) + generate_and_sleep(model, [0] * NUM_INPUT_TOKENS, sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 assert len( @@ -277,7 +287,7 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector, assert len(scheduler.request_finished.call_args.args[1]) == math.ceil( (NUM_INPUT_TOKENS + NUM_TOKENS) / BLOCK_SIZE) - model.generate([1] * NUM_INPUT_TOKENS, sampling_params) + generate_and_sleep(model, [1] * NUM_INPUT_TOKENS, sampling_params) # The initial computed position should be 0, since we haven't yet onboarded any blocks. assert scheduler.build_connector_meta.call_args_list[0].args[ @@ -306,7 +316,7 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, sampling_params = SamplingParams(max_tokens=BLOCK_SIZE, ignore_eos=True) - model.generate([0] * (CHUNK_SIZE * 2), sampling_params) + generate_and_sleep(model, [0] * (CHUNK_SIZE * 2), sampling_params) assert scheduler.update_state_after_alloc.call_count == 1 @@ -379,16 +389,16 @@ def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, scheduler.request_finished.return_value = False worker.get_finished.return_value = [], [] - result = prefill_worker.generate([0] * 48, - sampling_params=sampling_params, - disaggregated_params=disaggregated_params) + result = generate_and_sleep(prefill_worker, [0] * 48, + sampling_params=sampling_params, + disaggregated_params=disaggregated_params) gen_disagg_params = result.disaggregated_params gen_disagg_params.request_type = "generation_only" - result = decode_worker.generate([0] * 48, - sampling_params=sampling_params, - disaggregated_params=gen_disagg_params) + result = generate_and_sleep(decode_worker, [0] * 48, + sampling_params=sampling_params, + disaggregated_params=gen_disagg_params) assert scheduler.build_connector_meta.call_count == 1 From 9f4f95b622df48b9b22a1e6fa789659208125c2a Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Wed, 3 Dec 2025 10:59:53 -0800 Subject: [PATCH 2/5] Fix comment Signed-off-by: jthomson04 --- tests/integration/defs/llmapi/test_llm_api_connector.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index caebd104af33..68b5be99d8a4 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -64,9 +64,10 @@ def enforce_single_worker(monkeypatch): def generate_and_sleep(model, *args, **kwargs): - # When using the overlap scheduler, we generate an extra token. We want to be able to track calls to the connector made with this extra token. - # This is problematic since the final calls to the connector API occur after the final response has been returned. - # To get around this, we sleep for a short period of time to account for this additional time. + # Some KV connector API calls are made after a full response is returned. We want to be able to track these calls. + # However, we don't have any indication of when all the calls are complete. + # To compensate for this, we sleep between the generate call and the return of the outputs. + # TODO(jthomson04): Surely there's a better way to do this? outputs = model.generate(*args, **kwargs) time.sleep(1) return outputs From 45307471468f5c2bc3c3b4aca614367f9ff10bcf Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Thu, 4 Dec 2025 09:49:51 -0800 Subject: [PATCH 3/5] Fix test + add multi request test to ci Signed-off-by: jthomson04 --- tests/integration/defs/llmapi/test_llm_api_connector.py | 3 ++- tests/integration/test_lists/test-db/l0_a10.yml | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 68b5be99d8a4..5e74b6bb8128 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -14,6 +14,7 @@ # limitations under the License. import math +import time from unittest.mock import MagicMock, patch import pytest @@ -422,7 +423,7 @@ def test_connector_multi_request(enforce_single_worker, model_with_connector): model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=True, - kv_cache_config=KvCacheConfig(max_tokens=120)) + kv_cache_config=KvCacheConfig(max_tokens=144)) sampling_params = SamplingParams(ignore_eos=True, max_tokens=4) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 1b79caaec78e..3e31703bd67a 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -130,6 +130,7 @@ l0_a10: - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[False] - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[False] - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[True] + - llmapi/test_llm_api_connector.py::test_connector_multi_request - llmapi/test_llm_e2e.py::test_llmapi_load_engine_from_build_command[llama-llama-models/llama-7b-hf] # 5min - llmapi/test_llm_e2e.py::test_llmapi_build_command_parameters_align[llama-llama-models-v2/TinyLlama-1.1B-Chat-v1.0] - llmapi/test_llm_e2e.py::test_llmapi_load_engine_from_build_command_with_lora[llama-llama-models-v2/llama-v2-7b-hf] From cb82c18551d68e42b6019cd87d45e2fc01c67388 Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Thu, 4 Dec 2025 10:03:08 -0800 Subject: [PATCH 4/5] coderabbit Signed-off-by: jthomson04 --- .../defs/llmapi/test_llm_api_connector.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 5e74b6bb8128..f3053d73f138 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -165,10 +165,9 @@ def test_connector_async_onboard(enforce_single_worker, model_with_connector, worker.get_finished.side_effect = lambda finished_gen, load_async: ( finished_gen, load_async) - outputs = generate_and_sleep( - model, [ - "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua." - ], SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True)) + generate_and_sleep(model, [ + "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua." + ], SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True)) # Once for the initial poll, then once for each token. One extra token when using the overlap scheduler. assert worker.get_finished.call_count == NUM_TOKENS + 1 + int( @@ -398,9 +397,9 @@ def test_connector_disagg_prefill(enforce_single_worker, model_with_connector, gen_disagg_params = result.disaggregated_params gen_disagg_params.request_type = "generation_only" - result = generate_and_sleep(decode_worker, [0] * 48, - sampling_params=sampling_params, - disaggregated_params=gen_disagg_params) + generate_and_sleep(decode_worker, [0] * 48, + sampling_params=sampling_params, + disaggregated_params=gen_disagg_params) assert scheduler.build_connector_meta.call_count == 1 From 1c6a60e161fd7e2557b5bfce8a25184996cd8e0a Mon Sep 17 00:00:00 2001 From: jthomson04 Date: Thu, 4 Dec 2025 15:58:06 -0800 Subject: [PATCH 5/5] Move tests to pytorch runtime Signed-off-by: jthomson04 --- .../integration/test_lists/test-db/l0_a10.yml | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 7560409120e5..85ca756d2c4f 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -78,6 +78,19 @@ l0_a10: - unittest/llmapi/apps/test_chat_utils.py - unittest/llmapi/apps/test_tool_parsers.py - unittest/llmapi/apps/test_harmony_channel_validation.py + - llmapi/test_llm_api_connector.py::test_connector_simple[True] + - llmapi/test_llm_api_connector.py::test_connector_simple[False] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[True] + - llmapi/test_llm_api_connector.py::test_connector_async_onboard[False] + - llmapi/test_llm_api_connector.py::test_connector_async_save[True] + - llmapi/test_llm_api_connector.py::test_connector_async_save[False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[False] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[True] + - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[False] + - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[True] + - llmapi/test_llm_api_connector.py::test_connector_multi_request - condition: ranges: system_gpu_count: @@ -119,19 +132,6 @@ l0_a10: - examples/test_bert.py::test_llm_bert_general[compare_hf-enable_remove_input_padding-use_attention_plugin-enable_context_fmha-tp:1-pp:1-float16-BertModel-bert/bert-base-uncased] - unittest/trt/model/test_mistral.py - unittest/trt/model/test_llama.py - - llmapi/test_llm_api_connector.py::test_connector_simple[True] - - llmapi/test_llm_api_connector.py::test_connector_simple[False] - - llmapi/test_llm_api_connector.py::test_connector_async_onboard[True] - - llmapi/test_llm_api_connector.py::test_connector_async_onboard[False] - - llmapi/test_llm_api_connector.py::test_connector_async_save[True] - - llmapi/test_llm_api_connector.py::test_connector_async_save[False] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[True] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output[False] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[True] - - llmapi/test_llm_api_connector.py::test_connector_scheduler_output_chunked_context[False] - - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[False] - - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[True] - - llmapi/test_llm_api_connector.py::test_connector_multi_request - llmapi/test_llm_e2e.py::test_llmapi_load_engine_from_build_command[llama-llama-models/llama-7b-hf] # 5min - llmapi/test_llm_e2e.py::test_llmapi_build_command_parameters_align[llama-llama-models-v2/TinyLlama-1.1B-Chat-v1.0] - llmapi/test_llm_e2e.py::test_llmapi_load_engine_from_build_command_with_lora[llama-llama-models-v2/llama-v2-7b-hf]