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

config.progress_monitor = MyMonitor()      # Python

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

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()

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)

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

#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);

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