TraceStudio/server/app/nodes/trace_loader_nodes.py

"""
Trace 文件加载节点
用于从 Unreal Insights 导出和加载 Trace 数据
"""
import time
import polars as pl
from server.app.core.user_manager import CLOUD_ROOT
import server.app.utils.io as io
import subprocess
from pathlib import Path
from typing import Any, Dict, Optional
from server.app.core.node_base import (
    TraceNode,
    input_port,
    output_port,
    param,
    context_var,
    NodeType,
    CachePolicy
)
from server.app.core.node_registry import register_node
from server.app.services.agent_client import AgentClient
import yaml


def load_system_config():
    """加载系统配置"""
    config_path = Path(__file__).parent.parent.parent / "system_config.yaml"
    if config_path.exists():
        with open(config_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    return {}


def _agent_client() -> AgentClient:
    return AgentClient.from_env(container_cloud_root=CLOUD_ROOT)


@register_node
class TraceLoader_Metadata(TraceNode):
    """
    Trace 元数据加载器
    
    功能：
    1. 调用 UnrealInsights.exe 打开 .utrace 文件
    2. 执行 ExportMetadata 命令导出 CSV
    3. 读取 CSV 文件并加载到内存
    4. 返回 DataFrame 格式的数据表
    
    使用场景：
    - 分析 Trace 文件的元数据信息
    - 作为数据处理工作流的起点
    """
    
    CATEGORY = "Loader"
    DISPLAY_NAME = "Trace 元数据加载器"
    DESCRIPTION = "从 .utrace 文件导出并加载元数据"
    ICON = "📥"
    NODE_TYPE = NodeType.INPUT  # 输入节点（无输入端口）
    CACHE_POLICY = CachePolicy.DISK  # 使用磁盘缓存（元数据较大）

    @output_port("metadata", "DataTable", description="导出的元数据表")
    @param("utrace_file", "String", default="", description="Utrace 文件路径", widget="file", required=True)
    @param("tokens", "String", description="Token事件列表，以 | 进行分割")
    def process(self, inputs: Dict[str, Any], context: Optional[Dict] = None) -> Dict[str, Any]:
        # 获取参数
        tokens = self.get_param("tokens", "")
        utrace_file = self.get_param("utrace_file", "")
        output_csv = self.get_param("output_csv", "temp/metadata.csv")

        print("TraceLoader_Metadata: process called with utrace_file =", utrace_file, ", tokens =", tokens)
        if not utrace_file:
            raise ValueError("必须指定 utrace_file 参数")
        
        # 加载系统配置
        config = load_system_config()
        insights_path = config.get("unreal", {}).get("insights_path", "")
        
        # 构建完整路径
        # utrace 文件路径（用户目录下）
        if not insights_path or not Path(insights_path).exists():
            raise FileNotFoundError(
                f"UnrealInsights.exe 未找到：{insights_path}\n"
                f"请在 system_config.yaml 中配置正确的 insights_path"
            )
        
        utrace_full_path = CLOUD_ROOT / utrace_file
        
        if not utrace_full_path.exists():
            raise FileNotFoundError(f"Utrace 文件不存在：{utrace_full_path}")
        
        # 构建输出 CSV 完整路径（容器内路径 + 宿主机路径，宿主机路径给 Agent 用）
        csv_full_path = CLOUD_ROOT / output_csv
        agent = _agent_client()
        host_csv_full_path = agent.map_cloud_to_host(output_csv)

        # 确保输出目录存在（容器内），Agent 写宿主路径时也会自动创建
        csv_full_path.parent.mkdir(parents=True, exist_ok=True)
        io.safe_delete(csv_full_path)

        # 通过 Agent 调用宿主机 UnrealInsights
        host_utrace_full_path = agent.map_cloud_to_host(utrace_full_path)
        cmd_args = [
            f'-OpenTraceFile="{host_utrace_full_path}"',
            f'-ExecOnAnalysisCompleteCmd="TimingInsights.ExportMetadata {host_csv_full_path} -tokens={tokens}"'
        ]
        print("TraceLoader_Metadata: call agent with", cmd_args)
        resp = agent.run("unreal_insights", args=cmd_args, timeout=None)

        if not resp.get("success", False):
            raise RuntimeError(
                f"UnrealInsights 通过 Agent 执行失败 (code={resp.get('return_code')})\n"
                f"stdout: {resp.get('stdout','')}\n"
                f"stderr: {resp.get('stderr','')}"
            )

        if not csv_full_path.exists():
            raise FileNotFoundError(
                f"CSV 文件未生成：{csv_full_path}\n"
                f"Agent stdout: {resp.get('stdout','')}\n"
                f"Agent stderr: {resp.get('stderr','')}"
            )
        
        # 读取 CSV 文件
        df = pl.read_csv(csv_full_path)
        
        # 准备输出
        outputs = {
            "metadata": df,
            "file_path": str(csv_full_path)
        }
        
        #io.safe_delete(csv_full_path)
        return outputs


@register_node
class TraceLoader_Events(TraceNode):
    """
    Trace 事件加载器
    
    功能：
    导出 Trace 文件中的事件数据（Timing Events）
    """
    
    CATEGORY = "Loader"
    DISPLAY_NAME = "Trace 事件加载器"
    DESCRIPTION = "从 .utrace 文件导出并加载事件数据"
    ICON = "⏱️"
    NODE_TYPE = NodeType.INPUT
    CACHE_POLICY = CachePolicy.DISK
    
    @output_port("events", "DataTable", description="导出的事件数据表")
    @output_port("file_path", "String", description="导出的 CSV 文件路径")
    @param("utrace_file", "String", default="", description="Utrace 文件路径", widget="file", required=True)
    @param("output_csv", "String", default="temp/events.csv", description="输出 CSV 路径", required=True)
    @param("timeout", "Number", default=120, description="命令执行超时时间（秒）", min=10, max=600, step=10)
    @param("auto_cleanup", "Boolean", default=False, description="执行后自动删除 CSV 文件")
    @context_var("row_count", "Integer", description="事件数量")
    @context_var("column_count", "Integer", description="列数")
    @context_var("execution_time", "Number", description="执行时间（秒）")
    def process(self, inputs: Dict[str, Any], context: Optional[Dict] = None) -> Dict[str, Any]:
        start_time = time.time()
        
        # 获取参数
        utrace_file = self.get_param("utrace_file", "")
        output_csv = self.get_param("output_csv", "temp/events.csv")
        timeout = self.get_param("timeout", 120)
        auto_cleanup = self.get_param("auto_cleanup", False)
        
        if not utrace_file:
            raise ValueError("必须指定 utrace_file 参数")
        
        # 加载系统配置
        config = load_system_config()
        insights_path = config.get("unreal", {}).get("insights_path", "")
        cloud_root = Path(config.get("storage", {}).get("cloud_root", "./cloud"))
        
        user_id = (context or {}).get("user_id", "guest")
        
        # 构建路径
        if not insights_path or not Path(insights_path).exists():
            raise FileNotFoundError(f"UnrealInsights.exe 未找到：{insights_path}")
        
        if Path(utrace_file).is_absolute():
            utrace_full_path = Path(utrace_file)
        else:
            utrace_full_path = cloud_root / "users" / user_id / utrace_file
        
        if not utrace_full_path.exists():
            raise FileNotFoundError(f"Utrace 文件不存在：{utrace_full_path}")
        
        agent = _agent_client()
        if Path(output_csv).is_absolute():
            csv_full_path = Path(output_csv)
            host_csv_full_path = agent.map_cloud_to_host(output_csv)
        else:
            csv_full_path = cloud_root / output_csv
            host_csv_full_path = agent.map_cloud_to_host(output_csv)
        
        csv_full_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 通过 Agent 执行导出
        host_utrace_full_path = agent.map_cloud_to_host(utrace_full_path)
        cmd_args = [
            f'-OpenTraceFile="{host_utrace_full_path}"',
            f'-ExecOnAnalysisCompleteCmd="TimingInsights.ExportTimingEvents {host_csv_full_path}"'
        ]
        try:
            resp = agent.run("unreal_insights", args=cmd_args, timeout=timeout)
            
            if not resp.get("success", False):
                raise RuntimeError(
                    f"UnrealInsights 通过 Agent 执行失败 (code={resp.get('return_code')})\n"
                    f"stdout: {resp.get('stdout','')}\n"
                    f"stderr: {resp.get('stderr','')}"
                )
            
            # 等待文件生成
            wait_time = 0
            while not csv_full_path.exists() and wait_time < 5:
                time.sleep(0.5)
                wait_time += 0.5
            
            if not csv_full_path.exists():
                raise FileNotFoundError(
                    f"CSV 文件未生成：{csv_full_path}\n"
                    f"Agent stdout: {resp.get('stdout','')}\n"
                    f"Agent stderr: {resp.get('stderr','')}"
                )
            
            # 读取 CSV（使用 polars，保持后端一致）
            df = pl.read_csv(csv_full_path)
            execution_time = time.time() - start_time
            
            outputs = {
                "events": df,
                "file_path": str(csv_full_path)
            }
            
            result_context = {
                "row_count": len(df),
                "column_count": len(df.columns),
                "execution_time": round(execution_time, 2)
            }
            
            if auto_cleanup:
                try:
                    csv_full_path.unlink()
                    result_context["cleaned_up"] = True
                except Exception as e:
                    result_context["cleanup_error"] = str(e)
            
            return outputs
            
        except subprocess.TimeoutExpired:
            raise TimeoutError(f"命令执行超时（超过 {timeout} 秒）")
        except Exception as e:
            if csv_full_path.exists():
                try:
                    csv_full_path.unlink()
                except:
                    pass
            raise