148 lines
4.8 KiB
Python
Executable File
148 lines
4.8 KiB
Python
Executable File
import argparse
|
|
import json
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
from typing import AsyncGenerator
|
|
|
|
import uvicorn
|
|
from fastapi import BackgroundTasks, FastAPI, Request
|
|
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
|
from vllm.engine.arg_utils import AsyncEngineArgs
|
|
from vllm.engine.async_llm_engine import AsyncLLMEngine
|
|
from vllm.sampling_params import SamplingParams
|
|
from vllm.transformers_utils.tokenizer import get_tokenizer
|
|
from vllm.utils import random_uuid
|
|
|
|
TIMEOUT_KEEP_ALIVE = 5 # seconds.
|
|
TIMEOUT_TO_PREVENT_DEADLOCK = 1 # seconds.
|
|
app = FastAPI()
|
|
|
|
served_model = None
|
|
model_config = {}
|
|
startup_time = int(time.time())
|
|
|
|
|
|
# TODO: figure out ROPE scaling
|
|
# TODO: make sure error messages are returned in the response
|
|
|
|
def get_gpu_pstate():
|
|
cmd = ['nvidia-smi', '--query-gpu=pstate', '--format=csv']
|
|
output = subprocess.check_output(cmd).decode('utf-8')
|
|
lines = output.strip().split('\n')
|
|
if len(lines) > 1:
|
|
return int(lines[1].strip('P'))
|
|
else:
|
|
return None
|
|
|
|
|
|
@app.get("/info")
|
|
async def generate(request: Request) -> Response:
|
|
return JSONResponse({
|
|
'uptime': int(time.time() - startup_time),
|
|
'startup_time': startup_time,
|
|
'model': served_model,
|
|
'model_config': model_config,
|
|
'nvidia': {
|
|
'pstate': get_gpu_pstate()
|
|
}
|
|
})
|
|
|
|
|
|
@app.get("/model")
|
|
async def generate(request: Request) -> Response:
|
|
return JSONResponse({'model': served_model, 'timestamp': int(time.time())})
|
|
|
|
|
|
@app.post("/tokenize")
|
|
async def generate(request: Request) -> Response:
|
|
request_dict = await request.json()
|
|
to_tokenize = request_dict.get("input")
|
|
if not to_tokenize:
|
|
JSONResponse({'error': 'must have input field'}, status_code=400)
|
|
tokens = tokenizer.tokenize(to_tokenize)
|
|
response = {}
|
|
if request_dict.get("return", False):
|
|
response['tokens'] = tokens
|
|
response['length'] = len(tokens)
|
|
return JSONResponse(response)
|
|
|
|
|
|
@app.post("/generate")
|
|
async def generate(request: Request) -> Response:
|
|
"""Generate completion for the request.
|
|
|
|
The request should be a JSON object with the following fields:
|
|
- prompt: the prompt to use for the generation.
|
|
- stream: whether to stream the results or not.
|
|
- other fields: the sampling parameters (See `SamplingParams` for details).
|
|
"""
|
|
request_dict = await request.json()
|
|
prompt = request_dict.pop("prompt")
|
|
stream = request_dict.pop("stream", False)
|
|
sampling_params = SamplingParams(**request_dict)
|
|
request_id = random_uuid()
|
|
results_generator = engine.generate(prompt, sampling_params, request_id)
|
|
|
|
# Streaming case
|
|
async def stream_results() -> AsyncGenerator[bytes, None]:
|
|
async for request_output in results_generator:
|
|
prompt = request_output.prompt
|
|
text_outputs = [
|
|
prompt + output.text for output in request_output.outputs
|
|
]
|
|
ret = {"text": text_outputs}
|
|
yield (json.dumps(ret) + "\0").encode("utf-8")
|
|
|
|
async def abort_request() -> None:
|
|
await engine.abort(request_id)
|
|
|
|
if stream:
|
|
background_tasks = BackgroundTasks()
|
|
# Abort the request if the client disconnects.
|
|
background_tasks.add_task(abort_request)
|
|
return StreamingResponse(stream_results(), background=background_tasks)
|
|
|
|
# Non-streaming case
|
|
final_output = None
|
|
async for request_output in results_generator:
|
|
if await request.is_disconnected():
|
|
# Abort the request if the client disconnects.
|
|
await engine.abort(request_id)
|
|
return Response(status_code=499)
|
|
final_output = request_output
|
|
|
|
assert final_output is not None
|
|
prompt = final_output.prompt
|
|
text_outputs = [prompt + output.text for output in final_output.outputs]
|
|
ret = {"text": text_outputs}
|
|
return JSONResponse(ret)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", type=str, default="localhost")
|
|
parser.add_argument("--port", type=int, default=8000)
|
|
parser = AsyncEngineArgs.add_cli_args(parser)
|
|
args = parser.parse_args()
|
|
engine_args = AsyncEngineArgs.from_cli_args(args)
|
|
|
|
model_path = Path(args.model)
|
|
served_model = model_path.name
|
|
|
|
try:
|
|
model_config = json.loads((model_path / 'config.json').read_text())
|
|
except Exception as e:
|
|
print(f"Failed to load the model's config - {e.__class__.__name__}: {e}")
|
|
|
|
engine = AsyncLLMEngine.from_engine_args(engine_args)
|
|
tokenizer = get_tokenizer(engine_args.tokenizer,
|
|
tokenizer_mode=args.tokenizer_mode,
|
|
trust_remote_code=args.trust_remote_code)
|
|
|
|
uvicorn.run(app,
|
|
host=args.host,
|
|
port=args.port,
|
|
log_level="debug",
|
|
timeout_keep_alive=TIMEOUT_KEEP_ALIVE)
|