TensorRTの長時間エンジン構築を可視化・中止

NVIDIAは、TensorRTの長時間エンジン構築に進捗表示と中止操作を加えるIProgressMonitorの実装方法を解説しました。

要点

• PythonとC++でphase_start、step_complete、phase_finishを実装します。
• step_completeから偽を返すと、次の安全な区切りで構築中止を要求できます。
• 端末、IDE、HTTPサービス、エージェント実行基盤へ進捗を転送できます。

今回のブログ記事で語られていること

TensorRTのエンジン構築は、モデル規模、厳密な型指定、戦略探索、新しいGPUでの冷たいタイミングキャッシュによって数秒から数十分かかります。進捗が見えなければ、利用者やAIエージェントは処理が停止したのか、待つべきか、再試行すべきかを判断できません。GPU時間を使い続け、外側のセッションも停止したように見えます。

IProgressMonitorは構築器から呼び出される抽象基底で、段階開始、ステップ完了、段階終了を受け取ります。段階は親子関係を持つため、全体構築の中に戦略選択などの子段階を表示できます。TensorRT内部の複数スレッドから同じ監視器が呼ばれる可能性があり、実装はスレッド安全でなければなりません。監視器はIBuilderConfigへ設定し、進捗を外部表示へ流します。

中止はstep_completeの戻り値で要求します。Ctrl-Cや外側のイベントループから停止フラグを設定し、次のステップ境界で偽を返すと、TensorRTは新しい作業を始めず、開いている段階を終了します。ただし即時停止ではなく、長い一ステップの途中では待つ場合があります。導入時は、進捗更新の頻度、標準出力の扱い、複数スレッド、途中終了時の一時成果物、再実行の条件を試してください。公式のPythonとC++サンプルを基準に、自社のサービス監視へ接続するのが安全です。

サービスへ組み込む場合は、利用者に見せる百分率と内部の段階名を分け、同じモデル・設定の過去時間から残り時間を推定します。進捗通知自体が大量のログやロック競合を生まないよう、更新を間引き、監視器の例外が構築を壊さない実装にします。中止要求を受けた時刻と実際に停止した時刻も記録すると、長いステップの改善対象を見つけられます。

再実行では、破損したエンジンを公開しない原子的な保存と、タイミングキャッシュを再利用できる条件が重要です。ワーカー終了、GPU障害、利用者中止を区別し、同じ要求を自動再試行してよいかを状態ごとに決める必要があります。

今回のブログ記事が関係する人

TensorRTを組み込む推論基盤開発者、GPU費用を管理する担当、長時間処理を実行するIDEやエージェント基盤担当に関係します。

結局、今回のブログ記事をどう読むべきか

新しい推論性能ではなく、構築作業を運用可能にするAPI解説です。可視化と中止を入れることで、長いGPU処理を待つか止めるかを外側から判断できます。