run_all.py

#!/usr/bin/env python3
"""
SysStream 一键评测启动脚本（shard 版）

基于 datasets/final/shards_mini/ 下的 12 个分片并行评测。
每个 shard 对应一个独立 tmux session，命名为 shard_XX。

分片映射:
    shard_00           → proactive（S1-solvable）
    shard_01 ~ 03      → clue（S2-required）
    shard_04 ~ 11      → mixed（realtime / reasoning / sequential 混合）

五个子任务集: proactive, clue, realtime, reasoning, sequential

用法:
    python run_all.py              # 生成 config + 启动所有 tmux session
    python run_all.py --dry-run    # 只生成 config，打印命令，不启动 tmux

附着到某个 session:
    tmux attach -t shard_00
    tmux attach -t shard_04
    ...

列出所有 session:
    tmux ls
"""

import copy
import json
import os
import shutil
import subprocess
import sys
import time
from collections import Counter

import yaml

Job = tuple[str, str, list[str]]
GpuAssignments = dict[str, int]

ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
SRC_DIR = os.path.join(ROOT_DIR, "src")
if SRC_DIR not in sys.path:
    sys.path.insert(0, SRC_DIR)

from src.utils.data import get_result_value, load_benchmark_data, load_data

# ── 全局配置 ─────────────────────────────────────────────────────────────────
def _detect_conda_root() -> str:
    env_root = os.environ.get("CONDA_ROOT")
    if env_root:
        return env_root

    conda_bin = shutil.which("conda")
    if conda_bin:
        try:
            return subprocess.check_output(
                [conda_bin, "info", "--base"],
                text=True,
            ).strip()
        except Exception:
            pass

    return os.path.expanduser("~/miniconda3")


CONDA_ROOT = _detect_conda_root()
CONDA_ENV  = "SysStream"
CFG_DIR    = os.path.join(SRC_DIR, "configs")
BASE_CFG   = os.path.join(SRC_DIR, "config.yaml")

# ── 数据集定义 ────────────────────────────────────────────────────────────────
# 路径相对于 src/ 目录（eval_from_config.py 的工作目录）
SHARD_DIRS = {
    "mini": "../datasets/final/shards_mini",
    "full": "../datasets/final/shards",
    "memory": "../datasets/final/memory_ablation",
}
DATA_DIR   = "../datasets/data"

# (session_name, shard_file, benchmark_type)
# shard_00: proactive, shard_01-03: clue, shard_04-11: mixed

JOBS = [
    ("shard_00", "shard_00.json", "proactive"),
    ("shard_01", "shard_01.json", "clue"),
    ("shard_02", "shard_02.json", "clue"),
    ("shard_03", "shard_03.json", "clue"),
    ("shard_04", "shard_04.json", "mixed"),
    ("shard_05", "shard_05.json", "mixed"),
    ("shard_06", "shard_06.json", "mixed"),
    ("shard_07", "shard_07.json", "mixed"),
    ("shard_08", "shard_08.json", "mixed"),
    ("shard_09", "shard_09.json", "mixed"),
    ("shard_10", "shard_10.json", "mixed"),
    ("shard_11", "shard_11.json", "mixed"),
]

_EXPECTED_QUESTION_CACHE: dict[str, tuple[int, Counter[tuple[str, ...]]]] = {}


def _apply_no_memory_overrides(cfg: dict, *, context_clip_enabled: bool = True) -> None:
    """Configure a generated shard for offline frame-only evaluation without memory."""
    dataset_cfg = cfg.setdefault("dataset", {})
    system_cfg = cfg.setdefault("system", {})
    memory_cfg = cfg.setdefault("memory", {})
    stream_cfg = cfg.setdefault("stream", {})
    embed_cfg = system_cfg.setdefault("streaming_embed_cache", {})

    dataset_cfg["frame_output_dir"] = "../batch_frames"

    system_cfg["s1_input_mode"] = "frames"
    system_cfg["s2_input_mode"] = "frames"
    system_cfg["s2_memory_conf_threshold"] = 1.1
    system_cfg["low_conf_visible_to_s2"] = False
    system_cfg["s2_see_current_s1_answer"] = False

    embed_cfg["enabled"] = False
    embed_cfg["prefill_on_frame"] = False
    embed_cfg["prefill_s2"] = False
    embed_cfg["kv_reuse_enabled"] = False
    embed_cfg["context_clip_enabled"] = context_clip_enabled

    memory_cfg["short_term_window"] = 0
    memory_cfg["top_k_retrieve"] = 0
    memory_cfg["keyframe_retrieve_top_k"] = 0
    memory_cfg["persist"] = False
    memory_cfg["clear_on_startup"] = True

    stream_cfg["enabled"] = False


