| 14 | logger = setup_logger("llm_processor") |
| 15 | |
| 16 | class InstanceManager: |
| 17 | def __init__(self, ports, gpus, base_url, startup_delay=5): |
| 18 | self.ports = ports |
| 19 | self.gpus = gpus |
| 20 | self.base_url = base_url |
| 21 | self.instances = [] |
| 22 | self.lock = threading.Lock() |
| 23 | self.current_instance = 0 # 用于轮询策略 |
| 24 | |
| 25 | for port, gpu in zip(self.ports, self.gpus): |
| 26 | self.start_instance(gpu, port) |
| 27 | self.instances.append({"port": port, "load": 0}) |
| 28 | time.sleep(startup_delay) # 等待所有实例启动 |
| 29 | |
| 30 | def start_instance(self, num, port): |
| 31 | """启动ollama实例在特定GPU和端口上""" |
| 32 | # cmd = f"CUDA_VISIBLE_DEVICES={num} OLLAMA_HOST={self.base_url}:{port} ollama serve" |
| 33 | cmd = f"OLLAMA_HOST={self.base_url}:{port} ollama serve" |
| 34 | print("Running command:", cmd) |
| 35 | # subprocess.Popen(cmd, shell=True, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL) |
| 36 | |
| 37 | def get_available_instance(self): |
| 38 | """使用轮询策略获取一个可用的实例""" |
| 39 | with self.lock: |
| 40 | instance = self.instances[self.current_instance] |
| 41 | self.current_instance = (self.current_instance + 1) % len(self.instances) |
| 42 | return instance["port"] # 返回端口 |
| 43 | |
| 44 | def generate_text(self, prompt, model="qwen:72b", temperature=0): |
| 45 | """发送请求到选择的实例""" |
| 46 | port = self.get_available_instance() |
| 47 | base_url = f"{self.base_url}:{port}" |
| 48 | |
| 49 | response = requests.post( |
| 50 | f"{base_url}/api/generate", |
| 51 | json={"model": model, "prompt": prompt, "temperature": temperature}, |
| 52 | timeout=30 # 设置超时时间，避免无限等待 |
| 53 | ) |
| 54 | response.raise_for_status() |
| 55 | return response |
| 56 | |
| 57 | |
| 58 | |