Python과 C++에서 NVIDIA TensorRT 엔진 빌드를 가시화하고 취소 가능하게 만드는 개발자 도구와 기법

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}", end="")

#include <unordered_map>
#include <mutex>
#include <atomic>

class RichProgressMonitor : public nvinfer1::IProgressMonitor {
public:
  void phaseStart(char const* phaseName, char const* parentPhase, int32_t nbSteps) noexcept override {
    std::lock_guard g(mu_);
    phases_[phaseName] = {nbSteps, 0, parentPhase ? parentPhase : ""};
    render();
  }
  bool stepComplete(char const* phaseName, int32_t step) noexcept override {
    std::lock_guard 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 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;
};

// Attach with:
// auto config = std::unique_ptr( builder->createBuilderConfig());
// RichProgressMonitor monitor; config->setProgressMonitor(&monitor);