| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from copy import copy |
| 5 | from typing import cast |
| 6 | |
| 7 | from vllm.outputs import CompletionOutput |
| 8 | from vllm.sampling_params import RequestOutputKind, SamplingParams |
| 9 | from vllm.v1.engine import EngineCoreRequest |
| 10 | from vllm.v1.metrics.stats import IterationStats |
| 11 | |
| 12 | |
| 13 | class ParentRequest: |
| 14 | """Info, state & processing for parallel sampling request. |
| 15 | |
| 16 | Store parent request ID and sampling params. |
| 17 | Facilitate generating child request sampling params. |
| 18 | """ |
| 19 | |
| 20 | request_id: str |
| 21 | external_req_id: str |
| 22 | sampling_params: SamplingParams |
| 23 | |
| 24 | # To track the completion of child requests |
| 25 | child_requests: set[str] |
| 26 | |
| 27 | # To aggregate child completions when not streaming |
| 28 | output_aggregator: list[CompletionOutput] |
| 29 | |
| 30 | # To find the max number of generated tokens across all children |
| 31 | max_num_generation_tokens: int |
| 32 | |
| 33 | # To efficiently obtain child sampling params |
| 34 | cached_child_sampling_params: SamplingParams | None |
| 35 | |
| 36 | def __init__(self, request: EngineCoreRequest) -> None: |
| 37 | assert request.external_req_id is not None |
| 38 | sampling_params = request.params |
| 39 | self.request_id = request.request_id |
| 40 | self.external_req_id = request.external_req_id |
| 41 | self.sampling_params = sampling_params |
| 42 | |
| 43 | self.child_requests = set() |
| 44 | self.output_aggregator = ( |
| 45 | [cast(CompletionOutput, None)] * sampling_params.n |
| 46 | if (sampling_params.output_kind == RequestOutputKind.FINAL_ONLY) |
| 47 | else [] |
| 48 | ) |
| 49 | self.max_num_generation_tokens = 0 |
| 50 | self.cached_child_sampling_params = None |
| 51 | |
| 52 | def _get_child_sampling_params( |
| 53 | self, |
| 54 | index: int, |
| 55 | ) -> SamplingParams: |
| 56 | """Efficiently obtain child `sampling_params`. |
| 57 | |
| 58 | If `sampling_params.seed` is not `None` then |
| 59 | each child request requires a unique clone of |
| 60 | parent `sampling_params` with a unique seed. |
| 61 | |
| 62 | Args: |
| 63 | index: index within `n` child requests |
| 64 | |
| 65 | Returns: |
| 66 | Child `sampling_params` instance. |
| 67 | |
| 68 | """ |
| 69 | seed = self.sampling_params.seed |
| 70 | if self.cached_child_sampling_params: |
| 71 | # Reuse child sampling_params data structure |
| 72 | return self.cached_child_sampling_params |
| 73 | # Build child sampling_params |
| 74 | child_sampling_params = copy(self.sampling_params) |
| 75 | child_sampling_params.n = 1 |
| 76 | if seed is None: |
| 77 | # Cache child sampling_params for later reuse |
| 78 | self.cached_child_sampling_params = child_sampling_params |
| 79 | else: |
| 80 | # Each child gets a clone with a unique seed |
| 81 | child_sampling_params.seed = seed + index |
| 82 | return child_sampling_params |
| 83 | |
| 84 | def get_child_info(self, index: int) -> tuple[str, SamplingParams]: |
| 85 | """Get child request ID and sampling params. |
| 86 | |
| 87 | Args: |
| 88 | index: index within `n` child requests. |
| 89 | |
| 90 | Returns: |
| 91 | (request ID, sampling_params) tuple |
| 92 | |
| 93 | """ |
| 94 | child_req_id = f"{index}_{self.request_id}" |
| 95 | self.child_requests.add(child_req_id) |
| 96 | return child_req_id, self._get_child_sampling_params(index) |
| 97 | |
| 98 | @property |
| 99 | def n(self) -> int: |
| 100 | return self.sampling_params.n |
| 101 | |
| 102 | def get_outputs( |
| 103 | self, |
| 104 | child_request_id: str, |
| 105 | completion_output: CompletionOutput, |
| 106 | ) -> tuple[list[CompletionOutput], bool]: |
| 107 | already_finished_and_returned: bool = False |
| 108 | if completion_output.finished(): |
| 109 | if child_request_id in self.child_requests: |
| 110 | self.child_requests.remove(child_request_id) |
| 111 | else: |
| 112 | # child request ID is not available in child_requests |
| 113 | # which means the request had finished in previous |
| 114 | # batch step and returned to the client earlier |
| 115 | already_finished_and_returned = True |
| 116 | |
| 117 | if self.sampling_params.output_kind != RequestOutputKind.FINAL_ONLY: |
| 118 | # If streaming, just return the current output |
| 119 | # |
| 120 | # DO NOT output finished and already returned child request to client again |
| 121 | outputs = [] if already_finished_and_returned else [completion_output] |
| 122 | else: |
| 123 | # If not streaming, aggregate the n final outputs. |
| 124 | self.output_aggregator[completion_output.index] = completion_output |
| 125 | outputs = [] if self.child_requests else self.output_aggregator |
| 126 | |
| 127 | finished = not self.child_requests |
| 128 | return outputs, finished |
| 129 | |
| 130 | def observe_num_generation_tokens(self, num_generation_tokens: int): |
| 131 | self.max_num_generation_tokens = max( |
| 132 | num_generation_tokens, self.max_num_generation_tokens |
| 133 | ) |
| 134 | return self.max_num_generation_tokens |
| 135 | |
| 136 | @staticmethod |
| 137 | def observe_finished_request( |
| 138 | parent_req: "ParentRequest | None", |
| 139 | iteration_stats: IterationStats, |
| 140 | num_generation_tokens: int, |
| 141 | ): |
| 142 | n_param = parent_req.n if parent_req is not None else 1 |
| 143 | |
| 144 | if parent_req is not None: |
| 145 | num_generation_tokens = parent_req.observe_num_generation_tokens( |
| 146 | num_generation_tokens |
| 147 | ) |
| 148 | |
| 149 | # Child requests finished, we can now record to iteration stats |
| 150 | if parent_req is None or not parent_req.child_requests: |
| 151 | iteration_stats.max_num_generation_tokens_iter.append(num_generation_tokens) |
| 152 | iteration_stats.n_params_iter.append(n_param) |