Apache Geaflow推理框架Geaflow-infer 解析系列（五）环境上下文管理

第5章：环境上下文管理

章节导读

在虚拟环境创建和依赖安装完成后，我们需要一个数据容器 来持有所有的环境信息，并在不同的组件之间传递。这就是 InferEnvironmentContext 的使命。

本章将讲解：

• 如何组织和管理环境信息（路径、配置、标识）
• 如何生成进程标识（主机名+进程ID）
• 如何传递参数给 Python 进程（环境变量）

核心设计

单一责任:
  ✓ InferEnvironmentContext 只负责信息承载
  ✓ 不涉及文件 I/O、进程管理
  ✓ 不涉及数据交换、序列化

易于扩展:
  ✓ 需要新增参数？只需在 Context 中添加字段
  ✓ 需要生成新格式参数？只需添加新方法
  ✓ Python 端需要新参数？设置新环境变量即可

多实例隔离:
  ✓ roleNameIndex 确保每个实例独立
  ✓ 共享内存队列、虚拟环境目录都独立
  ✓ 同一机器上可运行任意数量的实例

---

5.1 InferEnvironmentContext 设计

设计目标

InferEnvironmentContext 的角色:
  ├─ 承载虚拟环境的所有信息（路径、配置、状态）
  ├─ 提供便捷方法生成命令参数
  ├─ 在 Java 各组件间传递环境信息
  └─ 与 Python 进程共享关键参数

核心数据结构

public class InferEnvironmentContext {
    
    // 1. 虚拟环境相关路径
    private final String virtualEnvDirectory;    // /tmp/geaflow_infer_env
    private final String inferFilesDirectory;    // /tmp/geaflow_infer_files
    private final String inferLibPath;           // /tmp/geaflow_infer_env/conda/lib
    
    // 2. 可执行程序路径
    private String pythonExec;                  // /tmp/.../conda/bin/python3
    private String inferScript;                 // /tmp/.../infer_server.py
    
    // 3. 进程标识（主机名:进程ID）
    private final String roleNameIndex;         // "hostname:12345"
    
    // 4. 配置
    private final Configuration configuration;  // GeaFlow 配置
    
    // 5. 初始化状态标记
    private Boolean envFinished;                // 虚拟环境是否初始化完成
    
    // 6. 常量定义（路径相对位置）
    private static final String LIB_PATH = "/conda/lib";
    private static final String INFER_SCRIPT_FILE = "/infer_server.py";
    private static final String PYTHON_EXEC = "/conda/bin/python3";
}

初始化流程

public InferEnvironmentContext(
    String virtualEnvDirectory,
    String pythonFilesDirectory,
    Configuration configuration) {
    
    // 1. 保存虚拟环境目录
    this.virtualEnvDirectory = virtualEnvDirectory;
    
    // 2. 保存 Python 文件目录
    this.inferFilesDirectory = pythonFilesDirectory;
    
    // 3. 计算库路径
    this.inferLibPath = virtualEnvDirectory + LIB_PATH;
    // 结果: /tmp/geaflow_infer_env/conda/lib
    
    // 4. 计算 Python 可执行文件路径
    this.pythonExec = virtualEnvDirectory + PYTHON_EXEC;
    // 结果: /tmp/geaflow_infer_env/conda/bin/python3
    
    // 5. 计算推理脚本路径
    this.inferScript = pythonFilesDirectory + INFER_SCRIPT_FILE;
    // 结果: /tmp/geaflow_infer_files/infer_server.py
    
    // 6. 保存配置
    this.configuration = configuration;
    
    // 7. 生成唯一的角色标识（最关键的一步）
    this.roleNameIndex = queryRoleNameIndex();
    // 结果: "my-host:12345" (hostname:pid)
    
    // 8. 初始化状态标记
    this.envFinished = false;  // 初始未完成
}

---

5.2 进程标识与角色管理

为什么需要进程标识？

场景: 同一台机器上运行多个 GeaFlow 应用实例

问题:
  实例A                    实例B
    PID: 1001                PID: 1002
    创建共享内存队列         创建共享内存队列
    "input_queue"            "input_queue"  ✗ 冲突！

解决方案: 为每个实例生成唯一标识
  实例A: roleNameIndex = "hostname:1001"
  实例B: roleNameIndex = "hostname:1002"
  
  共享内存队列:
    "input_queue_hostname_1001"
    "input_queue_hostname_1002"
  ✓ 不再冲突

进程标识生成

