NVIDIA TensorRT: monitorear y cancelar builds de motores

IProgressMonitor, disponible en NVIDIA TensorRT desde hace varias versiones, permite un seguimiento fino y seguro entre hilos del progreso, además de la cancelación durante la construcción de un motor. Es clave para manejar tiempos de build largos y evitar horas de GPU desperdiciadas.

La API exige sobrescribir tres métodos (phase_start, step_complete, phase_finish), tanto en Python como en C++, para seguir fases anidadas, actualizar el progreso y atender solicitudes de cancelación en los límites de cada paso. Al integrarla, el progreso puede reportarse en tiempo real a terminales, IDEs, servicios HTTP o entornos de agentes, desacoplando la lógica de build del renderizado de la aplicación.

Un build de TensorRT puede tardar desde segundos hasta varios minutos. Los modelos grandes con tipado estricto, una búsqueda profunda de tácticas y una caché de temporización fría en una GPU recién estrenada pueden dejar a desarrolladores, usuarios finales o agentes de IA mirando una terminal congelada, sin saber si conviene esperar, reintentar o matar el proceso. La mayoría de las integraciones de TensorRT no reportan nada durante un build ni ofrecen forma de abortar temprano. En un flujo de agente de larga duración, eso se traduce en horas de GPU tiradas y sesiones trabadas.

TensorRT provee IProgressMonitor, una API para resolver esto, presente en NvInfer.h desde hace varias versiones. Este tutorial recorre una implementación mínima y lista para usar en Python y C++, agrega un camino de cancelación que responde a Ctrl-C o a una señal de parada programática desde un bucle de eventos externo, y muestra dónde exponer el flujo de progreso resultante para que un IDE, un servicio o un entorno de agente lo aproveche.

Cada bloque de código de este artículo está tomado o modelado sobre dos ejemplos de código abierto mantenidos por NVIDIA:

• Python: samples/python/simple_progress_monitor/ (ResNet-50, red con tipado estricto).
• C++: samples/sampleProgressMonitor/ (MNIST).

¿Qué te entrega IProgressMonitor?

IProgressMonitor es una clase base abstracta que TensorRT invoca durante el build del motor. Uno crea una subclase y sobrescribe tres métodos. La forma es idéntica en Python y C++, solo cambia la escritura. Una fase cuyo parent_phase no es nulo queda anidada dentro de otra, de modo que el monitor ve un árbol de progreso en lugar de una lista plana. La implementación debe ser segura entre hilos, porque TensorRT puede llamar a la misma instancia del monitor desde varios hilos internos.

Se conecta el monitor al builder asignándolo en el IBuilderConfig. Es una sola llamada en cualquiera de los dos lenguajes:

config.progress_monitor = MyMonitor()      # Python

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

Lee el diagrama de arriba hacia abajo. El builder abre la fase Building Engine con phase_start, y luego abre Tactic Selection anidada dentro, con su parent_phase apuntando de vuelta a Building Engine. A medida que avanza el build, el builder llama a step_complete y tu monitor devuelve un booleano: true deja continuar y false solicita la cancelación. En la corrida mostrada, el monitor devuelve false en el paso 47 (el camino rojo de cancelación) y el builder deja de emitir pasos y se repliega, cerrando cada fase activa en orden inverso.

¿Qué construye este tutorial?

Muestra cómo implementar IProgressMonitor en Python y C++, agregar la cancelación mediante step_complete y enrutar las actualizaciones de progreso hacia una terminal, un IDE, un servicio o un entorno de agente.

Requisitos previos

• Una GPU NVIDIA.
• TensorRT (versión OSS actual) y sus bindings de Python, o un build de los ejemplos en C++.
• Python 3.10 o más nuevo (para la ruta de Python).
• Los datos de ejemplo de TensorRT: ResNet-50 ONNX para Python y MNIST ONNX para C++.
• Una terminal que soporte secuencias de escape ANSI. Cualquier shell moderno de Linux sirve, y Windows Terminal funciona si VT está habilitado.

1. Crear una subclase de IProgressMonitor en Python

La subclase es pequeña. Solo lleva registro de qué fases están activas y cuántos pasos contiene cada una.

import tensorrt as trt
from dataclasses import dataclass
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()

Dos cosas a notar. Primero, el Lock no es opcional: TensorRT llamará al monitor desde varios hilos internos, y renderizar desde un hilo que no es dueño del estado corromperá la vista. Segundo, step_complete es el único callback capaz de detener el build. phase_start devuelve None, así que no puedes rechazar una fase antes de que empiece: el punto de cancelación más temprano es el primer step_complete de esa fase.

2. Agregar un camino de cancelación

La cancelación es un agregado de tres líneas una vez que el monitor existe. Instala un manejador de SIGINT que active la bandera, y deja que step_complete la respete.

import signal

def install_cancel(monitor: RichProgressMonitor):
    def handler(signum, frame):
        monitor._cancelled = True
        print("\nCancelando el build de TensorRT en el proximo limite de paso...")
    signal.signal(signal.SIGINT, handler)

build_serialized_network() devuelve None al cancelar. El builder se repliega en el siguiente límite de paso, normalmente rápido pero no instantáneo, sobre todo dentro de un paso largo de búsqueda de tácticas. La misma bandera puede activarse desde cualquier ruta ajena a una señal, como un botón Stop de un IDE, un timeout de agente o un webhook de cancelación de CI.

3. El mismo patrón en C++

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 requestCancel() noexcept { cancelled_.store(true); }
};

Usar std::atomic<bool> para la bandera de cancelación importa, porque requestCancel() puede invocarse desde otro hilo o desde un manejador de señales. Todo lo demás refleja la versión en Python.

Dónde integrarlo en sistemas reales

El valor real aparece cuando el progreso deja de vivir solo en una terminal. Para equipos que arriendan GPU por hora en la nube, un build cancelable se traduce en ahorro directo: abortar una compilación fallida en el paso 47 en lugar de esperar a que termine evita pagar minutos de cómputo inútil. El mismo flujo alimenta paneles de un IDE, servicios HTTP que transmiten el avance o runtimes de agentes que deciden por sí solos cuándo cortar un build que se demora demasiado.