tb_real_dataset_agent_tree_structure.py

#!/usr/bin/env python3
import asyncio
import shutil
import time
import httpx
import json
import requests
import os
import re
import random
from typing import Dict, List, Optional, Tuple, Any, DefaultDict
from collections import defaultdict

from transformers import AutoTokenizer
from datasets import load_dataset, concatenate_datasets

from shell_router import run_batch_requests

# =========================
# CONFIG & DATA LOADING
# =========================

CONFIG_PATH = os.environ.get("CONFIG_PATH", "real_dataset_tb_configs.json")
CONFIG_NAME = os.path.splitext(os.path.basename(CONFIG_PATH))[0].strip(".json")

DATASET = os.environ.get("DATASET", "GSM8K").upper()  # "MMLU" | "GSM8K" | "AIME25" | "REAL_DATASET"
WARMUP = bool(int(os.environ.get("WARMUP", "0")))
RUN_BASELINE_ENV = bool(int(os.environ.get("RUN_BASELINE", "1")))
RUN_PROPOSED_ENV = bool(int(os.environ.get("RUN_PROPOSED", "1")))
NUM_SAMPLES = int(os.environ.get("NUM_SAMPLES", "1"))
CATEGORY = os.environ.get("CATEGORY", "math")
FIXED_TOKENS_DEFAULT = int(os.environ.get("FIXED_TOKENS", "4096"))

MMLUCategoryDict = {
    "math": ["college_mathematics", "abstract_algebra"],
    "computer_science": ["college_computer_science", "machine_learning"],
    "medicine": ["college_medicine", "professional_medicine", "clinical_knowledge"],
    "science": ["college_physics", "college_biology", "college_chemistry", "electrical_engineering"],
    "finance": ["high_school_macroeconomics", "high_school_microeconomics", "professional_accounting"]
}

def load_config(path: str) -> Dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

# ---------- Dataset loaders (kept) ----------
def load_mmlu_samples(categories: List[str], num_samples: int = 100, seed: int = 42):
    datasets_list = []
    for cfg in categories:
        ds = load_dataset('cais/mmlu', cfg, split='test', cache_dir='./llama_cache')
        datasets_list.append(ds)
    full = concatenate_datasets(datasets_list)
    return full.shuffle(seed=seed).select(range(min(num_samples, len(full)))) if num_samples > 0 else full

def load_gsm8k_samples(num_samples: int = 100, seed: int = 42):
    ds = load_dataset('openai/gsm8k', 'main', split='test', cache_dir='./llama_cache')
    ds = ds.shuffle(seed=seed)
    return ds.select(range(min(num_samples, len(ds)))) if num_samples > 0 else ds

def load_aime25_samples(num_samples: int = 100, seed: int = 42):
    i = load_dataset('opencompass/AIME2025', 'AIME2025-I', split='test', cache_dir='./llama_cache')
    ii = load_dataset('opencompass/AIME2025', 'AIME2025-II', split='test', cache_dir='./llama_cache')
    full = concatenate_datasets([i, ii]).shuffle(seed=42)
    return full.select(range(min(num_samples, len(full)))) if num_samples > 0 else full

def load_hmmt_feb_2025_samples(num_samples: int = 100, seed: int = 42):
    """
    HuggingFace: MathArena/hmmt_feb_2025 (split: train)
    Fields: problem (str), answer (str), problem_idx (int), problem_type (list[str])
    """
    ds = load_dataset('MathArena/hmmt_feb_2025', split='train', cache_dir='./llama_cache')
    ds = ds.shuffle(seed=seed)
    return ds.select(range(min(num_samples, len(ds)))) if num_samples > 0 else ds

def load_math500_samples(num_samples: int = 100, seed: int = 42):
    """
    HuggingFace: HuggingFaceH4/MATH-500 (split: test)
    Fields: problem (str), solution (str), answer (str), subject, level, unique_id
    """
    ds = load_dataset('HuggingFaceH4/MATH-500', split='test', cache_dir='./llama_cache')
    ds = ds.shuffle(seed=seed)
    return ds.select(range(min(num_samples, len(ds)))) if num_samples > 0 else ds


# --- NEW: GPQA-Diamond loader and canonizer ---

def _get_first_present(d: dict, keys: list, default=None):
    for k in keys:
        if k in d and d[k] is not None:
            return d[k]
    return default

