run_rca.py

import os
import pickle
import json
import time
from typing import List, Dict, Any
import torch
from torch_geometric.loader import DataLoader
from langchain_ollama import ChatOllama
from langchain.prompts import PromptTemplate

import graphbuilder
from faults import FaultType
from profile_agent.prompts import output_fault_functions
from profile_agent.prompts import profiling_compare_syscall_with_similarity
from utils import clean_llm_output
from faiss_retriever import FaissRetriever
from graphcl_model import GraphCLModel
from profile_dataset import ProfileDataset

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


llm = ChatOllama(
    model="qwen3:30b-a3b-64k",
    temperature=0,
    num_predict=16384,
)

def generate_rca_with_retry(comparison, similarity_results, service, max_retries=3):
    """
    生成根因分析，如果生成失败则重试
    
    Args:
        comparison: 比较结果
        similarity_results: 相似样例检索结果
        max_retries: 最大重试次数
    
    Returns:
        tuple: (raw_output, clean_output, elapsed_time)
    """
    for attempt in range(max_retries):
        # 记录开始时间
        start_time = time.time()
        
        # 生成分析
        llm_input = prompt.format(input=comparison, similarity_results=similarity_results, service=service, output_format_with_example=output_fault_functions)
        print(f"LLM input: {llm_input}")
        response = llm.invoke(llm_input)
        raw_output = response.content
        
        # 计算耗时
        end_time = time.time()
        elapsed_time = end_time - start_time
        
        # 清理LLM输出，移除<think></think>部分
        clean_output = clean_llm_output(raw_output)
        
        if clean_output:
            return raw_output, clean_output, elapsed_time
        
        print(f"根因分析生成失败，正在进行第{attempt+1}次重试...")
    
    return "", "", elapsed_time

def format_similarity_results(results):
    """
    格式化相似样例检索结果为字符串
    
    Args:
        results: 相似样例检索结果列表，每个元素包含meta和score信息
    
    Returns:
        str: 格式化后的字符串
    """
    formatted_results = []
    for i, result in enumerate(results, 1):  # 从1开始计数
        meta = result['meta']
        formatted_results.append(f"Top {i} Similar Case:")
        formatted_results.append(f"  Fault: {meta['fault']}")
        formatted_results.append(f"  Fault Description: {meta['fault_description']}")
        formatted_results.append(f"  Description: {meta['description']}")
        formatted_results.append("")
    
    return "\n".join(formatted_results)



services = [
    'adservice',
    'checkoutservice',
    'emailservice',
    'frontend',
    'recommendationservice',
]

fault_types = [
    FaultType.EPOLL_WAIT_DELAY,
    FaultType.FUTEX_DELAY,
    FaultType.READ_DELAY,
    FaultType.WRITE_DELAY,
]


prompt = PromptTemplate(
    template=profiling_compare_syscall_with_similarity,
    input_variables=["output_format_with_example", "input", "similarity_results", "service"]
)


# summary文件路径
result_dir = '4faults_2025-07-22_14-22-03'
base_dir = './results/' + result_dir

# 加载模型参数
model_params_path = f'{base_dir}/model_parameters.json'
with open(model_params_path, 'r') as f:
    model_params = json.load(f)
model_config = model_params['model_parameters']

# 加载summary, 包含训练和测试索引
summary_file = f'{base_dir}/all_services_summary.json'
with open(summary_file, 'r') as f:
    summary = json.load(f)

