()
| 835 | is_fetching = False |
| 836 | |
| 837 | def _fetch_request(): |
| 838 | try: |
| 839 | with self._pause_cond: |
| 840 | self._pause_cond.wait_for(lambda: not self.is_paused) |
| 841 | nonlocal is_fetching |
| 842 | num_prefill_batch = min( |
| 843 | int(self.resource_manager.available_batch()), |
| 844 | self.cfg.max_prefill_batch, |
| 845 | ) |
| 846 | |
| 847 | if self.cfg.scheduler_config.splitwise_role != "mixed": |
| 848 | max_num_batched_tokens = self.cfg.scheduler_config.max_num_batched_tokens |
| 849 | else: |
| 850 | max_num_batched_tokens = self.cfg.model_config.max_model_len |
| 851 | |
| 852 | available_blocks = self.cfg.cache_config.max_block_num_per_seq |
| 853 | tasks = self.scheduler.get_requests( |
| 854 | available_blocks=available_blocks, |
| 855 | block_size=self.cfg.cache_config.block_size, |
| 856 | reserved_output_blocks=0, # self.cfg.cache_config.enc_dec_block_num |
| 857 | max_num_batched_tokens=max_num_batched_tokens, |
| 858 | batch=num_prefill_batch, |
| 859 | ) |
| 860 | for task in tasks: |
| 861 | task.metrics.engine_get_req_time = time.time() |
| 862 | trace_print(LoggingEventName.REQUEST_QUEUE_END, task.request_id, getattr(task, "user", "")) |
| 863 | |
| 864 | if self.cfg.scheduler_config.splitwise_role == "decode": |
| 865 | # TODO: refine scheduler to remove this limitation |
| 866 | # Decode will process and schedule the request sent by prefill to engine, |
| 867 | # so the same request sent by the decode api server will be ignored |
| 868 | is_fetching = False |
| 869 | return |
| 870 | |
| 871 | if tasks: |
| 872 | self.llm_logger.debug( |
| 873 | f"Engine has fetched tasks from {self.scheduler.__class__.__name__}: {[task.request_id for task in tasks]}" |
| 874 | ) |
| 875 | |
| 876 | if self.cfg.scheduler_config.splitwise_role == "prefill": |
| 877 | for task in tasks: |
| 878 | # start async preprocess |
| 879 | self.resource_manager.apply_async_preprocess(task) |
| 880 | need_delete_tasks = [] |
| 881 | if envs.PREFILL_CONTINUOUS_REQUEST_DECODE_RESOURCES: |
| 882 | for task in tasks: |
| 883 | # assure can allocate block ids in P |
| 884 | while not self.resource_manager.preallocate_resource_in_p(task): |
| 885 | time.sleep(0.005) |
| 886 | self.llm_logger.debug( |
| 887 | f"P has allocated resources and then ask D resource for request: {task.request_id}" |
| 888 | ) |
| 889 | task.metrics.ask_decode_resource_start_time = time.time() |
| 890 | while True: |
| 891 | self.split_connector.send_splitwise_tasks([task], task.idx) |
| 892 | status, msg = self.split_connector.check_decode_allocated(task) |
| 893 | if not status: |
| 894 | self.llm_logger.warning( |
nothing calls this directly
no test coverage detected