427 lines
14 KiB
Python
427 lines
14 KiB
Python
|
import logging
|
||
|
|
import os
|
||
|
|
import pathlib
|
||
|
|
import shutil
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
import re
|
||
|
|
import time
|
||
|
|
import ast
|
||
|
|
from concurrent.futures import ThreadPoolExecutor
|
||
|
|
from hashlib import md5
|
||
|
|
from typing import Optional, Union, List, Tuple
|
||
|
|
from generic_escape import GenericEscape
|
||
|
|
|
||
|
|
from aios_kernel import AIStorage
|
||
|
|
|
||
|
|
try:
|
||
|
|
import docker
|
||
|
|
except ImportError:
|
||
|
|
docker = None
|
||
|
|
|
||
|
|
CODE_BLOCK_PATTERN = r"```[ \t]*(\w+)?[ \t]*\r?\n(.*?)\r?\n[ \t]*```"
|
||
|
|
UNKNOWN = "unknown"
|
||
|
|
TIMEOUT_MSG = "Timeout"
|
||
|
|
DEFAULT_TIMEOUT = 600
|
||
|
|
WIN32 = sys.platform == "win32"
|
||
|
|
PATH_SEPARATOR = WIN32 and "\\" or "/"
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
BUILT_IN_MODULES = set(
|
||
|
|
[
|
||
|
|
"sys",
|
||
|
|
"os",
|
||
|
|
"math",
|
||
|
|
"random",
|
||
|
|
"datetime",
|
||
|
|
"json",
|
||
|
|
"re",
|
||
|
|
"subprocess",
|
||
|
|
"time",
|
||
|
|
"threading",
|
||
|
|
"logging",
|
||
|
|
"collections",
|
||
|
|
"itertools",
|
||
|
|
"functools",
|
||
|
|
"operator",
|
||
|
|
"pathlib",
|
||
|
|
"shutil",
|
||
|
|
"tempfile",
|
||
|
|
"pickle",
|
||
|
|
"io",
|
||
|
|
"argparse",
|
||
|
|
"typing",
|
||
|
|
"unittest",
|
||
|
|
"contextlib",
|
||
|
|
"abc",
|
||
|
|
"heapq",
|
||
|
|
"bisect",
|
||
|
|
"copy",
|
||
|
|
"decimal",
|
||
|
|
"fractions",
|
||
|
|
"hashlib",
|
||
|
|
"secrets",
|
||
|
|
"statistics",
|
||
|
|
"difflib",
|
||
|
|
"doctest",
|
||
|
|
"enum",
|
||
|
|
"inspect",
|
||
|
|
"traceback",
|
||
|
|
"weakref",
|
||
|
|
"gc",
|
||
|
|
"mmap",
|
||
|
|
"msvcrt",
|
||
|
|
"winreg",
|
||
|
|
"array",
|
||
|
|
"audioop",
|
||
|
|
"binascii",
|
||
|
|
"cProfile",
|
||
|
|
"concurrent.futures",
|
||
|
|
"configparser",
|
||
|
|
"csv",
|
||
|
|
"ctypes",
|
||
|
|
"dateutil",
|
||
|
|
"dis",
|
||
|
|
"fnmatch",
|
||
|
|
"getopt",
|
||
|
|
"glob",
|
||
|
|
"gzip",
|
||
|
|
"pdb",
|
||
|
|
"pprint",
|
||
|
|
"profile",
|
||
|
|
"pstats",
|
||
|
|
"queue",
|
||
|
|
"socket",
|
||
|
|
"sqlite3",
|
||
|
|
"ssl",
|
||
|
|
"struct",
|
||
|
|
"tarfile",
|
||
|
|
"telnetlib",
|
||
|
|
"timeit",
|
||
|
|
"tokenize",
|
||
|
|
"uuid",
|
||
|
|
"xml",
|
||
|
|
"zipfile",
|
||
|
|
"zlib",
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def get_imports(code: str) -> List[str]:
|
||
|
|
root = ast.parse(code)
|
||
|
|
|
||
|
|
imports = []
|
||
|
|
for node in ast.iter_child_nodes(root):
|
||
|
|
if isinstance(node, ast.Import):
|
||
|
|
module_names = [alias.name for alias in node.names]
|
||
|
|
elif isinstance(node, ast.ImportFrom):
|
||
|
|
module_names = [node.module]
|
||
|
|
else:
|
||
|
|
continue
|
||
|
|
|
||
|
|
for name in module_names:
|
||
|
|
# Exclude built-in modules
|
||
|
|
if name not in BUILT_IN_MODULES:
|
||
|
|
imports.append(name)
|
||
|
|
|
||
|
|
return imports
|
||
|
|
|
||
|
|
|
||
|
|
def write_requirements(code: str, requirements_filepath: str):
|
||
|
|
imports = get_imports(code)
|
||
|
|
|
||
|
|
with open(requirements_filepath, "w") as file:
|
||
|
|
for module in imports:
|
||
|
|
file.write(module + "\n")
|
||
|
|
|
||
|
|
|
||
|
|
def _cmd(lang):
|
||
|
|
if lang.startswith("python") or lang in ["bash", "sh", "powershell"]:
|
||
|
|
return lang
|
||
|
|
if lang in ["shell"]:
|
||
|
|
return "sh"
|
||
|
|
if lang in ["ps1"]:
|
||
|
|
return "powershell"
|
||
|
|
raise NotImplementedError(f"{lang} not recognized in code execution")
|
||
|
|
|
||
|
|
|
||
|
|
def create_runner(code: str, timeout: int = 30) -> str:
|
||
|
|
"""
|
||
|
|
Create a Python script that runs the code and prints the output
|
||
|
|
"""
|
||
|
|
code = GenericEscape().escape(code)
|
||
|
|
# Create a runner script
|
||
|
|
runner = f"""
|
||
|
|
import os
|
||
|
|
import subprocess
|
||
|
|
|
||
|
|
my_env = os.environ.copy()
|
||
|
|
my_env["PYTHONIOENCODING"] = "utf-8"
|
||
|
|
|
||
|
|
process = subprocess.Popen(
|
||
|
|
f"python -i -q -u".split(),
|
||
|
|
stdin=subprocess.PIPE,
|
||
|
|
stdout=subprocess.PIPE,
|
||
|
|
stderr=subprocess.PIPE,
|
||
|
|
text=True,
|
||
|
|
bufsize=0,
|
||
|
|
universal_newlines=True,
|
||
|
|
env=my_env
|
||
|
|
)
|
||
|
|
|
||
|
|
process.stdin.write("{code}" + "\\n")
|
||
|
|
process.stdin.write("exit()\\n")
|
||
|
|
process.stdin.flush()
|
||
|
|
|
||
|
|
try:
|
||
|
|
process.wait({timeout})
|
||
|
|
except Exception as e:
|
||
|
|
process.terminate()
|
||
|
|
|
||
|
|
for line in iter(process.stdout.readline, ""):
|
||
|
|
print(line)
|
||
|
|
|
||
|
|
for line in iter(process.stderr.readline, ""):
|
||
|
|
if line.startswith(">>>"):
|
||
|
|
continue
|
||
|
|
print(line)
|
||
|
|
"""
|
||
|
|
return runner
|
||
|
|
|
||
|
|
|
||
|
|
def _run_cmd(cmd: [str], work_dir: str, timeout: int) -> str:
|
||
|
|
if WIN32:
|
||
|
|
logger.warning("SIGALRM is not supported on Windows. No timeout will be enforced.")
|
||
|
|
result = subprocess.run(
|
||
|
|
cmd,
|
||
|
|
cwd=work_dir,
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
||
|
|
future = executor.submit(
|
||
|
|
subprocess.run,
|
||
|
|
cmd,
|
||
|
|
cwd=work_dir,
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
)
|
||
|
|
result = future.result(timeout=timeout)
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def execute_code(
|
||
|
|
code: Optional[str] = None,
|
||
|
|
timeout: Optional[int] = None,
|
||
|
|
filename: Optional[str] = None,
|
||
|
|
work_dir: Optional[str] = None,
|
||
|
|
use_docker: Optional[Union[List[str], str, bool]] = None,
|
||
|
|
lang: Optional[str] = "python",
|
||
|
|
) -> Tuple[int, str]:
|
||
|
|
"""Execute code in a docker container.
|
||
|
|
This function is not tested on MacOS.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
code (Optional, str): The code to execute.
|
||
|
|
If None, the code from the file specified by filename will be executed.
|
||
|
|
Either code or filename must be provided.
|
||
|
|
timeout (Optional, int): The maximum execution time in seconds.
|
||
|
|
If None, a default timeout will be used. The default timeout is 600 seconds. On Windows, the timeout is not enforced when use_docker=False.
|
||
|
|
filename (Optional, str): The file name to save the code or where the code is stored when `code` is None.
|
||
|
|
If None, a file with a randomly generated name will be created.
|
||
|
|
The randomly generated file will be deleted after execution.
|
||
|
|
The file name must be a relative path. Relative paths are relative to the working directory.
|
||
|
|
work_dir (Optional, str): The working directory for the code execution.
|
||
|
|
If None, a default working directory will be used.
|
||
|
|
The default working directory is the "extensions" directory under
|
||
|
|
"path_to_autogen".
|
||
|
|
use_docker (Optional, list, str or bool): The docker image to use for code execution.
|
||
|
|
If a list or a str of image name(s) is provided, the code will be executed in a docker container
|
||
|
|
with the first image successfully pulled.
|
||
|
|
If None, False or empty, the code will be executed in the current environment.
|
||
|
|
Default is None, which will be converted into an empty list when docker package is available.
|
||
|
|
Expected behaviour:
|
||
|
|
- If `use_docker` is explicitly set to True and the docker package is available, the code will run in a Docker container.
|
||
|
|
- If `use_docker` is explicitly set to True but the Docker package is missing, an error will be raised.
|
||
|
|
- If `use_docker` is not set (i.e., left default to None) and the Docker package is not available, a warning will be displayed, but the code will run natively.
|
||
|
|
If the code is executed in the current environment,
|
||
|
|
the code must be trusted.
|
||
|
|
lang (Optional, str): The language of the code. Default is "python".
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
int: 0 if the code executes successfully.
|
||
|
|
str: The error message if the code fails to execute; the stdout otherwise.
|
||
|
|
"""
|
||
|
|
if all((code is None, filename is None)):
|
||
|
|
error_msg = f"Either {code=} or {filename=} must be provided."
|
||
|
|
logger.error(error_msg)
|
||
|
|
raise AssertionError(error_msg)
|
||
|
|
|
||
|
|
# Warn if use_docker was unspecified (or None), and cannot be provided (the default).
|
||
|
|
# In this case the current behavior is to fall back to run natively, but this behavior
|
||
|
|
# is subject to change.
|
||
|
|
if use_docker is None:
|
||
|
|
if docker is None:
|
||
|
|
use_docker = False
|
||
|
|
logger.warning(
|
||
|
|
"execute_code was called without specifying a value for use_docker. Since the python docker package is not available, code will be run natively. Note: this fallback behavior is subject to change"
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
# Default to true
|
||
|
|
use_docker = True
|
||
|
|
|
||
|
|
timeout = timeout or DEFAULT_TIMEOUT
|
||
|
|
original_filename = filename
|
||
|
|
if WIN32 and lang in ["sh", "shell"] and (not use_docker):
|
||
|
|
lang = "ps1"
|
||
|
|
if filename is None:
|
||
|
|
code_hash = md5(code.encode()).hexdigest()
|
||
|
|
# create a file with a automatically generated name
|
||
|
|
filename = f"tmp_code_{code_hash}.{'py' if lang.startswith('python') else lang}"
|
||
|
|
if work_dir is None:
|
||
|
|
WORKING_DIR = os.path.join(AIStorage.get_instance().get_myai_dir(), "tmp_code")
|
||
|
|
pathlib.Path(WORKING_DIR).mkdir(exist_ok=True)
|
||
|
|
work_dir = os.path.join(WORKING_DIR, code_hash)
|
||
|
|
pathlib.Path(work_dir).mkdir(exist_ok=True)
|
||
|
|
filepath = os.path.join(work_dir, filename)
|
||
|
|
file_dir = os.path.dirname(filepath)
|
||
|
|
os.makedirs(file_dir, exist_ok=True)
|
||
|
|
if code is not None:
|
||
|
|
write_requirements(code, os.path.join(file_dir, "requirements.txt"))
|
||
|
|
code = create_runner(code, 30)
|
||
|
|
with open(filepath, "w", encoding="utf-8") as fout:
|
||
|
|
fout.write(code)
|
||
|
|
|
||
|
|
|
||
|
|
# check if already running in a docker container
|
||
|
|
in_docker_container = os.path.exists("/.dockerenv")
|
||
|
|
if not use_docker or in_docker_container:
|
||
|
|
try:
|
||
|
|
env_cmd = ["python", "-m", "venv", os.path.join(file_dir, "venv")]
|
||
|
|
_run_cmd(env_cmd, file_dir, timeout)
|
||
|
|
if WIN32:
|
||
|
|
venv_path = os.path.join(file_dir, "venv", "Scripts")
|
||
|
|
else:
|
||
|
|
venv_path = os.path.join(file_dir, "venv", "bin")
|
||
|
|
pip_cmd = [os.path.join(venv_path, "python"), "-m", "pip", "install", "-r", "requirements.txt"]
|
||
|
|
_run_cmd(pip_cmd, file_dir, timeout)
|
||
|
|
# already running in a docker container
|
||
|
|
cmd = [
|
||
|
|
os.path.join(venv_path, "python"),
|
||
|
|
f".\\{filename}" if WIN32 else filename,
|
||
|
|
]
|
||
|
|
result = _run_cmd(cmd, file_dir, timeout)
|
||
|
|
except TimeoutError:
|
||
|
|
if original_filename is None:
|
||
|
|
shutil.rmtree(os.path.join(file_dir, "venv"))
|
||
|
|
os.remove(filepath)
|
||
|
|
os.remove(os.path.join(file_dir, "requirements.txt"))
|
||
|
|
try:
|
||
|
|
os.removedirs(file_dir)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
return 1, TIMEOUT_MSG
|
||
|
|
if original_filename is None:
|
||
|
|
shutil.rmtree(os.path.join(file_dir, "venv"))
|
||
|
|
os.remove(filepath)
|
||
|
|
os.remove(os.path.join(file_dir, "requirements.txt"))
|
||
|
|
try:
|
||
|
|
os.removedirs(file_dir)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
if result.returncode:
|
||
|
|
logs = result.stderr
|
||
|
|
if original_filename is None:
|
||
|
|
abs_path = str(pathlib.Path(filepath).absolute())
|
||
|
|
logs = logs.replace(str(abs_path), "").replace(filename, "")
|
||
|
|
else:
|
||
|
|
abs_path = str(pathlib.Path(work_dir).absolute()) + PATH_SEPARATOR
|
||
|
|
logs = logs.replace(str(abs_path), "")
|
||
|
|
else:
|
||
|
|
logs = result.stdout
|
||
|
|
return result.returncode, logs
|
||
|
|
|
||
|
|
# create a docker client
|
||
|
|
client = docker.from_env()
|
||
|
|
image_list = (
|
||
|
|
["python:3-alpine", "python:3", "python:3-windowsservercore"]
|
||
|
|
if use_docker is True
|
||
|
|
else [use_docker]
|
||
|
|
if isinstance(use_docker, str)
|
||
|
|
else use_docker
|
||
|
|
)
|
||
|
|
for image in image_list:
|
||
|
|
# check if the image exists
|
||
|
|
try:
|
||
|
|
client.images.get(image)
|
||
|
|
break
|
||
|
|
except docker.errors.ImageNotFound:
|
||
|
|
# pull the image
|
||
|
|
logger.info("Pulling image", image)
|
||
|
|
try:
|
||
|
|
client.images.pull(image, stream=True, decode=True)
|
||
|
|
break
|
||
|
|
except docker.errors.DockerException as e:
|
||
|
|
logger.error("Failed to pull image", image)
|
||
|
|
logger.exception(e)
|
||
|
|
# get a randomized str based on current time to wrap the exit code
|
||
|
|
exit_code_str = f"exitcode{time.time()}"
|
||
|
|
start_str = f'start{time.time()}'
|
||
|
|
abs_path = pathlib.Path(work_dir).absolute()
|
||
|
|
cmd = [
|
||
|
|
"sh",
|
||
|
|
"-c",
|
||
|
|
f"pip install --quiet -r requirements.txt; echo -n {start_str}; {_cmd(lang)} {filename}; exit_code=$?; echo -n {exit_code_str}; echo -n $exit_code; echo {exit_code_str};",
|
||
|
|
]
|
||
|
|
# create a docker container
|
||
|
|
container = client.containers.run(
|
||
|
|
image,
|
||
|
|
command=cmd,
|
||
|
|
working_dir="/workspace",
|
||
|
|
detach=True,
|
||
|
|
# get absolute path to the working directory
|
||
|
|
volumes={abs_path: {"bind": "/workspace", "mode": "rw"}},
|
||
|
|
)
|
||
|
|
start_time = time.time()
|
||
|
|
while container.status != "exited" and time.time() - start_time < timeout:
|
||
|
|
# Reload the container object
|
||
|
|
container.reload()
|
||
|
|
if container.status != "exited":
|
||
|
|
container.stop()
|
||
|
|
container.remove()
|
||
|
|
if original_filename is None:
|
||
|
|
os.remove(filepath)
|
||
|
|
return 1, TIMEOUT_MSG, image
|
||
|
|
# get the container logs
|
||
|
|
logs: str = container.logs().decode("utf-8").rstrip()
|
||
|
|
start_pos = logs.find(start_str)
|
||
|
|
if start_pos != -1:
|
||
|
|
logs = logs[start_pos + len(start_str):]
|
||
|
|
# # commit the image
|
||
|
|
# tag = filename.replace("/", "")
|
||
|
|
# container.commit(repository="python", tag=tag)
|
||
|
|
# remove the container
|
||
|
|
container.remove()
|
||
|
|
# check if the code executed successfully
|
||
|
|
exit_code = container.attrs["State"]["ExitCode"]
|
||
|
|
if exit_code == 0:
|
||
|
|
# extract the exit code from the logs
|
||
|
|
pattern = re.compile(f"{exit_code_str}(\\d+){exit_code_str}")
|
||
|
|
match = pattern.search(logs)
|
||
|
|
exit_code = 1 if match is None else int(match.group(1))
|
||
|
|
# remove the exit code from the logs
|
||
|
|
logs = logs if match is None else pattern.sub("", logs)
|
||
|
|
|
||
|
|
if original_filename is None:
|
||
|
|
os.remove(filepath)
|
||
|
|
os.remove(os.path.join(file_dir, "requirements.txt"))
|
||
|
|
os.removedirs(file_dir)
|
||
|
|
if exit_code:
|
||
|
|
logs = logs.replace(f"/workspace/{filename if original_filename is None else ''}", "")
|
||
|
|
# return the exit code, logs and image
|
||
|
|
return exit_code, logs
|
||
|
|
|