def _resolve_from_src(path: str | None) -> str:
    if not path:
        return ""
    if os.path.isabs(path):
        return path
    return os.path.normpath(os.path.join(SRC_DIR, path))


def _question_signature(video_item: dict, question: dict, question_index: int) -> tuple[str, ...]:
    qid = question.get("qid")
    if qid is not None:
        return ("qid", str(qid))

    video_key = str(
        video_item.get("video_id")
        or video_item.get("video_path")
        or video_item.get("source")
        or ""
    )
    timestamp = question.get("ask_timestamp")
    if timestamp is None:
        timestamp = question.get("time_sec")
    if timestamp is None:
        timestamp = question.get("time_stamp") or question.get("timestamp")
    question_text = " ".join(str(question.get("question", "")).split())
    return (
        "fallback",
        video_key,
        str(question_index),
        str(timestamp),
        question_text,
    )


def _collect_question_signatures(data: list[dict]) -> Counter[tuple[str, ...]]:
    signatures: Counter[tuple[str, ...]] = Counter()
    for item in data:
        questions = item.get("questions", [])
        for question_index, question in enumerate(questions):
            if not isinstance(question, dict):
                continue
            signatures[_question_signature(item, question, question_index)] += 1
    return signatures


def _load_expected_question_signatures(
        cfg_path: str,
) -> tuple[int, Counter[tuple[str, ...]]]:
    cached = _EXPECTED_QUESTION_CACHE.get(cfg_path)
    if cached is not None:
        return cached

    with open(cfg_path, encoding="utf-8") as f:
        cfg = yaml.safe_load(f)

    dataset_cfg = cfg.get("dataset", {}) if isinstance(cfg, dict) else {}
    data_file = _resolve_from_src(dataset_cfg.get("data_file", ""))
    if not data_file:
        raise KeyError("dataset.data_file is missing")

    benchmark = dataset_cfg.get("benchmark")
    data_dir = _resolve_from_src(dataset_cfg.get("data_dir", "")) or None
    if benchmark:
        expected_data = load_benchmark_data(
            data_file=data_file,
            benchmark_type=benchmark,
            data_dir=data_dir,
        )
    else:
        expected_data = load_data(data_file)

    expected_signatures = _collect_question_signatures(expected_data)
    cached = (sum(expected_signatures.values()), expected_signatures)
    _EXPECTED_QUESTION_CACHE[cfg_path] = cached
    return cached


def generate_configs(
        shard_dir: str,
        *,
        cross_gpu_mode: bool = False,
        max_samples: int | None = None,
        context_clip_enabled: bool = True,
        resume: bool = True,
        allow_missing_shards: bool = False,
        output_prefix: str = "eval_results",
        no_memory_mode: bool = False,
) -> list[Job]:
    """根据 base config 为每个 shard 生成独立 config 文件，并构建 job 列表。

    Returns:
        list of (session_name, config_path, extra_cli_args)
    """
    os.makedirs(CFG_DIR, exist_ok=True)
    with open(BASE_CFG, encoding="utf-8") as f:
        base = yaml.safe_load(f)

    results = []
    for session, shard_file, benchmark in JOBS:
        data_file = f"{shard_dir}/{shard_file}"
        if not os.path.exists(_resolve_from_src(data_file)):
            if allow_missing_shards:
                print(f"  {session:12s} -> skip missing {data_file}")
                continue
            raise FileNotFoundError(f"Missing shard file: {data_file}")

        cfg = copy.deepcopy(base)
        cfg["dataset"]["data_file"]  = data_file
        cfg["dataset"]["data_dir"]   = DATA_DIR
        cfg["dataset"]["benchmark"]  = benchmark
        cfg["dataset"]["max_samples"] = max_samples
        cfg["system"]["streaming_embed_cache"]["context_clip_enabled"] = context_clip_enabled
        cfg["memory"]["storage_dir"] = f"../memory_storage/{session}"
        cfg["output"]["output_file"] = f"{output_prefix}_{session}.json"
        cfg["output"]["resume"] = resume
        if no_memory_mode:
            _apply_no_memory_overrides(cfg, context_clip_enabled=context_clip_enabled)

        model_cfg = cfg.get("model", {})
        if isinstance(model_cfg, dict):
            if cross_gpu_mode:
                # cross-GPU 模式不设置 CUDA_VISIBLE_DEVICES，让后端按 device_map=auto 分布。
                model_cfg["force_single_gpu"] = False
            elif model_cfg.get("force_single_gpu", False):
                # per-GPU 模式会把每个 shard 绑到单张物理卡；会话内可见索引通常只有 0。
                model_cfg["force_single_gpu"] = True

        out_path = os.path.join(CFG_DIR, f"config_{session}.yaml")
        with open(out_path, "w", encoding="utf-8") as f:
            yaml.dump(cfg, f, default_flow_style=False, allow_unicode=True, sort_keys=False)

        print(f"  {session:12s} -> config_{session}.yaml  bench={benchmark}")
        results.append((session, out_path, []))

    return results


