(self)
| 13535 | self.metrics.requests_admitted_total += 1 |
| 13536 | |
| 13537 | def build_batch(self) -> List[CompletionScheduler.BatchItem]: |
| 13538 | prompt_requests = [ |
| 13539 | self.requests[request_id] |
| 13540 | for request_id in self.active_request_ids |
| 13541 | if self.requests[request_id].admitted |
| 13542 | and not self.requests[request_id].prompt_done |
| 13543 | and not self.requests[request_id].cancelled |
| 13544 | ] |
| 13545 | completions = [ |
| 13546 | completion |
| 13547 | for request_id in self.active_request_ids |
| 13548 | for completion in self.requests[request_id].completions |
| 13549 | if (completion.pending_input_tokens or completion.draft_tokens) |
| 13550 | and not completion.finished |
| 13551 | ] |
| 13552 | if not prompt_requests and not completions: |
| 13553 | return [] |
| 13554 | |
| 13555 | for request in prompt_requests: |
| 13556 | if request.prompt_cursor < len(request.prompt_tokens): |
| 13557 | segment = self._current_prompt_segment(request) |
| 13558 | if segment.kind != "text": |
| 13559 | token_count = min( |
| 13560 | segment.batch_rows |
| 13561 | if segment.media is not None and segment.media.non_causal |
| 13562 | else self._pending_tokens_length(request), |
| 13563 | self.model.n_batch, |
| 13564 | ) |
| 13565 | if segment.batch_rows > self.model.n_batch: |
| 13566 | if segment.media is not None and segment.media.non_causal: |
| 13567 | raise RuntimeError( |
| 13568 | "non-causal media prompt segment exceeds model.n_batch; " |
| 13569 | "increase n_batch" |
| 13570 | ) |
| 13571 | item, _ = self._build_pending_batch_item( |
| 13572 | request, |
| 13573 | token_count, |
| 13574 | 0, |
| 13575 | ) |
| 13576 | return [item] |
| 13577 | |
| 13578 | ordered_sequences = self._ordered_pending_sequences(prompt_requests, completions) |
| 13579 | allocations = self._allocate_pending_tokens_for_model(ordered_sequences) |
| 13580 | if ordered_sequences: |
| 13581 | self.sequence_round_robin += 1 |
| 13582 | |
| 13583 | items: List[CompletionScheduler.BatchItem] = [] |
| 13584 | output_index = 0 |
| 13585 | for source in ordered_sequences: |
| 13586 | token_count = allocations.get(self._pending_sequence_id(source), 0) |
| 13587 | if token_count <= 0: |
| 13588 | continue |
| 13589 | item, output_index = self._build_pending_batch_item( |
| 13590 | source, |
| 13591 | token_count, |
| 13592 | output_index, |
| 13593 | ) |
| 13594 | items.append(item) |
no test coverage detected