def _canonize_gpqa_row(row: Dict[str, Any]) -> Dict[str, Any]:
    """Return a unified row: {'question', 'choices', 'answer'} where answer is index into choices."""
    # Question
    q = _get_first_present(row, ["question", "Question", "prompt", "Prompt"])
    if not isinstance(q, str):
        q = str(q)

    # Correct answer (string)
    correct = _get_first_present(
        row,
        ["correct_answer", "Correct Answer", "answer", "Answer", "correct"],
    )
    if not isinstance(correct, str):
        correct = str(correct)

    # Incorrect answers (list[str] or 3 separate columns)
    inc = _get_first_present(
        row,
        ["incorrect_answers", "Incorrect Answers", "distractors", "options_incorrect", "choices_incorrect"],
        default=None
    )
    if isinstance(inc, str):
        # sometimes comma/semicolon separated
        parts = [x.strip() for x in re.split(r"[;,\n]\s*", inc) if x.strip()]
        inc = parts if parts else None

    # Fall back to the common 3 columns
    if not isinstance(inc, list):
        inc_cols = []
        for k in ["incorrect_answer1", "Incorrect Answer 1", "incorrect1",
                  "incorrect_answer2", "Incorrect Answer 2", "incorrect2",
                  "incorrect_answer3", "Incorrect Answer 3", "incorrect3"]:
            v = row.get(k)
            if isinstance(v, str) and v.strip():
                inc_cols.append(v.strip())
        inc = inc_cols if inc_cols else []

    # If dataset already provides all choices and a 'correct_idx'
    choices_full = _get_first_present(row, ["choices", "Choices", "options", "Options"], default=None)
    correct_idx = _get_first_present(row, ["correct_idx", "label", "answer_index", "Correct Index"], default=None)

    if isinstance(choices_full, list) and choices_full:
        choices = [str(c) for c in choices_full]
        if isinstance(correct_idx, int) and 0 <= correct_idx < len(choices):
            ans_idx = correct_idx
        else:
            # locate the correct string inside choices
            norm = lambda s: re.sub(r"\s+", " ", s).strip().lower()
            try:
                ans_idx = [norm(c) for c in choices].index(norm(correct))
            except ValueError:
                # As a last resort, prepend correct then others (rare mirror mismatch)
                rem = [c for c in choices if c != correct]
                choices = [correct] + rem
                ans_idx = 0
        return {"question": q, "choices": choices, "answer": ans_idx}

    # Build choices ourselves: [correct] + incorrects (then deterministic shuffle)
    choices = [correct] + [c for c in inc if isinstance(c, str)]
    # Deterministic per-question shuffle to avoid answer-position bias
    rnd = random.Random(hash(q) & 0xFFFFFFFF)
    rnd.shuffle(choices)
    ans_idx = choices.index(correct)
    return {"question": q, "choices": choices, "answer": ans_idx}

def load_gpqa_diamond_samples(num_samples: int = 100, seed: int = 42) -> List[Dict[str, Any]]:
    """
    Loads Idavidrein/gpqa with config 'gpqa_diamond'.
    Prefer 'test' split; fall back to 'train' if needed (mirrors may vary).
    """
    try:
        ds = load_dataset('Idavidrein/gpqa', 'gpqa_diamond', split='test', cache_dir='./llama_cache')
    except Exception:
        ds = load_dataset('Idavidrein/gpqa', 'gpqa_diamond', split='train', cache_dir='./llama_cache')
    ds = ds.shuffle(seed=seed)
    ds = ds.select(range(min(num_samples, len(ds)))) if num_samples > 0 else ds
    # Canonicalize rows
    canon = [_canonize_gpqa_row(r) for r in ds]
    return canon


def build_prompt_from_sample(dataset_name: str, sample: Dict[str, Any]) -> str:
    if dataset_name == "MMLU":
        question = sample['question']
        choices = sample.get('choices') or sample.get('answer_choices')
        choices_str = '[' + ', '.join(f'"{c}"' for c in choices) + ']'
        return f"\nQuestion: {question} \nChoices: {choices_str}"

    elif dataset_name in ("GSM8K", "AIME25"):
        return f"\nQuestion: {sample['question']}"

    # --- NEW: GPQA-Diamond ---
    elif dataset_name in ("GPQA", "GPQA_DIAMOND", "GPQA-DIAMOND"):
        question = sample['question']
        choices = sample['choices']
        choices_str = '[' + ', '.join(f'"{c}"' for c in choices) + ']'
        return f"\nQuestion: {question} \nChoices: {choices_str}"

    # --- NEW: HMMT Feb 2025 ---
    elif dataset_name in ("HMMT_FEB_2025", "HMMT2025", "HMMT"):
        q = sample.get("problem") or sample.get("question")
        return f"\nQuestion: {q}\n\nPlease return ONLY the final answer enclosed in \\boxed{{}}."
    # --- NEW: MATH-500 ---
    elif dataset_name in ("MATH500", "MATH-500"):
        q = sample.get("problem") or sample.get("question")
        return f"\nQuestion: {q}\n\nPlease return ONLY the final answer enclosed in \\boxed{{}}."

    else:
        q = sample.get("question", "What is 2+2? Put final answer in \\boxed{}.")
        return f"\nQuestion: {q}"


