让 TensorRT 引擎构建可观察、可取消：Python 与 C++ 实战

TensorRT 引擎构建可能耗时数秒，也可能长达数分钟。强类型的大模型、深度的策略（tactic）搜索，再加上全新 GPU SKU 上冷启动的计时缓存，都可能让开发者、终端用户或 AI 智能体盯着一个卡住的终端，不知道该继续等待、重试还是直接杀掉进程。大多数 NVIDIA TensorRT 集成在构建期间不会报告任何信息，也没有提前中止的手段。在长时间运行的智能体工作流中，这会变成白白浪费的 GPU 小时和卡死的会话。

TensorRT 提供了 IProgressMonitor，这是专门解决该问题的 API，并且已在 NvInfer.h 中存在了好几个版本。本教程将演示适用于 Python 和 C++ 的最小化即插即用实现，加入可响应 Ctrl-C 或外层事件循环发出的程序化停止信号的取消路径，并说明应在何处暴露由此产生的进度流，以便 IDE、服务或智能体运行时使用。

本文中的每一段代码都摘自或仿照 NVIDIA 维护的两个开源示例：

• Python：samples/python/simple_progress_monitor/（ResNet-50，强类型网络）
• C++：samples/sampleProgressMonitor/（MNIST）

IProgressMonitor 能为你提供什么

IProgressMonitor 是一个抽象基类，TensorRT 会在引擎构建期间调用它。你需要继承它并重写三个方法。Python 与 C++ 中的结构完全一致，只是写法不同。

表 1. IProgressMonitor 接口在 Python 与 C++ 中的对应关系。三个方法的语义完全相同，step_complete 是唯一一个其返回值会改变构建器行为的回调

如果某个阶段的 parent_phase 非空，则它嵌套在另一个阶段之内，因此监视器看到的是一棵进度树，而不是扁平的列表。该实现必须是线程安全的，因为 TensorRT 可能从多个内部线程调用同一个监视器实例。

通过在 IBuilderConfig 上设置监视器，将其接入构建器。两种语言都只需一次调用：

config.progress_monitor = MyMonitor()      # Python

config->setProgressMonitor(&myMonitor);     // C++

请自上而下阅读该图。构建器先用 phase_start 打开 Building Engine 阶段，然后在其内部打开嵌套的 Tactic Selection 阶段，其 parent_phase 指回 Building Engine。构建推进过程中，构建器会调用 step_complete（实线箭头），而你的监视器返回一个布尔值（虚线箭头）：true 表示让构建继续，false 表示请求取消。在图中所示的这次运行里，监视器在第 47 步返回 false，即红色的取消路径，构建器随即停止发出新的步骤并开始回退。它先对 Tactic Selection 提前调用 phase_finish，再对 Building Engine 调用，按相反顺序关闭所有活动阶段。

本教程要构建什么

本教程演示如何在 Python 和 C++ 中实现 IProgressMonitor，通过 step_complete 加入取消功能，并将进度更新路由到终端、IDE、服务或智能体运行时。

前置条件

• 一块 NVIDIA GPU。
• TensorRT（当前的开源版本）及其 Python 绑定，或者已构建好的 C++ 示例。
• Python 3.10 或更高版本（Python 路径）。
• TensorRT 示例数据：Python 使用 ResNet-50 ONNX，C++ 使用 MNIST ONNX。二者都随示例数据压缩包提供，或挂载在官方 NGC 容器的 /usr/src/tensorrt/data 下。
• 支持 ANSI 虚拟终端转义序列的终端。任何现代 Linux shell 都满足要求；在启用 VT 的情况下，Windows Terminal 也可以。

1. 在 Python 中继承 IProgressMonitor

这个子类很小，只需跟踪哪些阶段处于活动状态，以及每个阶段包含多少步骤。

import tensorrt as trt
from dataclasses import dataclass, field
from threading import Lock
 
@dataclass
class _PhaseState:
    num_steps: int
    current_step: int = 0
    parent: str | None = None
 
class RichProgressMonitor(trt.IProgressMonitor):
    def __init__(self):
        super().__init__()
        self._lock = Lock()
        self._phases: dict[str, _PhaseState] = {}
        self._cancelled = False
	     self._rendered_lines = 0
 
    def phase_start(self, phase_name, parent_phase, num_steps):
        with self._lock:
        	self._phases[phase_name] = _PhaseState(
            	num_steps=num_steps, parent=parent_phase
        	)
        	self._render()
 
    def step_complete(self, phase_name, step) -> bool:
        with self._lock:
        	if phase_name in self._phases:
                self._phases[phase_name].current_step = step
        	self._render()
        	return not self._cancelled
 
    def phase_finish(self, phase_name):
        with self._lock:
        	self._phases.pop(phase_name, None)
        	self._render()

