让长时间运行的 NVIDIA TensorRT 引擎在 Python 或 C++ 中构建可观察和可取消

构建 TensorRT 引擎可能需要几秒钟到几分钟的时间。大型强类型模型、深度策略搜索和全新 GPU SKU 上的冷定时缓存可能会让开发者、最终用户或 AI 智能体盯着冻结的终端,而不知道是否要等待、重试或终止该进程。大多数 NVIDIA TensorRT 集成在构建过程中不会报告任何内容,或者无法提前终止。在长时间运行的智能体工作流中,这会浪费 GPU 小时并导致会话卡顿。

图 1. 用于 TensorRT 引擎构建的实时嵌套进度条,由 IProgressMonitor 子类渲染
TensorRT 提供 IProgressMonitor (用于修复此问题的 API) ,并且多个版本都使用 NvInfer.h。本教程将介绍 Python 和 C++ 的最小插入式实现,添加响应 Ctrl-C 的取消路径或来自外部事件循环的编程停止信号,并展示在何处显示生成的进度流,以便 IDE、服务或代理运行时使用。
本文中的每个代码块均基于两个由 NVIDIA 维护的 OSS 示例进行提升或建模:
- Python:
samples/python/simple_progress_monitor/( ResNet-50,强类型网络) - C++:
samples/sampleProgressMonitor/( MNIST)
IProgressMonitor 为您提供的功能
IProgressMonitor 是 TensorRT 在引擎构建期间调用的抽象基类。您可以对其进行子类并覆盖三种方法。在 Python 和 C++ 中,形状相同;只有拼写不同。
| 概念 | Python 方法 | C++ 方法 | 您要做什么 |
|---|---|---|---|
| 阶段进入 | phase_start(phase_name, parent_phase, num_steps) |
phaseStart(phaseName, parentPhase, nbSteps) |
保留进度行并录制 num_steps。 |
| 阶段内步骤完成 | step_complete(phase_name, step) -> bool |
stepComplete(phaseName, step) -> bool |
前进吧。返回 False/ false 以取消构建。 |
| 阶段结束 | phase_finish(phase_name) |
phaseFinish(phaseName) |
拆掉这一行。 |
表 1. IProgressMonitor 接口镜像在 Python 和 C++ 上。这三种方法具有相同的语义,并且 step_complete 是唯一返回值会改变构建器行为的回调函数
parent_phase 为非 null 的阶段会嵌套在另一个阶段内,因此显示器会看到进度树,而非平面列表。实现必须是线程安全的,因为 TensorRT 可以从多个内部线程调用同一监视器实例。
通过在 IBuilderConfig 上进行设置,将显示器连接到构建器。这两种语言的单次调用:
|
|