for service in services:
    print(f"处理服务: {service}")
    
    # 加载模型
    model_path = f'{base_dir}/{service}/model_semi_{service}.pt'
    model = GraphCLModel(
        hidden_dim=model_config['hidden_dim'],
        dataset_num_features=model_config['dataset_num_features'],
        num_edge_features=model_config['num_edge_features'],
        num_gc_layers=model_config['num_gc_layers'],
        global_dim=model_config['global_dim'],
        multi_heads_num=model_config['multi_heads_num']
    ).to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))

    # 加载common图
    common_graph = graphbuilder.load_common_graph(service, generate=True)
    
    # 收集所有故障图数据
    all_fault_graphs = []
    all_fault_labels = []
    train_graphs = []
    train_labels = []
    test_graphs = []
    test_labels = []
    
    # 加载所有故障类型的图数据
    for fault_type in fault_types:
        print(f"收集故障类型: {fault_type} 的图数据")
        graph_dir = fault_type.path.format(service)
        graph_with_desc_dir = graph_dir.replace('data_fault', 'data_fault_description')
        if os.path.exists(graph_with_desc_dir):
            graph_dir = graph_with_desc_dir
        else:
            print(f"没有描述的故障图目录: {graph_with_desc_dir}")
        
        fault_graphs, fault_labels = graphbuilder.load_graph_from_glob_directory(
            graph_dir, 
            fault_type.code, 
            generate=False
        )
        
        # 只使用训练数据
        train_indices = summary[service][fault_type.name]['train_indices']
        train_graphs.extend([fault_graphs[i] for i in train_indices])
        train_labels.extend([fault_labels[i] for i in train_indices])

        test_indices = summary[service][fault_type.name]['test_indices']
        test_graphs.extend([fault_graphs[i] for i in test_indices])
        test_labels.extend([fault_labels[i] for i in test_indices])
        
        print(f"已收集 {len(train_graphs)} 个 {fault_type} 故障图")
    
    # 初始化向量检索器
    retriever = FaissRetriever(
        service_name=service,
        data_dir='./data_normal/5m/{}/'.format(service),
        index_dir='./vector_db/{}/'.format(service)
    )
    
    # 构建包含所有故障类型的索引
    print(f"为服务 {service} 构建索引，共有 {len(train_graphs)} 个图")
    retriever.build_index(model, generate=False, graph_list=train_graphs, label_list=train_labels)
    retriever.load_index()
    

        
    # 创建输出目录
    output_dir = base_dir.replace("results", "results_rca_with_function")
    if output_dir == base_dir:
        print(f"输出目录与故障图目录相同: {output_dir}")
        exit(1)
    
    os.makedirs(output_dir, exist_ok=True)

    dataset = ProfileDataset(graph_list=test_graphs, label_list=test_labels, embed=True)
    loader = DataLoader(dataset, batch_size=1, shuffle=False)
    query_emb, _ = model.encoder.get_embeddings(loader)


    for i, fault_graph in enumerate(test_graphs):
        # 确定输出文件路径
        output_filename = os.path.basename(fault_graph.filename) + '.json'
        output_file_path = os.path.join(output_dir, "5m", service, output_filename)
        
        # 检查输出文件是否已存在
        if os.path.exists(output_file_path):
            print(f"跳过图{i+1}: {fault_graph.filename} - 输出文件已存在")
            continue
        
        print(f"处理图{i+1}: {fault_graph.filename}")
        
        # 比较图
        comparison = graphbuilder.compare_call_graph(fault_graph, common_graph)
        
        
        # 检索相似样例
        similarity_results = retriever.search(query_emb[i], topk=3)[0]
        formatted_similarity = format_similarity_results(similarity_results)
        
        # 生成根因分析，如果失败则重试
        raw_output, clean_output, elapsed_time = generate_rca_with_retry(
            comparison, formatted_similarity, service
        )
        if clean_output == "":
            print(f"生成失败，进程退出")
            exit(1)
        
        print(f"LLM生成耗时: {elapsed_time:.2f}秒")
        print(f"LLM output: {clean_output}")
        
        # 解析JSON输出
        result = {'fault': fault_graph.fault, 'description': fault_graph.description, 'similarity_results': [x['meta']['fault'] for x in similarity_results]}
        try:
            rca_result = json.loads(clean_output)
            result['rca'] = rca_result
        except json.JSONDecodeError as e:
            print(f"JSON解析失败: {e}")
            continue
        
        # 保存结果前，确保父目录存在
        os.makedirs(os.path.dirname(output_file_path), exist_ok=True)
        with open(output_file_path, 'w') as f:
            json.dump(result, f, indent=2)
        
        print(f"已保存结果到: {output_file_path}")