有两点需要注意。第一，Lock 不是可选的。TensorRT 会从多个内部线程调用监视器，如果在不持有状态的线程中渲染，显示就会错乱。第二，step_complete 是唯一能够停止构建的回调。phase_start 返回 None，因此你无法在阶段开始之前拒绝它。最早的取消时机是该阶段的第一次 step_complete。

2. 使用虚拟终端转义序列渲染嵌套进度条

渲染器是最容易因环境而异的部分，因此本节只给出基本形态，并指向上游示例以获取生产级实现。模式如下：

def _render(self):
    # Order phases by nesting depth so children draw under parents.
    rows = sorted(
        self._phases.items(),
        key=lambda kv: (kv[1].parent or "", kv[0]),
    )
    # Move the cursor up by the number of lines the PREVIOUS render printed,
    # not the current row count — phases are added on nesting and removed on
    # phase_finish, so the two differ exactly when the tree changes shape.
    if self._rendered_lines:
        print(f"\x1b[{self._rendered_lines}A", end="")
    for name, st in rows:
        # step is a 0-based index in [0, num_steps); +1 turns it into a
        # completed count so the bar can actually reach 100%.
        done = min(st.current_step + 1, st.num_steps)
        pct = done / max(st.num_steps, 1)
        bar = "█" * int(40 * pct) + "·" * (40 - int(40 * pct))
        indent = "  " if st.parent else ""
        print(f"\x1b[2K{indent}{name:<28} [{bar}] {done}/{st.num_steps}")
    # Clear rows left behind when a phase finishes and the count shrinks.
    for _ in range(self._rendered_lines - len(rows)):
        print("\x1b[2K")
    self._rendered_lines = len(rows)

