@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 129 lines 4.24 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Sanity tests for ec_transfer_params protocol plumbing.
4
5 No running engine required.
6 """
7
8 from unittest.mock import MagicMock
9
10 import pytest
11
12 from vllm.entrypoints.openai.chat_completion.protocol import (
13 ChatCompletionRequest,
14 )
15 from vllm.outputs import CompletionOutput, RequestOutput
16 from vllm.sampling_params import SamplingParams
17 from vllm.v1.request import Request, RequestStatus
18
19 # tests.v1.core.utils transitively imports the example KV connector, which
20 # requires the compiled CUDA flash-attention extensions; skip cleanly rather
21 # than erroring at collection when those extensions aren't built.
22 _core_utils = pytest.importorskip(
23 "tests.v1.core.utils",
24 reason="requires compiled CUDA flash-attention extensions",
25 exc_type=ImportError,
26 )
27 create_scheduler = _core_utils.create_scheduler
28
29 EC_PARAMS: dict = {"mm_hash_abc": {"peer_host": "10.0.0.1", "peer_port": 5501}}
30
31
32 def test_ec_transfer_params_routed_to_sampling_params_extra_args():
33 """ec_transfer_params on the request must land in SamplingParams.extra_args."""
34 req = ChatCompletionRequest(
35 model="test-model",
36 messages=[{"role": "user", "content": "hi"}],
37 max_tokens=5,
38 ec_transfer_params=EC_PARAMS,
39 )
40 sp = req.to_sampling_params(max_tokens=5, default_sampling_params={})
41 assert sp.extra_args is not None
42 assert sp.extra_args.get("ec_transfer_params") == EC_PARAMS
43
44
45 def test_request_output_add_propagates_ec_transfer_params():
46 """RequestOutput.add() must carry ec_transfer_params forward to the caller."""
47
48 def _out(ec_params):
49 return RequestOutput(
50 request_id="r1",
51 prompt="p",
52 prompt_token_ids=[1],
53 prompt_logprobs=None,
54 outputs=[
55 CompletionOutput(
56 index=0,
57 text="",
58 token_ids=[],
59 cumulative_logprob=0.0,
60 logprobs=None,
61 finish_reason=None,
62 )
63 ],
64 finished=False,
65 ec_transfer_params=ec_params,
66 )
67
68 accumulated = _out(None)
69 accumulated.add(_out(EC_PARAMS), aggregate=True)
70 assert accumulated.ec_transfer_params == EC_PARAMS
71
72
73 def test_request_reads_ec_transfer_params_from_extra_args():
74 """v1 Request must pull ec_transfer_params out of SamplingParams.extra_args."""
75 sp = SamplingParams(extra_args={"ec_transfer_params": EC_PARAMS})
76 req = Request(
77 request_id="r1",
78 prompt_token_ids=[1, 2, 3],
79 sampling_params=sp,
80 pooling_params=None,
81 )
82 assert req.ec_transfer_params == EC_PARAMS
83
84
85 def test_free_request_calls_ec_connector_and_surfaces_params():
86 """_free_request must call ec_connector.request_finished() and return its params."""
87 sp = SamplingParams(max_tokens=1)
88 sp.update_from_generation_config({}, 50256)
89 request = Request(
90 request_id="test-req",
91 prompt_token_ids=[1, 2, 3],
92 sampling_params=sp,
93 pooling_params=None,
94 client_index=0,
95 )
96 scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer")
97 scheduler.add_request(request)
98 request.status = RequestStatus.FINISHED_STOPPED
99
100 mock_ec = MagicMock()
101 mock_ec.request_finished.return_value = (False, EC_PARAMS)
102 scheduler.ec_connector = mock_ec
103
104 kv_params, ec_params = scheduler._free_request(request)
105
106 mock_ec.request_finished.assert_called_once_with(request)
107 assert ec_params == EC_PARAMS
108 assert kv_params is None
109
110
111 def test_free_request_without_ec_connector_returns_none():
112 """When no EC connector is configured, ec_transfer_params must be None."""
113 sp = SamplingParams(max_tokens=1)
114 sp.update_from_generation_config({}, 50256)
115 request = Request(
116 request_id="test-req",
117 prompt_token_ids=[1, 2, 3],
118 sampling_params=sp,
119 pooling_params=None,
120 client_index=0,
121 )
122 scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer")
123 scheduler.add_request(request)
124 request.status = RequestStatus.FINISHED_STOPPED
125
126 kv_params, ec_params = scheduler._free_request(request)
127
128 assert ec_params is None
129 assert kv_params is None