使用 Rust 构建支持 Python 脚本执行的 Agent 框架虚拟运行环境

使用 Rust 构建支持 Python 脚本执行的 Agent 框架虚拟运行环境 | 极客日志

前言

为了增加 Agent 框架的扩展性和适用性，我们需要能够在流程节点中安全、高效地运行 Python 脚本。Python 作为 AI 领域的事实标准语言，拥有庞大的生态库，是构建智能体不可或缺的工具。

技术选型思考

1. 为什么选择 Python？

尽管从性能角度看，Lua 等脚本语言更为轻快，但在实际工程落地中，Python 占据了绝对主导地位。其丰富的机器学习、数据处理库（如 NumPy, Pandas, PyTorch）是其他语言难以替代的。虽然未来可能会考虑支持 JavaScript 或 WebAssembly，但现阶段 Python 是必须支持的核心能力。

2. 沙盒环境的必要性

直接在业务 Pod 内部执行 Python 脚本存在巨大风险：

环境一致性：不同服务器本地环境差异会导致脚本运行失败。
安全性：防止恶意脚本访问宿主机文件系统、网络资源或消耗过多系统资源。
隔离性：避免脚本崩溃影响主服务进程。

因此，必须构建一个独立的沙盒运行环境，通过容器化技术实现资源隔离和权限控制。

架构设计

整体架构采用 Rust 作为核心运行时，通过 PyO3 绑定 Python 解释器，对外提供 gRPC 接口。客户端发送请求，Rust 服务解析参数，在沙盒内调用 Python 函数，并将结果序列化返回。

[Client] --(gRPC)--> [Rust Runtime Service] --(PyO3)--> [Python VM]
                              |
                      (Docker Container)

核心实现

1. PyO3 绑定与运行时封装

使用 pyo3 库完成 Rust 与 Python 的交互。Python 入口必须是一个可导出的函数，支持指定 sys.path 以加载外部文件。

关键代码逻辑如下：

pub fn eval_function(&self, function_name: &str, args: Value) -> PyResult<Value> {
    wd_log::log_debug_ln!("eval_function -> {:?} args:{:?}", self, args);
    let Self { src, module_name, file_name, sys_path } = self;
    
    let file_name = file_name.clone().unwrap_or(format!("{}.py", module_name));
    
    Python::with_gil(move |py| {
        // 设置系统路径，支持挂载卷加载文件
        if let Some(path) = sys_path {
            let syspath: &PyList = py.import_bound("sys")?.getattr("path")?.extract()?;
            syspath.insert(0, &path)?;
        }

        // 根据源码类型加载模块
        let module = match src {
            ScriptSrc::ScriptCode(script) => PyModule::from_code_bound(
                py,
                script.as_str(),
                file_name.as_str(),
                module_name.as_str(),
            )?,
            ScriptSrc::ModuleName => PyModule::import_bound(py, module_name.as_str())?,
        };

        // 获取目标函数
        let function = module.getattr(function_name)?;
        
        // 转换输入参数为 Python 对象
        let input_obj = Self::value_to_py_object(py, args)?;
        let input_class = FunctionInput { data: input_obj };
        
        // 调用函数并获取输出
        let output_obj = function.call((input_class,), None)?.extract::<PyObject>()?;
        
        // 将 Python 对象转回通用 Value 格式
        let value = Self::py_object_to_value(output_obj, py)?;
        Ok(value)
    })
}

2. 入参绑定设计

为了增强 Python 脚本对 Rust 侧能力的访问，我们定义了包装类 FunctionInput。这允许在 Python 脚本中通过类实例访问数据，同时预留了扩展 Rust 原生方法的空间。

#[pyclass]
pub struct FunctionInput {
    #[pyo3(get, set)]
    data: PyObject,
}

3. gRPC 接口定义

对外暴露标准的 gRPC 接口，使用 Protocol Buffers 定义数据结构。入参类型采用 google.protobuf.Struct，最终转化为 Python 的字典对象。

syntax = "proto3";
package proto;

import "google/api/annotations.proto";
import "google/protobuf/struct.proto";

service PythonRuntimeService {
    rpc CallFunction(CallFunctionRequest) returns (CallFunctionResponse) {
        option (google.api.http) = {
            post: "/api/v1/function/{function_name}"
            body: "*"
        };
    }
}

enum SrcType {
    SRC_TYPE_SCRIPT = 0;
    SRC_TYPE_MODULE = 1;
}

