| 143 | print(item['entity_name'],item['parent']) |
| 144 | |
| 145 | class InstanceManager: |
| 146 | def __init__(self, url,ports, gpus, generate_model,startup_delay=30): |
| 147 | self.ports = ports |
| 148 | self.gpus = gpus |
| 149 | self.base_url=url |
| 150 | self.instances = [] |
| 151 | self.lock = threading.Lock() |
| 152 | self.current_instance = 0 # 用于轮询策略 |
| 153 | self.generate_model=generate_model |
| 154 | self.TOTAL_TOKEN_COST = 0 |
| 155 | self.TOTAL_API_CALL_COST = 0 |
| 156 | for port, gpu in zip(self.ports, self.gpus): |
| 157 | self.instances.append({"port": port, "load": 0}) |
| 158 | def reset_token_cost(self): |
| 159 | """重置总的token消耗和API调用次数""" |
| 160 | self.TOTAL_TOKEN_COST = 0 |
| 161 | self.TOTAL_API_CALL_COST = 0 |
| 162 | def get_tokens_cosumption(self): |
| 163 | |
| 164 | return self.TOTAL_TOKEN_COST, self.TOTAL_API_CALL_COST |
| 165 | |
| 166 | # time.sleep(startup_delay) # 等待所有实例启动 |
| 167 | |
| 168 | |
| 169 | |
| 170 | |
| 171 | def get_available_instance(self): |
| 172 | """使用轮询策略获取一个可用的实例""" |
| 173 | with self.lock: |
| 174 | instance = self.instances[self.current_instance] |
| 175 | self.current_instance = (self.current_instance + 1) % len(self.instances) |
| 176 | return instance["port"] # 返回端口 |
| 177 | |
| 178 | def generate_text(self, prompt,system_prompt=None, history_messages=[], **kwargs): |
| 179 | """发送请求到选择的实例""" |
| 180 | port = self.get_available_instance() |
| 181 | base_url = f"{self.base_url}:{port}/v1" |
| 182 | messages = [] |
| 183 | if system_prompt: |
| 184 | messages.append({"role": "system", "content": system_prompt}) |
| 185 | |
| 186 | # Get the cached response if having------------------- |
| 187 | |
| 188 | if len(history_messages)>1: |
| 189 | history_messages[0]['content'] = truncate_text(history_messages[0]['content'], max_tokens=3000) |
| 190 | history_messages[1]['content'] = truncate_text(history_messages[1]['content'], max_tokens=25000) |
| 191 | messages.extend(history_messages) |
| 192 | messages.append({"role": "user", "content": prompt}) |
| 193 | try: |
| 194 | cur_token_cost = len(tokenizer.encode(messages[0]['content'])) |
| 195 | if cur_token_cost>31000: |
| 196 | cur_token_cost = 31000 |
| 197 | messages[0]['content'] = truncate_text(messages[0]['content'], max_tokens=31000) |
| 198 | |
| 199 | # logging api call cost |
| 200 | self.TOTAL_API_CALL_COST += 1 |
| 201 | response = requests.post( |
| 202 | f"{base_url}/chat/completions", |
no outgoing calls
no test coverage detected