Refactor the code directory structure to better suit the current complexity
This commit is contained in:
Liu Zhicong
@@ -1,426 +0,0 @@
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user