airockchip / rknn-llm

Other
258 stars 26 forks source link

rknn_init failed #31

Open YanxingLiu opened 2 months ago

YanxingLiu commented 2 months ago

我尝试利用 rknn-llm 构建一个 OpenAI API兼容的server,由于我对 python 比较熟悉,因此我尝试通过 python 调用 librknnrt.so,并将librknnrt.so的一些函数进行了封装。 rkllm_createDefaultParam 可以正常工作,但是rkllm_init在进行初始化时会自行退出,并且没有任何提示信息,由于librknnrt.so是闭源的,我无法知道究竟发生了什么,我自行编写了一个和 rkllm_init具有相同声明的函数,它可以正常工作。(PS 驱动应该没有问题,rknn-llm 的 cpp examples 可以正常运行) 以下是我编写的代码:

# rknn_llm.py
import ctypes
import functools
import pathlib
import sys
import os

from typing import (
    Any,
    Callable,
    List,
    Union,
    NewType,
    Optional,
    TYPE_CHECKING,
    TypeVar,
    Generic,
)

# Load the library
def _load_shared_library(lib_base_name: str):
    # Construct the paths to the possible shared library names
    _base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__)))
    _lib_paths: List[pathlib.Path] = []
    _lib_paths += [
        _base_path /
        f"../vender/rknn-llm/rkllm-runtime/runtime/Linux/librkllm_api/aarch64/lib{lib_base_name}.so",
    ]

    cdll_args = dict()  # type: ignore

    # Try to load the shared library, handling potential errors
    for _lib_path in _lib_paths:
        if _lib_path.exists():
            try:
                return ctypes.CDLL(str(_lib_path), **cdll_args)  # type: ignore
            except Exception as e:
                raise RuntimeError(
                    f"Failed to load shared library '{_lib_path}': {e}")

    raise FileNotFoundError(
        f"Shared library with base name '{lib_base_name}' not found"
    )

_lib_base_name = "rkllmrt"

_lib = _load_shared_library(_lib_base_name)

F = TypeVar("F", bound=Callable[..., Any])

def ctypes_function_for_shared_library(lib: ctypes.CDLL):
    def ctypes_function(
        name: str, argtypes: List[Any], restype: Any, enabled: bool = True
    ):
        def decorator(f: F) -> F:
            if enabled:
                func = getattr(lib, name)
                func.argtypes = argtypes
                func.restype = restype
                functools.wraps(f)(func)
                return func
            else:
                return f

        return decorator

    return ctypes_function

ctypes_function = ctypes_function_for_shared_library(_lib)

LLMHandle = ctypes.c_void_p

LLM_RUN_NORMAL = 0
LLM_RUN_FINISH = 1
LLM_RUN_ERROR = 2

class RKLLMParam(ctypes.Structure):
    _fields_ = [
        ("modelPath", ctypes.c_char_p),
        ("target_platform", ctypes.c_char_p),

        ("num_npu_core", ctypes.c_int32),
        ("max_context_len", ctypes.c_int32),
        ("max_new_tokens", ctypes.c_int32),

        ("top_k", ctypes.c_int32),
        ("top_p", ctypes.c_float),
        ("temperature", ctypes.c_float),
        ("repeat_penalty", ctypes.c_float),
        ("frequency_penalty", ctypes.c_float),
        ("presence_penalty", ctypes.c_float),
        ("mirostat", ctypes.c_int32),
        ("mirostat_tau", ctypes.c_float),
        ("mirostat_eta", ctypes.c_float)
    ]

LLMResultCallback = ctypes.CFUNCTYPE(None, ctypes.c_char_p, ctypes.c_void_p, ctypes.c_int32)

@ctypes_function(
    "rkllm_createDefaultParam",
    [],
    RKLLMParam,
)
def rkllm_createDefaultParam():
    ...

@ctypes_function(
    "rkllm_init",
    [LLMHandle, RKLLMParam, LLMResultCallback],
    ctypes.c_int
)
def rkllm_init(handle: LLMHandle,
               param: RKLLMParam,
               callback: LLMResultCallback):
    ...

@ctypes_function(
    "rkllm_run",
    [LLMHandle, ctypes.c_char_p, ctypes.c_void_p],
    ctypes.c_int
)
def rkllm_run(handle: LLMHandle,
              prompt: ctypes.c_char_p,
              userdata: ctypes.c_void_p):
    ...

@ctypes_function(
    'rkllm_destroy',
    [LLMHandle],
    ctypes.c_int,
)
def rkllm_destroy(handle: LLMHandle):
    ...

读取模型函数

import argparse
import ctypes
import sys
import os

sys.path.append(os.getcwd())
import rkllm
from rkllm import rkllm_createDefaultParam, rkllm_init, rkllm_run, rkllm_destroy, LLMHandle

def callback(text, userdata, state):
    if (state == rkllm.LLM_RUN_NORMAL):
        print("llm_run_normal")
    elif (state == rkllm.LLM_RUN_FINISH):
        print("llm_run_finish")
    elif (state == rkllm.LLM_RUN_ERROR):
        print("llm_run_error")
    else:
        print("%s", text)

def parse_args():
    parser = argparse.ArgumentParser(description="RKLLM demo")
    parser.add_argument("--model", type=str, default="", help="model path")
    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = parse_args()
    print("rkllm init start")
    param = rkllm_createDefaultParam()
    param.modelPath = args.model.encode("utf-8")
    param.target_platform = "rk3588".encode("utf-8")
    param.num_npu_core = 3
    param.top_k = 1
    param.max_new_tokens = 256
    param.max_context_len = 512
    c_callback = ctypes.CFUNCTYPE(None, ctypes.c_char_p, ctypes.c_void_p, ctypes.c_int32)(callback)
    llm_handle = LLMHandle(None)
    res = rkllm_init(llm_handle, param, c_callback)
    print("rkllm init end")

debug 结果:

image

从图中可以看出,程序没有运行到rkllm init end 部分,通过单步调试发现程序在 rkllm_init 时就被 kill 掉了。 想知道是因为什么原因导致初始化失败

waydong commented 2 months ago

刚发布的v1.0.1版本里面个server的demo,有python调用的实现,可以参考下哈。rkllm_server_demo