record and allocate resources for the engine
| 26 | |
| 27 | |
| 28 | class ResourceManager: |
| 29 | """ |
| 30 | record and allocate resources for the engine |
| 31 | """ |
| 32 | |
| 33 | def __init__( |
| 34 | self, |
| 35 | max_num_seqs, |
| 36 | config, |
| 37 | tensor_parallel_size, |
| 38 | splitwise_role, |
| 39 | local_data_parallel_id=0, |
| 40 | ): |
| 41 | """ |
| 42 | Args: |
| 43 | cfg (Config): config object containing parameters for the engine |
| 44 | initialization |
| 45 | |
| 46 | Returns: |
| 47 | None |
| 48 | |
| 49 | Initializes the engine with the given configuration and sets up necessary |
| 50 | data structures to manage tasks and blocks. |
| 51 | """ |
| 52 | self.cfg = config.cache_config |
| 53 | self.max_num_seqs = max_num_seqs |
| 54 | self.stop_flags = [True] * max_num_seqs # flag set to true if the slot has not been taken |
| 55 | self.enable_prefix_cache = config.cache_config.enable_prefix_caching |
| 56 | self.cache_manager = PrefixCacheManager(config, tensor_parallel_size, splitwise_role, local_data_parallel_id) |
| 57 | self.tasks_list = [None] * max_num_seqs # task slots |
| 58 | self.req_dict = dict() |
| 59 | # current batch status of the engine |
| 60 | self.real_bsz = 0 |
| 61 | self.abort_req_ids_set = set() |
| 62 | llm_logger.info(f"{self.info()}") |
| 63 | main_process_metrics.max_batch_size.set(max_num_seqs) |
| 64 | |
| 65 | def reset_cache_config(self, cfg): |
| 66 | """ |
| 67 | reset cache config |
| 68 | """ |
| 69 | self.cfg = cfg |
| 70 | self.cache_manager.update_cache_config(cfg) |
| 71 | |
| 72 | def get_required_block_number(self, input_token_num): |
| 73 | """ |
| 74 | Calculate Block resources are needed |
| 75 | |
| 76 | Args: |
| 77 | input_token_num (int): input token number |
| 78 | |
| 79 | Returns: |
| 80 | int: block number |
| 81 | """ |
| 82 | block_num = (input_token_num + self.cfg.block_size - 1 + self.cfg.dec_token_num) // self.cfg.block_size |
| 83 | return block_num |
| 84 | |
| 85 | def get_encoder_block_number(self, input_token_num): |