Skip to content

Docker Code Exec delete temp files #6211

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
husseinmozannar merged 3 commits into from
Apr 6, 2025
Merged

Docker Code Exec delete temp files #6211

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
20ca7d8
docker delete temp
husseinmozannar Apr 4, 2025
cfa75f5
Merge branch 'main' into tmp_code_code
husseinmozannar Apr 4, 2025
5894de6
Merge branch 'main' into tmp_code_code
husseinmozannar Apr 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 41 additions & 27 deletions python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]):
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.")
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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,
)
46 changes: 46 additions & 0 deletions python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Loading