diff --git a/src/aios/storage/storage.py b/src/aios/storage/storage.py index fe0b3fd..11edd29 100644 --- a/src/aios/storage/storage.py +++ b/src/aios/storage/storage.py @@ -216,6 +216,13 @@ class AIStorage: ~/myai/ """ return Path.home() / "myai" + + def get_download_dir(self) -> str: + """ + download dir is the dir for user to store the files downloaded with the system. + ~/myai/download + """ + return f"{self.get_myai_dir()}/download" def get_db(self,app_name:str)->ResourceLocation: pass @@ -242,5 +249,3 @@ class AIStorage: except Exception as e: logger.error(f"open or create file {path} failed! {str(e)}") - - diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index 0ac7a81..b61b906 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/src/service/aios_shell/aios_shell.py @@ -28,6 +28,7 @@ sys.path.append(directory + '/../../') import proxy from aios import * +import local_compute_node_builder sys.path.append(directory + '/../../component/') @@ -402,41 +403,22 @@ class AIOS_Shell: async def handle_node_commands(self, args): show_text = FormattedText([("class:title", "sub command not support!\n" - "/node add llama $model_name $url\n" - "/node rm llama $model_name $url\n" + "/node add\n" + "/node rm $model_name $url\n" "/node list\n")]) if len(args) < 1: return show_text sub_cmd = args[0] if sub_cmd == "add": - if len(args) < 2: - return show_text - if args[1] == "llama": - if len(args) < 4: - return show_text - - model_name = args[2] - url = args[3] - ComputeNodeConfig.get_instance().add_node("llama", url, model_name) - ComputeNodeConfig.get_instance().save() - node = LocalLlama_ComputeNode(url, model_name) - node.start() - ComputeKernel.get_instance().add_compute_node(node) - else: - return show_text + await local_compute_node_builder.build(session, shell_style) elif sub_cmd == "rm": - if len(args) < 2: - return show_text - if args[1] == "llama": - if len(args) < 4: - return show_text - - model_name = args[2] - url = args[3] - ComputeNodeConfig.get_instance().remove_node("llama", url, model_name) - ComputeNodeConfig.get_instance().save() - else: + if len(args) < 3: return show_text + + model_name = args[1] + url = args[2] + ComputeNodeConfig.get_instance().remove_node("llama", url, model_name) + ComputeNodeConfig.get_instance().save() elif sub_cmd == "list": print_formatted_text(ComputeNodeConfig.get_instance().list()) @@ -785,8 +767,8 @@ async def main(): '/set_config $key', '/enable $feature', '/disable $feature', - '/node add llama $model_name $url', - '/node rm llama $model_name $url', + '/node add', + '/node rm $model_name $url', '/node list', '/show', '/exit', diff --git a/src/service/aios_shell/local_compute_node_builder/__init__.py b/src/service/aios_shell/local_compute_node_builder/__init__.py new file mode 100644 index 0000000..cb84ca9 --- /dev/null +++ b/src/service/aios_shell/local_compute_node_builder/__init__.py @@ -0,0 +1,30 @@ +from prompt_toolkit import PromptSession +from prompt_toolkit.styles import Style +from service.aios_shell.local_compute_node_builder.local_llama_node_builder import LocalLlamaNodeBuilder +from .local_compute_node_builder import BuilderState, LocalComputeNodeBuilder + +async def build(prompt_session: PromptSession, shell_style: Style) -> str or None: + # model_type = await prompt_session.prompt_async(f"Please select the node server type (default: llama.cpp):", style = shell_style) + + model_type = 'llama.cpp' + + state = BuilderState(prompt_session, shell_style) + + match model_type: + case 'llama.cpp': + builder = LocalLlamaNodeBuilder(state) + + while True: + param = builder.next_parameter() + if param is None: + return None + + value = await state.prompt_session.prompt_async(f"{state.last_result_prompt}{param.desc}:", style = state.shell_style) + if value: + value = value.strip() + + state.params[param.name] = value + url = await param.applier.apply(state, param.name, value) + + if url is not None: + return url \ No newline at end of file diff --git a/src/service/aios_shell/local_compute_node_builder/local_compute_node_builder.py b/src/service/aios_shell/local_compute_node_builder/local_compute_node_builder.py new file mode 100644 index 0000000..7d9bc3d --- /dev/null +++ b/src/service/aios_shell/local_compute_node_builder/local_compute_node_builder.py @@ -0,0 +1,39 @@ +from abc import abstractmethod + +from prompt_toolkit import PromptSession +from prompt_toolkit.styles import Style + +class BuilderState: + def __init__(self, prompt_session: PromptSession, shell_style: Style): + self.prompt_session = prompt_session + self.shell_style = shell_style + self.next_step = 0 + self.last_result_prompt = "" + self.params = {} + +# class ApplyResult: +# def __init__(self, next_step: any, url: str or None = None, result_prompt: str or None = None) -> None: +# self.next_step = next_step +# self.url = url +# self.result_prompt = result_prompt + + +class ParameterApplier: + @abstractmethod + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + pass + +class BuildParameter: + def __init__(self, name: str, applier: ParameterApplier, desc: str or None = None, default_value: str or None = None): + self.name = name + self.desc = desc + self.default_value = default_value + self.applier = applier + +class LocalComputeNodeBuilder: + def __init__(self, state: BuilderState) -> None: + self.state = state + + @abstractmethod + def next_parameter(self) -> BuildParameter or None: + pass \ No newline at end of file diff --git a/src/service/aios_shell/local_compute_node_builder/local_llama_node_builder.py b/src/service/aios_shell/local_compute_node_builder/local_llama_node_builder.py new file mode 100644 index 0000000..25b7b14 --- /dev/null +++ b/src/service/aios_shell/local_compute_node_builder/local_llama_node_builder.py @@ -0,0 +1,248 @@ +import os +import random +import subprocess +import requests + +from prompt_toolkit import print_formatted_text +from prompt_toolkit.shortcuts import ProgressBar +from prompt_toolkit.formatted_text import FormattedText +from aios_kernel.compute_kernel import ComputeKernel +from aios_kernel.compute_node_config import ComputeNodeConfig +from aios_kernel.local_llama_compute_node import LocalLlama_ComputeNode + +from aios_kernel.storage import AIStorage +from .local_compute_node_builder import BuildParameter, BuilderState, LocalComputeNodeBuilder, ParameterApplier + +class BuildParameterModelPath: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + if value: + if os.path.exists(value): + state.next_step += 2 + else: + print_formatted_text(FormattedText([("class:error", f"Model not exist at {value}")]), style = state.shell_style) + else: + state.next_step += 1 + + +class BuildParameterModelUrl: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + if value is None: + value = "1" + + url = value + recommend = _recommend_model_urls.get(value) + if recommend: + url = recommend["url"] + + save_path = f"{AIStorage.get_instance().get_download_dir()}/{url.split('/').pop()}" + + print_formatted_text(FormattedText([("class:prompt", f"Will save the model to {save_path}:\n")]), style = state.shell_style) + + try: + # get file size + response = requests.head(url) + file_size = int(response.headers.get('content-length', 0)) + + # start download + response = requests.get(url, stream=True) + + if response.status_code == 200: + with open(save_path, 'wb') as f, ProgressBar() as pb: + for data in pb(response.iter_content(1024), total = (file_size + 1023) // 1024): + f.write(data) + + print_formatted_text(FormattedText([("class:prompt", f"Download model success, save at: {save_path}\n")]), style = state.shell_style) + + state.params["model_path"] = save_path + state.next_step += 1 + else: + print_formatted_text(FormattedText([("class:error", f"Download model failed, error: {response.status_code}\nYou can retry it or select another one.")]), style = state.shell_style) + + except Exception as e: + print_formatted_text(FormattedText([("class:error", f"Download model failed: {e}\nYou can retry it or select another one.")]), style = state.shell_style) + +class ParameterNodeNameApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + value = value or os.path.basename(state.params["model_path"]) + state.params["node_name"] = value + state.next_step += 1 + +class ParameterPortApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + if value is None or value == "0": + value = str(random.randint(10000, 60000)) + + state.params["port"] = value + state.next_step += 1 + +class ParameterNGpuLayersApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + value = value or "83" + state.params["n_gpu_layers"] = value + state.next_step += 1 + +class ParameterNCtxApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + value = value or "4096" + state.params["n_ctx"] = value + state.next_step += 1 + +class ParameterChatFormatApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + value = value or "llama-2" + state.params["chat_format"] = value + state.next_step += 1 + +class ParameterExternParamsApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + extern_params = value + docker_image = "" + gpu_options = None + + if state.params["n_gpu_layers"] == "0": + docker_image = "ghcr.io/abetlen/llama-cpp-python:latest" + else: + gpu_options = "--gpus all" + llama_cpp_python_repo_url = "https://github.com/abetlen/llama-cpp-python.git" + download_path = AIStorage.get_instance().get_download_dir() + llama_cpp_python_path = download_path + "/llama-cpp-python" + + # update the `llama-cpp-python` + retry = True + while retry: + retry = False + result = None + if os.path.exists(llama_cpp_python_path): + result = subprocess.run(['git', 'pull'], cwd = llama_cpp_python_path, stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + else: + result = subprocess.run(['git', 'clone', llama_cpp_python_repo_url, download_path], stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + + if result.stderr: + while True: + sel = await state.prompt_session.prompt_async(f"Update 'llama-cpp-python' failed, you can press 'r' to retry, or 'c' to continue with the current version.", style = state.shell_style) + if sel == 'r': + retry = True + break + elif sel == 'c': + break + else: + pass # Select again + else: + break + + # build the image + docker_image = 'llama-cpp-python-cuda' + retry = True + while retry: + retry = False + result = subprocess.run(['docker', 'rmi', docker_image], stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + result = subprocess.run(['docker', 'build', '-t', docker_image, f"{llama_cpp_python_path}/docker/cuda_simple/"], stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + + if result.stderr: + while True: + sel = await state.prompt_session.prompt_async(f"Build the image failed, you can press 'r' to retry, or 'c' to continue with the current version.", style = state.shell_style) + if sel == 'r': + retry = True + break + elif sel == 'c': + break + else: + pass # Select again + else: + break + + retry = True + while True: + retry = False + run_options = ['docker', 'run', '-d'] + + if gpu_options: + run_options.append(gpu_options) + + run_options.extend([ + '-p', f"{state.params['port']}:8000", + '-v', f"{os.path.dirname(state.params['model_path'])}:/models", '-e', f"MODEL=/models/{os.path.basename(state.params['model_path'])}", + 'llama-cpp-python-cuda', + 'python3', '-m', 'llama_cpp.server', + '--n_gpu_layers', state.params["n_gpu_layers"], + '--n_ctx', state.params["n_ctx"], + '--chat_format', state.params["chat_format"], + ]) + + if extern_params: + run_options.extend(extern_params.split(' ')) + + result = subprocess.run(run_options, stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + + if result.stderr: + while True: + sel = await state.prompt_session.prompt_async(f"Start the node service failed, you can press 'r' to retry, or 'a' to abort.", style = state.shell_style) + if sel == 'r': + retry = True + break + elif sel == 'a': + break + else: + pass # Select again + else: + local_url = f'http://localhost:{state.params["port"]}' + foreign_url = 'http://{your-host-address}:' + state.params["port"] + model_name = state.params['node_name'] + + ComputeNodeConfig.get_instance().add_node("llama", local_url, model_name) + ComputeNodeConfig.get_instance().save() + node = LocalLlama_ComputeNode(local_url, model_name) + node.start() + ComputeKernel.get_instance().add_compute_node(node) + + print_formatted_text(FormattedText([( + "class:prompt", + f""" + Congratulations! The node ({model_name}) service successed. + You can access it with follow url: + {local_url} + And 'http://{foreign_url}' in other computers. + Now you can refer it in agents as `llm_model_name={model_name}` + """ + )]), style = state.shell_style) + break + +_recommend_model_urls = { + "1": { + "model": "Llama-2-70B-chat", + "url": "https://huggingface.co/TheBloke/Llama-2-70B-chat-GGUF/resolve/main/llama-2-70b-chat.Q4_0.gguf" + }, + "2": { + "model": "Llama-2-13B-chat", + "url": "https://huggingface.co/TheBloke/Llama-2-13B-chat-GGUF/resolve/main/llama-2-13b-chat.Q4_0.gguf" + }, + "3": { + "model": "Llama-2-7B-Chat-GGUF", + "url": "https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/blob/main/llama-2-7b-chat.Q4_K_M.gguf" + }, +} + +_recommend_model_url_table_str = map(lambda id, info: f"\t{id}\t{info['model']}\t{info['url']}\n", _recommend_model_urls) + +_params = [ + BuildParameter("model_path", BuildParameterModelPath(), "Please input the model file path (Press 'Enter' if you need to download it)"), + BuildParameter("model_url", BuildParameterModelUrl(), + f""" + Please input the url to download the model, or you can input the 'ID' in the follow table to select one: + ID\tmodel\turl + {_recommend_model_url_table_str} + Please input (default: Llama-2-70B-chat) + """ + ), + BuildParameter("node_name", ParameterNodeNameApplier(), "Please input name for your node, and you can set it in 'llm_model_name' of 'agent.toml' (default: the name of the model file)"), + BuildParameter("port", ParameterPortApplier(), "Please input the port which the node server will listen on (default: random)"), + BuildParameter("n_gpu_layers", ParameterNGpuLayersApplier(), "Please input layers offload to GPU (<=83 for Llama, 0 for CPU only, default: 83)"), + BuildParameter("n_ctx", ParameterNCtxApplier(), "Please input the content limit (default: 4096)"), + BuildParameter("chat_format", ParameterChatFormatApplier(), "Please input the chat format (default: llama-2)"), + BuildParameter("extern_params", ParameterExternParamsApplier(), "Please input other parameters refer to 'llama-cpp-python'(https://github.com/abetlen/llama-cpp-python), press 'Enter' to ignore it"), +] + +class LocalLlamaNodeBuilder(LocalComputeNodeBuilder): + def next_parameter(self) -> BuildParameter or None: + if self.state.next_step < len(_params): + return _params[self.state.next_step]