fastchat/serve/dashinfer_worker.py

"""
A model worker that executes the model based on dash-infer.

See documentations at docs/dashinfer_integration.md
"""

import argparse
import asyncio
import copy
import json
import os
import subprocess
from typing import List

from fastapi import FastAPI, Request, BackgroundTasks
from fastapi.responses import StreamingResponse, JSONResponse
import uvicorn
from dashinfer.helper import EngineHelper, ConfigManager

from fastchat.constants import ErrorCode, SERVER_ERROR_MSG
from fastchat.serve.base_model_worker import BaseModelWorker
from fastchat.serve.model_worker import (
    logger,
    worker_id,
)
from fastchat.utils import build_logger, get_context_length, is_partial_stop


app = FastAPI()


def download_model(model_id, revision):
    source = "huggingface"
    if os.environ.get("FASTCHAT_USE_MODELSCOPE", "False").lower() == "true":
        source = "modelscope"

    logger.info(f"Downloading model {model_id} (revision: {revision}) from {source}")
    if source == "modelscope":
        from modelscope import snapshot_download

        model_dir = snapshot_download(model_id, revision=revision)
    elif source == "huggingface":
        from huggingface_hub import snapshot_download

        model_dir = snapshot_download(repo_id=model_id)
    else:
        raise ValueError("Unknown source")

    logger.info(f"Save model to path {model_dir}")

    return model_dir


