打造安全沙盒！如何完美执行AI生成代码？
- Rifx.Online
- Programming , Technology , Machine Learning
- 12 Feb, 2025
在构建AI代理时，特别是那些动态生成和执行分析代码的代理，一个主要的关注点是安全性和稳定性。允许代码在生产服务器或本地机器上任意执行可能会带来风险。代码可能无限运行，消耗过多资源，甚至可能危及系统安全。
解决方案是沙盒化或容器化代码执行的环境。这确保了与主机系统的隔离，一致的依赖关系，以及对在后台运行的任何代码更可控的生命周期。
在这篇博客文章中，我们将探讨：
- 为什么沙盒化是必要的。
- 设置基于FastAPI + Jupyter的沙盒的逐步指南。
- 相关的Docker配置和FastAPI服务器代码。
- 编写临时Python “exec”包装器时的陷阱。
- 在容器化Jupyter内核中需要注意的可能问题。
为什么您需要一个沙盒执行环境
安全性
当代码可以由 AI 代理动态生成时，您面临着恶意代码注入的风险，或者代码试图访问本地文件或不应该访问的网络资源。沙盒技术强制代码在与其他服务或主机操作系统分开的受限环境中运行。
资源控制
沙盒技术允许您对 CPU、内存和运行时间施加资源限制。例如，如果用户的代码陷入无限循环，您可以终止容器或内核，而不会影响主机。
更简单的依赖管理
不同的 AI 或数据分析任务可能对 Python 库有冲突的要求。通过使用 Docker 容器，您可以隔离每个工作流的依赖关系。这确保了安装或升级一个库不会破坏另一个工作流。
一致的环境
一旦您为沙盒环境创建了 Docker 镜像，每个用户或工作流都将拥有相同的环境。不再有“但它在我的机器上可以工作”的情况。
状态保持性
在代理工作流中，您通常希望 AI 在后续迭代中仅纠正错误代码。在这种情况下，您希望保留命名空间，这就像执行 jupyter notebook 的另一个单元格。
设置基于 FastAPI + Jupyter 的沙箱服务
为了说明如何构建一个沙箱代码执行环境，我们将使用：
- FastAPI 作为网络服务器框架。
- Jupyter 作为创建 Python 内核和运行代码的后端。
- Docker 来容器化所有内容。
以下是我们的项目结构（简化版）：
├── Dockerfile
├── fastapi_jupyter_server.py
├── test_api.py
└── ...mounted folders(data , jupyter session etc.)
Docker 配置和 FastAPI 服务器代码
Dockerfile
这是设置我们沙箱环境的 Dockerfile，使用 Python 3.10-slim。它安装系统依赖项、用于数据分析的 Python 库，以及 FastAPI 和 Jupyter 所需的包。
## Use a lightweight Python base image
FROM python:3.10-slim
## Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential \
python3-pip \
&& apt-get clean && rm -rf /var/lib/apt/lists/*
## Install Python dependencies
RUN pip install python-multipart fastapi uvicorn jupyter-client nbformat ipykernel
RUN python3 -m ipykernel install --user
RUN pip install pandas numpy matplotlib scipy seaborn scikit-learn pyarrow tabulate openpyxl xlrd
## Create necessary directories
RUN mkdir -p /mnt/data /mnt/jupyter_sessions /workspace
## Set environment variables for mounted volumes
ENV DATA_DIR=/mnt/data
ENV JUPYTER_SESSIONS_DIR=/mnt/jupyter_sessions
## Copy the FastAPI server script into the container
COPY fastapi_jupyter_server.py /workspace/fastapi_jupyter_api.py
## Set the working directory
WORKDIR /workspace
## Expose the FastAPI server port
EXPOSE 5000
## Use Python to start the FastAPI server
CMD ["python3.10", "-m", "uvicorn", "fastapi_jupyter_api:app", "--host", "0.0.0.0", "--port", "5000"]
此 Docker 设置的关键点：
- 隔离：容器将代码执行与您的本地机器隔离。
- 挂载卷：我们定义
/mnt/data和/mnt/jupyter_sessions用于外部数据和存储 Jupyter 笔记本会话。 - 暴露端口：我们在容器内暴露端口
5000，以便可以将其转发到外部（例如，-p 5002:5000）。
FastAPI + Jupyter 服务器
在 fastapi_jupyter_server.py 中，我们定义了以下端点：
- 启动新会话（启动 Jupyter 内核，创建一个新的笔记本文件）。
- 在用户内核中执行代码。
- 安装包（尽管应谨慎使用，因为这可能导致环境漂移）。
- 重置或结束会话。
from fastapi import FastAPI, UploadFile, Form, HTTPException
import subprocess
import os
import queue
import asyncio
from jupyter_client import KernelManager
import nbformat
from nbformat.v4 import new_notebook
import time
from typing import Dict, Optional
app = FastAPI()
BASE_FOLDER = "/mnt/data"
SESSIONS_FOLDER = "/mnt/jupyter_sessions"
...
(…请查看 GitHub Repo 获取完整代码。)
JupyterController 类
此类：
- 管理 Jupyter 内核的生命周期。
- 处理笔记本文件的创建。
- 提供执行代码和捕获输出的方法。
- 如果内核崩溃或无响应，则重置内核。
- 清理资源并在完成后删除笔记本。
我们依赖 jupyter_client.KernelManager 来启动/停止内核。这是运行完整 Jupyter 服务器的轻量级替代方案。
class JupyterController:
def __init__(self, folder_path):
self.folder_path = folder_path
self.notebook_path = None
self.kernel_manager = None
self.kernel_client = None
self._kernel_ready = False
# Create notebook, start kernel, wait for readiness
async def create_notebook(self, notebook_name):
os.makedirs(self.folder_path, exist_ok=True)
self.notebook_path = os.path.join(self.folder_path, f"{notebook_name}.ipynb")
nb = new_notebook()
with open(self.notebook_path, "w") as f:
nbformat.write(nb, f)
self.kernel_manager = KernelManager()
self.kernel_manager.start_kernel()
self.kernel_client = self.kernel_manager.client()
self.kernel_client.start_channels()
# Wait until kernel is up
await self._wait_for_kernel_ready()
self._clear_output_queue()
return self.notebook_path
async def execute_code(self, code):
# …
pass
async def reset_kernel(self):
# …
pass
def cleanup(self):
# …
pass
FastAPI 端点
端点是标准的：
/start_session– 创建一个新的 Jupyter 会话。/execute– 在用户的会话中运行用户提交的代码。/install_package– 安装一个包/reset– 重置 Jupyter 内核。/end_session– 清理并结束用户的会话。
测试 API
test_api.py 显示了如何调用 /start_session 和 /execute 端点。
import requests
import json
def execute_code(user_id, code):
url = "http://localhost:5002/execute"
payload = {"user_id": user_id, "code": code}
response = requests.post(url, json=payload)
return response.json()
def start_session(user_id):
url = "http://localhost:5002/start_session"
response = requests.post(url, data={"user_id": user_id})
return response.json()
## Example usage:
if __name__ == "__main__":
user_id = "user_test"
session_result = start_session(user_id)
print(session_result)
# Execute some code
result = execute_code(user_id, 'x = 42\nprint(f"x = {x}")')
print(result)
运行您的容器
1. 构建 Docker 镜像
docker build -t fastapi-jupyter-api .
2. 运行容器
docker run -d -p 5002:5000 \
-v $(pwd)/data:/mnt/data \
-v $(pwd)/jupyter_sessions:/mnt/jupyter_sessions \
fastapi-jupyter-api
这将在 http://localhost:5002 上发布 FastAPI 应用程序，并挂载两个卷用于数据和笔记本会话。
使用 python test_api.py 测试端点。
Ad-Hoc “exec” 包装器的陷阱
在没有正式沙箱的情况下，开发人员有时会将 Python 的 exec 嵌入到一个函数中（如下片段所示）以捕获输出和可视化：
def execute_analysis(code: str) -> Dict[str, Any]:
output = io.StringIO()
result = None
plots_info = []
output_string = ""
# ...
local_namespace = {
'pd': pd,
'agent_state': self.agent_state,
'np': __import__('numpy'),
'plt': __import__('matplotlib.pyplot'),
'sns': __import__('seaborn'),
'stats': __import__('scipy.stats')
}
with redirect_stdout(output):
exec(code, local_namespace)
# If 'result' is defined in code, capture it
result = local_namespace.get('result', None)
# Get the captured output
output_string = output.getvalue()
# Check for generated plots
plot_files = sorted(self.output_dir.glob("output_*.png"))
for i, plot_path in enumerate(plot_files):
try:
# Read the plot file and convert to base64
with open(plot_path, 'rb') as img_file:
plot_data = base64.b64encode(img_file.read()).decode('utf-8')
plot_info = {
'plot_number': i + 1,
'file_name': plot_path.name,
'file_path': str(plot_path),
'base64_data': plot_data
}
plots_info.append(plot_info)
except Exception as e:
print(f"Error processing plot {plot_path}: {str(e)}")
plots_info.append({
'plot_number': i + 1,
'file_name': plot_path.name,
'error': str(e)
})
return {
'printed_output': output_string,
'result_object': result,
# ...
}
缺点：
- 安全性：任意代码可以读取或写入文件、网络资源或使用系统调用。
- 缺乏资源控制：如果代码触发内存泄漏或无限循环，您只能终止整个 Python 进程或依赖复杂的超时机制。
- 依赖冲突：此代码可能导入与主环境冲突的库。
- 没有真正的隔离：如果恶意或有缺陷的代码在同一进程中运行，它可能会崩溃或干扰整个系统。
对于博客文章或演示，您可以包含一个简化版本，但在生产环境中，您确实希望使用基于 Docker 或类似的沙箱来安全、稳健地执行代码。
FastAPI + Jupyter 沙盒代码中的潜在问题
即使使用 Docker，仍然有一些细微之处需要注意：
内核稳定性
有时 Jupyter 内核可能会挂起。我们的代码中包含一个 _wait_for_kernel_ready() 方法来轮询内核是否准备就绪。如果失败，您可能需要更强大的检查或强制重启。我使用 time.sleep() 来等待，然后再发送代码进行执行。
会话清理
如果您不主动结束会话，过期的会话可能会累积。我们使用后台任务 (cleanup_inactive_sessions()) 定期删除超过 1 小时的会话。根据您的需要进行调整。
包安装
/install_package 端点可能会导致“依赖漂移”，如果多个用户或进程不断安装随机库。考虑在生产环境中限制或移除此端点。
文件传输
您需要将代码将要分析的文件传输到已挂载在 Docker 容器上的目录。如果后端服务器和您的沙盒环境不在同一位置，网络速度可能会导致延迟。
资源限制
仅使用 Docker 并不会自动限制内存或 CPU 使用率。如果您需要严格的资源隔离，请确保通过 Docker 或您的 Kubernetes/OpenShift 平台配置容器资源限制。
安全隐患
根据您的威胁模型，以根权限运行 Docker 容器并不足以保护您免受所有容器突破漏洞的影响。为了更好的安全性，考虑使用无根 Docker 或更严格的运行时配置。
尽管有这些警告，在容器内运行代码和短暂的 Jupyter 内核比在主机环境中调用 exec 安全得多（且更易于管理）。
架构的视觉概述
以下是请求流动的简化图示：
- 客户端发送请求以启动或与会话进行交互。
- FastAPI 将工作委托给
JupyterController。 JupyterController启动或重置底层的 Jupyter 内核。- 内核执行代码并返回输出。
- 这些输出通过 FastAPI 响应返回给用户。
结论
一个 沙箱代码执行环境 对于 AI 代理工作流程是 必不可少的。它提供了安全性、资源控制、环境一致性以及动态代码执行的稳定接口。虽然使用 Python 的 exec 进行快速代码执行看起来很诱人，但这种方法会使系统面临严重风险，并使调试和依赖管理变得更加困难。
通过 Docker + FastAPI + Jupyter，您可以实现：
- 进程级隔离 通过容器。
- 程序化控制 通过 Jupyter 内核。
- 灵活的代码执行，具有良好的错误处理、超时和会话管理。
欢迎查看包含完整源代码的 GitHub 仓库。
感谢您的阅读！ 如果您有问题或想分享改进意见，请随时在 GitHub 上留言或提出问题。
Github Repo : https://github.com/anukriti-ranjan/sandboxed-jupyter-code-exec
