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

Dark terminal showing four nested ASCII progress bars in NVIDIA green, labeled "Building Engine," "Tactic Selection," "Timing Cache Warmup," and "Kernel Autotune." Each row is indented to indicate nesting under its parent phase.

图 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 前进吧。返回 Falsefalse 以取消构建。
阶段结束 phase_finish(phase_name) phaseFinish(phaseName) 拆掉这一行。

表 1. IProgressMonitor 接口镜像在 Python 和 C++ 上。这三种方法具有相同的语义,并且 step_complete 是唯一返回值会改变构建器行为的回调函数

parent_phase 为非 null 的阶段会嵌套在另一个阶段内,因此显示器会看到进度树,而非平面列表。实现必须是线程安全的,因为 TensorRT 可以从多个内部线程调用同一监视器实例。

通过在 IBuilderConfig 上进行设置,将显示器连接到构建器。这两种语言的单次调用:

config.progress_monitor = MyMonitor()      # Python

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

UML sequence diagram between the TensorRT Builder and a user IProgressMonitor subclass showing nested phase_start calls, step_complete returning true (continue) or false (cancel), and a red-highlighted cancel path where the builder unwinds by calling phase_finish early on every active phase.

图 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

子类很小。它仅追踪哪些阶段处于活动状态,以及每个阶段包含多少步长。

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。构建器在下一步边界时展开,通常是快速的,但不是瞬间的,尤其是在漫长的战术搜索步骤中。

应用应向用户公布取消延迟。在放松窗口中显示一条简单的“Cancelling …” (取消 …) 消息会大有益。

您可以从任何非信号路径 (例如 IDE Stop 按钮、智能体超时或 CI cancel 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> 对于“Cancel” (取消) 标志很重要,因为 requestCancel() 可能会从其他线程或信号处理器调用。其他所有内容都与 Python 版本一致。

在实际系统中连接到何处

Three-layer architecture diagram with the TensorRT Builder at the top, IProgressMonitor in the middle, and four application sinks below: terminal renderer, IDE extension, FastAPI service, and agent runtime. A dashed red arrow shows the cancel signal flowing from a sink back through the monitor to the builder.

图 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 访问最终将崩溃或撕裂。

开始使用

运行此端到端的最快方法是:

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 的 monitor 类替换为上述版本,或在现有类周围附加取消处理程序。samples/sampleProgressMonitor/ 中提供了 C++ 等效函数。

对于较大的系统,下一步是将终端渲染器替换为应用已使用的传输,例如语言服务器协议通知、服务器发送事件或结构化工具调用块。IProgressMonitor 将成为 TensorRT 构建进度转换为应用进度模型的点。

了解详情

有关更多信息,请参阅以下资源:

Logo

分享最新的 NVIDIA AI Software 资源以及活动/会议信息,精选收录AI相关技术内容,欢迎大家加入社区并参与讨论。

更多推荐