import docker, sys, os, time, requests, psutil, json import subprocess import webbrowser from docker.types import Mount, DeviceRequest from loguru import logger src_dir = os.path.join( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) sys.path.append(src_dir) from configs.model_config import USE_FASTCHAT, JUPYTER_WORK_PATH from configs.server_config import ( NO_REMOTE_API, SANDBOX_SERVER, SANDBOX_IMAGE_NAME, SANDBOX_CONTRAINER_NAME, WEBUI_SERVER, API_SERVER, SDFILE_API_SERVER, CONTRAINER_NAME, IMAGE_NAME, DOCKER_SERVICE, DEFAULT_BIND_HOST, NEBULA_GRAPH_SERVER, SANDBOX_DO_REMOTE ) import platform system_name = platform.system() USE_TTY = system_name in ["Windows"] def check_process(content: str, lang: str = None, do_stop=False): '''process-not-exist is true, process-exist is false''' for process in psutil.process_iter(["pid", "name", "cmdline"]): # check process name contains "jupyter" and port=xx if content in str(process.info["cmdline"]).lower(): logger.info(f"content, {process.info}") # 关闭进程 if do_stop: process.terminate() return True return False return True def check_docker(client, container_name, do_stop=False): '''container-not-exist is true, container-exist is false''' if client is None: return True for i in client.containers.list(all=True): if i.name == container_name: if do_stop: container = i if container_name == CONTRAINER_NAME and i.status == 'running': # wrap up db logger.info(f'inside {container_name}') # cp nebula data res = container.exec_run('''sh chatbot/coagent/utils/nebula_cp.sh''') logger.info(f'cp res={res}') # stop nebula service res = container.exec_run('''/usr/local/nebula/scripts/nebula.service stop all''') logger.info(f'stop res={res}') container.stop() container.remove() return True return False return True def start_docker(client, script_shs, ports, image_name, container_name, mounts=None, network=None): container = client.containers.run( image=image_name, command="bash", mounts=mounts, name=container_name, mem_limit="8g", # device_requests=[DeviceRequest(count=-1, capabilities=[['gpu']])], # network_mode="host", ports=ports, stdin_open=True, detach=True, tty=USE_TTY, network=network, ) logger.info(f"docker id: {container.id[:10]}") # 启动notebook for script_sh in script_shs: if USE_FASTCHAT and "llm_api" in script_sh: logger.debug(script_sh) response = container.exec_run(["sh", "-c", script_sh]) logger.debug(response) elif "llm_api" not in script_sh: logger.debug(script_sh) response = container.exec_run(["sh", "-c", script_sh]) logger.debug(response) return container ######################################### ############# 开始启动服务 ############### ######################################### network_name ='my_network' def start_sandbox_service(network_name ='my_network'): mount = Mount( type='bind', source=os.path.join(src_dir, "jupyter_work"), target='/home/user/chatbot/jupyter_work', read_only=False # 如果需要只读访问,将此选项设置为True ) mounts = [mount] # 沙盒的启动与服务的启动是独立的 if SANDBOX_DO_REMOTE: client = docker.from_env() networks = client.networks.list() if any([network_name==i.attrs["Name"] for i in networks]): network = client.networks.get(network_name) else: network = client.networks.create('my_network', driver='bridge') # 启动容器 logger.info("start container sandbox service") JUPYTER_WORK_PATH = "/home/user/chatbot/jupyter_work" script_shs = [f"cd /home/user/chatbot/jupyter_work && nohup jupyter-notebook --NotebookApp.token=mytoken --port=5050 --allow-root --ip=0.0.0.0 --notebook-dir={JUPYTER_WORK_PATH} --no-browser --ServerApp.disable_check_xsrf=True &"] ports = {f"{SANDBOX_SERVER['docker_port']}/tcp": f"{SANDBOX_SERVER['port']}/tcp"} if check_docker(client, SANDBOX_CONTRAINER_NAME, do_stop=True, ): container = start_docker(client, script_shs, ports, SANDBOX_IMAGE_NAME, SANDBOX_CONTRAINER_NAME, mounts=mounts, network=network_name) # 判断notebook是否启动 time.sleep(5) retry_nums = 3 while retry_nums>0: logger.info(f"http://localhost:{SANDBOX_SERVER['port']}") response = requests.get(f"http://localhost:{SANDBOX_SERVER['port']}", timeout=270) if response.status_code == 200: logger.info("container & notebook init success") break else: retry_nums -= 1 logger.info(client.containers.list()) logger.info("wait container running ...") time.sleep(5) else: try: client = docker.from_env() except: client = None check_docker(client, SANDBOX_CONTRAINER_NAME, do_stop=True, ) logger.info("start local sandbox service") def start_api_service(sandbox_host=DEFAULT_BIND_HOST): # 启动service的容器 if DOCKER_SERVICE: client = docker.from_env() logger.info("start container service") check_process("api.py", do_stop=True) check_process("llm_api.py", do_stop=True) check_process("sdfile_api.py", do_stop=True) check_process("webui.py", do_stop=True) mount = Mount( type='bind', source=src_dir, target='/home/user/chatbot/', read_only=False # 如果需要只读访问,将此选项设置为True ) ports={ f"{API_SERVER['docker_port']}/tcp": f"{API_SERVER['port']}/tcp", f"{WEBUI_SERVER['docker_port']}/tcp": f"{WEBUI_SERVER['port']}/tcp", f"{SDFILE_API_SERVER['docker_port']}/tcp": f"{SDFILE_API_SERVER['port']}/tcp", f"{NEBULA_GRAPH_SERVER['docker_port']}/tcp": f"{NEBULA_GRAPH_SERVER['port']}/tcp" } # mounts = [mount, mount_database, mount_code_database] mounts = [mount] script_shs = [ "mkdir -p /home/user/chatbot/logs", ''' if [ -d "/home/user/chatbot/data/nebula_data/data/meta" ]; then cp -r /home/user/chatbot/data/nebula_data/data /usr/local/nebula/ fi ''', "/usr/local/nebula/scripts/nebula.service start all", "/usr/local/nebula/scripts/nebula.service status all", "sleep 2", '''curl -X PUT -H "Content-Type: application/json" -d'{"heartbeat_interval_secs":"2"}' -s "http://127.0.0.1:19559/flags"''', '''curl -X PUT -H "Content-Type: application/json" -d'{"heartbeat_interval_secs":"2"}' -s "http://127.0.0.1:19669/flags"''', '''curl -X PUT -H "Content-Type: application/json" -d'{"heartbeat_interval_secs":"2"}' -s "http://127.0.0.1:19779/flags"''', "pip install zdatafront-sdk-python==0.1.2 -i https://artifacts.antgroup-inc.cn/simple", "pip install jieba", "pip install duckduckgo-search", "pip install codefuse-muagent", "nohup python chatbot/examples/sdfile_api.py > /home/user/chatbot/logs/sdfile_api.log 2>&1 &", f"export DUCKDUCKGO_PROXY=socks5://host.docker.internal:13659 && export SANDBOX_HOST={sandbox_host} &&\ nohup python chatbot/examples/api.py > /home/user/chatbot/logs/api.log 2>&1 &", "nohup python chatbot/examples/llm_api.py > /home/user/llm.log 2>&1 &", f"export DUCKDUCKGO_PROXY=socks5://host.docker.internal:13659 && export SANDBOX_HOST={sandbox_host} &&\ cd chatbot/examples && nohup streamlit run webui.py > /home/user/chatbot/logs/start_webui.log 2>&1 &" ] if check_docker(client, CONTRAINER_NAME, do_stop=True): container = start_docker(client, script_shs, ports, IMAGE_NAME, CONTRAINER_NAME, mounts, network=network_name) logger.info("You can open http://localhost:8501 to use chatbot!") webbrowser.open("http://localhost:8501") else: logger.info("start local service") # 关闭之前启动的docker 服务 # check_docker(client, CONTRAINER_NAME, do_stop=True, ) api_sh = "nohup python api.py > ../logs/api.log 2>&1 &" sdfile_sh = "nohup python sdfile_api.py > ../logs/sdfile_api.log 2>&1 &" notebook_sh = f"nohup jupyter-notebook --NotebookApp.token=mytoken --port={SANDBOX_SERVER['port']} --allow-root --ip=0.0.0.0 --notebook-dir={JUPYTER_WORK_PATH} --no-browser --ServerApp.disable_check_xsrf=True > ../logs/sandbox.log 2>&1 &" llm_sh = "nohup python llm_api.py > ../logs/llm_api.log 2>&1 &" webui_sh = "streamlit run webui.py" if USE_TTY else "streamlit run webui.py" # if SANDBOX_DO_REMOTE and check_process("jupyter-notebook --NotebookApp"): # logger.debug(f"{notebook_sh}") # subprocess.Popen(notebook_sh, shell=True) # if not NO_REMOTE_API and check_process("api.py"): subprocess.Popen(api_sh, shell=True) # if USE_FASTCHAT and check_process("llm_api.py"): subprocess.Popen(llm_sh, shell=True) # if check_process("sdfile_api.py"): subprocess.Popen(sdfile_sh, shell=True) subprocess.Popen(webui_sh, shell=True) logger.info("You can please open http://localhost:8501 to use chatbot!") def start_main(): global SANDBOX_DO_REMOTE, DOCKER_SERVICE SANDBOX_DO_REMOTE = SANDBOX_DO_REMOTE if os.environ.get("SANDBOX_DO_REMOTE") is None else json.loads(os.environ.get("SANDBOX_DO_REMOTE")) DOCKER_SERVICE = DOCKER_SERVICE if os.environ.get("DOCKER_SERVICE") is None else json.loads(os.environ.get("DOCKER_SERVICE")) start_sandbox_service() sandbox_host = DEFAULT_BIND_HOST if SANDBOX_DO_REMOTE: client = docker.from_env() containers = client.containers.list(all=True) for container in containers: container_a_info = client.containers.get(container.id) if container_a_info.name == SANDBOX_CONTRAINER_NAME: container1_networks = container.attrs['NetworkSettings']['Networks'] sandbox_host = container1_networks.get(network_name)["IPAddress"] break start_api_service(sandbox_host) if __name__ == "__main__": start_main()