def tmux_session_exists(session: str) -> bool:
    ret = subprocess.run(["tmux", "has-session", "-t", session],
                         capture_output=True)
    return ret.returncode == 0


def _gpu_instance_limit(
        gpu: int,
        max_instances_per_gpu: int | dict[int, int] | None,
) -> int:
    """Return the configured live-session limit for one GPU; 0 disables it."""
    if max_instances_per_gpu is None:
        return 0
    if isinstance(max_instances_per_gpu, dict):
        value = max_instances_per_gpu.get(gpu, 0)
    else:
        value = max_instances_per_gpu
    try:
        return max(0, int(value))
    except (TypeError, ValueError):
        return 0


def _active_sessions_on_gpu(
        jobs: list[Job],
        gpu: int,
        assigned_gpus: GpuAssignments,
) -> list[str]:
    """List known shard sessions currently alive on the given GPU."""
    active = []
    for session, _cfg, _extra in jobs:
        if assigned_gpus.get(session) == gpu and tmux_session_exists(session):
            active.append(session)
    return active


def _select_available_gpu(
        gpu_pool: list[int],
        known_jobs: list[Job],
        assigned_gpus: GpuAssignments,
        max_instances_per_gpu: int | dict[int, int] | None,
) -> int | None:
    """Pick the least-loaded GPU that still has an available instance slot."""
    best_gpu = None
    best_load = None
    for gpu in gpu_pool:
        active = _active_sessions_on_gpu(known_jobs, gpu, assigned_gpus)
        limit = _gpu_instance_limit(gpu, max_instances_per_gpu)
        if limit > 0 and len(active) >= limit:
            continue
        if best_load is None or len(active) < best_load:
            best_gpu = gpu
            best_load = len(active)
    return best_gpu


def _gpu_wait_summary(
        gpu_pool: list[int],
        known_jobs: list[Job],
        assigned_gpus: GpuAssignments,
        max_instances_per_gpu: int | dict[int, int] | None,
) -> str:
    parts = []
    for gpu in gpu_pool:
        active = _active_sessions_on_gpu(known_jobs, gpu, assigned_gpus)
        limit = _gpu_instance_limit(gpu, max_instances_per_gpu)
        limit_text = str(limit) if limit > 0 else "unlimited"
        sessions = ",".join(active) if active else "-"
        parts.append(f"gpu={gpu} active={len(active)}/{limit_text} sessions=[{sessions}]")
    return "; ".join(parts)


def _wait_for_available_gpu(
        gpu_pool: list[int],
        known_jobs: list[Job],
        assigned_gpus: GpuAssignments,
        max_instances_per_gpu: int | dict[int, int] | None,
        wait_interval: int,
) -> int:
    """Block before tmux creation until any GPU has one available slot."""
    wait_interval = max(1, int(wait_interval))
    while True:
        gpu = _select_available_gpu(
            gpu_pool,
            known_jobs,
            assigned_gpus,
            max_instances_per_gpu,
        )
        if gpu is not None:
            return gpu
        print(
            "  WAIT   all GPU slots full before launching next tmux: "
            f"{_gpu_wait_summary(gpu_pool, known_jobs, assigned_gpus, max_instances_per_gpu)}",
            flush=True,
        )
        time.sleep(wait_interval)