private String queryRoleNameIndex() {
    try {
        // Step 1: 获取主机名
        InetAddress address = InetAddress.getLocalHost();
        String hostName = address.getHostName();
        // 结果: "my-MacBook-Pro.local" 或 "server-01"
        
        // Step 2: 获取进程ID
        RuntimeMXBean runtime = ManagementFactory.getRuntimeMXBean();
        String name = runtime.getName();
        // 结果: "12345@hostname" (pid@hostname)
        
        // Step 3: 解析出 PID
        int processId = Integer.parseInt(
            name.substring(0, name.indexOf("@"))  // "@" 前的部分
        );
        // 结果: 12345
        
        // Step 4: 组合生成标识
        return hostName + HOST_SEPARATOR + processId;
        // 结果: "hostname:12345"
        
    } catch (Exception e) {
        throw new GeaflowRuntimeException(
            "生成角色标识失败", e);
    }
}

进程标识的应用

生成的 roleNameIndex = "hostname:12345"

应用 1: 共享内存队列命名
  input_queue_id = "input_hostname_12345"
  output_queue_id = "output_hostname_12345"

应用 2: 日志标记
  [hostname:12345] 开始初始化虚拟环境
  [hostname:12345] Python 进程已启动
  [hostname:12345] 推理完成

应用 3: 调试和监控
  ps aux | grep java
  # 看到 PID 12345 对应的就是这个 GeaFlow 实例

应用 4: 环境变量传递给 Python
  export GEAFLOW_ROLE_NAME_INDEX="hostname:12345"
  python infer_server.py

多个进程间的隔离

┌─────────────────────────────────────────────────────────┐
│              Linux 进程表                                │
├─────────────────────────────────────────────────────────┤
│ PID   │ 角色标识           │ 共享内存队列               │
├───────┼──────────────────┼────────────────────────────┤
│ 1001  │ hostname:1001    │ input_hostname_1001       │
│       │ (Java 进程 A)    │ output_hostname_1001      │
├───────┼──────────────────┼────────────────────────────┤
│ 1002  │ hostname:1002    │ input_hostname_1002       │
│       │ (Java 进程 B)    │ output_hostname_1002      │
├───────┼──────────────────┼────────────────────────────┤
│ 1003  │ hostname:1003    │ input_hostname_1003       │
│       │ (Java 进程 C)    │ output_hostname_1003      │
└─────────────────────────────────────────────────────────┘

benefits:
  ✓ 每个进程有独立的队列
  ✓ Python 推理进程知道自己的身份
  ✓ 日志和监控可以追踪到具体进程

---

5.3 参数传递机制

参数传递的两种方式

方式 1: 环境变量（推荐）

// Java 端：设置环境变量
ProcessBuilder pb = new ProcessBuilder(pythonExec, inferScript);
Map<String, String> env = pb.environment();

// 1. 虚拟环境路径
env.put("VIRTUAL_ENV", virtualEnvDirectory);

// 2. 库路径（动态库加载）
env.put("LD_LIBRARY_PATH", inferLibPath);

// 3. Python 模块查找路径
env.put("PYTHONPATH", inferFilesDirectory);

// 4. 共享内存队列 ID
env.put("input_queue_shm_id", "input_hostname_12345");
env.put("output_queue_shm_id", "output_hostname_12345");

// 5. 用户自定义类名
env.put("tfClassName", "my.custom.Transform");

// 6. 角色标识
env.put("GEAFLOW_ROLE_NAME_INDEX", "hostname:12345");

// 启动进程
Process process = pb.start();

方式 2: 命令行参数（备选）

// 虽然不推荐，但也可以用命令行参数

ProcessBuilder pb = new ProcessBuilder(
    pythonExec,
    inferScript,
    "--input_queue_shm_id=input_hostname_12345",
    "--output_queue_shm_id=output_hostname_12345",
    "--tfClassName=my.custom.Transform",
    "--virtual_env=" + virtualEnvDirectory
);

// Python 端需要解析命令行参数
import sys
args = sys.argv[1:]
# 解析 --key=value 格式

Python 端的接收

# infer_server.py

import os
import sys

# Step 1: 读取环境变量
virtual_env = os.environ.get('VIRTUAL_ENV')
input_queue_id = os.environ.get('input_queue_shm_id')
output_queue_id = os.environ.get('output_queue_shm_id')
tf_class_name = os.environ.get('tfClassName')
role_index = os.environ.get('GEAFLOW_ROLE_NAME_INDEX')