def _normalize_numeric(s: str) -> str:
    return re.sub(r"[,\s]", "", str(s)).strip()

def gold_from_sample(dataset_name: str, sample: Dict[str, Any]) -> str:
    if dataset_name == "MMLU":
        gold_idx = sample['answer']
        choices = sample.get('choices') or sample.get('answer_choices')
        return str(choices[gold_idx]).strip()

    elif dataset_name == "GSM8K":
        ans = sample['answer']
        m = re.search(r"####\s*([^\n]+)", ans)
        final = m.group(1).strip() if m else ans.strip()
        return _normalize_numeric(final)

    elif dataset_name == "AIME25":
        return _normalize_numeric(sample['answer'])

    elif dataset_name in ("GPQA", "GPQA_DIAMOND", "GPQA-DIAMOND"):
        # Canonicalized sample: 'choices' + 'answer' (index)
        choices = sample['choices']
        idx = int(sample['answer'])
        return str(choices[idx]).strip()

    elif dataset_name in ("HMMT_FEB_2025", "HMMT2025", "HMMT"):
        # HMMT answers are final-answer strings; normalize numerics
        return _normalize_numeric(sample.get("answer", ""))

    elif dataset_name in ("MATH500", "MATH-500"):
        # MATH-500 'answer' is a final string (can be numeric or text)
        return _normalize_numeric(sample.get("answer", ""))

    else:
        return _normalize_numeric(sample.get("answer", ""))


def extract_boxed(text: str) -> str:
    pattern = r"\\boxed\{(.*?)\}"
    matches = re.findall(pattern, text)
    return matches[-1].strip() if matches else "No answer found"

# =========================
# CONFIG → GRAPH
# =========================

class AgentSpec:
    def __init__(self, layer_name: str, agent_id: str, model_key: str,
                 prefix_template: Optional[str], suffix_template: Optional[str],
                 depend_on: Optional[List[str]], sampling_params: Dict[str, Any]):
        self.layer_name = layer_name
        self.agent_id = agent_id
        self.model_key = model_key  # key into config["model_dict"]
        self.prefix_template = prefix_template or ""
        self.suffix_template = suffix_template or ""
        self.depend_on = depend_on or []
        self.sampling_params = sampling_params or {}

def build_agent_graph(cfg: Dict[str, Any]) -> Tuple[List[str], Dict[str, AgentSpec]]:
    layers_order: List[str] = []
    agent_map: Dict[str, AgentSpec] = {}
    for layer in cfg.get("layers", []):
        lname = layer["layer_name"]
        layers_order.append(lname)
        for a in layer.get("agents", []):
            # In your JSON, this is typically called "model" or "model_name".
            model_key = a.get("model_path") or a.get("model") or a.get("model_name")
            if not model_key:
                raise ValueError(f"Agent {a.get('agent_id')} missing 'model' key")
            spec = AgentSpec(
                layer_name=lname,
                agent_id=a["agent_id"],
                model_key=model_key,
                prefix_template=cfg["prompt_templates"].get(a.get("prefix_template", "default")),
                suffix_template=cfg["prompt_templates"].get(a.get("suffix_template", "default")),
                depend_on=a.get("depend_on") or [],
                sampling_params=a.get("sampling_params", {}),
            )
            agent_map[spec.agent_id] = spec
    return layers_order, agent_map

def substitute_q(template: str, question: str) -> str:
    return template.replace("{QUESTION}", question)

def build_prompt_for_agent_baseline(agent: AgentSpec, question: str, dep_texts: List[str]) -> str:
    prefix = substitute_q(agent.prefix_template, question)
    suffix = substitute_q(agent.suffix_template, question)
    blocks = []
    if dep_texts:
        for t in dep_texts:
            blocks.append(t)
    joined = ("\n\n".join(blocks) + "\n\n") if blocks else ""
    return f"{prefix}\n\n{joined}{suffix}".strip()

