第09章 - 边缘推理与端云协同部署

# edge_budget.py —— 把现场约束折算成可执行的配置上限
# 用法：python3 edge_budget.py（内置示例场景，改参数后可直接用于选型评审）
# 注意：所有系数都是工程经验值，必须用目标板实测校准后才能写进交付文档

from dataclasses import dataclass


@dataclass
class Device:
    name: str
    supply_w: float          # 供电上限：POE、电池放电能力、机巢供电三者取最小值
    board_base_w: float      # 板卡空载功耗（SoC + 内存 + 载板）
    peripheral_w: float      # 相机、风扇、4G、SSD 等外设峰值
    peak_tops: float         # 标称 AI 算力，以厂商数据手册为准
    sustained_ratio: float   # 持续算力系数：被动散热 0.5~0.7，主动 0.8~0.95
    nvdec_channels: int      # 硬件解码通道数，以 SoC 规格为准


def power_headroom(dev: Device, startup_ratio: float = 1.5, rail_v: float = 12.0) -> dict:
    """按 70% 负载率反推可用功耗，并估算启动涌流。"""
    budget = dev.supply_w * 0.7
    static = dev.board_base_w + dev.peripheral_w
    usable = budget - static
    if usable <= 0:
        return {
            "ok": False,
            "usable_w": round(usable, 1),
            "advice": "供电不足：减外设、换更高功率电源，或改用低功耗档",
        }
    return {
        "ok": True,
        "usable_w": round(usable, 1),
        "sustained_tops": round(dev.peak_tops * dev.sustained_ratio, 1),
        "startup_current_a": round(dev.supply_w * startup_ratio / rail_v, 2),
    }