def _instance_limit(max_instances: int | None) -> int:
    """Return a process-count limit; 0 disables it."""
    try:
        return max(0, int(max_instances or 0))
    except (TypeError, ValueError):
        return 0


def _active_sessions(jobs: list[Job]) -> list[str]:
    """List known shard sessions currently alive, regardless of GPU."""
    return [
        session
        for session, _cfg, _extra in jobs
        if tmux_session_exists(session)
    ]


def _cross_gpu_slot_available(
        jobs: list[Job],
        max_instances: int | None,
) -> bool:
    limit = _instance_limit(max_instances)
    return limit <= 0 or len(_active_sessions(jobs)) < limit


def _cross_gpu_wait_summary(
        jobs: list[Job],
        max_instances: int | None,
) -> str:
    active = _active_sessions(jobs)
    limit = _instance_limit(max_instances)
    limit_text = str(limit) if limit > 0 else "unlimited"
    sessions = ",".join(active) if active else "-"
    return f"active={len(active)}/{limit_text} sessions=[{sessions}]"


def _wait_for_cross_gpu_slot(
        known_jobs: list[Job],
        max_instances: int | None,
        wait_interval: int,
) -> None:
    """Block before tmux creation until the global cross-GPU slot is available."""
    wait_interval = max(1, int(wait_interval))
    while not _cross_gpu_slot_available(known_jobs, max_instances):
        print(
            "  WAIT   cross-GPU slots full before launching next tmux: "
            f"{_cross_gpu_wait_summary(known_jobs, max_instances)}",
            flush=True,
        )
        time.sleep(wait_interval)


def launch_tmux(
        jobs: list[Job],
        launch_delay: int = 0,
        *,
        all_jobs: list[Job] | None = None,
        gpu_pool: list[int] | None = None,
        assigned_gpus: GpuAssignments | None = None,
        max_instances_per_gpu: int | dict[int, int] | None = None,
        cross_gpu_mode: bool = False,
        cross_gpu_max_instances: int | None = None,
        gpu_wait_interval: int = 30,
        block_on_gpu_limit: bool = True,
) -> list[str]:
    """每个 benchmark 开启一个独立的 tmux session。

    Args:
        launch_delay: 每次启动后等待的秒数，避免多进程同时加载模型导致 OOM。
                      推荐值 90~120 s（7B 模型从磁盘加载并转到 GPU 约需 60~90 s）。
        all_jobs: 用于统计同 GPU 已存活 session 的完整 job 列表。
        gpu_pool: 可动态分配的物理 GPU 列表。
        assigned_gpus: 当前进程内 session -> GPU 的动态分配表。
        max_instances_per_gpu: 每张 GPU 允许同时存活的 shard session 数；0/None 表示不限制。
        cross_gpu_mode: True 时不设置 CUDA_VISIBLE_DEVICES，允许单实例看到所有 GPU。
        cross_gpu_max_instances: cross-GPU 模式下允许同时存活的全局 shard session 数；0/None 表示不限制。
        gpu_wait_interval: 同卡 session 数达到上限时的等待轮询间隔。
        block_on_gpu_limit: True 时等到有 slot；False 时本轮跳过，交给监控循环重试。
    """
    conda_init = (
        f"source {CONDA_ROOT}/etc/profile.d/conda.sh "
        f"&& conda activate {CONDA_ENV}"
    )
    known_jobs = all_jobs or jobs
    pool = list(gpu_pool or [0])
    assignments = assigned_gpus if assigned_gpus is not None else {}

    launched = 0
    launched_sessions: list[str] = []
    for name, cfg_path, extra_args in jobs:
        rel_cfg = os.path.relpath(cfg_path, SRC_DIR)
        extra_str = " ".join(extra_args)

        if tmux_session_exists(name):
            print(f"  SKIP   session '{name}' (already exists)")
            continue

        if cross_gpu_mode:
            if block_on_gpu_limit:
                _wait_for_cross_gpu_slot(
                    known_jobs,
                    cross_gpu_max_instances,
                    gpu_wait_interval,
                )
                slot_available = True
            else:
                slot_available = _cross_gpu_slot_available(
                    known_jobs,
                    cross_gpu_max_instances,
                )
                if not slot_available:
                    print(
                        f"  WAIT   cross-GPU active limit before launching '{name}': "
                        f"{_cross_gpu_wait_summary(known_jobs, cross_gpu_max_instances)}",
                        flush=True,
                    )
                    continue
            gpu = None
            gpu_label = "ALL"
            cuda_prefix = ""
        else:
            if block_on_gpu_limit:
                gpu = _wait_for_available_gpu(
                    pool,
                    known_jobs,
                    assignments,
                    max_instances_per_gpu,
                    gpu_wait_interval,
                )
                slot_available = True
            else:
                gpu = _select_available_gpu(
                    pool,
                    known_jobs,
                    assignments,
                    max_instances_per_gpu,
                )
                slot_available = gpu is not None
                if not slot_available:
                    print(
                        f"  WAIT   no GPU slot before launching '{name}': "
                        f"{_gpu_wait_summary(pool, known_jobs, assignments, max_instances_per_gpu)}",
                        flush=True,
                    )
                    continue
            gpu_label = str(gpu)
            cuda_prefix = f"CUDA_VISIBLE_DEVICES={gpu} "
        if not slot_available:
            continue
        if tmux_session_exists(name):
            print(f"  SKIP   session '{name}' (already exists)")
            continue

        cmd = (
            f"{conda_init} "
            f"&& cd {SRC_DIR} "
            f"&& {cuda_prefix}python eval_from_config.py --config {rel_cfg} {extra_str}; "
            f"exit"
        )
        subprocess.run(
            ["tmux", "new-session", "-d", "-s", name, "-x", "220", "-y", "50"],
            check=True,
        )
        subprocess.run(
            ["tmux", "send-keys", "-t", name, cmd, "Enter"],
            check=True,
        )
        if not cross_gpu_mode and gpu is not None:
            assignments[name] = int(gpu)
        print(f"  LAUNCH session='{name}'  gpu={gpu_label}  config={rel_cfg}  {extra_str}")
        launched += 1
        launched_sessions.append(name)
        if launch_delay > 0 and launched < len(jobs):
            print(f"  WAIT   {launch_delay}s before next launch (avoid OOM during model load)...")
            time.sleep(launch_delay)
    return launched_sessions


