make code execution async (#219)

* make code execution async

* python 3.10 does not support asyncio.timeout()

* make code execution cancellable

* make code execution async

* python 3.10 does not support asyncio.timeout()

* make code execution cancellable

* make entire callstack for code_executor async

* Update python/src/agnext/components/code_executor/_impl/local_commandline_code_executor.py

Co-authored-by: Jack Gerrits <jackgerrits@users.noreply.github.com>

* fix variable description

* remove unnecessary code

* fix usage of execute_code_blocks

* fix usage of execute_code_blocks

---------

Co-authored-by: Jack Gerrits <jackgerrits@users.noreply.github.com>
Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
This commit is contained in:
peterychang
2024-07-26 18:37:34 -04:00
committed by GitHub
parent 53343972f0
commit 070a97ceaa
8 changed files with 100 additions and 82 deletions

View File

@@ -12,15 +12,15 @@ UNIX_SHELLS = ["bash", "sh", "shell"]
WINDOWS_SHELLS = ["ps1", "pwsh", "powershell"]
PYTHON_VARIANTS = ["python", "Python", "py"]
def test_execute_code() -> None:
@pytest.mark.asyncio
async def test_execute_code() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
executor = LocalCommandLineCodeExecutor(work_dir=temp_dir)
# Test single code block.
code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")]
code_result = executor.execute_code_blocks(code_blocks)
code_result = await executor.execute_code_blocks(code_blocks)
assert code_result.exit_code == 0 and "hello world!" in code_result.output and code_result.code_file is not None
# Test multiple code blocks.
@@ -28,7 +28,7 @@ def test_execute_code() -> None:
CodeBlock(code="import sys; print('hello world!')", language="python"),
CodeBlock(code="a = 100 + 100; print(a)", language="python"),
]
code_result = executor.execute_code_blocks(code_blocks)
code_result = await executor.execute_code_blocks(code_blocks)
assert (
code_result.exit_code == 0
and "hello world!" in code_result.output
@@ -39,13 +39,13 @@ def test_execute_code() -> None:
# Test bash script.
if sys.platform not in ["win32"]:
code_blocks = [CodeBlock(code="echo 'hello world!'", language="bash")]
code_result = executor.execute_code_blocks(code_blocks)
code_result = await executor.execute_code_blocks(code_blocks)
assert code_result.exit_code == 0 and "hello world!" in code_result.output and code_result.code_file is not None
# Test running code.
file_lines = ["import sys", "print('hello world!')", "a = 100 + 100", "print(a)"]
code_blocks = [CodeBlock(code="\n".join(file_lines), language="python")]
code_result = executor.execute_code_blocks(code_blocks)
code_result = await executor.execute_code_blocks(code_blocks)
assert (
code_result.exit_code == 0
and "hello world!" in code_result.output
@@ -59,12 +59,12 @@ def test_execute_code() -> None:
for file_line, code_line in zip(file_lines, code_lines):
assert file_line.strip() == code_line.strip()
def test_commandline_code_executor_timeout() -> None:
@pytest.mark.asyncio
async def test_commandline_code_executor_timeout() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
executor = LocalCommandLineCodeExecutor(timeout=1, work_dir=temp_dir)
code_blocks = [CodeBlock(code="import time; time.sleep(10); print('hello world!')", language="python")]
code_result = executor.execute_code_blocks(code_blocks)
code_result = await executor.execute_code_blocks(code_blocks)
assert code_result.exit_code and "Timeout" in code_result.output
@@ -75,18 +75,18 @@ def test_local_commandline_code_executor_restart() -> None:
def test_invalid_relative_path() -> None:
@pytest.mark.asyncio
async def test_invalid_relative_path() -> None:
executor = LocalCommandLineCodeExecutor()
code = """# filename: /tmp/test.py
print("hello world")
"""
result = executor.execute_code_blocks([CodeBlock(code=code, language="python")])
result = await executor.execute_code_blocks([CodeBlock(code=code, language="python")])
assert result.exit_code == 1 and "Filename is not in the workspace" in result.output
def test_valid_relative_path() -> None:
@pytest.mark.asyncio
async def test_valid_relative_path() -> None:
with tempfile.TemporaryDirectory() as temp_dir_str:
temp_dir = Path(temp_dir_str)
executor = LocalCommandLineCodeExecutor(work_dir=temp_dir)
@@ -94,7 +94,7 @@ def test_valid_relative_path() -> None:
print("hello world")
"""
result = executor.execute_code_blocks([CodeBlock(code=code, language="python")])
result = await executor.execute_code_blocks([CodeBlock(code=code, language="python")])
assert result.exit_code == 0
assert "hello world" in result.output
assert result.code_file is not None

View File

@@ -47,7 +47,8 @@ def function_missing_reqs() -> "polars.DataFrame":
return polars.DataFrame()
def test_can_load_function_with_reqs() -> None:
@pytest.mark.asyncio
async def test_can_load_function_with_reqs() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
executor = LocalCommandLineCodeExecutor(
work_dir=temp_dir, functions=[load_data]
@@ -59,7 +60,7 @@ import polars
data = load_data()
print(data['name'][0])"""
result = executor.execute_code_blocks(
result = await executor.execute_code_blocks(
code_blocks=[
CodeBlock(language="python", code=code),
]
@@ -68,7 +69,8 @@ print(data['name'][0])"""
assert result.exit_code == 0
def test_can_load_function() -> None:
@pytest.mark.asyncio
async def test_can_load_function() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
executor = LocalCommandLineCodeExecutor(
work_dir=temp_dir, functions=[add_two_numbers]
@@ -76,7 +78,7 @@ def test_can_load_function() -> None:
code = f"""from {executor.functions_module} import add_two_numbers
print(add_two_numbers(1, 2))"""
result = executor.execute_code_blocks(
result = await executor.execute_code_blocks(
code_blocks=[
CodeBlock(language="python", code=code),
]
@@ -85,7 +87,8 @@ print(add_two_numbers(1, 2))"""
assert result.exit_code == 0
def test_fails_for_function_incorrect_import() -> None:
@pytest.mark.asyncio
async def test_fails_for_function_incorrect_import() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
executor = LocalCommandLineCodeExecutor(
work_dir=temp_dir, functions=[function_incorrect_import]
@@ -94,14 +97,15 @@ def test_fails_for_function_incorrect_import() -> None:
function_incorrect_import()"""
with pytest.raises(ValueError):
executor.execute_code_blocks(
await executor.execute_code_blocks(
code_blocks=[
CodeBlock(language="python", code=code),
]
)
def test_fails_for_function_incorrect_dep() -> None:
@pytest.mark.asyncio
async def test_fails_for_function_incorrect_dep() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
executor = LocalCommandLineCodeExecutor(
work_dir=temp_dir, functions=[function_incorrect_dep]
@@ -110,7 +114,7 @@ def test_fails_for_function_incorrect_dep() -> None:
function_incorrect_dep()"""
with pytest.raises(ValueError):
executor.execute_code_blocks(
await executor.execute_code_blocks(
code_blocks=[
CodeBlock(language="python", code=code),
]
@@ -152,7 +156,8 @@ def add_two_numbers(a: int, b: int) -> int:
)
def test_can_load_str_function_with_reqs() -> None:
@pytest.mark.asyncio
async def test_can_load_str_function_with_reqs() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
func = FunctionWithRequirements.from_str(
'''
@@ -166,7 +171,7 @@ def add_two_numbers(a: int, b: int) -> int:
code = f"""from {executor.functions_module} import add_two_numbers
print(add_two_numbers(1, 2))"""
result = executor.execute_code_blocks(
result = await executor.execute_code_blocks(
code_blocks=[
CodeBlock(language="python", code=code),
]
@@ -187,7 +192,8 @@ invaliddef add_two_numbers(a: int, b: int) -> int:
)
def test_cant_run_broken_str_function_with_reqs() -> None:
@pytest.mark.asyncio
async def test_cant_run_broken_str_function_with_reqs() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
func = FunctionWithRequirements.from_str(
'''
@@ -201,7 +207,7 @@ def add_two_numbers(a: int, b: int) -> int:
code = f"""from {executor.functions_module} import add_two_numbers
print(add_two_numbers(object(), False))"""
result = executor.execute_code_blocks(
result = await executor.execute_code_blocks(
code_blocks=[
CodeBlock(language="python", code=code),
]