图 2. TensorRT 在构建过程中驱动的回调序列,取消路径以红色突出显示
从上到下阅读图表。构建器使用 phase_start 打开构建引擎阶段,然后打开内嵌的战术选择,其中 parent_phase 指向构建引擎。随着构建过程的进行,构建器调用 step_complete (实体箭头) ,您的显示器将返回一个布尔值 (虚线箭头) :true 允许构建继续,false 则会请求取消。在此处显示的运行中,显示器在步骤 47 (即红色取消路径) 返回 false,构建器停止发出新步骤并展开。它在战术选择和构建引擎中提前调用 phase_finish,以反向顺序关闭每个活动阶段。
本教程构建的内容
本教程展示了如何在 Python 和 C++ 中实现 IProgressMonitor,通过 step_complete 添加取消,以及将进度更新路由到终端、IDE、服务或代理运行时。
预备知识
- 一个 NVIDIA GPU。
- TensorRT (当前 OSS 版本) 及其 Python 绑定,或 C++ 示例的构建。
- Python 3.10 或更高版本 ( Python 路径) 。
- TensorRT 样本数据:ResNet-50 ONNX (适用于 Python) 和 MNIST ONNX (适用于 C++) 。两者都随附样本数据存档,或安装在官方 NGC 容器中的
/usr/src/tensorrt/data下。 - 支持 ANSI 虚拟终端转义的终端。任何现代 Linux shell 都符合要求;如果启用了 VT,Windows 终端可以正常运行。
1. Python 中的子类 IProgressMonitor
子类很小。它仅追踪哪些阶段处于活动状态,以及每个阶段包含多少步长。
|
|
有两点需要注意。首先,Lock 不是可选项。TensorRT 将从多个内部线程调用监视器,而从不拥有状态的线程进行渲染将导致显示器撕裂。其次,step_complete 是唯一可以停止构建的回调函数。phase_start 返回 None,因此您无法在相位开始前拒绝该相位。最早的取消点是该相位的第一个 step_complete。
2. 使用虚拟终端转义渲染嵌套进度条
渲染器是因环境而异的部分,因此本节给出形状并指向用于生产级实现的上游示例。模式为:
|
|
上游 simple_progress_monitor.py 渲染相同的形状,并改进了颜色和宽度处理。转义序列 \x1b[NA 将光标向上移动 N 行,然后 \x1b[2K 清除一行。第一次渲染调用会写入空行;后续调用会将其覆盖到位。
连接此监视器时,请勿将标准输出重定向到文件或管道。转义码将逐字写入日志,使其不可读。对于非终端汇,将 _render() 替换为结构化发射器。
3. 添加取消路径
显示器存在后,取消是三行附加内容。安装可翻转标志的 SIGINT 处理程序,然后让 step_complete 遵守该标志。
|
|
连接显示器并运行构建器:
|
|
build_serialized_network() 在取消时返回 None。构建器在下一步边界时展开,通常是快速的,但不是瞬间的,尤其是在漫长的战术搜索步骤中。
应用应向用户公布取消延迟。在放松窗口中显示一条简单的“Cancelling …” (取消 …) 消息会大有益。
您可以从任何非信号路径 (例如 IDE Stop 按钮、智能体超时或 CI cancel webhook) 设置相同的标志。设置 monitor._cancelled = True,构建将在下一步边界处中止。
4。C++ 中的模式相同
|
|
以相同的方式连接:
|
|
std::atomic<bool> 对于“Cancel” (取消) 标志很重要,因为 requestCancel() 可能会从其他线程或信号处理器调用。其他所有内容都与 Python 版本一致。
在实际系统中连接到何处

图 3. IProgressMonitor 是构建器表面和应用表面之间的单一集成点
从 agent runtime 中提取取消箭头以确定具体程度,但相同的机制适用于每个汇点。终端的 Ctrl-C、IDE Stop 按钮、HTTP cancel webhook 或智能体 timeout 都会翻转相同的 monitor._cancelled 标志,并且取消将在下一个 step_complete 返回时生效。在真实系统中的连接位置
终端是最简单的例子。有趣的集成会在其他地方取得进展:
- IDE 扩展:覆盖
_render()以在语言服务器协议中发出$/progress通知,或在协议中发出同等的window/showProgress通知。每个阶段成为一个进度令牌;step_complete()成为报告消息;phase_finish()成为结束。 - FastAPI/ HTTP 服务:在后台线程上运行构建,并让
_render()将条目推送到请求处理程序通过服务器发送事件排空的asyncio.Queue中。客户端将获得一个实时串流;cancel hook 只是一个调用monitor.requestCancel()的POST /builds/{id}/cancel。 - 代理工具调用:在每个阶段转换 (
{"phase": ..., "step": ..., "total": ...}) 中,向工具调用流发送一个结构化数据块。智能体运行时在用户可见追踪中进行渲染,而相同的requestCancel()hook 是智能体超时在构建超出预算时调用的数据。这种模式对于智能体运行时也很重要。长时间运行的构建需要是可观察和可取消的,以便智能体能够报告进度、执行时间预算,并彻底停止。
在这三种情况下,IProgressMonitor 都是正确的边界。高于此级别的内容 (渲染、流式传输、传输) 是应用级的;低于此级别的内容 (策略计时、内核选择) 是构建者的业务。
要处理的边缘案例
这些行为是集成错误的常见来源:
- 连接终端渲染器时,请勿重定向
stdout。转义序列会污染日志。对于非交互式汇点,将渲染器替换为结构化发射器。 phase_start()无法取消。返回None。最早的取消点是该阶段的第一个step_complete()。如果用户在很长的phase_start()期间取消,构建将继续进行,直到第一步边界。phase_finish()可能会在所有num_steps报告之前发射。这可能会在错误恢复、构建器内部短路或step_complete()返回False期间发生。将其视为权威的相位结束信号;不要假设current_step == num_steps。- 取消延迟是有限的,但不是零。构建器在检查返回值之前完成当前步骤。长时间的战术搜索步骤可以将其推向数秒到几十秒的范围。
- 线程安全性是必需的。从多个构建器线程调用同一监视器实例;从
_render()进行的未被分析的dict或unordered_map访问最终将崩溃或撕裂。
开始使用
运行此端到端的最快方法是:
|
|
这将启动 ResNet-50 引擎的实时动画构建。将 simple_progress_monitor.py 的 monitor 类替换为上述版本,或在现有类周围附加取消处理程序。samples/sampleProgressMonitor/ 中提供了 C++ 等效函数。
对于较大的系统,下一步是将终端渲染器替换为应用已使用的传输,例如语言服务器协议通知、服务器发送事件或结构化工具调用块。IProgressMonitor 将成为 TensorRT 构建进度转换为应用进度模型的点。
了解详情
有关更多信息,请参阅以下资源:
更多推荐





所有评论(0)