def build_template_for_agent_proposed(agent: AgentSpec, question: str) -> str:
    prefix = substitute_q(agent.prefix_template, question)
    suffix = substitute_q(agent.suffix_template, question)
    denom_construction = ((f"\nResponse from {dep}: \n{{DENOM}}\n\n" for dep in agent.depend_on) if agent.depend_on else "")
    denom_chunk = "".join(str(item) for item in denom_construction)
    return f"{prefix}\n\n{denom_chunk}{suffix}".strip()

# =========================
# ENDPOINTS & TOKENIZER
# =========================

def endpoints_for_agent(model_dict: Dict[str, Dict[str, str]], model_key: str) -> Tuple[str, str]:
    """
    Returns (shell_url, sglang_url) for this agent based on its model_key.
    model_dict example:
      {
        "llama31_8b": {"shell_url": "...:8100/generate", "sglang_url": "...:8000/generate"},
        "qwen2_7b":   {"shell_url": "...:9100/generate", "sglang_url": "...:9000/generate"}
      }
    """
    if not isinstance(model_dict, dict) or model_key not in model_dict:
        raise KeyError(f"Model '{model_key}' not found in model_dict.")
    entry = model_dict[model_key]
    shell_url = entry.get("shell_url")
    sglang_url = entry.get("sglang_url")
    if not shell_url or not sglang_url:
        raise KeyError(f"Model '{model_key}' missing 'shell_url' or 'sglang_url' in model_dict.")
    return shell_url, sglang_url

def sampling_with_defaults(sp: Dict[str, Any], fixed_max: int) -> Dict[str, Any]:
    out = {
        "temperature": 0.7,
        "max_new_tokens": fixed_max,
        "top_p": 0.8,
        "top_k": 20,
        "min_p": 0.0,
        "regex": [r'\\boxed\{(?P<ans>(?:\\.|[^\\}])*)\}'],
    }
    out.update({k: v for k, v in sp.items() if v is not None})
    return out

# =========================
# STREAMING (BASELINE)
# =========================

async def run_baseline_one_request_and_collect(prompt: str, sglang_url: str, tokenizer, samp: Dict[str, Any]) -> Dict[str, Any]:
    payload = {
        "text": prompt,
        "sampling_params": {
            "temperature": samp.get("temperature", 0.7),
            "max_new_tokens": samp.get("max_new_tokens", 64),
            "top_p": samp.get("top_p", 0.8),
            "top_k": samp.get("top_k", 20),
            "min_p": samp.get("min_p", 0.0),
        },
        "stream": True,
        "return_logprobs": False,
    }

    input_tokens = len(tokenizer.encode(prompt, add_special_tokens=False))
    t0 = time.perf_counter()
    ttft = None
    final_text = ""
    prev_accum = ""
    summary = None
    e2e_server = None

    async with httpx.AsyncClient(timeout=None) as client:
        async with client.stream("POST", sglang_url, json=payload) as resp:
            resp.raise_for_status()
            async for raw in resp.aiter_lines():
                if not raw or not raw.startswith("data:"):
                    continue
                if raw.strip() == "data: [DONE]":
                    break
                try:
                    obj = json.loads(raw[len("data:"):].strip())
                except Exception:
                    continue
                if obj.get("event") == "summary":
                    summary = obj
                    continue

                new_text = obj.get("text")
                delta = obj.get("delta") or obj.get("output")
                has_output = (isinstance(new_text, str) and new_text) or (isinstance(delta, str) and delta)
                if has_output and ttft is None:
                    ttft = time.perf_counter() - t0

                if isinstance(new_text, str) and new_text:
                    if len(new_text) < len(prev_accum):
                        prev_accum = ""
                        final_text = ""
                    if len(new_text) > len(prev_accum):
                        tail = new_text[len(prev_accum):]
                        final_text += tail
                        prev_accum = new_text
                elif isinstance(delta, str) and delta:
                    final_text += delta

                if obj.get("e2e_latency") is not None:
                    e2e_server = obj["e2e_latency"]

    if ttft is None:
        ttft = 0.0
    e2e = e2e_server if e2e_server is not None else (time.perf_counter() - t0)
    generated_tokens = len(tokenizer.encode(final_text, add_special_tokens=False))

    if summary is not None:
        ttft_s = (summary.get("ttft_ms", ttft * 1000.0) / 1000.0)
        e2e_s = (summary.get("e2e_ms", e2e * 1000.0) / 1000.0)
        return {"text": final_text, "ttft": ttft_s, "e2e": e2e_s,
                "input_tokens": input_tokens, "generated_tokens": generated_tokens}
    else:
        return {"text": final_text, "ttft": ttft, "e2e": e2e,
                "input_tokens": input_tokens, "generated_tokens": generated_tokens}

