| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from copy import copy |
| 5 | |
| 6 | from vllm import SamplingParams |
| 7 | from vllm.outputs import CompletionOutput |
| 8 | from vllm.sampling_params import RequestOutputKind |
| 9 | from vllm.v1.engine import EngineCoreRequest |
| 10 | from vllm.v1.engine.parallel_sampling import ParentRequest |
| 11 | |
| 12 | |
| 13 | def test_parent_request_to_output_stream() -> None: |
| 14 | parent_request = ParentRequest(make_request(SamplingParams(n=2))) |
| 15 | parent_request.child_requests = {"child_id_0", "child_id_1"} |
| 16 | output_0 = CompletionOutput( |
| 17 | index=0, text="child 0", token_ids=[], cumulative_logprob=None, logprobs=None |
| 18 | ) |
| 19 | output_1 = CompletionOutput( |
| 20 | index=1, text="child 1", token_ids=[], cumulative_logprob=None, logprobs=None |
| 21 | ) |
| 22 | # Request not finished |
| 23 | assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0) |
| 24 | assert ([output_1], False) == parent_request.get_outputs("child_id_1", output_1) |
| 25 | assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0) |
| 26 | assert ([output_1], False) == parent_request.get_outputs("child_id_1", output_1) |
| 27 | |
| 28 | # output_1 finished |
| 29 | output_1.finish_reason = "ended" |
| 30 | assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0) |
| 31 | assert ([output_1], False) == parent_request.get_outputs("child_id_1", output_1) |
| 32 | # Finished output_1 had already returned, DO NOT returned again |
| 33 | assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0) |
| 34 | assert parent_request.get_outputs("child_id_1", output_1) == ([], False) |
| 35 | |
| 36 | # output_0 finished |
| 37 | output_0.finish_reason = "ended" |
| 38 | assert ([output_0], True) == parent_request.get_outputs("child_id_0", output_0) |
| 39 | assert parent_request.get_outputs("child_id_1", output_1) == ([], True) |
| 40 | # Finished output_0 had already returned, DO NOT returned again |
| 41 | assert parent_request.get_outputs("child_id_0", output_0) == ([], True) |
| 42 | assert parent_request.get_outputs("child_id_1", output_1) == ([], True) |
| 43 | |
| 44 | |
| 45 | def test_parent_request_to_output_final_only() -> None: |
| 46 | parent_request = ParentRequest( |
| 47 | make_request(SamplingParams(n=2, output_kind=RequestOutputKind.FINAL_ONLY)) |
| 48 | ) |
| 49 | parent_request.child_requests = {"child_id_0", "child_id_1"} |
| 50 | output_0 = CompletionOutput( |
| 51 | index=0, text="child 0", token_ids=[], cumulative_logprob=None, logprobs=None |
| 52 | ) |
| 53 | output_1 = CompletionOutput( |
| 54 | index=1, text="child 1", token_ids=[], cumulative_logprob=None, logprobs=None |
| 55 | ) |
| 56 | # Request not finished, return nothing |
| 57 | assert parent_request.get_outputs("child_id_0", output_0) == ([], False) |
| 58 | assert parent_request.get_outputs("child_id_1", output_1) == ([], False) |
| 59 | # output_1 finished, but outputs won't be returned until all child requests finished |
| 60 | output_1.finish_reason = "ended" |
| 61 | assert parent_request.get_outputs("child_id_0", output_0) == ([], False) |
| 62 | assert parent_request.get_outputs("child_id_1", output_1) == ([], False) |
| 63 | # output_0 finished, as all child requests finished, the output would be returned |
| 64 | output_0.finish_reason = "ended" |
| 65 | assert ([output_0, output_1], True) == parent_request.get_outputs( |
| 66 | "child_id_0", output_0 |
| 67 | ) |
| 68 | assert ([output_0, output_1], True) == parent_request.get_outputs( |
| 69 | "child_id_1", output_1 |
| 70 | ) |
| 71 | |
| 72 | |
| 73 | def test_parallel_sampling_child_requests_preserve_session_id() -> None: |
| 74 | request = make_request(SamplingParams(n=2)) |
| 75 | request.session_id = "session-1" |
| 76 | parent_request = ParentRequest(request) |
| 77 | |
| 78 | for idx in range(parent_request.n): |
| 79 | request_id, child_params = parent_request.get_child_info(idx) |
| 80 | child_request = request if idx == parent_request.n - 1 else copy(request) |
| 81 | child_request.request_id = request_id |
| 82 | child_request.sampling_params = child_params |
| 83 | |
| 84 | assert child_request.session_id == "session-1" |
| 85 | |
| 86 | |
| 87 | def make_request(sampling_params: SamplingParams) -> EngineCoreRequest: |
| 88 | return EngineCoreRequest( |
| 89 | request_id="parent_id", |
| 90 | external_req_id="ext_parent_id", |
| 91 | prompt_token_ids=None, |
| 92 | mm_features=None, |
| 93 | sampling_params=sampling_params, |
| 94 | pooling_params=None, |
| 95 | arrival_time=0.0, |
| 96 | lora_request=None, |
| 97 | cache_salt=None, |
| 98 | data_parallel_rank=None, |
| 99 | ) |