gradio-web/process_session.py

"""
多进程会话管理模块

本模块实现了多进程架构，将每个用户的 OracleSession 运行在独立的工作进程中。
这样可以确保重计算任务不会阻塞主进程，多个用户可以并发使用系统。

架构说明：
1. ProcessSessionProxy: 主进程中的代理类，提供与 OracleSession 相同的接口
2. session_worker_loop: 工作进程中的循环函数，运行实际的 OracleSession
3. 进程间通信：通过 multiprocessing.Queue 进行命令和结果的传递
4. 视频帧同步：工作进程产生的新帧通过 stream_queue 推送到主进程，由后台线程同步到代理的本地缓存
"""
import logging
import multiprocessing
import queue
import threading
import time
import sys
import os

# 添加父目录到路径（逻辑复制自 oracle_logic.py）
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
src_dir = os.path.join(parent_dir, "src")
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)
if src_dir not in sys.path:
    sys.path.insert(0, src_dir)

from oracle_logic import OracleSession, DEFAULT_DATASET_ROOT

# Import ScrewPlanFailure for exception handling
try:
    from robomme.robomme_env.utils.planner_fail_safe import ScrewPlanFailure
except ImportError:
    # Fallback if import fails
    ScrewPlanFailure = RuntimeError

# Custom exception for screw plan failures (to be caught in gradio_callbacks)
class ScrewPlanFailureError(RuntimeError):
    """Exception raised when screw plan fails, to be caught and displayed via gr.Info popup"""
    pass

# 定义命令常量
CMD_LOAD_EPISODE = "load_episode"
CMD_UPDATE_OBSERVATION = "update_observation"
CMD_GET_PIL_IMAGE = "get_pil_image"
CMD_EXECUTE_ACTION = "execute_action"
CMD_GET_REFERENCE_ACTION = "get_reference_action"
CMD_CLOSE = "close"
LOGGER = logging.getLogger("robomme.process_session")


def _setup_worker_logging():
    level_name = os.getenv("LOG_LEVEL", "DEBUG").upper()
    level = getattr(logging, level_name, logging.DEBUG)
    logging.basicConfig(
        level=level,
        format=(
            "%(asctime)s | %(levelname)s | %(name)s | "
            "pid=%(process)d tid=%(threadName)s | %(message)s"
        ),
        stream=sys.stdout,
        force=True,
    )
    for noisy_logger in [
        "asyncio",
        "httpx",
        "httpcore",
        "urllib3",
        "matplotlib",
        "PIL",
        "h5py",
        "trimesh",
        "toppra",
    ]:
        logging.getLogger(noisy_logger).setLevel(logging.WARNING)
    logging.getLogger("robomme").setLevel(logging.DEBUG)
    LOGGER.debug("worker logging initialized level=%s", level_name)

def _sanitize_options(options):
    """
    清理选项数据，移除不可序列化的项（如 'solve' 函数）
    
    在跨进程通信时，需要确保所有数据都可以被 pickle 序列化。
    raw_solve_options 中包含的 'solve' 函数无法序列化，需要移除。
    'available' 字段可能是复杂对象，需要转换为简单的布尔值。
    
    Args:
        options: 原始选项列表
        
    Returns:
        list: 清理后的选项列表
    """
    clean_opts = []
    if not options:
        return clean_opts
    for opt in options:
        clean_opt = opt.copy()
        if "solve" in clean_opt:
            del clean_opt["solve"]
        if "available" in clean_opt:
            # Only keep truthiness for UI logic
            clean_opt["available"] = bool(clean_opt["available"])
        clean_opts.append(clean_opt)
    return clean_opts