message CallFunctionRequest {
    SrcType src = 1;
    optional string script_code = 2; // 当 src 为 SCRIPT 时必填
    string module_name = 3;
    optional string file_name = 4;   // 默认为 module_name.py
    optional string sys_path = 5;    // 自定义搜索路径
    string function_name = 6;
    google.protobuf.Struct function_input = 7;
}

message CallFunctionResponse {
    int32 code = 1;      // 0 表示成功
    string msg = 2;
    optional google.protobuf.Struct output = 3;
}

安全与沙盒策略

在生产环境中，安全是重中之重。虽然示例代码中简化了部分安全逻辑，但实际部署必须包含以下措施：

资源限制：通过 Docker 配置 CPU 和内存上限（Cgroups），防止脚本死循环耗尽资源。
网络隔离：禁止沙盒内的 Python 脚本访问外网，仅允许访问必要的内部服务。
文件系统权限：挂载只读卷，限制脚本只能读取特定目录，禁止写入敏感路径。
模块白名单：在 Python 启动时禁用危险模块（如 os, subprocess, socket），仅保留业务所需功能。
超时控制：设置脚本执行的最大时间，超过时限强制终止进程。

Docker 镜像构建

由于 PyO3 依赖 libpython 和 Rust 编译环境，直接构建镜像较为复杂。建议使用多阶段构建（Multi-stage Build）来优化镜像体积。

构建阶段问题与解决

Alpine 镜像问题：Rust 官方 Alpine 镜像基于 musl libc，可能导致动态链接库缺失。建议优先使用 Debian 或 Ubuntu 基础镜像以确保兼容性。
跨平台编译：PyO3 编译时绑定了特定 Python 版本信息，建议在同一架构下完成编译和运行，避免跨平台兼容性问题。

测试与验证

单元测试

使用 Cargo Test 进行后端逻辑验证，确保参数解析和函数调用正确无误。

cargo test tests::run_server -- --nocapture

集成测试

通过模拟 gRPC 请求验证完整链路。以下是文件执行模式的测试用例：

请求示例：

{
  "src": 1,
  "module_name": "sys_info",
  "sys_path": "./custom_plugin",
  "function_name": "get_system_info",
  "function_input": {
    "fields": {
      "hello": { "kind": "world" }
    }
  }
}

响应示例：

{
  "code": 0,
  "msg": "success",
  "output": {
    "fields": {
      "cpu_count": { "numberValue": 12, "kind": "numberValue" },
      "memory_total": { "numberValue": 32768, "kind": "numberValue" },
      "os_type": { "stringValue": "Linux", "kind": "stringValue" }
    }
  }
}

文本脚本测试

支持直接传入 Python 代码字符串，无需预编译模块。这对于快速原型开发和临时任务非常有用。在 Agent 编排界面中，用户可直接编辑 Python 代码块并即时调试。

总结与展望

本文详细介绍了如何使用 Rust 构建一个支持 Python 脚本执行的 Agent 框架虚拟运行环境。通过 PyO3 实现语言绑定，结合 gRPC 提供标准化接口，并利用 Docker 实现环境隔离。

主要优势：

高性能：Rust 保证了运行时的高效与稳定。
安全性：沙盒机制有效隔离了潜在风险。
易用性：保留了 Python 丰富的生态库，降低了开发门槛。

后续计划：

完善安全模块，实现更细粒度的权限控制。
探索 WASM 支持，进一步降低运行时开销。
增加更多编程语言插件，提升框架通用性。

通过合理的权衡与设计，该方案能够有效支撑复杂的 Agent 业务流程，为智能化应用提供坚实的底层基础设施。

使用 Rust 构建支持 Python 脚本执行的 Agent 框架虚拟运行环境

前言

技术选型思考

1. 为什么选择 Python？

2. 沙盒环境的必要性

架构设计

核心实现

1. PyO3 绑定与运行时封装

2. 入参绑定设计

3. gRPC 接口定义

安全与沙盒策略

Docker 镜像构建

构建阶段问题与解决

推荐 Dockerfile 结构

测试与验证

单元测试

集成测试

文本脚本测试

总结与展望

更多推荐文章

相关免费在线工具

使用 Rust 构建支持 Python 脚本执行的 Agent 框架虚拟运行环境

前言

技术选型思考

1. 为什么选择 Python？

2. 沙盒环境的必要性

架构设计

核心实现

1. PyO3 绑定与运行时封装

2. 入参绑定设计

3. gRPC 接口定义

安全与沙盒策略

Docker 镜像构建

构建阶段问题与解决

推荐 Dockerfile 结构

测试与验证

单元测试

集成测试

文本脚本测试

总结与展望

微信扫一扫，关注极客日志

更多推荐文章

相关免费在线工具