def stream_limit(dev: Device, sustained_tops: float, per_stream_tops: float, reserve: float = 0.25) -> int:
    """并发上限取「算力允许」与「解码通道允许」的较小值。"""
    by_compute = int(sustained_tops * (1 - reserve) // per_stream_tops)
    return max(1, min(by_compute, dev.nvdec_channels))


def latency_budget(frame_interval_ms: float, decode_ms: float, preprocess_ms: float,
                   queue_ms: float, postproc_ms: float) -> dict:
    """反推留给推理的时延预算；结果为负说明配置本身不可行。"""
    inference_ms = frame_interval_ms - (decode_ms + preprocess_ms + queue_ms + postproc_ms)
    return {
        "inference_ms": round(inference_ms, 1),
        "ok": inference_ms > 0,
        "advice": "按此预算压模型与后处理" if inference_ms > 0 else "预算为负：先降分辨率或抽帧，再谈模型",
    }


if __name__ == "__main__":
    dev = Device("Orin NX 16G（示例）", 48.0, 9.0, 6.0, 100.0, 0.7, 2)
    p = power_headroom(dev)
    print(p)
    if p["ok"]:
        print("并发上限（路）:", stream_limit(dev, p["sustained_tops"], per_stream_tops=18.0))
    print(latency_budget(40.0, 6.0, 3.0, 4.0, 3.0))

# export_onnx.py —— 固定空间尺寸、动态 batch 的导出模板
import torch
from ultralytics import YOLO

PT = "runs/detect/edge/weights/best.pt"
OUT = "edge.onnx"
IMGSZ = 640

wrapper = YOLO(PT).model.float().eval()
dummy = torch.rand(1, 3, IMGSZ, IMGSZ)
with torch.no_grad():
    ref = wrapper(dummy)
ref = (ref[0] if isinstance(ref, (list, tuple)) else ref).numpy()

torch.onnx.export(
    wrapper,
    dummy,
    OUT,
    opset_version=17,
    input_names=["images"],
    output_names=["output0"],
    dynamic_axes={"images": {0: "batch"}, "output0": {0: "batch"}},
    do_constant_folding=True,
)
print("exported:", OUT, "ref max:", float(ref.max()))

# build_engine.py —— ONNX → TensorRT 11 引擎（FP16 / INT8 二选一）
# 依赖：tensorrt、pycuda、numpy；Jetson 镜像通常已带 tensorrt，pycuda 用 apt 或 pip 安装
# 版本要求：TensorRT 版本必须与 JetPack 7.2.x 配套，跨版本引擎不能混用
import argparse
import os

import numpy as np
import pycuda.autoinit  # noqa: F401  初始化 CUDA context
import pycuda.driver as cuda
import tensorrt as trt


class EntropyCalibrator(trt.IInt8EntropyCalibrator2):
    """INT8 熵校准：校准图必须是真实帧，且按现场分布采样。"""

    def __init__(self, paths, batch_size, cache_path):
        super().__init__()
        self.paths = paths
        self.batch_size = batch_size
        self.cache_path = cache_path
        self.index = 0
        self.device_buffer = None

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.index + self.batch_size > len(self.paths):
            return None
        chunk = self.paths[self.index:self.index + self.batch_size]
        batch = np.stack([np.load(p) for p in chunk]).astype(np.float32)
        self.index += self.batch_size
        if self.device_buffer is None:
            self.device_buffer = cuda.mem_alloc(batch.nbytes)
        cuda.memcpy_htod(self.device_buffer, batch)
        return [int(self.device_buffer)]

    def read_calibration_cache(self):
        if os.path.exists(self.cache_path):
            with open(self.cache_path, "rb") as f:
                return f.read()
        return None

    def write_calibration_cache(self, cache):
        with open(self.cache_path, "wb") as f:
            f.write(cache)


def build(onnx_path, engine_path, precision, batch, imgsz, calib_dir=None):
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)
    with open(onnx_path, "rb") as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print("parse error:", parser.get_error(i))
            raise SystemExit("ONNX 解析失败：先补插件或改等价结构")

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)

    input_name = network.get_input(0).name
    profile = builder.create_optimization_profile()
    profile.set_shape(input_name,
                      (1, 3, imgsz, imgsz),
                      (batch, 3, imgsz, imgsz),
                      (batch, 3, imgsz, imgsz))
    config.add_optimization_profile(profile)

    if precision == "fp16":
        config.set_flag(trt.BuilderFlag.FP16)
    elif precision == "int8":
        config.set_flag(trt.BuilderFlag.INT8)
        config.set_flag(trt.BuilderFlag.FP16)
        if not calib_dir:
            raise SystemExit("INT8 必须提供真实帧校准集目录")
        paths = sorted(os.path.join(calib_dir, n) for n in os.listdir(calib_dir) if n.endswith(".npy"))
        config.int8_calibrator = EntropyCalibrator(paths, batch, "calib.cache")

    serialized = builder.build_serialized_network(network, config)
    if serialized is None:
        raise SystemExit("构建失败：查 trtexec 日志里不支持的算子与显存是否够")
    with open(engine_path, "wb") as f:
        f.write(serialized)
    print("engine saved:", engine_path)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--onnx", required=True)
    ap.add_argument("--engine", required=True)
    ap.add_argument("--precision", choices=["fp16", "int8"], default="fp16")
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--imgsz", type=int, default=640)
    ap.add_argument("--calib-dir")
    args = ap.parse_args()
    build(args.onnx, args.engine, args.precision, args.batch, args.imgsz, args.calib_dir)

# precision_regression.py —— 张量级回归：引擎 vs ONNX，逐图比对原始输出
import argparse
import glob

import numpy as np
import onnxruntime as ort
import pycuda.driver as cuda
import tensorrt as trt


def load_engine(path):
    logger = trt.Logger(trt.Logger.ERROR)
    with open(path, "rb") as f:
        engine = trt.Runtime(logger).deserialize_cuda_engine(f.read())
    if engine is None:
        raise SystemExit("引擎反序列化失败：TensorRT 版本或 GPU 架构不匹配")
    return engine


