让长时间运行的 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