上游的 simple_progress_monitor.py 渲染出相同的形态，并改进了颜色与宽度的处理。转义序列 \x1b[NA 将光标上移 N 行，\x1b[2K 清除一行。第一次渲染调用会写入空白行，之后的调用则原地覆盖这些行。

挂载此监视器后，不要把标准输出重定向到文件或管道，否则转义码会被原样写入日志，使其无法阅读。对于非终端的输出目标，请把 _render() 替换为结构化的输出器。

3. 加入取消路径

监视器就绪后，取消功能只需增加三行代码。安装一个翻转标志位的 SIGINT 处理器，再让 step_complete 遵从该标志。

import signal

def install_cancel(monitor: RichProgressMonitor):
    def handler(signum, frame):
        monitor._cancelled = True
        print("\nCancelling TensorRT build at next step boundary...")

    signal.signal(signal.SIGINT, handler)

接入监视器并运行构建器：

builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
)

parser = trt.OnnxParser(network, TRT_LOGGER)

with open(onnx_path, "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()

monitor = RichProgressMonitor()
config.progress_monitor = monitor

install_cancel(monitor)

serialized = builder.build_serialized_network(network, config)

if serialized is None:
    if monitor._cancelled:
        print("Build cancelled cleanly.")
    else:
        print("Build failed.")

取消时 build_serialized_network() 返回 None。构建器会在下一个步骤边界回退，通常很快，但并非瞬时完成，尤其是在耗时较长的策略搜索步骤中。

应用应当向用户展示取消延迟。在回退期间显示一条简单的“正在取消……”提示就很有帮助。

同一个标志位可以通过任何非信号路径设置，例如 IDE 的停止按钮、智能体超时或 CI 的取消 webhook。只要设置 monitor._cancelled = True，构建就会在下一个步骤边界中止。

4. C++ 中的相同模式

#include <NvInfer.h>
#include <atomic>
#include <mutex>
#include <unordered_map>

class RichProgressMonitor : public nvinfer1::IProgressMonitor {
public:
    void phaseStart(char const* phaseName,
                    char const* parentPhase,
                    int32_t nbSteps) noexcept override {
        std::lock_guard<std::mutex> g(mu_);
        phases_[phaseName] = {nbSteps, 0, parentPhase ? parentPhase : ""};
        render();
    }

    bool stepComplete(char const* phaseName,
                      int32_t step) noexcept override {
        std::lock_guard<std::mutex> g(mu_);
        auto it = phases_.find(phaseName);
        if (it != phases_.end())
            it->second.current = step;
        render();
        return !cancelled_.load();
    }

    void phaseFinish(char const* phaseName) noexcept override {
        std::lock_guard<std::mutex> g(mu_);
        phases_.erase(phaseName);
        render();
    }

    void requestCancel() noexcept {
        cancelled_.store(true);
    }

private:
    struct Phase {
        int32_t nbSteps;
        int32_t current;
        std::string parent;
    };

    std::mutex mu_;
    std::unordered_map<std::string, Phase> phases_;
    std::atomic<bool> cancelled_{false};

    void render() noexcept;
};

以同样的方式挂载：

auto config =
    std::unique_ptr<nvinfer1::IBuilderConfig>(
        builder->createBuilderConfig());

RichProgressMonitor monitor;

config->setProgressMonitor(&monitor);

取消标志使用 std::atomic<bool> 很重要，因为 requestCancel() 可能从另一个线程或信号处理器中被调用。其余部分都与 Python 版本一致。

在真实系统中的接入位置

图 3. IProgressMonitor 是构建器与应用各界面之间唯一的集成点

为了具体说明，取消箭头是从 agent runtime 画出的，但同样的机制适用于每一种输出目标。终端中的 Ctrl-C、IDE 的停止按钮、HTTP 取消 webhook 或智能体超时，都会翻转同一个 monitor._cancelled 标志，取消在下一次 step_complete 返回时生效。在真实系统中的接入位置

终端是最简单的情形。真正有意思的集成是把进度路由到别处：

• IDE 扩展：重写 _render()，发出语言服务器协议（Language Server Protocol）中的 $/progress 通知，或协议中等价的 window/showProgress。每个阶段对应一个进度令牌；step_complete() 变为一条报告消息；phase_finish() 变为结束。
• FastAPI / HTTP 服务：在后台线程中运行构建，并让 _render() 把条目推入一个 asyncio.Queue，由请求处理器通过 Server-Sent Events 取出并发送。客户端得到实时流；取消钩子就是一个调用 monitor.requestCancel() 的 POST /builds/{id}/cancel。
• 智能体工具调用：每次阶段切换（{"phase": ..., "step": ..., "total": ...}）向工具调用流中发出一个结构化数据块。智能体运行时会把它渲染到用户可见的轨迹中，而当构建超出预算时，智能体超时调用的正是同一个 requestCancel() 钩子。这一模式对智能体运行时同样重要。长时间运行的构建需要可观察、可取消，这样智能体才能汇报进度、执行时间预算并干净地停止。

在这三种情形下，IProgressMonitor 都是合适的边界。它之上的一切（渲染、流式传输、传输层）属于应用层；它之下的一切（策略计时、内核选择）则是构建器自己的事。

需要处理的边界情况

以下这些行为是集成缺陷的常见来源：

• 挂载终端渲染器时，不要重定向 stdout。转义序列会污染日志。对于非交互式输出目标，请把渲染器换成结构化的输出器。
• phase_start() 无法取消，它返回 None。最早的取消点是该阶段的第一次 step_complete()。如果用户在耗时较长的 phase_start() 期间取消，构建会一直持续到第一个步骤边界。
• phase_finish() 可能在所有 num_steps 都上报之前就触发。这可能发生在错误恢复、构建器内部短路，或 step_complete() 返回 False 的时候。请将其视为阶段结束的权威信号；不要假定 current_step == num_steps。
• 取消延迟有上限，但不为零。构建器会先完成当前步骤，再检查返回值。耗时较长的策略搜索步骤可能把这一延迟推高到数秒乃至数十秒。
• 线程安全是必需的。同一个监视器实例会被多个构建器线程调用；来自 _render() 的未加保护的 dict 或 unordered_map 访问最终会导致崩溃或显示错乱。

快速上手

端到端运行的最快方式是：

git clone --depth 1 https://github.com/NVIDIA/TensorRT.git
cd TensorRT/samples/python/simple_progress_monitor
python3 simple_progress_monitor.py

这会启动一个 ResNet-50 引擎的实时动画构建。将 simple_progress_monitor.py 的监视器类替换为上文的版本，或在现有类外围附加一个取消处理器。C++ 的等价实现见 samples/sampleProgressMonitor/。

对于更大的系统，下一步的正确做法是把终端渲染器替换为应用已在使用的传输方式，例如语言服务器协议通知、Server-Sent Events 或结构化的工具调用数据块。IProgressMonitor 就成为把 TensorRT 构建进度转换为应用自身进度模型的那个节点。

了解更多

更多信息请参阅以下资源：

• TensorRT Python API 文档：IProgressMonitor
• TensorRT C++ API 文档：nvinfer1::IProgressMonitor
• Python simple_progress_monitor 示例
• C++ sampleProgressMonitor 示例
• TensorRT GitHub 发布版本与示例数据