def run_engine(engine, x):
    context = engine.create_execution_context()
    in_name = engine.get_tensor_name(0)
    context.set_input_shape(in_name, x.shape)
    out_name = engine.get_tensor_name(1)
    out_shape = tuple(context.get_tensor_shape(out_name))
    out = np.empty(out_shape, dtype=np.float32)
    d_in = cuda.mem_alloc(x.nbytes)
    d_out = cuda.mem_alloc(out.nbytes)
    cuda.memcpy_htod(d_in, x)
    context.set_tensor_address(in_name, int(d_in))
    context.set_tensor_address(out_name, int(d_out))
    stream = cuda.Stream()
    context.execute_async_v3(stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(out, d_out, stream)
    stream.synchronize()
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--engine", required=True)
    ap.add_argument("--onnx", required=True)
    ap.add_argument("--npy-dir", required=True)
    ap.add_argument("--tol", type=float, default=1e-2)
    args = ap.parse_args()

    engine = load_engine(args.engine)
    sess = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
    in_name = sess.get_inputs()[0].name
    files = sorted(glob.glob(args.npy_dir + "/*.npy"))
    if not files:
        raise SystemExit("没有找到 .npy 样本")

    diffs = []
    for p in files:
        x = np.load(p).astype(np.float32)
        if x.ndim == 3:
            x = x[None, ...]
        a = run_engine(engine, np.ascontiguousarray(x))
        b = sess.run(None, {in_name: x})[0]
        diffs.append(float(np.abs(a - b).max()))

    diffs = np.asarray(diffs)
    print("samples:", len(diffs), "max:", round(float(diffs.max()), 6),
          "p95:", round(float(np.percentile(diffs, 95)), 6), "tol:", args.tol)
    if diffs.max() > args.tol:
        raise SystemExit("张量级超差：先查预处理、输出节点与量化精度策略")


if __name__ == "__main__":
    main()

# pipeline_mp.py —— 多路视频流水线骨架（抓帧 → 有界队列 → 凑批推理 → 结果消费）
# 说明：骨架用 OpenCV 抓帧便于本地复现；上板后把 capture_worker 换成 GStreamer/NVDEC 管道，
#       把 engine.infer 换成 TensorRT execute_async_v3，其余结构不变。
import multiprocessing as mp
import time

import cv2
import numpy as np

FRAME_Q = 4          # 每路帧队列上限：满了丢最旧帧
RESULT_Q = 64        # 结果队列上限：满了丢结果并计数，绝不阻塞推理
BATCH_WAIT_MS = 10   # 凑批等待上限
BATCH_MAX = 8
URLS = ["rtsp://cam-01/live", "rtsp://cam-02/live", "rtsp://cam-03/live", "rtsp://cam-04/live"]


def capture_worker(url, out_q, every_n=2, drop_counter=None):
    cap = cv2.VideoCapture(url)
    idx = 0
    while True:
        ok, frame = cap.read()
        if not ok:
            cap.release()
            time.sleep(1.0)          # 断流重连退避，不要忙轮询
            cap = cv2.VideoCapture(url)
            continue
        idx += 1
        if idx % every_n:
            continue                 # 抽帧：25 fps 抽到 12.5 fps，推理压力减半
        item = (time.time(), frame)
        try:
            out_q.put_nowait(item)
        except Exception:
            try:
                out_q.get_nowait()   # 保新弃旧
            except Exception:
                pass
            try:
                out_q.put_nowait(item)
            except Exception:
                pass
            if drop_counter is not None:
                drop_counter.value += 1


def inference_worker(qs, result_q, engine, batch_wait_ms=BATCH_WAIT_MS):
    pending = []
    while True:
        for q in qs:
            try:
                pending.append(q.get_nowait())
            except Exception:
                continue
        if not pending:
            time.sleep(batch_wait_ms / 1000.0)
            continue
        batch, pending = pending[:BATCH_MAX], pending[BATCH_MAX:]
        t0 = time.time()
        frames = np.stack([cv2.resize(f, (640, 640)) for _, f in batch])
        out = engine.infer(frames)       # 目标板：TensorRT 动态 batch 引擎
        infer_ms = (time.time() - t0) * 1000
        for (ts, _), boxes in zip(batch, out):
            try:
                result_q.put_nowait((ts, boxes, infer_ms, len(batch)))
            except Exception:
                pass                     # 结果堆积时丢结果，保推理节拍


if __name__ == "__main__":
    ctx = mp.get_context("spawn")
    qs = [ctx.Queue(maxsize=FRAME_Q) for _ in URLS]
    result_q = ctx.Queue(maxsize=RESULT_Q)
    drop = ctx.Value("i", 0)
    procs = [ctx.Process(target=capture_worker, args=(u, q, 2, drop), daemon=True)
             for u, q in zip(URLS, qs)]
    procs.append(ctx.Process(target=inference_worker, args=(qs, result_q, engine), daemon=True))
    for p in procs:
        p.start()
    while True:
        ts, boxes, infer_ms, batch_n = result_q.get()
        e2e_ms = (time.time() - ts) * 1000
        # 这四行就是最小埋点集合：端到端时延、推理时延、批大小、丢帧累计
        print(f"e2e={e2e_ms:.0f}ms infer={infer_ms:.0f}ms batch={batch_n} drops={drop.value}")

# event_uploader.py —— 事件包落盘 + 断网补传（至少一次语义，靠 event_id 去重）
import hashlib
import json
import os
import time
import uuid

import cv2

SPOOL = "/data/edge/spool"


def pack_event(device_id, camera_id, frame, boxes, ts=None):
    """把一次告警打包成可上传的事件目录：image.jpg + meta.json"""
    os.makedirs(SPOOL, exist_ok=True)
    event_id = uuid.uuid4().hex
    event_dir = os.path.join(SPOOL, event_id)
    os.makedirs(event_dir, exist_ok=True)

    image_path = os.path.join(event_dir, "image.jpg")
    ok, buf = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 88])
    if not ok:
        raise RuntimeError("JPEG 编码失败")
    with open(image_path, "wb") as f:
        f.write(buf.tobytes())
    digest = hashlib.sha256(buf.tobytes()).hexdigest()

    meta = {
        "event_id": event_id,
        "device_id": device_id,
        "camera_id": camera_id,
        "captured_at": ts or time.time(),
        "boxes": boxes,
        "sha256": digest,
        "schema": "edge-event/1",
    }
    with open(os.path.join(event_dir, "meta.json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False)
    return event_dir


def flush_spool(upload_fn, limit=50):
    """按时间顺序补传；上传成功后删目录，失败保留等下次。"""
    if not os.path.isdir(SPOOL):
        return 0
    entries = [os.path.join(SPOOL, d) for d in os.listdir(SPOOL)]
    dirs = sorted(entries, key=os.path.getmtime)
    sent = 0
    for d in dirs[:limit]:
        try:
            upload_fn(d)                  # 幂等：服务端按 event_id 去重
        except Exception:
            break                         # 网络未恢复，保持顺序，先不跳过
        for name in os.listdir(d):
            os.remove(os.path.join(d, name))
        os.rmdir(d)
        sent += 1
    return sent

# telemetry.py —— 最小指标采集与上报（1 分钟一包，断网保留本地）
import json
import os
import time
import urllib.request

METRICS_PATH = "/data/edge/metrics.jsonl"
ENDPOINT = os.environ.get("EDGE_TELEMETRY_URL", "https://cloud.example.com/v1/edge/metrics")


def parse_tegrastats(line):
    """从 tegrastats 一行里取 RAM、GPU 频率、结温；字段缺失就跳过。"""
    parts = line.split()
    out = {}
    for i, tok in enumerate(parts[:-1]):
        if tok == "RAM":
            cur, total = parts[i + 1].replace("MB", "").split("/")
            out["ram_mb"] = int(cur)
            out["ram_total_mb"] = int(total)
        elif tok == "GR3D_FREQ":
            out["gpu_pct"] = int(parts[i + 1].replace("%", ""))
        elif tok.startswith("tj@"):
            out["tj_c"] = float(tok.replace("tj@", "").replace("C", ""))
    return out


def snapshot(device_id, pipeline_stats, tegrastats_line):
    rec = {"ts": int(time.time() * 1000), "device_id": device_id, "schema": "edge-metrics/1"}
    rec.update(pipeline_stats)        # queue_len / drop_rate / infer_p95_ms / alarm_delay_ms
    rec.update(parse_tegrastats(tegrastats_line))
    with open(METRICS_PATH, "a", encoding="utf-8") as f:
        f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    return rec


def flush():
    if not os.path.exists(METRICS_PATH):
        return
    with open(METRICS_PATH, encoding="utf-8") as f:
        rows = [line for line in f if line.strip()]
    if not rows:
        return
    body = ("[" + ",".join(rows) + "]").encode("utf-8")
    req = urllib.request.Request(ENDPOINT, data=body, headers={"Content-Type": "application/json"})
    try:
        urllib.request.urlopen(req, timeout=5).read()
        os.remove(METRICS_PATH)       # 成功才清空，失败保留等下一分钟
    except Exception:
        pass


if __name__ == "__main__":
    while True:
        stats = {"queue_len": 0, "drop_rate": 0.0, "infer_p95_ms": 0, "alarm_delay_ms": 0}
        with open("/tmp/tegrastats.txt", encoding="utf-8") as f:
            last = f.read().splitlines()[-1]
        snapshot("edge-001", stats, last)
        flush()
        time.sleep(60)