wangtao
/
autogen
mirror of https://github.com/microsoft/autogen.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Run LocalCommandLineCodeExecutor within venv (#3977) 

* Run LocalCommandLineCodeExecutor within venv

* Remove create_virtual_env func and add docstring

* Add explanation for LocalCommandLineExecutor docstring example

* Enhance docstring example explanation

---------

Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
This commit is contained in:
Gerardo Moreno 2024-10-29 08:17:34 -07:00 committed by GitHub
parent eb4b1f856e
commit 93733dbd65
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 176 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

72
python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/command-line-code-executors.ipynb
View File

@ -136,6 +136,76 @@
" )\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Local within a Virtual Environment\n",
"\n",
"If you want the code to run within a virtual environment created as part of the application’s setup, you can specify a directory for the newly created environment and pass its context to {py:class}`~autogen_core.components.code_executor.LocalCommandLineCodeExecutor`. This setup allows the executor to use the specified virtual environment consistently throughout the application's lifetime, ensuring isolated dependencies and a controlled runtime environment."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"CommandLineCodeResult(exit_code=0, output='', code_file='/Users/gziz/Dev/autogen/python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/coding/tmp_code_d2a7db48799db3cc785156a11a38822a45c19f3956f02ec69b92e4169ecbf2ca.bash')"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import venv\n",
"from pathlib import Path\n",
"\n",
"from autogen_core.base import CancellationToken\n",
"from autogen_core.components.code_executor import CodeBlock, LocalCommandLineCodeExecutor\n",
"\n",
"work_dir = Path(\"coding\")\n",
"work_dir.mkdir(exist_ok=True)\n",
"\n",
"venv_dir = work_dir / \".venv\"\n",
"venv_builder = venv.EnvBuilder(with_pip=True)\n",
"venv_builder.create(venv_dir)\n",
"venv_context = venv_builder.ensure_directories(venv_dir)\n",
"\n",
"local_executor = LocalCommandLineCodeExecutor(work_dir=work_dir, virtual_env_context=venv_context)\n",
"await local_executor.execute_code_blocks(\n",
" code_blocks=[\n",
" CodeBlock(language=\"bash\", code=\"pip install matplotlib\"),\n",
" ],\n",
" cancellation_token=CancellationToken(),\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As we can see, the code has executed successfully, and the installation has been isolated to the newly created virtual environment, without affecting our global environment."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
@ -154,7 +224,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
"version": "3.12.4"
}
},
"nbformat": 4,

57
python/packages/autogen-core/src/autogen_core/components/code_executor/_impl/local_commandline_code_executor.py
View File