print(f"[{role_index}] 推理服务启动")
print(f"虚拟环境: {virtual_env}")
print(f"输入队列: {input_queue_id}")
print(f"输出队列: {output_queue_id}")

# Step 2: 连接共享内存队列
input_queue = DataExchangeQueue(input_queue_id)
output_queue = DataExchangeQueue(output_queue_id)

# Step 3: 加载用户类
from importlib import import_module
module_path, class_name = tf_class_name.rsplit('.', 1)
module = import_module(module_path)
TransformClass = getattr(module, class_name)

# Step 4: 创建转换实例
transformer = TransformClass()

# Step 5: 监听队列，接收推理请求
while True:
    # 从输入队列读取数据
    input_bytes = input_queue.get(timeout=1000)
    
    if input_bytes is None:
        continue
    
    # 反序列化
    input_obj = unpickle(input_bytes)
    
    # 执行推理
    result = transformer.transform(input_obj)
    
    # 序列化结果
    output_bytes = pickle(result)
    
    # 写入输出队列
    output_queue.put(output_bytes)

参数的作用列表

参数名	值示例	作用	由谁设置
VIRTUAL_ENV	/tmp/env	Python 虚拟环境位置	InferEnvironmentContext
LD_LIBRARY_PATH	/tmp/env/lib	动态库查找路径	InferEnvironmentContext
PYTHONPATH	/tmp/files	Python 模块查找路径	InferEnvironmentContext
input_queue_shm_id	input_host_1001	输入队列标识	DataExchangeContext
output_queue_shm_id	output_host_1001	输出队列标识	DataExchangeContext
tfClassName	my.Transform	用户类全名	用户指定
GEAFLOW_ROLE_NAME_INDEX	hostname:1001	进程角色标识	InferEnvironmentContext

---

5.4 InferEnvironmentContext 的便捷方法

参数生成方法

public class InferEnvironmentContext {
    
    /**
     * 生成 TensorFlow 类名参数
     */
    public String getInferTFClassNameParam(String udfClassName) {
        return TF_CLASSNAME_KEY + udfClassName;
        // 结果: "--tfClassName=my.custom.Transform"
    }
    
    /**
     * 生成输入队列参数
     */
    public String getInferShareMemoryInputParam(String queueId) {
        return SHARE_MEMORY_INPUT_KEY + queueId;
        // 结果: "--input_queue_shm_id=input_hostname_1001"
    }
    
    /**
     * 生成输出队列参数
     */
    public String getInferShareMemoryOutputParam(String queueId) {
        return SHARE_MEMORY_OUTPUT_KEY + queueId;
        // 结果: "--output_queue_shm_id=output_hostname_1001"
    }
    
    /**
     * 获取 Python 可执行文件路径
     */
    public String getPythonExec() {
        return pythonExec;
        // 结果: /tmp/geaflow_infer_env/conda/bin/python3
    }
    
    /**
     * 获取推理脚本路径
     */
    public String getInferScript() {
        return inferScript;
        // 结果: /tmp/geaflow_infer_files/infer_server.py
    }
    
    /**
     * 获取角色标识
     */
    public String getRoleNameIndex() {
        return roleNameIndex;
        // 结果: "hostname:12345"
    }
}

使用示例

// 在 InferTaskRunImpl 中的使用
public class InferTaskRunImpl implements InferTaskRun {
    
    public void run(List<String> script) {
        // 获取环境上下文
        InferEnvironmentContext ctx = 
            inferEnvironmentContext;
        
        // 构建完整的命令
        List<String> command = new ArrayList<>();
        command.add(ctx.getPythonExec());
        command.add(ctx.getInferScript());
        
        // 添加参数（使用便捷方法）
        command.add(ctx.getInferTFClassNameParam(userClassName));
        command.add(ctx.getInferShareMemoryInputParam(inputQueueId));
        command.add(ctx.getInferShareMemoryOutputParam(outputQueueId));
        
        LOGGER.info("[{}] 执行命令: {}", 
            ctx.getRoleNameIndex(),  // 日志标记
            String.join(" ", command));
        
        // 启动进程
        ProcessBuilder pb = new ProcessBuilder(command);
        Map<String, String> env = pb.environment();
        env.put("LD_LIBRARY_PATH", inferLibPath);
        env.put("PYTHONPATH", ctx.inferFilesDirectory);
        
        Process process = pb.start();
    }
}

---

参考资源

• Python os.environ 文档