From 20ca7d8499e21b90b8d3ff6274af8c548b7968e5 Mon Sep 17 00:00:00 2001 From: Hussein Mozannar Date: Fri, 4 Apr 2025 23:07:57 +0000 Subject: [PATCH] docker delete temp --- .../docker/_docker_code_executor.py | 68 +++++++++++-------- .../test_docker_commandline_code_executor.py | 46 +++++++++++++ 2 files changed, 87 insertions(+), 27 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py index 51ac4f200efa..fda23da0b349 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py @@ -78,6 +78,7 @@ class DockerCommandLineCodeExecutorConfig(BaseModel): extra_volumes: Dict[str, Dict[str, str]] = {} extra_hosts: Dict[str, str] = {} init_command: Optional[str] = None + delete_tmp_files: bool = False class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCodeExecutorConfig]): @@ -124,6 +125,7 @@ class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCod Example: extra_hosts = {"kubernetes.docker.internal": "host-gateway"} init_command (Optional[str], optional): A shell command to run before each shell operation execution. Defaults to None. Example: init_command="kubectl config use-context docker-hub" + delete_tmp_files (bool, optional): If true, will delete temporary files after execution. Defaults to False. .. note:: Using the current directory (".") as working directory is deprecated. Using it will raise a deprecation warning. @@ -172,6 +174,7 @@ def __init__( extra_volumes: Optional[Dict[str, Dict[str, str]]] = None, extra_hosts: Optional[Dict[str, str]] = None, init_command: Optional[str] = None, + delete_tmp_files: bool = False, ): if timeout < 1: raise ValueError("Timeout must be greater than or equal to 1.") @@ -225,6 +228,7 @@ def __init__( self._extra_volumes = extra_volumes if extra_volumes is not None else {} self._extra_hosts = extra_hosts if extra_hosts is not None else {} self._init_command = init_command + self._delete_tmp_files = delete_tmp_files # Setup could take some time so we intentionally wait for the first code block to do it. if len(functions) > 0: @@ -311,33 +315,41 @@ async def _execute_code_dont_check_setup( outputs: List[str] = [] files: List[Path] = [] last_exit_code = 0 - for code_block in code_blocks: - lang = code_block.language.lower() - code = silence_pip(code_block.code, lang) - - # Check if there is a filename comment - try: - filename = get_file_name_from_content(code, self.work_dir) - except ValueError: - outputs.append("Filename is not in the workspace") - last_exit_code = 1 - break - - if not filename: - filename = f"tmp_code_{sha256(code.encode()).hexdigest()}.{lang}" - - code_path = self.work_dir / filename - with code_path.open("w", encoding="utf-8") as fout: - fout.write(code) - files.append(code_path) - - command = ["timeout", str(self._timeout), lang_to_cmd(lang), filename] - - output, exit_code = await self._execute_command(command, cancellation_token) - outputs.append(output) - last_exit_code = exit_code - if exit_code != 0: - break + try: + for code_block in code_blocks: + lang = code_block.language.lower() + code = silence_pip(code_block.code, lang) + + # Check if there is a filename comment + try: + filename = get_file_name_from_content(code, self.work_dir) + except ValueError: + outputs.append("Filename is not in the workspace") + last_exit_code = 1 + break + + if not filename: + filename = f"tmp_code_{sha256(code.encode()).hexdigest()}.{lang}" + + code_path = self.work_dir / filename + with code_path.open("w", encoding="utf-8") as fout: + fout.write(code) + files.append(code_path) + + command = ["timeout", str(self._timeout), lang_to_cmd(lang), filename] + + output, exit_code = await self._execute_command(command, cancellation_token) + outputs.append(output) + last_exit_code = exit_code + if exit_code != 0: + break + finally: + if self._delete_tmp_files: + for file in files: + try: + file.unlink() + except (OSError, FileNotFoundError): + pass code_file = str(files[0]) if files else None return CommandLineCodeResult(exit_code=last_exit_code, output="".join(outputs), code_file=code_file) @@ -512,6 +524,7 @@ def _to_config(self) -> DockerCommandLineCodeExecutorConfig: extra_volumes=self._extra_volumes, extra_hosts=self._extra_hosts, init_command=self._init_command, + delete_tmp_files=self._delete_tmp_files, ) @classmethod @@ -531,4 +544,5 @@ def _from_config(cls, config: DockerCommandLineCodeExecutorConfig) -> Self: extra_volumes=config.extra_volumes, extra_hosts=config.extra_hosts, init_command=config.init_command, + delete_tmp_files=config.delete_tmp_files, ) diff --git a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py index dfaa1b99c3df..17bb2e7a67a9 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py @@ -315,3 +315,49 @@ async def test_directory_creation_cleanup() -> None: await executor.stop() assert not Path(directory).exists() + + +@pytest.mark.asyncio +async def test_delete_tmp_files() -> None: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + + with tempfile.TemporaryDirectory() as temp_dir: + # Test with delete_tmp_files=False (default) + async with DockerCommandLineCodeExecutor(work_dir=temp_dir) as executor: + cancellation_token = CancellationToken() + code_blocks = [CodeBlock(code="print('test output')", language="python")] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code == 0 + assert result.code_file is not None + # Verify file exists after execution + assert Path(result.code_file).exists() + + # Test with delete_tmp_files=True + async with DockerCommandLineCodeExecutor(work_dir=temp_dir, delete_tmp_files=True) as executor: + cancellation_token = CancellationToken() + code_blocks = [CodeBlock(code="print('test output')", language="python")] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code == 0 + assert result.code_file is not None + # Verify file is deleted after execution + assert not Path(result.code_file).exists() + + # Test with multiple code blocks + code_blocks = [ + CodeBlock(code="print('first block')", language="python"), + CodeBlock(code="print('second block')", language="python"), + ] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code == 0 + assert result.code_file is not None + # Verify files are deleted after execution + assert not Path(result.code_file).exists() + + # Test deletion with execution error + code_blocks = [CodeBlock(code="raise Exception('test error')", language="python")] + result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert result.exit_code != 0 + assert result.code_file is not None + # Verify file is deleted even after error + assert not Path(result.code_file).exists()