# =========================
# BASELINE: LAYERED EXEC (per-agent endpoints)
# =========================

async def run_baseline_layers(question: str,
                              layers_order: List[str],
                              agent_map: Dict[str, AgentSpec],
                              model_dict: Dict[str, Dict[str, str]],
                              tokenizer,
                              fixed_tokens_fallback: int, timestamp, warmup: bool=False) -> Dict[str, Any]:
    results: Dict[str, Dict[str, Any]] = {}

    layer_agents: Dict[str, List[AgentSpec]] = {}
    for spec in agent_map.values():
        layer_agents.setdefault(spec.layer_name, []).append(spec)

    t_wall0 = time.time()
    r1, r2 = [], []

    for lname in layers_order:
        specs = layer_agents.get(lname, [])
        coros = []
        for spec in specs:
            deps = spec.depend_on or []
            dep_texts = [results[d]["text"] for d in deps] if deps else []
            prompt = build_prompt_for_agent_baseline(spec, question, dep_texts)
            samp = sampling_with_defaults(spec.sampling_params, fixed_tokens_fallback)
            _, sglang_url = endpoints_for_agent(model_dict, spec.model_key)
            coros.append(run_baseline_one_request_and_collect(prompt, sglang_url, tokenizer, samp))

        outs = await asyncio.gather(*coros)
        for spec, out in zip(specs, outs):
            out["final_answer"] = extract_boxed(out["text"])
            results[spec.agent_id] = out

        if lname == "layer0":
            r1 = outs
        elif lname == "layer1":
            r2.extend([[o] for o in outs])

    wall = time.time() - t_wall0
    last_layer = layers_order[-1]
    last_agents = layer_agents.get(last_layer, [])
    final_outs = [results[a.agent_id] for a in last_agents] if last_agents else []

    total_ttft = sum(o["ttft"] for o in r1) + sum(o["ttft"] for group in r2 for o in group) + sum(o["ttft"] for o in final_outs)
    total_e2e = sum(o["e2e"] for o in r1) + sum(o["e2e"] for group in r2 for o in group) + sum(o["e2e"] for o in final_outs)

    data = {}
    data["round1"] = r1
    if r2:
        data["round2"] = r2
    data["final"] = {"outs": final_outs, "total_e2e_time": wall}
    data["summary"] = {"total_ttft": total_ttft, "total_e2e": total_e2e, "wall_time": wall}

    if not warmup:
        with open(f"{timestamp}_tb_agent_tree_structure_metrics_baseline.json", "w") as f:
            json.dump(data, f, indent=4)

    return {"wall_time": wall, "results": results, "final_outs": final_outs} if not warmup else {}

# =========================
# PROPOSED: TEMPLATED BATCH (per-agent endpoints)
# =========================