def check_results_completeness(jobs: list[Job]) -> list[str]:
    """检查每个 shard 的结果文件完整性，返回不完整的 session 名列表。"""
    def _is_terminal_skip(question: dict) -> bool:
        eval_mode = get_result_value(question, "eval_mode")
        return isinstance(eval_mode, str) and eval_mode.startswith("skipped_")

    def _has_answer(question: dict) -> bool:
        if _is_terminal_skip(question):
            return True
        value = get_result_value(question, "s1_response")
        if value is None:
            return False
        if isinstance(value, str):
            return bool(value.strip())
        return True

    incomplete = []
    for session, cfg_path, extra_args in jobs:
        try:
            with open(cfg_path, encoding="utf-8") as f:
                cfg = yaml.safe_load(f)
            output_cfg = cfg.get("output", {}) if isinstance(cfg, dict) else {}
            results_dir = _resolve_from_src(output_cfg.get("results_dir", "../results"))
            result_file = os.path.join(
                results_dir,
                output_cfg.get("output_file", f"eval_results_{session}.json"),
            )
            expected_total_q, expected_signatures = _load_expected_question_signatures(cfg_path)
        except (OSError, KeyError, ValueError, TypeError, yaml.YAMLError, json.JSONDecodeError) as e:
            print(f"  {session:12s}  ERROR reading expected shard: {e}")
            incomplete.append(session)
            continue

        if not os.path.isfile(result_file):
            print(f"  {session:12s}  MISSING result file")
            incomplete.append(session)
            continue
        try:
            with open(result_file, encoding="utf-8") as f:
                data = json.load(f)
            total_q = sum(len(item.get("questions", [])) for item in data)
            result_signatures = _collect_question_signatures(data)

            if total_q != expected_total_q or result_signatures != expected_signatures:
                print(
                    f"  {session:12s}  MISMATCH current shard/result file "
                    f"({total_q}/{expected_total_q} questions)"
                )
                incomplete.append(session)
                continue

            answered = sum(
                sum(1 for q in item.get("questions", []) if _has_answer(q))
                for item in data
            )
            if answered < total_q:
                print(f"  {session:12s}  INCOMPLETE {answered}/{total_q} questions")
                incomplete.append(session)
            else:
                print(f"  {session:12s}  OK {answered}/{total_q}")
        except (json.JSONDecodeError, KeyError, AttributeError, TypeError) as e:
            print(f"  {session:12s}  ERROR reading result: {e}")
            incomplete.append(session)
    return incomplete


