import time, os, docker, requests, json, uuid, subprocess, time, asyncio, aiohttp, re, traceback import psutil from typing import List, Optional, Union from loguru import logger from websockets.sync.client import connect as ws_connect_sync from websockets.client import connect as ws_connect from websocket import create_connection from websockets.client import WebSocketClientProtocol, ClientConnection from websockets.exceptions import ConnectionClosedError from configs.server_config import SANDBOX_SERVER from configs.model_config import JUPYTER_WORK_PATH from .basebox import BaseBox, CodeBoxResponse, CodeBoxStatus class PyCodeBox(BaseBox): enter_status: bool = False def __init__( self, remote_url: str = "", remote_ip: str = SANDBOX_SERVER["host"], remote_port: str = SANDBOX_SERVER["port"], token: str = "mytoken", do_code_exe: bool = False, do_remote: bool = False, do_check_net: bool = True, ): super().__init__(remote_url, remote_ip, remote_port, token, do_code_exe, do_remote) self.enter_status = True self.do_check_net = do_check_net asyncio.run(self.astart()) # logger.info(f"""remote_url: {self.remote_url}, # remote_ip: {self.remote_ip}, # remote_port: {self.remote_port}""") def decode_code_from_text(self, text: str) -> str: pattern = r'```.*?```' code_blocks = re.findall(pattern, text, re.DOTALL) code_text: str = "\n".join([block.strip('`') for block in code_blocks]) code_text = code_text[6:] if code_text.startswith("python") else code_text code_text = code_text.replace("python\n", "").replace("code", "") return code_text def run( self, code_text: Optional[str] = None, file_path: Optional[os.PathLike] = None, retry = 3, ) -> CodeBoxResponse: if not code_text and not file_path: return CodeBoxResponse( code_exe_response="Code or file_path must be specifieds!", code_text=code_text, code_exe_type="text", code_exe_status=502, do_code_exe=self.do_code_exe, ) if code_text and file_path: return CodeBoxResponse( code_exe_response="Can only specify code or the file to read_from!", code_text=code_text, code_exe_type="text", code_exe_status=502, do_code_exe=self.do_code_exe, ) if file_path: with open(file_path, "r", encoding="utf-8") as f: code_text = f.read() # run code in jupyter kernel if retry <= 0: raise RuntimeError("Could not connect to kernel") if not self.ws: raise RuntimeError("Jupyter not running. Make sure to start it first") # logger.debug(f"code_text: {len(code_text)}, {code_text}") self.ws.send( json.dumps( { "header": { "msg_id": (msg_id := uuid.uuid4().hex), "msg_type": "execute_request", }, "parent_header": {}, "metadata": {}, "content": { "code": code_text, "silent": False, # True，则内核会执行代码，但不会发送执行结果（如输出） "store_history": True, # True，则执行的代码会被记录在交互式环境的历史记录中 "user_expressions": {}, "allow_stdin": False, # True，允许代码执行时接受用户输入 "stop_on_error": True, # True，当执行中遇到错误时，后续代码将不会继续执行。 }, "channel": "shell", "buffers": [], } ) ) result = "" while True: try: if isinstance(self.ws, WebSocketClientProtocol): raise RuntimeError("Mixing asyncio and sync code is not supported") received_msg = json.loads(self.ws.recv()) except ConnectionClosedError: # logger.debug("box start, ConnectionClosedError!!!") self.start() return self.run(code_text, file_path, retry - 1) if ( received_msg["header"]["msg_type"] == "stream" and received_msg["parent_header"]["msg_id"] == msg_id ): msg = received_msg["content"]["text"].strip() if "Requirement already satisfied:" in msg: continue result += msg + "\n" elif ( received_msg["header"]["msg_type"] == "execute_result" and received_msg["parent_header"]["msg_id"] == msg_id ): result += received_msg["content"]["data"]["text/plain"].strip() + "\n" elif received_msg["header"]["msg_type"] == "display_data": if "image/png" in received_msg["content"]["data"]: return CodeBoxResponse( code_exe_type="image/png", code_text=code_text, code_exe_response=received_msg["content"]["data"]["image/png"], code_exe_status=200, do_code_exe=self.do_code_exe ) if "text/plain" in received_msg["content"]["data"]: return CodeBoxResponse( code_exe_type="text", code_text=code_text, code_exe_response=received_msg["content"]["data"]["text/plain"], code_exe_status=200, do_code_exe=self.do_code_exe ) return CodeBoxResponse( code_exe_type="error", code_text=code_text, code_exe_response=received_msg["content"]["data"]["text/plain"], code_exe_status=420, do_code_exe=self.do_code_exe ) elif ( received_msg["header"]["msg_type"] == "status" and received_msg["parent_header"]["msg_id"] == msg_id and received_msg["content"]["execution_state"] == "idle" ): if len(result) > 500: result = "[...]\n" + result[-500:] return CodeBoxResponse( code_exe_type="text", code_text=code_text, code_exe_response=result or "Code run successfully (no output)", code_exe_status=200, do_code_exe=self.do_code_exe ) elif ( received_msg["header"]["msg_type"] == "error" and received_msg["parent_header"]["msg_id"] == msg_id ): error = ( f"{received_msg['content']['ename']}: " f"{received_msg['content']['evalue']}" ) return CodeBoxResponse( code_exe_type="error", code_text=code_text, code_exe_response=error, code_exe_status=500, do_code_exe=self.do_code_exe ) def _get_kernelid(self, ) -> None: headers = {"Authorization": f'Token {self.token}', 'token': self.token} response = requests.get(f"{self.kernel_url}?token={self.token}", headers=headers) if len(response.json()) > 0: self.kernel_id = response.json()[0]["id"] else: response = requests.post(f"{self.kernel_url}?token={self.token}", headers=headers) self.kernel_id = response.json()["id"] if self.kernel_id is None: raise Exception("Could not start kernel") async def _aget_kernelid(self, ) -> None: headers = {"Authorization": f'Token {self.token}', 'token': self.token} response = requests.get(f"{self.kernel_url}?token={self.token}", headers=headers) async with aiohttp.ClientSession() as session: async with session.get(f"{self.kernel_url}?token={self.token}", headers=headers) as resp: if len(await resp.json()) > 0: self.kernel_id = (await resp.json())[0]["id"] else: async with session.post(f"{self.kernel_url}?token={self.token}", headers=headers) as response: self.kernel_id = (await response.json())["id"] # if len(response.json()) > 0: # self.kernel_id = response.json()[0]["id"] # else: # response = requests.post(f"{self.kernel_url}?token={self.token}", headers=headers) # self.kernel_id = response.json()["id"] # if self.kernel_id is None: # raise Exception("Could not start kernel") def _check_connect(self, ) -> bool: if self.kernel_url == "": return False try: response = requests.get(f"{self.kernel_url}?token={self.token}", timeout=270) return response.status_code == 200 except requests.exceptions.ConnectionError: return False async def _acheck_connect(self, ) -> bool: if self.kernel_url == "": return False try: async with aiohttp.ClientSession() as session: async with session.get(f"{self.kernel_url}?token={self.token}", timeout=270) as resp: return resp.status == 200 except aiohttp.ClientConnectorError: pass except aiohttp.ServerDisconnectedError: pass def _check_port(self, ) -> bool: try: response = requests.get(f"{self.remote_ip}:{self.remote_port}", timeout=270) logger.warning(f"Port is conflict, please check your codebox's port {self.remote_port}") return response.status_code == 200 except requests.exceptions.ConnectionError: return False async def _acheck_port(self, ) -> bool: try: async with aiohttp.ClientSession() as session: async with session.get(f"{self.remote_ip}:{self.remote_port}", timeout=270) as resp: logger.warning(f"Port is conflict, please check your codebox's port {self.remote_port}") return resp.status == 200 except aiohttp.ClientConnectorError: pass except aiohttp.ServerDisconnectedError: pass def _check_connect_success(self, retry_nums: int = 5) -> bool: if not self.do_check_net: return True while retry_nums > 0: try: connect_status = self._check_connect() if connect_status: logger.info(f"{self.remote_url} connection success") return True except requests.exceptions.ConnectionError: logger.info(f"{self.remote_url} connection fail") retry_nums -= 1 time.sleep(5) raise BaseException(f"can't connect to {self.remote_url}") async def _acheck_connect_success(self, retry_nums: int = 5) -> bool: if not self.do_check_net: return True while retry_nums > 0: try: connect_status = await self._acheck_connect() if connect_status: logger.info(f"{self.remote_url} connection success") return True except requests.exceptions.ConnectionError: logger.info(f"{self.remote_url} connection fail") retry_nums -= 1 time.sleep(5) raise BaseException(f"can't connect to {self.remote_url}") def start(self, ): '''判断是从外部service执行还是内部启动notebook执行''' self.jupyter = None if self.do_remote: # TODO自动检测日期,并重启容器 self.kernel_url = self.remote_url + "/api/kernels" self._check_connect_success() self._get_kernelid() # logger.debug(self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}") self.wc_url = self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}" headers = {"Authorization": f'Token {self.token}', 'token': self.token} self.ws = create_connection(self.wc_url, headers=headers) else: # TODO 自动检测本地接口 port_status = self._check_port() connect_status = self._check_connect() logger.info(f"port_status: {port_status}, connect_status: {connect_status}") if port_status and not connect_status: raise BaseException(f"Port is conflict, please check your codebox's port {self.remote_port}") if not connect_status: self.jupyter = subprocess.run( [ "jupyer", "notebnook", f"--NotebookApp.token={self.token}", f"--port={self.remote_port}", "--no-browser", "--ServerApp.disable_check_xsrf=True", "--notebook-dir={JUPYTER_WORK_PATH}" ], stderr=subprocess.PIPE, stdin=subprocess.PIPE, stdout=subprocess.PIPE, ) self.kernel_url = self.remote_url + "/api/kernels" self.do_check_net = True self._check_connect_success() self._get_kernelid() # logger.debug(self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}") self.wc_url = self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}" headers = {"Authorization": f'Token {self.token}', 'token': self.token} self.ws = create_connection(self.wc_url, headers=headers) async def astart(self, ): '''判断是从外部service执行还是内部启动notebook执行''' self.jupyter = None if self.do_remote: # TODO自动检测日期,并重启容器 self.kernel_url = self.remote_url + "/api/kernels" await self._acheck_connect_success() await self._aget_kernelid() self.wc_url = self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}" headers = {"Authorization": f'Token {self.token}', 'token': self.token} self.ws = create_connection(self.wc_url, headers=headers) else: # TODO 自动检测本地接口 port_status = await self._acheck_port() self.kernel_url = self.remote_url + "/api/kernels" connect_status = await self._acheck_connect() logger.info(f"port_status: {port_status}, connect_status: {connect_status}") if port_status and not connect_status: raise BaseException(f"Port is conflict, please check your codebox's port {self.remote_port}") if not connect_status: self.jupyter = subprocess.Popen( [ "jupyter", "notebook", f"--NotebookApp.token={self.token}", f"--port={self.remote_port}", "--no-browser", "--ServerApp.disable_check_xsrf=True" ], stderr=subprocess.PIPE, stdin=subprocess.PIPE, stdout=subprocess.PIPE, ) self.kernel_url = self.remote_url + "/api/kernels" self.do_check_net = True await self._acheck_connect_success() await self._aget_kernelid() self.wc_url = self.kernel_url.replace("http", "ws") + f"/{self.kernel_id}/channels?token={self.token}" headers = {"Authorization": f'Token {self.token}', 'token': self.token} self.ws = create_connection(self.wc_url, headers=headers) def status(self,) -> CodeBoxStatus: if not self.kernel_id: self._get_kernelid() return CodeBoxStatus( status="running" if self.kernel_id and requests.get(self.kernel_url, timeout=270).status_code == 200 else "stopped" ) async def astatus(self,) -> CodeBoxStatus: if not self.kernel_id: await self._aget_kernelid() return CodeBoxStatus( status="running" if self.kernel_id and requests.get(self.kernel_url, timeout=270).status_code == 200 else "stopped" ) def restart(self, ) -> CodeBoxStatus: return CodeBoxStatus(status="restared") def stop(self, ) -> CodeBoxStatus: try: if self.jupyter is not None: for process in psutil.process_iter(["pid", "name", "cmdline"]): # 检查进程名是否包含"jupyter" if f'port={self.remote_port}' in str(process.info["cmdline"]).lower() and \ "jupyter" in process.info['name'].lower(): logger.warning(f'port={self.remote_port}, {process.info}') # 关闭进程 process.terminate() self.jupyter = None except Exception as e: logger.error(traceback.format_exc()) if self.ws is not None: try: if self.ws is not None: self.ws.close() else: loop = asyncio.new_event_loop() loop.run_until_complete(self.ws.close()) except Exception as e: logger.error(traceback.format_exc()) self.ws = None # return CodeBoxStatus(status="stopped") def __del__(self): self.stop()