class DashInferWorker(BaseModelWorker):
    def __init__(
        self,
        controller_addr: str,
        worker_addr: str,
        worker_id: str,
        model_path: str,
        model_names: List[str],
        limit_worker_concurrency: int,
        revision: str,
        no_register: bool,
        config: json,
        conv_template: str,
    ):
        super().__init__(
            controller_addr,
            worker_addr,
            worker_id,
            model_path,
            model_names,
            limit_worker_concurrency,
            conv_template,
        )

        logger.info(
            f"Loading the model {self.model_names} on worker {worker_id}, worker type: dash-infer worker..."
        )
        # check if model_path is existed at local path
        if not os.path.exists(model_path):
            model_path = download_model(model_path, revision)
        engine_helper = EngineHelper(config)
        engine_helper.init_tokenizer(model_path)
        engine_helper.convert_model(model_path)
        engine_helper.init_engine()

        self.context_len = engine_helper.engine_config["engine_max_length"]
        self.tokenizer = engine_helper.tokenizer
        self.engine_helper = engine_helper

        if not no_register:
            self.init_heart_beat()

    async def generate_stream(self, params):
        self.call_ct += 1

        context = params.pop("prompt")
        temperature = params.get("temperature")
        top_k = params.get("top_k")
        top_p = params.get("top_p")
        repetition_penalty = params.get("repetition_penalty")
        presence_penalty = params.get("presence_penalty")
        max_new_tokens = params.get("max_new_tokens")
        stop_token_ids = params.get("stop_token_ids") or []
        if self.tokenizer.eos_token_id is not None:
            stop_token_ids.append(self.tokenizer.eos_token_id)
        seed = params.get("seed")
        echo = params.get("echo", True)
        logprobs = params.get("logprobs")
        # not supported parameters
        frequency_penalty = params.get("frequency_penalty")
        stop = params.get("stop")
        use_beam_search = params.get("use_beam_search", False)
        best_of = params.get("best_of", None)

        gen_cfg = copy.deepcopy(self.engine_helper.default_gen_cfg) or dict()
        if temperature is not None:
            gen_cfg["temperature"] = float(temperature)
        if top_k is not None:
            dashinfer_style_top_k = 0 if int(top_k) == -1 else int(top_k)
            gen_cfg["top_k"] = dashinfer_style_top_k
        if top_p is not None:
            gen_cfg["top_p"] = float(top_p)
        if repetition_penalty is not None:
            gen_cfg["repetition_penalty"] = float(repetition_penalty)
        if presence_penalty is not None:
            gen_cfg["presence_penalty"] = float(presence_penalty)
        if len(stop_token_ids) != 0:
            dashinfer_style_stop_token_ids = [[id] for id in set(stop_token_ids)]
            logger.info(
                f"dashinfer_style_stop_token_ids = {dashinfer_style_stop_token_ids}"
            )
            gen_cfg["stop_words_ids"] = dashinfer_style_stop_token_ids
        if seed is not None:
            gen_cfg["seed"] = int(seed)
        if logprobs is not None:
            gen_cfg["logprobs"] = True
            gen_cfg["top_logprobs"] = int(logprobs)
        if frequency_penalty is not None:
            logger.warning(
                "dashinfer worker does not support `frequency_penalty` parameter"
            )
        if stop is not None:
            logger.warning("dashinfer worker does not support `stop` parameter")
        if use_beam_search == True:
            logger.warning(
                "dashinfer worker does not support `use_beam_search` parameter"
            )
        if best_of is not None:
            logger.warning("dashinfer worker does not support `best_of` parameter")

        logger.info(
            f"dashinfer engine helper creates request with context: {context}, gen_cfg: {gen_cfg}"
        )

        request_list = self.engine_helper.create_request([context], gen_cfg=[gen_cfg])

        engine_req = request_list[0]

        # check if prompt tokens exceed the max_tokens
        max_tokens = (
            gen_cfg["max_length"]
            if max_new_tokens is None
            else engine_req.in_tokens_len + max_new_tokens
        )
        if engine_req.in_tokens_len > max_tokens:
            ret = {
                "text": f"This model's maximum generated tokens include context are {max_tokens}, However, your context resulted in {engine_req.in_tokens_len} tokens",
                "error_code": ErrorCode.CONTEXT_OVERFLOW,
            }
            yield json.dumps(ret).encode() + b"\0"
        else:
            gen_cfg["max_length"] = int(max_tokens)
            logger.info(
                f"dashinfer is going to process one request in stream mode: {engine_req}"
            )
            results_generator = self.engine_helper.process_one_request_stream(
                engine_req
            )

            try:
                for generate_text in results_generator:
                    if echo:
                        output_text = context + generate_text
                    else:
                        output_text = generate_text
                    prompt_tokens = engine_req.in_tokens_len
                    completion_tokens = engine_req.out_tokens_len
                    ret = {
                        "text": output_text,
                        "error_code": 0,
                        "usage": {
                            "prompt_tokens": prompt_tokens,
                            "completion_tokens": completion_tokens,
                            "total_tokens": prompt_tokens + completion_tokens,
                        },
                    }
                    yield (json.dumps(ret) + "\0").encode()
            except Exception as e:
                ret = {
                    "text": f"{SERVER_ERROR_MSG}\n\n({e})",
                    "error_code": ErrorCode.INTERNAL_ERROR,
                }
                yield json.dumps(ret).encode() + b"\0"

    async def generate(self, params):
        async for x in self.generate_stream(params):
            pass
        return json.loads(x[:-1].decode())


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(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 worker.generate(params)
    release_worker_semaphore()
    return JSONResponse(output)


@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}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="localhost")
    parser.add_argument("--port", type=int, default=21002)
    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
    parser.add_argument(
        "--controller-address", type=str, default="http://localhost:21001"
    )
    parser.add_argument("--model-path", type=str, default="qwen/Qwen-7B-Chat")
    parser.add_argument(
        "--model-names",
        type=lambda s: s.split(","),
        help="Optional display comma separated names",
    )

    parser.add_argument("--limit-worker-concurrency", type=int, default=1024)
    parser.add_argument("--no-register", action="store_true")
    parser.add_argument(
        "--revision",
        type=str,
        default="main",
        help="Hugging Face Hub model revision identifier",
    )
    parser.add_argument(
        "--conv-template", type=str, default=None, help="Conversation prompt template."
    )
    parser.add_argument(
        "config_file",
        metavar="config-file",
        type=str,
        default="config_qwen_v10_7b.json",
        help="A model config file which dash-inferread",
    )

    args = parser.parse_args()
    config = ConfigManager.get_config_from_json(args.config_file)

    cmd = f"pip show dashinfer | grep 'Location' | cut -d ' ' -f 2"
    package_location = subprocess.run(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, text=True
    )
    package_location = package_location.stdout.strip()
    os.environ["AS_DAEMON_PATH"] = package_location + "/dashinfer/allspark/bin"
    os.environ["AS_NUMA_NUM"] = str(len(config["device_ids"]))
    os.environ["AS_NUMA_OFFSET"] = str(config["device_ids"][0])
    worker = DashInferWorker(
        args.controller_address,
        args.worker_address,
        worker_id,
        args.model_path,
        args.model_names,
        args.limit_worker_concurrency,
        args.revision,
        args.no_register,
        config,
        args.conv_template,
    )
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")