async def run_proposed_pipelined(question: str,
                                 layers_order: List[str],
                                 agent_map: Dict[str, AgentSpec],
                                 model_dict: Dict[str, Dict[str, str]],
                                 fixed_tokens_fallback: int) -> List[Dict[str, Any]]:
    """
    If all agents share the same shell_url, call run_batch_requests once.
    Otherwise, execute layer-by-layer and within each layer group by shell_url.
    Dependencies only go from earlier to later layers, so this preserves correctness.
    """
    # Organize agents by layer
    layer_agents: Dict[str, List[AgentSpec]] = {}
    for spec in agent_map.values():
        layer_agents.setdefault(spec.layer_name, []).append(spec)

    # Detect if there's a single shell_url
    shell_urls = set(endpoints_for_agent(model_dict, spec.model_key)[0] for spec in agent_map.values())
    single_shell = (len(shell_urls) == 1)
    global_shell_url = next(iter(shell_urls)) if single_shell else None

    def make_independent(specs: List[AgentSpec]) -> List[Dict[str, Any]]:
        reqs = []
        for spec in specs:
            samp = sampling_with_defaults(spec.sampling_params, fixed_tokens_fallback)
            text0 = substitute_q(spec.prefix_template or "", question)
            if spec.suffix_template:
                text0 = f"{text0}\n\n{substitute_q(spec.suffix_template, question)}"
            shell_url, _ = endpoints_for_agent(model_dict, spec.model_key)
            req = {
                "name": spec.agent_id,
                "text": text0,
                "agent_id": spec.agent_id,
                "start_delay_s": 0.0,
                "sampling_params": samp,
            }
            # annotate endpoint for multi-shell mode
            if not single_shell:
                req["shell_url"] = shell_url
            reqs.append(req)
        return reqs

    def make_dependent(specs: List[AgentSpec]) -> List[Dict[str, Any]]:
        reqs = []
        for spec in specs:
            if not spec.depend_on:
                continue
            samp = sampling_with_defaults(spec.sampling_params, fixed_tokens_fallback)
            template_text = build_template_for_agent_proposed(spec, question)
            shell_url, _ = endpoints_for_agent(model_dict, spec.model_key)
            req = {
                "name": spec.agent_id,
                "template_text": template_text,
                "agent_id": spec.agent_id,
                "depends_on": list(spec.depend_on),
                "start_delay_s": 0.5,
                "sampling_params": samp,
            }
            if not single_shell:
                req["shell_url"] = shell_url
            reqs.append(req)
        return reqs

    t0 = time.time()
    outs_all: List[Dict[str, Any]] = []

    if single_shell:
        # Build everything at once
        independent_reqs = make_independent(layer_agents.get("layer0", []))
        dependent_reqs: List[Dict[str, Any]] = []
        # append higher layers as dependent (respect order)
        for lname in layers_order:
            if lname == "layer0":
                continue
            dependent_reqs.extend(make_dependent(layer_agents.get(lname, [])))
        outs = await run_batch_requests(
            global_shell_url,
            independent_reqs=independent_reqs,
            dependent_reqs=dependent_reqs,
            default_stagger_s=0.5,
        )
        for resp in outs:
            resp["final_answer"] = extract_boxed(resp.get("output", resp.get("text", "")))
        outs_all.extend(outs)
    else:
        # Multi-shell: run per layer, grouping by shell_url
        # Layer 0: independent only
        layer0 = layer_agents.get("layer0", [])
        groups: DefaultDict[str, List[Dict[str, Any]]] = defaultdict(list)
        for req in make_independent(layer0):
            groups[req["shell_url"]].append(req)

        for shell_url, reqs in groups.items():
            out = await run_batch_requests(
                shell_url,
                independent_reqs=reqs,
                dependent_reqs=[],
                default_stagger_s=0.5,
            )
            for resp in out:
                resp["final_answer"] = extract_boxed(resp.get("output", resp.get("text", "")))
            outs_all.extend(out)

        # Higher layers: dependent only (their deps are in earlier layers, already finished)
        for lname in layers_order:
            if lname == "layer0":
                continue
            groups_dep: DefaultDict[str, List[Dict[str, Any]]] = defaultdict(list)
            for req in make_dependent(layer_agents.get(lname, [])):
                groups_dep[req["shell_url"]].append(req)

            for shell_url, reqs in groups_dep.items():
                out = await run_batch_requests(
                    shell_url,
                    independent_reqs=[],  # these are dependent; router fills {DENOM} from previous results saved server-side
                    dependent_reqs=reqs,
                    default_stagger_s=0.5,
                )
                for resp in out:
                    resp["final_answer"] = extract_boxed(resp.get("output", resp.get("text", "")))
                outs_all.extend(out)

    dt = time.time() - t0
    outs_all.append({
        "total_ttft_ms": sum(resp.get("summary", {}).get("ttft_ms", 0) for resp in outs_all if isinstance(resp, dict)),
        "total_e2e_time": dt
    })
    return outs_all

# =========================
# MAIN
# =========================

