使用Rust+pyo3调用tensorboardX.SummaryWriter时无限生成进程的问题

原因分析 #

tensorboardX的SummaryWriter依赖tensorflow，而tensorflow在嵌入Python环境（比如通过PyO3从Rust调用）下的初始化逻辑，会尝试重新执行当前进程（你的Rust二进制程序）来完成环境配置。这就导致每次调用SummaryWriter构造函数都会触发程序重启，进而形成无限进程循环。

核心诱因是：tensorflow未检测到合法的PYTHONEXECUTABLE环境变量，会默认使用当前进程路径（Rust程序）重新启动，而非调用Python解释器。

修复方案 #

通过以下两步阻止进程重启：

1. 显式设置PYTHONEXECUTABLE环境变量，指向你的虚拟环境Python解释器路径
2. 修改Python的sys.argv，模拟正常Python脚本的执行参数，避免库的环境误判

修改后的可运行代码 #

use pyo3::prelude::*;

fn main() {
    println!("Start");
    pyo3::prepare_freethreaded_python();
    let x: PyResult<PyObject> = Python::with_gil(|py| {
        // 1. 修改sys.argv，模拟正常Python脚本执行
        let sys = PyModule::import(py, "sys")?;
        sys.setattr("argv", vec!["dummy_script.py"].into_py(py))?;

        // 2. 设置关键环境变量，指定Python解释器路径
        let os = PyModule::import(py, "os")?;
        // 替换为你的虚拟环境Python路径
        os.call_method1("environ.setdefault", ("PYTHONEXECUTABLE", "/path/to/venv/bin/python"))?;

        // 初始化tensorboardX SummaryWriter
        let tbx = PyModule::import(py, "tensorboardX")?;
        // 可指定log_dir避免临时目录问题，可选
        let x = tbx.getattr("SummaryWriter")?.call1(("log_dir", "./tb_logs"))?;
        Ok(x.to_object(py))
    });
    let x = x.unwrap();
    println!("End ({:?})", x);
}

额外注意事项 #

• 确保.cargo/config.toml中的PYTHONPATH正确指向虚拟环境的site-packages目录
• 如果仍有问题，可尝试添加tensorflow的日志屏蔽配置，减少干扰：

  os.call_method1("environ.setdefault", ("TF_CPP_MIN_LOG_LEVEL", "3"))?;
  

内容的提问来源于stack exchange，提问作者Zanapher