用Python或C++讓NVIDIA TensorRT引擎構建過程可觀測、可取消

TensorRT引擎的構建過程可能從幾秒到數分鐘不等。對于大型強類型模型、深度策略搜索以及全新GPU SKU上冷啟動的計時緩存，開發者、終端用戶或AI智能體往往只能面對一個卡住的終端，完全不知道該繼續等待、重試還是直接終止進程。大多數NVIDIA TensorRT集成在構建過程中既不報告任何進度，也不提供提前中止的方式。在長時運行的智能體工作流中，這會造成GPU算力的浪費和會話的阻塞。

TensorRT提供了IProgressMonitor接口來解決這一問題，該API已在NvInfer.h中存在多個版本。本文將介紹一種適用于Python和C++的最簡集成實現，并新增響應Ctrl-C或外部事件循環程序化停止信號的取消路徑，同時展示如何將進度流暴露給IDE、服務或智能體運行時使用。

IProgressMonitor接口簡介

IProgressMonitor是TensorRT在引擎構建過程中調用的抽象基類。開發者需要繼承該類并重寫三個方法，其結構在Python和C++中完全一致，只是語法寫法略有差異。

當某個階段的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"和"Building Engine"調用phase_finish，按相反順序關閉所有活躍階段。

環境準備

使用本文所述功能，需滿足以下條件：

一塊NVIDIA GPU；

TensorRT（當前OSS版本）及其Python綁定，或已編譯的C++示例；

Python 3.10及以上版本（Python路徑）；

TensorRT示例數據：Python使用ResNet-50 ONNX，C++使用MNIST ONNX，二者均包含在示例數據壓縮包中，或掛載于NGC官方容器的/usr/src/tensorrt/data目錄下；

支持ANSI虛擬終端轉義序列的終端，任何現代Linux Shell均可，Windows Terminal需啟用VT支持。

Python實現

以下是一個精簡的子類實現，僅追蹤當前活躍階段及各階段的步驟數量：

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。

進度渲染實現

渲染邏輯因使用環境不同而存在較大差異，其核心模式如下：

def _render(self):

rows = sorted(

self._phases.items(),

key=lambda kv: (kv[1].parent or "", kv[0]),

if self._rendered_lines:

print(f"\\x1b[{self._rendered_lines}A", end="")

for name, st in rows:

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

for _ in range(self._rendered_lines - len(rows)):

print("\\x1b[2K")

self._rendered_lines = len(rows)

注意：當此監控器處于附加狀態時，請勿將stdout重定向至文件或管道，否則轉義碼會被直接寫入日志，導致日志不可讀。對于非終端輸出目標，應將_render()替換為結構化發送器。

添加取消功能

取消功能的添加只需三行代碼，安裝SIGINT處理器將標志位置為True，再由step_complete響應該標志：

import signal

def install_cancel(monitor: RichProgressMonitor):

def handler(signum, frame):

monitor._cancelled = True

print("\Cancelling TensorRT build at next step boundary...")

signal.signal(signal.SIGINT, handler)

綁定監控器并運行構建器：

builder = trt.Builder(TRT_LOGGER)

network = builder.create_network(

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，構建即在下一個步驟邊界處中止。

C++實現

#include

#include

#include

#include

class RichProgressMonitor : public nvinfer1::IProgressMonitor {

public:

void phaseStart(char const* phaseName,

char const* parentPhase,

int32_t nbSteps) noexcept override {

std::lock_guardg(mu_);

phases_[phaseName] = {nbSteps, 0, parentPhase ? parentPhase : ""};

render();

bool stepComplete(char const* phaseName,

int32_t step) noexcept override {

std::lock_guardg(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_guardg(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_mapphases_;

std::atomiccancelled_{false};

void render() noexcept;

綁定方式與Python一致：

auto config =

std::unique_ptr(

builder->createBuilderConfig());

RichProgressMonitor monitor;

config->setProgressMonitor(&monitor);

取消標志使用std::atomic至關重要，因為requestCancel()可能從另一個線程或信號處理器中調用。其余部分與Python版本完全對應。

真實系統中的集成方式

IProgressMonitor是構建器與應用層之間唯一的集成點，取消操作的觸發來源可以是終端Ctrl-C、IDE停止按鈕、HTTP取消Webhook或智能體超時，所有路徑都通過同一個monitor._cancelled標志生效。

以下是幾種常見的集成場景：

IDE擴展：重寫_render()以發送語言服務器協議中的$/progress通知，每個階段對應一個進度Token，step_complete()對應report消息，phase_finish()對應end消息。

FastAPI/HTTP服務：在后臺線程運行構建，_render()將條目推入asyncio.Queue，請求處理器通過Server-Sent Events消費，客戶端獲得實時流，取消通過POST /builds/{id}/cancel調用monitor.requestCancel()實現。

智能體工具調用：每次階段切換時向工具調用流發送一條結構化數據塊，格式為{"phase": ..., "step": ..., "total": ...}，智能體運行時在用戶可見的追蹤中渲染，同一個requestCancel()鉤子供智能體超時時調用。該模式對智能體運行時尤為重要，長時間運行的構建需要可觀測、可取消，使智能體能夠匯報進度、執行時間預算并干凈地停止。

在所有上述場景中，IProgressMonitor都是正確的邊界：其上（渲染、流式傳輸、傳輸協議）屬于應用層，其下（策略計時、內核選擇）屬于構建器內部。

常見集成問題

以下是容易引發集成Bug的典型情形：

不要在終端渲染器處于附加狀態時重定向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/。

對于更大規模的系統，下一步應將終端渲染器替換為應用程序已有的傳輸方式，例如語言服務器協議通知、服務器推送事件或結構化工具調用數據塊。IProgressMonitor將成為TensorRT構建進度轉化為應用程序進度模型的接入點。

Q&A

Q1：IProgressMonitor接口有哪些核心方法，各自的作用是什么？

A：IProgressMonitor有三個核心方法：phase_start在一個構建階段開始時被調用，傳入階段名稱、父階段和步驟總數，返回None；step_complete在每個步驟完成時被調用，返回布爾值，返回true表示繼續構建，返回false則請求取消；phase_finish在階段結束時被調用。由于TensorRT會從多個內部線程調用監控器，三個方法的實現必須保證線程安全。

Q2：TensorRT引擎構建的取消操作為什么會有延遲？

A：取消操作的延遲來自構建器的內部機制。構建器不會立即響應取消信號，而是在完成當前步驟后才檢查step_complete的返回值。如果正處于耗時較長的策略搜索步驟中，這個延遲可能達到數秒甚至數十秒。因此取消不是即時生效的，建議應用程序在觸發取消后向用戶顯示"正在取消..."等提示，說明系統正處于回退窗口期。

Q3：如何在FastAPI服務中集成TensorRT構建進度監控？

A：在FastAPI服務中，可以在后臺線程運行TensorRT引擎構建，同時重寫監控器的_render()方法，將每次進度更新推入asyncio.Queue。請求處理器通過Server-Sent Events持續消費隊列中的數據，客戶端即可獲得實時進度流。取消功能通過新增一個POST /builds/{id}/cancel接口實現，該接口調用monitor.requestCancel()將取消標志置為True，構建器在下一個步驟邊界處響應并停止。