if __name__ == "__main__":
    random.seed(42)

    # 1) Load config
    cfg = load_config(CONFIG_PATH)

    RUN_BASELINE = cfg.get("run_baseline", True) and RUN_BASELINE_ENV
    RUN_PROPOSED = cfg.get("run_proposed", True) and RUN_PROPOSED_ENV
    WARMUP = cfg.get("warmup", WARMUP)

    # 2) Build agent graph; tokenizer from the first model (for token counts)
    layers_order, agent_map = build_agent_graph(cfg)
    first_agent = next(iter(agent_map.values()))
    tokenizer_model_key = first_agent.model_key
    # For tokenizer, we need a HF repo id. If your model_key is not a HF repo,
    # optionally provide cfg["tokenizer_repo"] to specify it.
    tokenizer_repo = cfg.get("tokenizer_repo", tokenizer_model_key)
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_repo, use_fast=True)

    # 3) Dataset
    dataset_name = cfg.get("dataset", DATASET).upper()

    if dataset_name == "MMLU":
        cats = MMLUCategoryDict.get(CATEGORY, ["math"])
        ds = load_mmlu_samples(cats, num_samples=NUM_SAMPLES, seed=42)
        samples = list(ds); gold_available = True
    elif dataset_name == "GSM8K":
        ds = load_gsm8k_samples(num_samples=NUM_SAMPLES, seed=42)
        samples = list(ds); gold_available = True
    elif dataset_name == "AIME25":
        ds = load_aime25_samples(num_samples=NUM_SAMPLES, seed=42)
        samples = list(ds); gold_available = True
    elif dataset_name in ("GPQA", "GPQA_DIAMOND", "GPQA-DIAMOND"):
        samples = load_gpqa_diamond_samples(num_samples=NUM_SAMPLES, seed=42)
        gold_available = True
    elif dataset_name in ("HMMT_FEB_2025", "HMMT2025", "HMMT"):
        ds = load_hmmt_feb_2025_samples(num_samples=NUM_SAMPLES, seed=42)
        samples = list(ds); gold_available = True
    elif dataset_name in ("MATH500", "MATH-500"):
        ds = load_math500_samples(num_samples=NUM_SAMPLES, seed=42)
        samples = list(ds); gold_available = True

    else:
        samples = [{"question": "What is 2+2? Put final answer in \\boxed{}."}]
        gold_available = False

    corrects_metric = 0
    metadata_metric = []
    total_ttft_metric = 0.0
    total_e2e_metric = 0.0

    corrects_metric_baseline = 0
    metadata_metric_baseline = []
    total_ttft_metric_baseline = 0.0
    total_e2e_metric_baseline = 0.0

    timestamp = int(time.time())
    dataset_tag = dataset_name.lower()

    if WARMUP:
        print("=== Warmup Run ===")
        # Randomly pick one sample for warmup
        samp = random.choice(samples)
        question = build_prompt_from_sample(dataset_name, samp)

        # Run baseline once
        try:
            for port in (30000, 30001, 30002, 31000, 31001, 31002):
                try:
                    requests.post(f"http://localhost:{port}/flush_cache", timeout=2)
                except Exception:
                    pass
        except Exception:
            print("Warning: Potential error in flushing cache...skip")
        time.sleep(2)
        asyncio.run(run_baseline_layers(
            question=question,
            layers_order=layers_order,
            agent_map=agent_map,
            model_dict=cfg.get("model_dict", {}),
            tokenizer=tokenizer,
            fixed_tokens_fallback=FIXED_TOKENS_DEFAULT,
            timestamp=timestamp,
            warmup=True,
        ))
        print("Warmup completed.")
        WARMUP = False

    # 4) Iterate samples
    for idx, samp in enumerate(samples, start=1):
        question = build_prompt_from_sample(dataset_name, samp)

        # ===== Baseline =====
        if RUN_BASELINE:
            print("\n=== Baseline: Sequential layer-by-layer ===")
            try:
                for port in (30000, 30001, 30002, 31000, 31001, 31002):
                    try:
                        requests.post(f"http://localhost:{port}/flush_cache", timeout=2)
                    except Exception:
                        pass
            except Exception:
                print("Warning: Potential error in flushing cache...skip")
            time.sleep(2)

            asyncio.run(run_baseline_layers(
                question=question,
                layers_order=layers_order,
                agent_map=agent_map,
                model_dict=cfg.get("model_dict", {}),
                tokenizer=tokenizer,
                fixed_tokens_fallback=FIXED_TOKENS_DEFAULT,
                timestamp=timestamp,
            ))

            with open(f"{timestamp}_tb_agent_tree_structure_metrics_baseline.json", "r") as f:
                results = json.load(f)

            total_ttft_metric_baseline += results["summary"]["total_ttft"]
            total_e2e_metric_baseline += results["summary"]["wall_time"]

            finals = results["final"]["outs"]
            pred = finals[0].get("final_answer", extract_boxed(finals[0].get("text", ""))) if finals else ""
            if gold_available:
                gold = gold_from_sample(dataset_name, samp)
                if _normalize_numeric(pred) == gold or pred == gold:
                    corrects_metric_baseline += 1
            else:
                gold = ""

            metadata_metric_baseline.append({
                'index': idx,
                'question': samp.get('question') or samp.get('problem'),
                'prediction': pred,
                'gold': gold,
                'correct': _normalize_numeric(pred) == gold if gold_available else None
            })

            current_accuracy = (corrects_metric_baseline / idx * 100.0) if gold_available else None
            summary = {
                'total': len(samples), 'current': idx, 'correct': corrects_metric_baseline,
                'accuracy': current_accuracy, 'total_ttft': total_ttft_metric_baseline,
                'total_e2e': total_e2e_metric_baseline
            }
            out_file = f"{timestamp}_tb_summary_agent_tree_structure_baseline_results.json"
            with open(out_file, 'w', encoding='utf-8') as f:
                json.dump({'summary': summary, 'details': metadata_metric_baseline}, f, indent=2)
            print(f"Baseline results written to {out_file}")

            os.rename(f"{timestamp}_tb_agent_tree_structure_metrics_baseline.json", f"{timestamp}_tb_agent_tree_structure_metrics_baseline_sample_{idx}.json")
            run_dir = f"results/{CONFIG_NAME}/{timestamp}_{dataset_tag}_tree_structure"
            os.makedirs(f"{run_dir}/baseline", exist_ok=True)
            shutil.move(f"{timestamp}_tb_agent_tree_structure_metrics_baseline_sample_{idx}.json", f"{run_dir}/baseline/")

        # ===== Proposed =====
        if RUN_PROPOSED:
            print("\n=== Proposed: Prefill Overlapping ===")
            try:
                for port in (30000, 30001, 30002, 31000, 31001, 31002):
                    try:
                        requests.post(f"http://localhost:{port}/flush_cache", timeout=2)
                    except Exception:
                        pass
            except Exception:
                print("Warning: Potential error in flushing cache...skip")
            time.sleep(2)

            outs = asyncio.run(run_proposed_pipelined(
                question=question,
                layers_order=layers_order,
                agent_map=agent_map,
                model_dict=cfg.get("model_dict", {}),
                fixed_tokens_fallback=FIXED_TOKENS_DEFAULT,
            ))

            with open(f"{timestamp}_tb_agent_tree_structure_metrics.json", "w") as f:
                json.dump(outs, f, indent=4)

            total_ttft_metric += outs[-1]["total_ttft_ms"] / 1000.0
            total_e2e_metric += outs[-1]["total_e2e_time"]

            # Heuristic: the last real response before the totals
            pred = ""
            for item in reversed(outs[:-1]):
                if isinstance(item, dict) and ("output" in item or "text" in item):
                    pred = item.get("final_answer", extract_boxed(item.get("output", item.get("text", ""))))
                    break

            if gold_available:
                gold = gold_from_sample(dataset_name, samp)
                if _normalize_numeric(pred) == gold:
                    corrects_metric += 1
            else:
                gold = ""

            metadata_metric.append({
                'index': idx,
                'question': samp.get('question') or samp.get('problem'),
                'prediction': pred,
                'gold': gold,
                'correct': _normalize_numeric(pred) == gold if gold_available else None
            })

            current_accuracy = (corrects_metric / idx * 100.0) if gold_available else None
            summary = {
                'total': len(samples), 'current': idx, 'correct': corrects_metric,
                'accuracy': current_accuracy, 'total_ttft': total_ttft_metric,
                'total_e2e': total_e2e_metric
            }
            out_file = f"{timestamp}_tb_summary_agent_tree_structure_results.json"
            with open(out_file, 'w', encoding='utf-8') as f:
                json.dump({'summary': summary, 'details': metadata_metric}, f, indent=2)
            print(f"Proposed results written to {out_file}")

            os.rename(f"{timestamp}_tb_agent_tree_structure_metrics.json", f"{timestamp}_tb_agent_tree_structure_metrics_sample_{idx}.json")
            run_dir = f"results/{CONFIG_NAME}/{timestamp}_{dataset_tag}_tree_structure"
            os.makedirs(f"{run_dir}/proposed", exist_ok=True)
            shutil.move(f"{timestamp}_tb_agent_tree_structure_metrics_sample_{idx}.json", f"{run_dir}/proposed/")

    run_dir = f"results/{CONFIG_NAME}/{timestamp}_{dataset_tag}_tree_structure"
    if RUN_PROPOSED:
        os.makedirs(run_dir, exist_ok=True)
        if os.path.exists(f"{timestamp}_tb_summary_agent_tree_structure_results.json"):
            shutil.move(f"{timestamp}_tb_summary_agent_tree_structure_results.json", run_dir)
    if RUN_BASELINE:
        os.makedirs(run_dir, exist_ok=True)
        if os.path.exists(f"{timestamp}_tb_summary_agent_tree_structure_baseline_results.json"):
            shutil.move(f"{timestamp}_tb_summary_agent_tree_structure_baseline_results.json", run_dir)

    print("All done.")