def session_worker_loop(cmd_queue, result_queue, stream_queue, dataset_root, gui_render):
    """
    工作进程主循环
    
    此函数在工作进程中运行，负责：
    1. 初始化 OracleSession 实例
    2. 监听来自主进程的命令（通过 cmd_queue）
    3. 执行命令并返回结果（通过 result_queue）
    4. 监控视频帧变化，将新帧推送到流队列（通过 stream_queue）
    5. 处理异常和清理资源
    
    Args:
        cmd_queue: 命令队列，主进程发送命令到此队列
        result_queue: 结果队列，工作进程返回命令执行结果到此队列
        stream_queue: 流队列，工作进程推送新视频帧到此队列
        dataset_root: 数据集根目录路径
        gui_render: 是否使用GUI渲染模式
    """
    _setup_worker_logging()
    session = None
    try:
        LOGGER.info(
            "worker loop starting dataset_root=%s gui_render=%s",
            dataset_root,
            gui_render,
        )
        session = OracleSession(dataset_root=dataset_root, gui_render=gui_render)
        session.stream_frame_callback = lambda frames: stream_queue.put({"base": frames, "wrist": []})
        LOGGER.info("worker OracleSession initialized")
        
        while True:
            try:
                # Check for commands
                cmd_data = cmd_queue.get(timeout=0.1)
            except queue.Empty:
                continue
                
            cmd = cmd_data["cmd"]
            args = cmd_data.get("args", [])
            kwargs = cmd_data.get("kwargs", {})
            LOGGER.debug("worker received cmd=%s args=%s kwargs_keys=%s", cmd, len(args), list(kwargs.keys()))
            
            if cmd == CMD_CLOSE:
                if session:
                    session.close()
                LOGGER.info("worker received close command, exiting")
                break
            
            elif cmd == CMD_LOAD_EPISODE:
                # 加载环境episode
                res = session.load_episode(*args, **kwargs)
                LOGGER.info(
                    "worker load_episode env=%s episode=%s result_msg=%s",
                    getattr(session, "env_id", None),
                    getattr(session, "episode_idx", None),
                    res[1] if isinstance(res, tuple) and len(res) > 1 else None,
                )
                
                # 更新帧索引跟踪（用于增量同步）
                session.last_base_frame_idx = len(session.base_frames)
                session.last_wrist_frame_idx = len(session.wrist_frames)
                
                # 获取演示状态（从 DemonstrationWrapper 获取）
                is_demonstration = False
                if session.env:
                    is_demonstration = getattr(session.env, 'current_task_demonstration', False)
                
                # 构建状态更新（完整同步，因为这是加载操作）
                state_update = {
                    "env_id": session.env_id,
                    "episode_idx": session.episode_idx,
                    "language_goal": session.language_goal,
                    "difficulty": session.difficulty,
                    "seed": session.seed,
                    "demonstration_frames": session.demonstration_frames,
                    "last_execution_frames": [],
                    "base_frames": session.base_frames,  # 加载时完整同步
                    "wrist_frames": session.wrist_frames,  # 加载时完整同步
                    "available_options": session.available_options,
                    "raw_solve_options": _sanitize_options(session.raw_solve_options),
                    "seg_vis": session.seg_vis,
                    "is_demonstration": is_demonstration,
                    "non_demonstration_task_length": session.non_demonstration_task_length  # 同步非demonstration任务长度
                }
                result_queue.put({"status": "success", "result": res, "state": state_update})
                
            elif cmd == CMD_EXECUTE_ACTION:
                # 执行动作（重计算任务）
                execute_base_start = len(session.base_frames)
                try:
                    res = session.execute_action(*args, **kwargs)
                    LOGGER.info(
                        "worker execute_action done env=%s episode=%s done=%s",
                        getattr(session, "env_id", None),
                        getattr(session, "episode_idx", None),
                        res[2] if isinstance(res, tuple) and len(res) > 2 else None,
                    )
                except ScrewPlanFailure as e:
                    # 捕获 ScrewPlanFailure 并作为特殊状态传递到主进程，用于显示弹窗
                    LOGGER.warning("worker screw_plan_failure: %s", e)
                    result_queue.put({"status": "screw_plan_failure", "message": str(e)})
                    continue
                except Exception as e:
                    # 捕获所有其他异常并传递到主进程，用于显示弹窗
                    LOGGER.exception("worker execution_error")
                    result_queue.put({"status": "execution_error", "message": str(e)})
                    continue
                
                # 增量帧同步：只发送新增的帧
                new_base = session.base_frames[session.last_base_frame_idx:]
                new_wrist = session.wrist_frames[session.last_wrist_frame_idx:]
                streamed_count = int(getattr(session, "_execute_streamed_frame_count", 0) or 0)
                # Frames already pushed by stream_frame_callback during solve() should not be sent twice.
                if streamed_count > 0 and new_base:
                    if streamed_count >= len(new_base):
                        new_base = []
                    else:
                        new_base = new_base[streamed_count:]
                
                # 更新帧索引
                session.last_base_frame_idx = len(session.base_frames)
                session.last_wrist_frame_idx = len(session.wrist_frames)
                
                # 如果有新帧，推送到流队列
                if new_base or new_wrist:
                    stream_queue.put({"base": new_base, "wrist": new_wrist})
                LOGGER.debug(
                    "worker execute_action streamed frames base=%s wrist=%s",
                    len(new_base),
                    len(new_wrist),
                )

                # 获取演示状态（从 DemonstrationWrapper 获取）
                is_demonstration = False
                if session.env:
                    is_demonstration = getattr(session.env, 'current_task_demonstration', False)

                execution_frames = session.base_frames[execute_base_start:]

                # 构建状态更新（只更新选项和分割视图，帧通过流队列同步）
                state_update = {
                    "last_execution_frames": execution_frames,
                    "available_options": session.available_options,
                    "raw_solve_options": _sanitize_options(session.raw_solve_options),
                    "seg_vis": session.seg_vis,
                    "is_demonstration": is_demonstration
                }
                result_queue.put({"status": "success", "result": res, "state": state_update})

            elif cmd == CMD_GET_PIL_IMAGE:
                res = session.get_pil_image(*args, **kwargs)
                result_queue.put({"status": "success", "result": res})
                
            elif cmd == CMD_UPDATE_OBSERVATION:
                # 更新观察（获取当前环境状态）
                res = session.update_observation(*args, **kwargs)
                
                # 增量帧同步
                new_base = session.base_frames[session.last_base_frame_idx:]
                new_wrist = session.wrist_frames[session.last_wrist_frame_idx:]
                
                # 更新帧索引
                session.last_base_frame_idx = len(session.base_frames)
                session.last_wrist_frame_idx = len(session.wrist_frames)

                # 如果有新帧，推送到流队列
                if new_base or new_wrist:
                    stream_queue.put({"base": new_base, "wrist": new_wrist})
                LOGGER.debug(
                    "worker update_observation streamed frames base=%s wrist=%s",
                    len(new_base),
                    len(new_wrist),
                )
                
                # 获取演示状态（从 DemonstrationWrapper 获取）
                is_demonstration = False
                if session.env:
                    is_demonstration = getattr(session.env, 'current_task_demonstration', False)
                    
                # 构建状态更新
                state_update = {
                    "available_options": session.available_options,
                    "raw_solve_options": _sanitize_options(session.raw_solve_options),
                    "seg_vis": session.seg_vis,
                    "is_demonstration": is_demonstration
                }
                result_queue.put({"status": "success", "result": res, "state": state_update})

            elif cmd == CMD_GET_REFERENCE_ACTION:
                res = session.get_reference_action(*args, **kwargs)
                LOGGER.debug("worker get_reference_action ok=%s", bool(res.get("ok")) if isinstance(res, dict) else None)
                result_queue.put({"status": "success", "result": res})
                
            else:
                LOGGER.error("worker unknown command=%s", cmd)
                result_queue.put({"status": "error", "message": f"Unknown command: {cmd}"})
                
    except Exception as e:
        LOGGER.exception("worker fatal error")
        result_queue.put({"status": "fatal", "message": str(e)})