@ -3,12 +3,14 @@
import asyncio
import logging
import os
import sys
import warnings
from hashlib import sha256
from pathlib import Path
from string import Template
from typing import Any, Callable, ClassVar, List, Sequence, Union
from types import SimpleNamespace
from typing import Any, Callable, ClassVar, List, Optional, Sequence, Union
from typing_extensions import ParamSpec
@ -54,6 +56,36 @@ class LocalCommandLineCodeExecutor(CodeExecutor):
directory is the current directory ".".
functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list.
functions_module (str, optional): The name of the module that will be created to store the functions. Defaults to "functions".
virtual_env_context (Optional[SimpleNamespace], optional): The virtual environment context. Defaults to None.
Example:
How to use `LocalCommandLineCodeExecutor` with a virtual environment different from the one used to run the autogen application:
Set up a virtual environment using the `venv` module, and pass its context to the initializer of `LocalCommandLineCodeExecutor`. This way, the executor will run code within the new environment.
.. code-block:: python
import venv
from pathlib import Path
from autogen_core.base import CancellationToken
from autogen_core.components.code_executor import CodeBlock, LocalCommandLineCodeExecutor
work_dir = Path("coding")
work_dir.mkdir(exist_ok=True)
venv_dir = work_dir / ".venv"
venv_builder = venv.EnvBuilder(with_pip=True)
venv_builder.create(venv_dir)
venv_context = venv_builder.ensure_directories(venv_dir)
local_executor = LocalCommandLineCodeExecutor(work_dir=work_dir, virtual_env_context=venv_context)
await local_executor.execute_code_blocks(
code_blocks=[
CodeBlock(language="bash", code="pip install matplotlib"),
],
cancellation_token=CancellationToken(),
)
"""
@ -86,6 +118,7 @@ $functions"""
]
] = [],
functions_module: str = "functions",
virtual_env_context: Optional[SimpleNamespace] = None,
):
if timeout < 1:
raise ValueError("Timeout must be greater than or equal to 1.")
@ -110,6 +143,8 @@ $functions"""
else:
self._setup_functions_complete = True
self._virtual_env_context: Optional[SimpleNamespace] = virtual_env_context
def format_functions_for_prompt(self, prompt_template: str = FUNCTION_PROMPT_TEMPLATE) -> str:
"""(Experimental) Format the functions for a prompt.
@ -164,9 +199,14 @@ $functions"""
cmd_args = ["-m", "pip", "install"]
cmd_args.extend(required_packages)
if self._virtual_env_context:
py_executable = self._virtual_env_context.env_exe
else:
py_executable = sys.executable
task = asyncio.create_task(
asyncio.create_subprocess_exec(
sys.executable,
py_executable,
*cmd_args,
cwd=self._work_dir,
stdout=asyncio.subprocess.PIPE,
@ -253,7 +293,17 @@ $functions"""
f.write(code)
file_names.append(written_file)
program = sys.executable if lang.startswith("python") else lang_to_cmd(lang)
env = os.environ.copy()
if self._virtual_env_context:
virtual_env_exe_abs_path = os.path.abspath(self._virtual_env_context.env_exe)
virtual_env_bin_abs_path = os.path.abspath(self._virtual_env_context.bin_path)
env["PATH"] = f"{virtual_env_bin_abs_path}{os.pathsep}{env['PATH']}"
program = virtual_env_exe_abs_path if lang.startswith("python") else lang_to_cmd(lang)
else:
program = sys.executable if lang.startswith("python") else lang_to_cmd(lang)
# Wrap in a task to make it cancellable
task = asyncio.create_task(
asyncio.create_subprocess_exec(
@ -262,6 +312,7 @@ $functions"""
cwd=self._work_dir,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
)
)
cancellation_token.link_future(task)

51
python/packages/autogen-core/tests/execution/test_commandline_code_executor.py
View File

@ -2,8 +2,11 @@
# Credit to original authors
import asyncio
import os
import shutil
import sys
import tempfile
import venv
from pathlib import Path
from typing import AsyncGenerator, TypeAlias
@ -143,3 +146,51 @@ print("hello world")
assert "test.py" in result.code_file
assert (temp_dir / Path("test.py")).resolve() == Path(result.code_file).resolve()
assert (temp_dir / Path("test.py")).exists()
@pytest.mark.asyncio
async def test_local_executor_with_custom_venv() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
env_builder = venv.EnvBuilder(with_pip=True)
env_builder.create(temp_dir)
env_builder_context = env_builder.ensure_directories(temp_dir)
executor = LocalCommandLineCodeExecutor(work_dir=temp_dir, virtual_env_context=env_builder_context)
code_blocks = [
# https://stackoverflow.com/questions/1871549/how-to-determine-if-python-is-running-inside-a-virtualenv
CodeBlock(code="import sys; print(sys.prefix != sys.base_prefix)", language="python"),
]
cancellation_token = CancellationToken()
result = await executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token)
assert result.exit_code == 0
assert result.output.strip() == "True"
@pytest.mark.asyncio
async def test_local_executor_with_custom_venv_in_local_relative_path() -> None:
relative_folder_path = "tmp_dir"
try:
if not os.path.isdir(relative_folder_path):
os.mkdir(relative_folder_path)
env_path = os.path.join(relative_folder_path, ".venv")
env_builder = venv.EnvBuilder(with_pip=True)
env_builder.create(env_path)
env_builder_context = env_builder.ensure_directories(env_path)
executor = LocalCommandLineCodeExecutor(work_dir=relative_folder_path, virtual_env_context=env_builder_context)
code_blocks = [
CodeBlock(code="import sys; print(sys.executable)", language="python"),
]
cancellation_token = CancellationToken()
result = await executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token)
assert result.exit_code == 0
# Check if the expected venv has been used
bin_path = os.path.abspath(env_builder_context.bin_path)
assert Path(result.output.strip()).parent.samefile(bin_path)
finally:
if os.path.isdir(relative_folder_path):
shutil.rmtree(relative_folder_path)