#!/usr/bin/env python3
"""采样设备时钟/温度/功耗，跨后端归一化成同一组列。

为什么要有它：六轮 A/B 会声称单侧噪声在 0.01% 量级。降频和热漂移足以移动比这大
得多的延迟，而且在 run.log 里不留任何痕迹——报告照样干净。NV 上复现 PR763 的那次
全程 `nvidia-smi -l 5` 写 clocks.csv，S5000 这边一直没有对应的东西。

输出每设备每次采样一行 TSV：

    timestamp  index  sm_clock_mhz  mem_clock_mhz  temp_c  power_w  util_pct

列在 cuda 与 musa 上完全一致，所以两边的 campaign 可以并排读。采样器不存在时安静
退出：拿不到遥测是缺一份证据，不是让整批测试失败的理由。

用法：
    python telemetry.py --backend musa --device 7 --interval 5 --run > telemetry.tsv
    python telemetry.py --backend cuda            # 只打印将要执行的采样命令
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
import time
from typing import List, Optional

FIELDS = ("index", "sm_clock_mhz", "mem_clock_mhz", "temp_c", "power_w", "util_pct")


def resolve_device(backend, explicit: Optional[str]) -> str:
    """显式设备优先，否则从后端自己的可见性变量取。

    单值才用得上：`MUSA_VISIBLE_DEVICES=7` 指一张卡，`0,1,2` 指一组，后者不传
    `-i`、采全部，比挑第一张更诚实。
    """
    if explicit:
        return explicit.strip()
    for key in getattr(backend, "visibility_env_keys", ()):
        raw = os.environ.get(key, "").strip()
        if raw and "," not in raw:
            return raw
    return ""


def sample_once(backend, argv: List[str]) -> List[dict]:
    try:
        out = subprocess.run(argv, capture_output=True, text=True, timeout=30)
    except (OSError, subprocess.SubprocessError):
        return []
    return backend.parse_telemetry(out.stdout or "")


def run(backend, device: str, interval: float, out=sys.stdout) -> int:
    argv = backend.telemetry_command(device)
    if not argv:
        print(f"[telemetry] backend={backend.name} 没有采样器；本批不记时钟/温度",
              file=sys.stderr, flush=True)
        return 0
    if not sample_once(backend, argv):
        # 第一次就取不到，说明工具缺失或没权限。不要留一个空文件让人以为采过了
        print(f"[telemetry] {argv[0]} 取不到读数；本批不记时钟/温度",
              file=sys.stderr, flush=True)
        return 0
    print("timestamp\t" + "\t".join(FIELDS), file=out, flush=True)
    while True:
        ts = time.strftime("%Y-%m-%dT%H:%M:%S")
        for row in sample_once(backend, argv):
            print(ts + "\t" + "\t".join(row.get(f, "") for f in FIELDS),
                  file=out, flush=True)
        time.sleep(interval)


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"))
    p.add_argument("--device", default=os.environ.get("ZL_DEVICE", ""))
    p.add_argument("--interval", type=float, default=5.0)
    p.add_argument("--run", action="store_true",
                   help="持续采样直到被杀；不给则只打印采样命令")
    args = p.parse_args(argv)

    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from backends import get_backend
    backend = get_backend(args.backend)
    device = resolve_device(backend, args.device)
    if not args.run:
        print(" ".join(backend.telemetry_command(device)))
        return 0
    try:
        return run(backend, device, args.interval)
    except KeyboardInterrupt:
        return 0
    except BrokenPipeError:
        # 读端先走（管给 head，或 campaign 结束时收掉）。这是正常终止，不该在
        # telemetry.err 里留一段 traceback 让人以为遥测本身坏了。按 CPython 的
        # 惯用做法把 stdout 指到 devnull——只吞异常不够，解释器退出时还会 flush
        # 一次，再抛一条 "Exception ignored" 到 stderr。
        try:
            devnull = os.open(os.devnull, os.O_WRONLY)
            os.dup2(devnull, sys.stdout.fileno())
        except Exception:
            pass
        return 0


if __name__ == "__main__":
    raise SystemExit(main())