From b2c621bd18f3b966fe4bc4015b54456fe5c2f889 Mon Sep 17 00:00:00 2001 From: Yazhou Cao Date: Mon, 22 Jul 2024 17:01:32 +0800 Subject: [PATCH] Add back code --- vision_agent/utils/execute.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vision_agent/utils/execute.py b/vision_agent/utils/execute.py index b689bab2..3dc571c2 100644 --- a/vision_agent/utils/execute.py +++ b/vision_agent/utils/execute.py @@ -462,6 +462,7 @@ def close(self, *args: Any, **kwargs: Any) -> None: ) def restart_kernel(self) -> None: + self._check_sandbox_liveness() self.interpreter.notebook.restart_kernel() @tenacity.retry( @@ -473,6 +474,7 @@ def restart_kernel(self) -> None: after=tenacity.after_log(_LOGGER, logging.INFO), ) def exec_cell(self, code: str) -> Execution: + self._check_sandbox_liveness() self.interpreter.set_timeout(_SESSION_TIMEOUT) # Extend the life of the sandbox try: _LOGGER.info(f"Start code execution in remote sandbox {self.interpreter.sandbox_id}. Timeout: {_SESSION_TIMEOUT}. Code hash: {hash(code)}") @@ -485,6 +487,7 @@ def exec_cell(self, code: str) -> Execution: ) from e def upload_file(self, file: Union[str, Path]) -> str: + self._check_sandbox_liveness() file_name = Path(file).name remote_path = f"/home/user/{file_name}" with open(file, "rb") as f: @@ -493,6 +496,7 @@ def upload_file(self, file: Union[str, Path]) -> str: return remote_path def download_file(self, file_path: str) -> Path: + self._check_sandbox_liveness() with tempfile.NamedTemporaryFile(mode="w+b", delete=False) as file: file.write(self.interpreter.files.read(path=file_path, format="bytes")) _LOGGER.info(f"File ({file_path}) is downloaded to: {file.name}") @@ -514,7 +518,7 @@ def _check_sandbox_liveness(self) -> None: @staticmethod def _new_e2b_interpreter_impl(*args, **kwargs) -> E2BCodeInterpreterImpl: # type: ignore template_name = os.environ.get("E2B_TEMPLATE_NAME", "va-sandbox") - _LOGGER.info(f"Creating a new E2BCodeInterpreter instance using template: {template_name}") + _LOGGER.info(f"Creating a new E2BCodeInterpreter using template: {template_name}") return E2BCodeInterpreterImpl(template=template_name, *args, **kwargs)