def filter_incomplete_jobs(
        jobs: list[Job],
) -> tuple[list[Job], list[str]]:
    """按结果完整性过滤待启动的 shard。"""
    incomplete_sessions = set(check_results_completeness(jobs))
    pending_jobs = [job for job in jobs if job[0] in incomplete_sessions]
    complete_sessions = [job[0] for job in jobs if job[0] not in incomplete_sessions]
    return pending_jobs, complete_sessions


def _fmt_duration(seconds: float) -> str:
    """将秒数格式化为 HH:MM:SS。"""
    h, rem = divmod(int(seconds), 3600)
    m, s = divmod(rem, 60)
    return f"{h:02d}:{m:02d}:{s:02d}"


def monitor_and_restart(
        jobs: list[Job],
        poll_interval: int = 30,
        max_retries: int = 3,
        launch_delay: int = 3,
        gpu_pool: list[int] | None = None,
        assigned_gpus: GpuAssignments | None = None,
        max_instances_per_gpu: int | dict[int, int] | None = None,
        cross_gpu_mode: bool = False,
        cross_gpu_max_instances: int | None = None,
        gpu_wait_interval: int = 30,
) -> None:
    """监控 tmux session 状态，异常退出自动重启，全部结束后检查完整性。

    Args:
        poll_interval: 轮询间隔（秒）
        max_retries: 每个 shard 最大重试次数
        launch_delay: 重启时的启动间隔（秒）
        gpu_pool: 可动态分配的物理 GPU 列表。
        assigned_gpus: 当前进程内 session -> GPU 的动态分配表。
        max_instances_per_gpu: 每张 GPU 允许同时存活的 shard session 数。
        cross_gpu_mode: True 时不设置 CUDA_VISIBLE_DEVICES，允许单实例看到所有 GPU。
        cross_gpu_max_instances: cross-GPU 模式下允许同时存活的全局 shard session 数。
        gpu_wait_interval: 同卡 session 数达到上限时的等待轮询间隔。
    """
    pool = list(gpu_pool or [0])
    assignments = assigned_gpus if assigned_gpus is not None else {}
    job_map = {session: (session, cfg, extra) for session, cfg, extra in jobs}
    retry_counts: dict[str, int] = {session: 0 for session, *_ in jobs}
    finished: set[str] = set()
    started: set[str] = {
        session for session, *_ in jobs
        if tmux_session_exists(session)
    }

    # 计时与重启记录
    start_times: dict[str, float] = {session: time.time() for session, *_ in jobs}
    elapsed: dict[str, float] = {}
    restart_log: list[str] = []

    print(f"\n==> [3/3] Monitoring {len(jobs)} sessions (poll every {poll_interval}s, max retries={max_retries}) ...")
    print()

    while len(finished) < len(jobs):
        time.sleep(poll_interval)

        for session, cfg_path, extra_args in jobs:
            if session in finished:
                continue

            if tmux_session_exists(session):
                # session 还在运行
                started.add(session)
                continue

            now = time.time()

            if session not in started:
                launched = launch_tmux(
                    [job_map[session]],
                    launch_delay=0,
                    all_jobs=jobs,
                    gpu_pool=pool,
                    assigned_gpus=assignments,
                    max_instances_per_gpu=max_instances_per_gpu,
                    cross_gpu_mode=cross_gpu_mode,
                    cross_gpu_max_instances=cross_gpu_max_instances,
                    gpu_wait_interval=gpu_wait_interval,
                    block_on_gpu_limit=False,
                )
                if launched:
                    started.add(session)
                    start_times[session] = now
                    if launch_delay > 0:
                        time.sleep(launch_delay)
                continue

            # session 不存在 → 已退出（正常或异常）
            # 检查结果是否完整
            incomplete = check_results_completeness([job_map[session]])
            if not incomplete:
                elapsed[session] = now - start_times[session]
                print(f"  DONE   {session} completed in {_fmt_duration(elapsed[session])}")
                finished.add(session)
                continue

            # 结果不完整 → 需要重启
            if retry_counts[session] >= max_retries:
                elapsed[session] = now - start_times[session]
                print(f"  GIVE UP {session} after {max_retries} retries ({_fmt_duration(elapsed[session])})")
                finished.add(session)
                continue

            launched = launch_tmux(
                [job_map[session]],
                launch_delay=0,
                all_jobs=jobs,
                gpu_pool=pool,
                assigned_gpus=assignments,
                max_instances_per_gpu=max_instances_per_gpu,
                cross_gpu_mode=cross_gpu_mode,
                cross_gpu_max_instances=cross_gpu_max_instances,
                gpu_wait_interval=gpu_wait_interval,
                block_on_gpu_limit=False,
            )
            if not launched:
                continue

            retry_counts[session] += 1
            attempt = retry_counts[session]
            ts = time.strftime("%Y-%m-%d %H:%M:%S")
            msg = f"[{ts}] {session} restart #{attempt}/{max_retries}"
            restart_log.append(msg)
            print(f"  RESTART {session} (attempt {attempt}/{max_retries}) ...")
            start_times[session] = now  # 重置计时
            if launch_delay > 0:
                time.sleep(launch_delay)

    total_elapsed = time.time() - min(start_times.values())

    # 最终报告
    print()
    print("=" * 60)
    print("==> Final Report")
    print("=" * 60)

    # 用时统计
    print("\n  Time per shard:")
    for session, *_ in jobs:
        dur = elapsed.get(session)
        retries = retry_counts[session]
        tag = f"  (retried {retries}x)" if retries else ""
        if dur is not None:
            print(f"    {session:12s}  {_fmt_duration(dur)}{tag}")
        else:
            print(f"    {session:12s}  N/A{tag}")
    print(f"\n  Total wall time: {_fmt_duration(total_elapsed)}")

    # 重启记录
    if restart_log:
        print(f"\n  Restart log ({len(restart_log)} events):")
        for entry in restart_log:
            print(f"    {entry}")

    # 完整性检查
    print("\n  Completeness:")
    incomplete = check_results_completeness(jobs)
    if incomplete:
        print(f"\n  WARNING: {len(incomplete)} shard(s) still incomplete: {', '.join(incomplete)}")
    else:
        print("\n  ALL shards completed successfully!")
    print("=" * 60)


