Handle a single response from await_response worker.
(
self, input: Union["PostprocWorker.Input", "ResponseWrapper"]
)
| 121 | tokenizer=tokenizer) |
| 122 | |
| 123 | async def _handle_input( |
| 124 | self, input: Union["PostprocWorker.Input", "ResponseWrapper"] |
| 125 | ) -> [Any, Optional[dict[str, float]]]: |
| 126 | ''' Handle a single response from await_response worker. ''' |
| 127 | if input.rsp.result.context_logits is not None or \ |
| 128 | input.rsp.result.generation_logits is not None: |
| 129 | raise ValueError( |
| 130 | "Context logits or generation logits are not supposed to be " |
| 131 | "sent to postprocessing workers.") |
| 132 | |
| 133 | with nvtx_range_debug("handle_input", |
| 134 | color="yellow", |
| 135 | category="Postproc"): |
| 136 | req_id = input.rsp.client_id |
| 137 | if req_id not in self._records: |
| 138 | # TODO: support variant creation later |
| 139 | self._records[req_id] = self._record_creator( |
| 140 | input, self._tokenizer) |
| 141 | |
| 142 | record = self._records[req_id] |
| 143 | record._handle_response(input.rsp) # inplace |
| 144 | # Left the result_handler determine the final output dtype. |
| 145 | # NOTE: This will change the CompletionOutput._postprocess_result |
| 146 | metrics_dict = record.metrics_dict |
| 147 | perf_metrics = None |
| 148 | disaggregated_params = None |
| 149 | if record.outputs: |
| 150 | perf_metrics = record.outputs[0].request_perf_metrics |
| 151 | disaggregated_params = record.outputs[0].disaggregated_params |
| 152 | if postproc_params := record.postproc_params: |
| 153 | result_handler, args = postproc_params.post_processor, postproc_params.postproc_args |
| 154 | args.tokenizer = self._tokenizer |
| 155 | out = result_handler(record, args) |
| 156 | else: |
| 157 | # This should only be called in streaming mode, and each time it |
| 158 | # produces a single output. |
| 159 | out = record.outputs[0] |
| 160 | |
| 161 | # TODO: Keep only the diff token_ids and text in streaming mode when |
| 162 | # result_handler is not set |
| 163 | return out, metrics_dict, perf_metrics, disaggregated_params |
| 164 | |
| 165 | async def _batched_put(self): |
| 166 | ''' Batched IPC send. ''' |
no test coverage detected