From fd9f2416ff672e32ee0d5a8aa4251e9ec3795662 Mon Sep 17 00:00:00 2001 From: lidanyang Date: Tue, 9 Jan 2024 16:14:44 +0800 Subject: [PATCH 1/2] add timeout and retry when code execution --- metagpt/actions/execute_code.py | 57 +++++++++++++++++----- tests/metagpt/actions/test_execute_code.py | 9 ++++ 2 files changed, 54 insertions(+), 12 deletions(-) diff --git a/metagpt/actions/execute_code.py b/metagpt/actions/execute_code.py index 6e4a6fd6e..ab8019e23 100644 --- a/metagpt/actions/execute_code.py +++ b/metagpt/actions/execute_code.py @@ -12,6 +12,8 @@ import re import nbformat from nbclient import NotebookClient +from nbclient.exceptions import DeadKernelError, CellTimeoutError +from nbformat import NotebookNode from nbformat.v4 import new_code_cell, new_output from rich.console import Console from rich.syntax import Syntax @@ -46,13 +48,23 @@ class ExecuteCode(ABC): class ExecutePyCode(ExecuteCode, Action): """execute code, return result to llm, and display it.""" - def __init__(self, name: str = "python_executor", context=None, llm=None, nb=None): + def __init__( + self, + name: str = "python_executor", + context=None, + llm=None, + nb=None, + timeout: int = 600, + max_tries: int = 3 + ): super().__init__(name, context, llm) if nb is None: self.nb = nbformat.v4.new_notebook() else: self.nb = nb - self.nb_client = NotebookClient(self.nb) + self.timeout = timeout + self.max_tries = max_tries + self.nb_client = NotebookClient(self.nb, timeout=self.timeout) self.console = Console() self.interaction = "ipython" if self.is_ipython() else "terminal" @@ -69,7 +81,8 @@ class ExecutePyCode(ExecuteCode, Action): async def reset(self): """reset NotebookClient""" await self.terminate() - self.nb_client = NotebookClient(self.nb) + await self.build() + self.nb_client = NotebookClient(self.nb, timeout=self.timeout) def add_code_cell(self, code): self.nb.cells.append(new_code_cell(source=code)) @@ -160,6 +173,19 @@ class ExecutePyCode(ExecuteCode, Action): return code, language + async def run_cell(self, cell: NotebookNode, cell_index: int) -> Tuple[bool, str]: + """set timeout for run code""" + try: + await self.nb_client.async_execute_cell(cell, cell_index) + return True, "" + except CellTimeoutError: + return False, "TimeoutError" + except DeadKernelError: + await self.reset() + return False, "DeadKernelError" + except Exception as e: + return False, f"{traceback.format_exc()}" + async def run(self, code: Union[str, Dict, Message], language: str = "python") -> Tuple[str, bool]: code, language = self._process_code(code, language) @@ -168,19 +194,26 @@ class ExecutePyCode(ExecuteCode, Action): if language == "python": # add code to the notebook self.add_code_cell(code=code) - try: + + tries = 0 + success = False + outputs = "" + while tries < self.max_tries and not success: # build code executor await self.build() # run code - # TODO: add max_tries for run code. cell_index = len(self.nb.cells) - 1 - await self.nb_client.async_execute_cell(self.nb.cells[-1], cell_index) - outputs = self.parse_outputs(self.nb.cells[-1].outputs) - success = True - except Exception as e: - outputs = traceback.format_exc() - success = False - return truncate(remove_escape_and_color_codes(outputs)), success + success, error_message = await self.run_cell(self.nb.cells[-1], cell_index) + + if success: + outputs = self.parse_outputs(self.nb.cells[-1].outputs) + else: + tries += 1 + + if success: + return truncate(remove_escape_and_color_codes(outputs)), True + else: + return error_message, False else: # TODO: markdown raise NotImplementedError(f"Not support this code type : {language}, Only support code!") diff --git a/tests/metagpt/actions/test_execute_code.py b/tests/metagpt/actions/test_execute_code.py index 95f883e12..8340272e4 100644 --- a/tests/metagpt/actions/test_execute_code.py +++ b/tests/metagpt/actions/test_execute_code.py @@ -88,3 +88,12 @@ def test_truncate(): assert truncate(output) == output output = "hello world" assert truncate(output, 5) == "Truncated to show only the last 5 characters\nworld" + + +@pytest.mark.asyncio +async def test_run_with_timeout(): + pi = ExecutePyCode(timeout=1) + code = "import time; time.sleep(2)" + message, success = await pi.run(code) + assert not success + assert message == "TimeoutError" From 3eee6eff8c7b2112cede5084cbe8b8fd81c4190b Mon Sep 17 00:00:00 2001 From: lidanyang Date: Tue, 9 Jan 2024 17:45:46 +0800 Subject: [PATCH 2/2] drop retry --- metagpt/actions/execute_code.py | 21 ++++++--------------- 1 file changed, 6 insertions(+), 15 deletions(-) diff --git a/metagpt/actions/execute_code.py b/metagpt/actions/execute_code.py index ab8019e23..d192ca79a 100644 --- a/metagpt/actions/execute_code.py +++ b/metagpt/actions/execute_code.py @@ -55,7 +55,6 @@ class ExecutePyCode(ExecuteCode, Action): llm=None, nb=None, timeout: int = 600, - max_tries: int = 3 ): super().__init__(name, context, llm) if nb is None: @@ -63,7 +62,6 @@ class ExecutePyCode(ExecuteCode, Action): else: self.nb = nb self.timeout = timeout - self.max_tries = max_tries self.nb_client = NotebookClient(self.nb, timeout=self.timeout) self.console = Console() self.interaction = "ipython" if self.is_ipython() else "terminal" @@ -195,22 +193,15 @@ class ExecutePyCode(ExecuteCode, Action): # add code to the notebook self.add_code_cell(code=code) - tries = 0 - success = False - outputs = "" - while tries < self.max_tries and not success: - # build code executor - await self.build() - # run code - cell_index = len(self.nb.cells) - 1 - success, error_message = await self.run_cell(self.nb.cells[-1], cell_index) + # build code executor + await self.build() - if success: - outputs = self.parse_outputs(self.nb.cells[-1].outputs) - else: - tries += 1 + # run code + cell_index = len(self.nb.cells) - 1 + success, error_message = await self.run_cell(self.nb.cells[-1], cell_index) if success: + outputs = self.parse_outputs(self.nb.cells[-1].outputs) return truncate(remove_escape_and_color_codes(outputs)), True else: return error_message, False