def main() -> None:
    # ── 变量参数（直接修改这里即可）──────────────────────────────────────────
    dataset = "full"           # "mini" / "full" / "memory"
    eval_mode = "default"        # "default" / "no_memory"
    dry_run = False            # True: 只生成 config，不启动 tmux
    launch_delay = 3           # 两次启动间隔（秒）
    monitor = True             # True: 监控 session 状态并自动重启
    poll_interval = 60         # 监控轮询间隔（秒）
    max_retries = 5            # 每个 shard 最大重试次数
    cross_gpu_mode = False      # True: 不按单 GPU 分配；每个实例可看到全部 GPU
    cross_gpu_max_instances = 1 # cross_gpu_mode=True 时的全局实例上限；0/None 不限制
    gpu_pool = [0]     # 可动态分配的物理 GPU；启动时从有空位的 GPU 中选择
    max_instances_per_gpu = 2   # 每张 GPU 的 shard tmux session 上限；也可写 {0: 2, 1: 2}，0/None 不限制
    gpu_wait_interval = 120      # 达到同卡 session 上限时，等待多久后重新检查（秒）
    config_max_samples = None    # 覆盖 dataset.max_samples；None 表示不限制
    config_context_clip_enabled = True  # 覆盖 context_clip_enabled
    config_resume = True         # 覆盖 output.resume
    # ────────────────────────────────────────────────────────────────────────

    if dataset not in SHARD_DIRS:
        raise ValueError(f"Unknown dataset={dataset!r}; choose from {sorted(SHARD_DIRS)}")
    shard_dir = SHARD_DIRS[dataset]
    allow_missing_shards = dataset == "memory"
    no_memory_mode = eval_mode == "no_memory"
    if eval_mode not in {"default", "no_memory"}:
        raise ValueError("eval_mode must be 'default' or 'no_memory'")
    if no_memory_mode:
        output_prefix = (
            f"eval_results_{dataset}_no_memory"
            if dataset != "full"
            else "eval_results_no_memory"
        )
    else:
        output_prefix = "eval_results_memory" if dataset == "memory" else "eval_results"

    print(f"==> [1/2] Generating configs in {CFG_DIR}  (dataset={dataset}, mode={eval_mode}) ...")
    print(
        "    Config overrides: "
        f"max_samples={config_max_samples}, "
        f"context_clip_enabled={config_context_clip_enabled}, "
        f"resume={config_resume}, "
        f"output_prefix={output_prefix}"
    )
    jobs = generate_configs(
        shard_dir,
        cross_gpu_mode=cross_gpu_mode,
        max_samples=config_max_samples,
        context_clip_enabled=config_context_clip_enabled,
        resume=config_resume,
        allow_missing_shards=allow_missing_shards,
        output_prefix=output_prefix,
        no_memory_mode=no_memory_mode,
    )
    print()

    print("==> [2/3] Checking existing result completeness before launch ...")
    jobs_to_run, complete_sessions = filter_incomplete_jobs(jobs)
    if complete_sessions:
        print()
        print(f"  SKIP   completed shards ({len(complete_sessions)}): {', '.join(complete_sessions)}")

    if not jobs_to_run:
        print()
        print("All shards are already complete. Nothing to launch.")
        return

    print()
    print(f"  PENDING shards to launch: {len(jobs_to_run)}")
    if cross_gpu_mode:
        cross_limit = _instance_limit(cross_gpu_max_instances)
        cross_limit_text = str(cross_limit) if cross_limit > 0 else "unlimited"
        print(f"  GPU mode: cross-GPU  max active instances={cross_limit_text}")
    else:
        print(f"  GPU pool: {gpu_pool}")
    if not cross_gpu_mode and max_instances_per_gpu:
        print(f"  GPU session limit: max {max_instances_per_gpu} active shard(s) per GPU")

    if not shutil.which("tmux") and not dry_run:
        print("ERROR: tmux not found. Install tmux or use dry_run=True.", file=sys.stderr)
        sys.exit(1)
    conda_init_script = os.path.join(CONDA_ROOT, "etc", "profile.d", "conda.sh")
    if not os.path.exists(conda_init_script) and not dry_run:
        print(f"ERROR: conda init script not found: {conda_init_script}", file=sys.stderr)
        print("Set CONDA_ROOT to your conda base path, or make `conda info --base` available.", file=sys.stderr)
        sys.exit(1)

    if dry_run:
        print("==> [3/3] Dry-run: skipping tmux launch.")
        print("    Commands to run manually:")
        for name, cfg_path, extra_args in jobs_to_run:
            rel_cfg = os.path.relpath(cfg_path, SRC_DIR)
            extra_str = " ".join(extra_args)
            if cross_gpu_mode:
                print(f"      python eval_from_config.py --config {rel_cfg} {extra_str}")
            else:
                print(f"      CUDA_VISIBLE_DEVICES=<auto:{gpu_pool}> python eval_from_config.py --config {rel_cfg} {extra_str}")
        return

    print("==> [3/3] Launching tmux sessions ...")
    gpu_assignments: GpuAssignments = {}
    launch_tmux(
        jobs_to_run,
        launch_delay=launch_delay,
        all_jobs=jobs,
        gpu_pool=gpu_pool,
        assigned_gpus=gpu_assignments,
        max_instances_per_gpu=max_instances_per_gpu,
        cross_gpu_mode=cross_gpu_mode,
        cross_gpu_max_instances=cross_gpu_max_instances,
        gpu_wait_interval=gpu_wait_interval,
        block_on_gpu_limit=not monitor,
    )

    if not monitor:
        print()
        print("Sessions launched. Monitoring disabled.")
        print()
        print("List all sessions:   tmux ls")
        print("Attach to one:       tmux attach -t <name>")
        print("Detach:              Ctrl-b + d")
        return

    monitor_and_restart(
        jobs_to_run,
        poll_interval=poll_interval,
        max_retries=max_retries,
        launch_delay=launch_delay,
        gpu_pool=gpu_pool,
        assigned_gpus=gpu_assignments,
        max_instances_per_gpu=max_instances_per_gpu,
        cross_gpu_mode=cross_gpu_mode,
        cross_gpu_max_instances=cross_gpu_max_instances,
        gpu_wait_interval=gpu_wait_interval,
    )


if __name__ == "__main__":
    main()