|
import asyncio |
|
import threading |
|
import time |
|
from typing import List |
|
|
|
from fastapi import FastAPI, Request, BackgroundTasks |
|
from fastapi.responses import StreamingResponse, JSONResponse |
|
import requests |
|
|
|
from src.constants import WORKER_HEART_BEAT_INTERVAL |
|
from src.conversation import Conversation |
|
from src.utils import pretty_print_semaphore, build_logger |
|
|
|
|
|
worker = None |
|
logger = None |
|
|
|
app = FastAPI() |
|
|
|
|
|
def heart_beat_worker(obj): |
|
while True: |
|
time.sleep(WORKER_HEART_BEAT_INTERVAL) |
|
obj.send_heart_beat() |
|
|
|
|
|
class BaseModelWorker: |
|
def __init__( |
|
self, |
|
controller_addr: str, |
|
worker_addr: str, |
|
worker_id: str, |
|
model_path: str, |
|
model_names: List[str], |
|
limit_worker_concurrency: int, |
|
conv_template: str = None, |
|
multimodal: bool = False, |
|
): |
|
global logger, worker |
|
|
|
self.controller_addr = controller_addr |
|
self.worker_addr = worker_addr |
|
self.worker_id = worker_id |
|
if model_path.endswith("/"): |
|
model_path = model_path[:-1] |
|
self.model_names = model_names or [model_path.split("/")[-1]] |
|
self.limit_worker_concurrency = limit_worker_concurrency |
|
self.conv = self.make_conv_template(conv_template, model_path) |
|
self.conv.sep_style = int(self.conv.sep_style) |
|
self.multimodal = multimodal |
|
self.tokenizer = None |
|
self.context_len = None |
|
self.call_ct = 0 |
|
self.semaphore = None |
|
|
|
self.heart_beat_thread = None |
|
|
|
if logger is None: |
|
logger = build_logger("model_worker", f"model_worker_{self.worker_id}.log") |
|
if worker is None: |
|
worker = self |
|
|
|
def make_conv_template( |
|
self, |
|
conv_template: str = None, |
|
model_path: str = None, |
|
) -> Conversation: |
|
""" |
|
can be overrided to costomize the conversation template for different model workers. |
|
""" |
|
from fastchat.conversation import get_conv_template |
|
from fastchat.model.model_adapter import get_conversation_template |
|
|
|
if conv_template: |
|
conv = get_conv_template(conv_template) |
|
else: |
|
conv = get_conversation_template(model_path) |
|
return conv |
|
|
|
def init_heart_beat(self): |
|
self.register_to_controller() |
|
self.heart_beat_thread = threading.Thread( |
|
target=heart_beat_worker, |
|
args=(self,), |
|
daemon=True, |
|
) |
|
self.heart_beat_thread.start() |
|
|
|
def register_to_controller(self): |
|
logger.info("Register to controller") |
|
|
|
url = self.controller_addr + "/register_worker" |
|
data = { |
|
"worker_name": self.worker_addr, |
|
"check_heart_beat": True, |
|
"worker_status": self.get_status(), |
|
"multimodal": self.multimodal, |
|
} |
|
r = requests.post(url, json=data) |
|
assert r.status_code == 200 |
|
|
|
def send_heart_beat(self): |
|
logger.info( |
|
f"Send heart beat. Models: {self.model_names}. " |
|
f"Semaphore: {pretty_print_semaphore(self.semaphore)}. " |
|
f"call_ct: {self.call_ct}. " |
|
f"worker_id: {self.worker_id}. " |
|
) |
|
|
|
url = self.controller_addr + "/receive_heart_beat" |
|
|
|
while True: |
|
try: |
|
ret = requests.post( |
|
url, |
|
json={ |
|
"worker_name": self.worker_addr, |
|
"queue_length": self.get_queue_length(), |
|
}, |
|
timeout=5, |
|
) |
|
exist = ret.json()["exist"] |
|
break |
|
except (requests.exceptions.RequestException, KeyError) as e: |
|
logger.error(f"heart beat error: {e}") |
|
time.sleep(5) |
|
|
|
if not exist: |
|
self.register_to_controller() |
|
|
|
def get_queue_length(self): |
|
if self.semaphore is None: |
|
return 0 |
|
else: |
|
sempahore_value = ( |
|
self.semaphore._value |
|
if self.semaphore._value is not None |
|
else self.limit_worker_concurrency |
|
) |
|
waiter_count = ( |
|
0 if self.semaphore._waiters is None else len(self.semaphore._waiters) |
|
) |
|
return self.limit_worker_concurrency - sempahore_value + waiter_count |
|
|
|
def get_status(self): |
|
return { |
|
"model_names": self.model_names, |
|
"speed": 1, |
|
"queue_length": self.get_queue_length(), |
|
} |
|
|
|
def count_token(self, params): |
|
prompt = params["prompt"] |
|
|
|
try: |
|
input_ids = self.tokenizer(prompt).input_ids |
|
input_echo_len = len(input_ids) |
|
except TypeError: |
|
input_echo_len = self.tokenizer.num_tokens(prompt) |
|
|
|
ret = { |
|
"count": input_echo_len, |
|
"error_code": 0, |
|
} |
|
return ret |
|
|
|
def get_conv_template(self): |
|
return {"conv": self.conv} |
|
|
|
def generate_stream_gate(self, params): |
|
raise NotImplementedError |
|
|
|
def generate_gate(self, params): |
|
raise NotImplementedError |
|
|
|
def get_embeddings(self, params): |
|
raise NotImplementedError |
|
|
|
|
|
def release_worker_semaphore(): |
|
worker.semaphore.release() |
|
|
|
|
|
def acquire_worker_semaphore(): |
|
if worker.semaphore is None: |
|
worker.semaphore = asyncio.Semaphore(worker.limit_worker_concurrency) |
|
return worker.semaphore.acquire() |
|
|
|
|
|
def create_background_tasks(): |
|
background_tasks = BackgroundTasks() |
|
background_tasks.add_task(release_worker_semaphore) |
|
return background_tasks |
|
|
|
|
|
@app.post("/worker_generate_stream") |
|
async def api_generate_stream(request: Request): |
|
params = await request.json() |
|
await acquire_worker_semaphore() |
|
generator = worker.generate_stream_gate(params) |
|
background_tasks = create_background_tasks() |
|
return StreamingResponse(generator, background=background_tasks) |
|
|
|
|
|
@app.post("/worker_generate") |
|
async def api_generate(request: Request): |
|
params = await request.json() |
|
await acquire_worker_semaphore() |
|
output = await asyncio.to_thread(worker.generate_gate, params) |
|
release_worker_semaphore() |
|
return JSONResponse(output) |
|
|
|
|
|
@app.post("/worker_get_embeddings") |
|
async def api_get_embeddings(request: Request): |
|
params = await request.json() |
|
await acquire_worker_semaphore() |
|
embedding = worker.get_embeddings(params) |
|
release_worker_semaphore() |
|
return JSONResponse(content=embedding) |
|
|
|
|
|
@app.post("/worker_get_status") |
|
async def api_get_status(request: Request): |
|
return worker.get_status() |
|
|
|
|
|
@app.post("/count_token") |
|
async def api_count_token(request: Request): |
|
params = await request.json() |
|
return worker.count_token(params) |
|
|
|
|
|
@app.post("/worker_get_conv_template") |
|
async def api_get_conv(request: Request): |
|
return worker.get_conv_template() |
|
|
|
|
|
@app.post("/model_details") |
|
async def api_model_details(request: Request): |
|
return {"context_length": worker.context_len} |
|
|