class ProcessSessionProxy:
    """
    进程会话代理类
    
    此类在主进程中运行，提供与 OracleSession 相同的接口。
    所有方法调用都会被转发到工作进程中的实际 OracleSession 实例。
    
    主要功能：
    1. 启动和管理工作进程
    2. 通过队列与工作进程通信
    3. 维护本地状态缓存（从工作进程同步）
    4. 后台线程实时同步视频帧
    """
    
    def __init__(self, dataset_root=DEFAULT_DATASET_ROOT, gui_render=False):
        """
        初始化代理对象
        
        Args:
            dataset_root: 数据集根目录路径
            gui_render: 是否使用GUI渲染模式
        """
        # 使用 spawn 上下文以获得更清晰的进程隔离
        ctx = multiprocessing.get_context("spawn")
        
        # 创建进程间通信队列
        self.cmd_queue = ctx.Queue()      # 命令队列：主进程 -> 工作进程
        self.result_queue = ctx.Queue()   # 结果队列：工作进程 -> 主进程
        self.stream_queue = ctx.Queue()    # 流队列：工作进程 -> 主进程（视频帧）
        
        # 启动工作进程
        self.process = ctx.Process(
            target=session_worker_loop,
            args=(self.cmd_queue, self.result_queue, self.stream_queue, dataset_root, gui_render),
            daemon=True
        )
        self.process.start()
        LOGGER.info(
            "ProcessSessionProxy started worker pid=%s dataset_root=%s gui_render=%s",
            self.process.pid,
            dataset_root,
            gui_render,
        )
        
        # 本地状态缓存（从工作进程同步）
        self.env_id = None
        self.episode_idx = None
        self.language_goal = ""
        self.difficulty = None
        self.seed = None
        self.demonstration_frames = []
        self.last_execution_frames = []
        self.base_frames = []  # 由后台同步线程持续更新
        self.wrist_frames = []  # 由后台同步线程持续更新
        self.available_options = []
        self.raw_solve_options = []
        self.seg_vis = None
        self.is_demonstration = False  # 演示模式标志
        self.non_demonstration_task_length = None  # 从工作进程同步的非demonstration任务长度
        
        # 帧同步线程：从流队列接收新帧并更新本地缓存
        self.stop_sync = False
        self.sync_thread = threading.Thread(target=self._sync_loop, daemon=True)
        self.sync_thread.start()

    def _sync_loop(self):
        """
        后台线程循环：从流队列消费视频帧并更新本地缓存
        
        此线程持续运行，实时接收工作进程推送的新视频帧，
        并将其追加到本地的 base_frames 和 wrist_frames 列表中。
        UI 刷新逻辑会直接从代理的本地缓存读取帧数据。
        """
        while not self.stop_sync:
            try:
                # Use a short timeout to check stop_sync frequently
                frames = self.stream_queue.get(timeout=0.1)
                new_base = frames.get("base", [])
                new_wrist = frames.get("wrist", [])
                
                # Append to local lists
                if new_base:
                    self.base_frames.extend(new_base)
                if new_wrist:
                    self.wrist_frames.extend(new_wrist)
            except queue.Empty:
                continue
            except Exception:
                LOGGER.exception("ProcessSessionProxy sync loop crashed")
                break
    
    def _send_cmd(self, cmd, *args, **kwargs):
        """
        发送命令到工作进程并等待结果
        
        Args:
            cmd: 命令名称
            *args: 位置参数
            **kwargs: 关键字参数
            
        Returns:
            命令执行结果
            
        Raises:
            RuntimeError: 工作进程返回错误或致命错误
            TimeoutError: 工作进程超时（600秒）
        """
        # 发送命令到工作进程
        start_ts = time.time()
        LOGGER.debug(
            "proxy send cmd=%s pid=%s args=%s kwargs_keys=%s",
            cmd,
            self.process.pid,
            len(args),
            list(kwargs.keys()),
        )
        self.cmd_queue.put({"cmd": cmd, "args": args, "kwargs": kwargs})
        try:
            # 等待结果（重任务如加载/执行可能需要较长时间，设置600秒超时）
            res = self.result_queue.get(timeout=600) 
            elapsed_ms = int((time.time() - start_ts) * 1000)
            LOGGER.debug(
                "proxy recv cmd=%s pid=%s status=%s elapsed_ms=%s",
                cmd,
                self.process.pid,
                res.get("status"),
                elapsed_ms,
            )
            
            # 检查错误状态并转换为异常，以便在 gradio_callbacks 中捕获并显示弹窗
            if res.get("status") == "screw_plan_failure":
                raise ScrewPlanFailureError(f"screw plan failed: {res.get('message', 'unknown error')}")
            if res.get("status") == "execution_error":
                raise RuntimeError(f"Execution error: {res.get('message', 'unknown error')}")
            if res.get("status") == "fatal":
                raise RuntimeError(f"工作进程致命错误: {res.get('message')}")
            if res.get("status") == "error":
                raise RuntimeError(f"命令执行错误: {res.get('message')}")
            
            # 更新本地状态缓存（如果工作进程返回了状态更新）
            if "state" in res:
                state = res["state"]
                for k, v in state.items():
                    if k in ["base_frames", "wrist_frames"]:
                        # 对于帧数据：只有在显式发送时才替换（如加载时）
                        # 否则由同步循环处理增量更新
                        if v is not None:
                             setattr(self, k, v)
                    else:
                        # 其他状态直接更新
                        setattr(self, k, v)
                        
            return res.get("result")
        except queue.Empty:
            LOGGER.error("proxy command timeout cmd=%s pid=%s", cmd, self.process.pid)
            raise TimeoutError("工作进程超时")

    def load_episode(self, env_id, episode_idx):
        """
        加载环境episode（在工作进程中执行）
        
        Args:
            env_id: 环境ID
            episode_idx: episode索引
            
        Returns:
            tuple: (PIL.Image, str) 图像和状态消息
        """
        return self._send_cmd(CMD_LOAD_EPISODE, env_id, episode_idx)

    def execute_action(self, action_idx, click_coords):
        """
        执行动作（在工作进程中执行，重计算任务）
        
        Args:
            action_idx: 动作索引
            click_coords: 点击坐标 (x, y) 或 None
            
        Returns:
            tuple: (PIL.Image, str, bool) 图像、状态消息、是否完成
        """
        self.last_execution_frames = []
        return self._send_cmd(CMD_EXECUTE_ACTION, action_idx, click_coords)
        
    def get_pil_image(self, use_segmented=True):
        """
        获取PIL图像（在工作进程中执行）
        
        Args:
            use_segmented: 是否使用分割视图
            
        Returns:
            PIL.Image: 图像对象
        """
        return self._send_cmd(CMD_GET_PIL_IMAGE, use_segmented=use_segmented)
        
    def update_observation(self, use_segmentation=True):
        """
        更新观察（在工作进程中执行）
        
        Args:
            use_segmentation: 是否使用分割视图
            
        Returns:
            tuple: (PIL.Image, str) 图像和状态消息
        """
        return self._send_cmd(CMD_UPDATE_OBSERVATION, use_segmentation=use_segmentation)

    def get_reference_action(self):
        """
        获取当前步参考动作与坐标（在工作进程中执行）

        Returns:
            dict: 参考动作结果
        """
        return self._send_cmd(CMD_GET_REFERENCE_ACTION)
        
    def close(self):
        """
        关闭代理并清理资源
        
        此方法会：
        1. 停止帧同步线程
        2. 发送关闭命令到工作进程
        3. 等待工作进程优雅退出（最多1秒）
        4. 如果进程仍在运行，强制终止
        """
        self.stop_sync = True
        try:
            self.cmd_queue.put({"cmd": CMD_CLOSE})
            LOGGER.debug("proxy close command sent pid=%s", self.process.pid)
        except Exception:
            LOGGER.exception("proxy failed to send close command pid=%s", self.process.pid)
            pass
        # 等待工作进程优雅退出
        self.process.join(timeout=1)
        if self.process.is_alive():
            LOGGER.warning("proxy worker still alive after join; terminating pid=%s", self.process.pid)
            self.process.terminate()
        LOGGER.info("ProcessSessionProxy closed pid=%s", self.process.pid)