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

构建一个 TensorRT 引擎可能需要几秒钟到几分钟不等。大型强类型模型、深度策略搜索，以及在全新 GPU SKU 上冷启动的计时缓存，可能会让开发者、最终用户或 AI 代理盯着一个冻结的终端，不知道是该等待、重试还是终止进程。大多数 NVIDIA TensorRT 集成在构建过程中不报告任何信息，或者不提供提前中止的方法。在长时间运行的代理工作流中，这会变成浪费的 GPU 小时数和卡住的会话。

图 1. TensorRT 引擎构建的实时嵌套进度条，由 IProgressMonitor 子类渲染

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++

图 2. TensorRT 在构建期间驱动的回调序列，取消路径以红色高亮显示

从上到下阅读该图。builder 使用 phase_start 打开 Building Engine 阶段

，然后在其内部打开嵌套的 Tactic Selection，其 parent_phase

指回 Building Engine。随着构建的进行，builder 调用 step_complete

（实线箭头），你的监视器返回一个布尔值（虚线箭头）：true

让构建继续，false

请求取消。在此处显示的运行中，监视器在第 47 步返回 false

，这就是红色的取消路径，builder 停止发出新步骤并展开。它在 Tactic Selection 上提前调用 phase_finish

，然后在 Building Engine 上调用，按相反顺序关闭每个活动阶段。

本教程构建的内容

本教程展示如何在 Python 和 C++ 中实现 IProgressMonitor

，通过 step_complete 添加取消

，并将进度更新路由到终端、IDE、服务或代理运行时。

先决条件

• 一块 NVIDIA GPU。
• TensorRT（当前 OSS 版本）及其 Python 绑定，或 C++ 示例的构建。
• Python 3.10 或更高版本（Python 路径）。
• TensorRT 示例数据：用于 Python 的 ResNet-50 ONNX 和用于 C++ 的 MNIST ONNX。两者都随示例数据存档提供，或挂载在 /usr/src/tensorrt/data

在官方 NGC 容器中。 - 支持 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):
# 按嵌套深度对阶段排序，使子阶段绘制在父阶段下方。
rows = sorted(
self._phases.items(),
key=lambda kv: (kv[1].parent or "", kv[0]),
)
# 将光标上移 PREVIOUS 渲染所打印的行数，
# 而不是当前的行数——阶段在嵌套时添加、在
# phase_finish 时移除，因此两者仅在树结构发生变化时不同。
if self._rendered_lines:
print(f"\x1b[{self._rendered_lines}A", end="")
for name, st in rows:
# step 是 [0, num_steps) 范围内的从 0 开始的索引；+1 将其转换为
# 已完成计数，这样进度条才能真正达到 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}")
# 清除阶段完成且计数减少时遗留的行。
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

清除一行。第一次渲染调用写入空白行；后续调用会就地覆盖它们。

当此监视器被附加时，不要将 stdout 重定向到文件或管道。转义码将被逐字写入日志，使其无法阅读。对于非终端接收端，请将 _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()

以在语言服务器协议中发出 $/progress

通知，或等效的 window/showProgress

。每个阶段变成一个进度令牌；step_complete()

变成报告消息；phase_finish()

变成结束。FastAPI / HTTP 服务：在后台线程上运行构建，并让 _render()

将条目推入 asyncio.Queue

，请求处理程序通过服务器发送事件排空该队列。客户端获得实时流；取消钩子只是一个 POST /builds/{id}/cancel

，它调用 monitor.requestCancel()

。代理工具调用：在每个阶段转换时发出一个结构化块（{"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/

中找到。

对于更大的系统，下一步正确的做法是将终端渲染器替换为应用程序已经使用的传输方式，例如 Language Server Protocol 通知、服务器发送事件或结构化工具调用块。IProgressMonitor

成为将 TensorRT 构建进度转换为应用程序进度模型的切入点。

了解更多

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