| """Shared utilities for distributed clean-test evaluation.""" |
|
|
| import gzip |
| import hashlib |
| import json |
| import os |
| import random |
| import shutil |
| import tempfile |
| import time |
| from collections import Counter, defaultdict |
| from datetime import timedelta |
|
|
| import dgl |
| import torch |
| import torch.distributed as dist |
| import torch.nn.functional as F |
|
|
| from common import EDGE_TYPE_MAPPING, build_paper_mtl_augmented_graph |
|
|
|
|
| TASKS = ("ret", "jumptable", "indirectcall", "tailcall") |
|
|
|
|
| def distributed_setup(): |
| rank = int(os.environ.get("RANK", "0")) |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| world_size = int(os.environ.get("WORLD_SIZE", "1")) |
| if world_size > 1: |
| backend = os.environ.get("ICFLOWNET_DIST_BACKEND", "nccl") |
| dist.init_process_group(backend=backend, timeout=timedelta(hours=12)) |
| torch.cuda.set_device(local_rank) |
| return rank, local_rank, world_size, torch.device("cuda", local_rank) |
|
|
|
|
| def cleanup_distributed(world_size): |
| if world_size > 1 and dist.is_initialized(): |
| dist.destroy_process_group() |
|
|
|
|
| def all_reduce(tensor, world_size): |
| if world_size > 1: |
| if dist.get_backend() == "gloo" and tensor.is_cuda: |
| cpu_tensor = tensor.cpu() |
| dist.all_reduce(cpu_tensor, op=dist.ReduceOp.SUM) |
| tensor.copy_(cpu_tensor) |
| else: |
| dist.all_reduce(tensor, op=dist.ReduceOp.SUM) |
| return tensor |
|
|
|
|
| def seed_everything(seed, rank): |
| random.seed(seed + rank) |
| torch.manual_seed(seed + rank) |
| torch.cuda.manual_seed_all(seed + rank) |
|
|
|
|
| def load_json(path): |
| with open(path, "r") as handle: |
| return json.load(handle) |
|
|
|
|
| def read_jsonl(path): |
| rows = [] |
| with open(path, "r") as handle: |
| for line in handle: |
| line = line.strip() |
| if line: |
| rows.append(json.loads(line)) |
| return rows |
|
|
|
|
| def load_compressed_graph(path, use_func_rel=False): |
| cache_dir = os.environ.get("ICFLOWNET_CLEAN_GRAPH_CACHE_DIR") |
| cached_path = None |
| if cache_dir: |
| os.makedirs(cache_dir, exist_ok=True) |
| path_digest = hashlib.sha256(os.path.abspath(path).encode("utf-8")).hexdigest()[:16] |
| graph_name = os.path.basename(path) |
| if graph_name.endswith(".gz"): |
| graph_name = graph_name[:-3] |
| cached_path = os.path.join(cache_dir, path_digest + "_" + graph_name) |
| if os.path.exists(cached_path): |
| graph = dgl.load_graphs(cached_path)[0][0] |
| if not use_func_rel and "code2funchead" in graph.etypes: |
| edge_ids = graph.edges(form="eid", etype="code2funchead") |
| graph = dgl.remove_edges(graph, edge_ids, etype="code2funchead") |
| return graph |
|
|
| with tempfile.NamedTemporaryFile( |
| prefix="icflownet_eval_", |
| suffix=".graph", |
| dir=cache_dir, |
| delete=False, |
| ) as temporary: |
| temporary_path = temporary.name |
| try: |
| with gzip.open(path, "rb") as source, open(temporary_path, "wb") as destination: |
| shutil.copyfileobj(source, destination) |
| if cached_path: |
| os.replace(temporary_path, cached_path) |
| temporary_path = cached_path |
| graph = dgl.load_graphs(temporary_path)[0][0] |
| if not use_func_rel and "code2funchead" in graph.etypes: |
| edge_ids = graph.edges(form="eid", etype="code2funchead") |
| graph = dgl.remove_edges(graph, edge_ids, etype="code2funchead") |
| return graph |
| finally: |
| if not cached_path and os.path.exists(temporary_path): |
| os.remove(temporary_path) |
|
|
|
|
| def discard_unused_full_code_features(graph): |
| """Keep pooled code features when both pooled and full tensors are stored.""" |
| if "code" in graph.ntypes: |
| code_data = graph.nodes["code"].data |
| if "featmean" in code_data and "feat" in code_data: |
| del code_data["feat"] |
| return graph |
|
|
|
|
| def prepare_clean_records(clean_root, resource_root, scope="overall"): |
| if scope not in ("overall", "long_range"): |
| raise ValueError("Unsupported clean-test scope: {}".format(scope)) |
|
|
| records = read_jsonl(os.path.join(clean_root, "binary_graph_mapping.jsonl")) |
| prepared = [] |
| pair_totals = Counter() |
| task_totals = Counter() |
| zero_pair_records = [] |
|
|
| for record in records: |
| binary_key = record["binary_key"] |
| resource_dir = os.path.join(resource_root, binary_key) |
| lookup_path = os.path.join(resource_dir, binary_key + "_nodelookup.json") |
| nidtoaddr_path = os.path.join(resource_dir, binary_key + "_nidtoaddr.json") |
| hubmeta_path = os.path.join(resource_dir, binary_key + "_hubmeta.json") |
| graphstats_path = os.path.join(resource_dir, binary_key + "_graphstats.json") |
| graph_path = os.path.join(clean_root, record["graph"]) |
| required = (lookup_path, hubmeta_path, graphstats_path, graph_path) |
| missing = [path for path in required if not os.path.exists(path)] |
| if missing: |
| raise RuntimeError( |
| "Missing clean-test resources for {}: {}".format(binary_key, missing) |
| ) |
|
|
| graphstats = load_json(graphstats_path) |
| expected_code_nodes = int(graphstats["code_nodes"]) |
| |
| |
| |
| |
| |
| if os.path.exists(nidtoaddr_path): |
| nidtoaddr = load_json(nidtoaddr_path) |
| code_nidtoaddr = nidtoaddr.get("code") |
| else: |
| code_nidtoaddr = None |
| if isinstance(code_nidtoaddr, dict): |
| lookup = { |
| str(address): int(node_id) |
| for node_id, address in code_nidtoaddr.items() |
| } |
| else: |
| raw_lookup = load_json(lookup_path) |
| lookup = { |
| str(address): int(node_id) |
| for address, node_id in raw_lookup.items() |
| } |
| if set(lookup.values()) != set(range(expected_code_nodes)): |
| raise RuntimeError( |
| "Node lookup does not exactly cover graph code ids for {}.".format( |
| binary_key |
| ) |
| ) |
|
|
| task_rows = defaultdict(list) |
| record_pair_count = 0 |
| for kind in (scope + "_positive", scope + "_negative"): |
| rows = read_jsonl(os.path.join(clean_root, record[kind])) |
| pair_totals[kind] += len(rows) |
| record_pair_count += len(rows) |
| for row in rows: |
| src_key = str(row["src"]) |
| dst_key = str(row["dst"]) |
| if src_key not in lookup or dst_key not in lookup: |
| raise RuntimeError( |
| "Unmapped clean pair in {}: {}".format(binary_key, row) |
| ) |
| task = row["task"] |
| task_rows[task].append( |
| { |
| "src": lookup[src_key], |
| "dst": lookup[dst_key], |
| "label": int(row["label"]), |
| } |
| ) |
| task_totals[(task, int(row["label"]))] += 1 |
|
|
| if record_pair_count == 0: |
| zero_pair_records.append( |
| {"graph_id": str(record["graph_id"]), "binary_key": binary_key} |
| ) |
|
|
| prepared.append( |
| { |
| "binary_key": binary_key, |
| "graph_id": str(record["graph_id"]), |
| "graph_path": graph_path, |
| "candidate_meta": load_json(hubmeta_path), |
| "task_rows": dict(task_rows), |
| "expected_code_nodes": expected_code_nodes, |
| "expected_data_nodes": int(graphstats["data_nodes"]), |
| } |
| ) |
|
|
| return prepared, { |
| "records": len(prepared), |
| "pair_totals": dict(pair_totals), |
| "nonzero_pair_records": len(prepared) - len(zero_pair_records), |
| "zero_pair_records": zero_pair_records, |
| "task_label_totals": { |
| "{}:{}".format(task, label): count |
| for (task, label), count in sorted(task_totals.items()) |
| }, |
| } |
|
|
|
|
| def confusion_metrics(tn, fp, fn, tp): |
| precision = tp / (tp + fp) if tp + fp else 0.0 |
| recall = tp / (tp + fn) if tp + fn else 0.0 |
| f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0 |
| return { |
| "precision": precision, |
| "recall": recall, |
| "f1": f1, |
| "support": tp + fn, |
| "pairs": tn + fp + fn + tp, |
| "tn": tn, |
| "fp": fp, |
| "fn": fn, |
| "tp": tp, |
| "confusion_matrix": [[tn, fp], [fn, tp]], |
| } |
|
|
|
|
| def evaluate_clean_full( |
| ddp_model, |
| clean_records, |
| device, |
| rank, |
| world_size, |
| output_dir, |
| log_every, |
| amp_enabled, |
| ): |
| module = ddp_model.module |
| module.eval() |
| local_records = clean_records[rank::world_size] |
| confusions = torch.zeros((len(TASKS), 2, 2), dtype=torch.long, device=device) |
| loss_sums = torch.zeros(len(TASKS), dtype=torch.float64, device=device) |
| pair_counts = torch.zeros(len(TASKS), dtype=torch.float64, device=device) |
| type_confusion = torch.zeros((4, 4), dtype=torch.long, device=device) |
| evaluated = torch.zeros((), dtype=torch.long, device=device) |
| local_details = [] |
| start_time = time.time() |
|
|
| for item_index, item in enumerate(local_records, start=1): |
| combined_pairs = [] |
| task_slices = {} |
| actual_types = [] |
| actual_labels = [] |
| for task in TASKS: |
| rows = item["task_rows"].get(task, []) |
| if not rows: |
| continue |
| start = len(combined_pairs) |
| combined_pairs.extend((row["src"], row["dst"]) for row in rows) |
| actual_labels.extend(row["label"] for row in rows) |
| actual_types.extend([EDGE_TYPE_MAPPING[task]] * len(rows)) |
| task_slices[task] = (start, len(combined_pairs), rows) |
|
|
| if not combined_pairs: |
| local_details.append( |
| { |
| "binary": item["binary_key"], |
| "graph_id": item["graph_id"], |
| "pairs": 0, |
| "tasks": {}, |
| "metric_contribution": "none_no_labeled_pairs", |
| } |
| ) |
| evaluated += 1 |
| _log_progress(rank, item_index, len(local_records), log_every, start_time) |
| continue |
|
|
| base_graph = discard_unused_full_code_features( |
| load_compressed_graph(item["graph_path"]) |
| ) |
| if ( |
| base_graph.num_nodes("code") != item["expected_code_nodes"] |
| or base_graph.num_nodes("data") != item["expected_data_nodes"] |
| ): |
| raise RuntimeError( |
| "Clean graph node count mismatch for {}.".format(item["binary_key"]) |
| ) |
| model_graph = build_paper_mtl_augmented_graph( |
| base_graph, |
| item["candidate_meta"], |
| use_gch=True, |
| use_gdh=True, |
| task_aware_routing=True, |
| gdh_radius=2, |
| ) |
| graph_gpu = model_graph.to(device) |
| pairs = torch.tensor(combined_pairs, dtype=torch.long, device=device) |
| labels = torch.tensor(actual_labels, dtype=torch.float32, device=device) |
| types = torch.tensor(actual_types, dtype=torch.long, device=device) |
|
|
| with torch.no_grad(), torch.autocast( |
| device_type=device.type, dtype=torch.float16, enabled=amp_enabled |
| ): |
| task_logits, edge_type_logits = module( |
| graph_gpu, pairs, candidate_meta=None |
| ) |
|
|
| detail = { |
| "binary": item["binary_key"], |
| "graph_id": item["graph_id"], |
| "pairs": len(combined_pairs), |
| "tasks": {}, |
| } |
| for task_index, task in enumerate(TASKS): |
| if task not in task_slices: |
| continue |
| start, end, rows = task_slices[task] |
| logits = task_logits[task].reshape(-1)[start:end] |
| task_labels = labels[start:end] |
| scores = torch.sigmoid(logits) |
| predictions = scores >= 0.5 |
| positives = task_labels == 1 |
| negatives = ~positives |
| tn = int(((~predictions) & negatives).sum().item()) |
| fp = int((predictions & negatives).sum().item()) |
| fn = int(((~predictions) & positives).sum().item()) |
| tp = int((predictions & positives).sum().item()) |
| confusions[task_index, 0, 0] += tn |
| confusions[task_index, 0, 1] += fp |
| confusions[task_index, 1, 0] += fn |
| confusions[task_index, 1, 1] += tp |
| loss_sums[task_index] += F.binary_cross_entropy_with_logits( |
| logits, task_labels, reduction="sum" |
| ).double() |
| pair_counts[task_index] += len(rows) |
| task_detail = { |
| "pairs": len(rows), |
| "confusion_matrix": [[tn, fp], [fn, tp]], |
| } |
| detail["tasks"][task] = task_detail |
|
|
| positive_mask = labels == 1 |
| predicted_types = edge_type_logits[positive_mask].argmax(dim=1) |
| true_types = types[positive_mask] |
| for true_type, predicted_type in zip( |
| true_types.tolist(), predicted_types.tolist() |
| ): |
| type_confusion[true_type, predicted_type] += 1 |
|
|
| local_details.append(detail) |
| evaluated += 1 |
| del graph_gpu, base_graph, model_graph, pairs, labels, types |
| del task_logits, edge_type_logits |
| if device.type == "cuda": |
| torch.cuda.empty_cache() |
| _log_progress(rank, item_index, len(local_records), log_every, start_time) |
|
|
| detail_dir = os.path.join(output_dir, "clean_details") |
| os.makedirs(detail_dir, exist_ok=True) |
| detail_path = os.path.join(detail_dir, "rank_{}.jsonl".format(rank)) |
| with open(detail_path, "w") as handle: |
| for detail in local_details: |
| handle.write(json.dumps(detail) + "\n") |
|
|
| for tensor in (confusions, loss_sums, pair_counts, type_confusion, evaluated): |
| all_reduce(tensor, world_size) |
|
|
| confusions = confusions.cpu() |
| loss_sums = loss_sums.cpu() |
| pair_counts = pair_counts.cpu() |
| task_metrics = {} |
| for task_index, task in enumerate(TASKS): |
| tn = int(confusions[task_index, 0, 0].item()) |
| fp = int(confusions[task_index, 0, 1].item()) |
| fn = int(confusions[task_index, 1, 0].item()) |
| tp = int(confusions[task_index, 1, 1].item()) |
| task_metrics[task] = confusion_metrics(tn, fp, fn, tp) |
| task_metrics[task]["loss"] = float( |
| loss_sums[task_index].item() / max(pair_counts[task_index].item(), 1) |
| ) |
|
|
| type_confusion = type_confusion.cpu() |
| type_total = int(type_confusion.sum().item()) |
| type_correct = int(type_confusion.diag().sum().item()) |
| return { |
| "evaluated_binaries": int(evaluated.item()), |
| "task_metrics": task_metrics, |
| "macro_f1": sum(value["f1"] for value in task_metrics.values()) / len(TASKS), |
| "edge_type": { |
| "class_order": ["ret", "jumptable", "tailcall", "indirectcall"], |
| "confusion_matrix": type_confusion.tolist(), |
| "accuracy": type_correct / type_total if type_total else 0.0, |
| "positive_pairs": type_total, |
| }, |
| "detail_rank_files": [ |
| os.path.join(detail_dir, "rank_{}.jsonl".format(worker_rank)) |
| for worker_rank in range(world_size) |
| ], |
| "elapsed_seconds": time.time() - start_time, |
| } |
|
|
|
|
| def _log_progress(rank, current, total, log_every, start_time): |
| if current == 1 or current % log_every == 0 or current == total: |
| print( |
| "[clean][rank {}] graphs={}/{} elapsed={:.1f}s".format( |
| rank, current, total, time.time() - start_time |
| ), |
| flush=True, |
| ) |
|
|
