Kkit.llm_utils.llamacpp_wrapper_server

  1import subprocess
  2import json
  3import time
  4from fastapi import FastAPI, HTTPException
  5from pydantic import BaseModel
  6from typing import Dict, Any
  7import uvicorn
  8from contextlib import asynccontextmanager
  9import threading
 10import argparse
 11from typing import List
 12from fastapi.logger import logger
 13from logging import StreamHandler, Formatter
 14import logging
 15
 16
 17server_lock = threading.Lock()
 18llama_server = None
 19
 20class ServerConfig(BaseModel):
 21    model_name: str
 22    configs: Dict[str, Any]
 23    server_path: List[str] = ["./llama-server"]
 24
 25class SwitchRequest(BaseModel):
 26    new_model_name: str
 27    new_configs: Dict[str, Any]
 28
 29class LlamaServer:
 30    def __init__(self, model_name: str, configs: Dict[str, Any], server_path: List[str]):
 31        self.server_path = server_path
 32        self.model_name = model_name
 33        self.configs = configs.copy()
 34        self.process = None
 35
 36    def _convert_configs_to_args(self) -> list:
 37        args = []
 38        for key, value in self.configs.items():
 39            if key.startswith("-"):
 40                arg_name = key
 41            else:
 42                arg_name = f"--{key.replace('_', '-')}"
 43            
 44            if isinstance(value, bool):
 45                if value:
 46                    args.append(arg_name)
 47            elif value == '':
 48                args.extend([arg_name])
 49            else:
 50                args.extend([arg_name, str(value)])
 51        return args
 52
 53    def start(self, log_file: str = "llama_server.log"):
 54        if self.process and self.process.poll() is None:
 55            raise RuntimeError("Server is already running")
 56        cmd = (self.server_path).copy()
 57        if app.state.llama_cpp_or_vllm == "llama_cpp":
 58            cmd.extend(["--model", self.model_name])
 59        elif app.state.llama_cpp_or_vllm == "vllm":
 60            cmd.extend([self.model_name])
 61        cmd.extend(self._convert_configs_to_args())
 62        print(cmd)
 63
 64        try:
 65            with open(log_file, "w") as log_f:
 66                self.process = subprocess.Popen(
 67                    cmd,
 68                    stdout=log_f,
 69                    stderr=log_f,
 70                    text=True,
 71                )
 72            time.sleep(1)
 73            if self.process.poll() is not None:
 74                raise RuntimeError(f"Server failed to start. Exit code: {self.process.returncode}")
 75        except Exception as e:
 76            self.process = None
 77            raise RuntimeError(f"Error starting server: {str(e)}")
 78
 79    def stop(self):
 80        if self.process:
 81            try:
 82                self.process.terminate()
 83                self.process.wait(timeout=5)
 84            except subprocess.TimeoutExpired:
 85                self.process.kill()
 86            finally:
 87                self.process = None
 88
 89    def switch(self, new_model_name: str, new_configs: Dict[str, Any]):
 90        self.stop()
 91        self.model_name = new_model_name
 92        self.configs = new_configs.copy()
 93        self.start(app.state.log_file)
 94
 95    def is_running(self) -> bool:
 96        return self.process is not None and self.process.poll() is None
 97
 98@asynccontextmanager
 99async def lifespan(app: FastAPI):
100    logger.info(f"Run in {app.state.llama_cpp_or_vllm} mode...")
101    yield
102    logger.info("Stopping server...")
103    global llama_server
104    if llama_server and llama_server.is_running():
105        llama_server.stop()
106
107app = FastAPI(lifespan=lifespan)
108
109@app.post("/wrapper_start")
110def start_server(config: ServerConfig):
111    global llama_server
112    with server_lock:
113        if llama_server and llama_server.is_running():
114            raise HTTPException(status_code=400, detail="Server is already running")
115        
116        try:
117            llama_server = LlamaServer(
118                model_name=config.model_name,
119                configs=config.configs,
120                server_path=config.server_path
121            )
122            llama_server.start(app.state.log_file)
123            return {
124                "status": "success",
125                "message": f"Server started with model: {config.model_name}",
126                "model": config.model_name,
127                "config": config.configs
128            }
129        except Exception as e:
130            raise HTTPException(status_code=500, detail=str(e))
131
132@app.post("/wrapper_switch")
133def switch_model(request: SwitchRequest):
134    global llama_server
135    with server_lock:
136        if not llama_server or not llama_server.is_running():
137            raise HTTPException(status_code=400, detail="Server is not running")
138        
139        try:
140            llama_server.switch(
141                new_model_name=request.new_model_name,
142                new_configs=request.new_configs
143            )
144            return {
145                "status": "success",
146                "message": f"Switched to model: {request.new_model_name}",
147                "new_model": request.new_model_name,
148                "new_config": request.new_configs
149            }
150        except Exception as e:
151            raise HTTPException(status_code=500, detail=str(e))
152
153@app.post("/wrapper_stop")
154def stop_server():
155    global llama_server
156    with server_lock:
157        if not llama_server or not llama_server.is_running():
158            raise HTTPException(status_code=400, detail="Server is not running")
159        
160        try:
161            llama_server.stop()
162            return {"status": "success", "message": "Server stopped"}
163        except Exception as e:
164            raise HTTPException(status_code=500, detail=str(e))
165
166@app.get("/wrapper_status")
167def get_status():
168    global llama_server
169    status = {
170        "is_running": False,
171        "current_model": None,
172        "current_config": None
173    }
174    
175    if llama_server and llama_server.is_running():
176        status.update({
177            "is_running": True,
178            "current_model": llama_server.model_name,
179            "current_config": llama_server.configs
180        })
181    
182    return status
183
184def main():
185    parser = argparse.ArgumentParser()
186    parser.add_argument("--host", type=str, default="0.0.0.0", help="Server listening host")
187    parser.add_argument("--port", type=int, default=8001, help="Server listening port")
188    parser.add_argument("--log_file", type=str, default="./server.log", help="Log file path")
189    parser.add_argument("--llama_cpp_or_vllm", "-lv", default="vllm", help="using `vllm` or `llama_cpp` as backend")
190    args = parser.parse_args()
191    args.llama_cpp_or_vllm in ["vllm", "llama_cpp"]
192    app.state.log_file = args.log_file
193    app.state.llama_cpp_or_vllm = args.llama_cpp_or_vllm
194    handler = StreamHandler()
195    formater = Formatter("%(levelname)s:     %(message)s")
196    handler.setFormatter(formater)
197    logger.setLevel(logging.INFO)
198    logger.addHandler(handler)
199    uvicorn.run(app, host="0.0.0.0", port=args.port)
200
201if __name__ == "__main__":
202    main()