From 1b65a94747ab9294e11770629be3e20debf7972b Mon Sep 17 00:00:00 2001 From: yifeif-nv Date: Fri, 28 Aug 2026 10:55:48 -0700 Subject: [PATCH] fix(whisper): harden runner initialization Release partially initialized CUDA resources safely and import TP graph operations without relying on lazy package state during isolated CPU tests. Signed-off-by: yifeif-nv --- .../families/whisper/debug_runner.py | 33 +++++++++++-------- .../families/whisper/decoder_tp_builder.py | 5 +-- src/runtime/models/whisper/MODEL.toml | 2 +- .../whisper/test_whisper_builder_engine.py | 7 ++-- 4 files changed, 27 insertions(+), 20 deletions(-) diff --git a/python/tensorrt_model_connect/families/whisper/debug_runner.py b/python/tensorrt_model_connect/families/whisper/debug_runner.py index 59a6fc1f5b..2b10c58d96 100644 --- a/python/tensorrt_model_connect/families/whisper/debug_runner.py +++ b/python/tensorrt_model_connect/families/whisper/debug_runner.py @@ -317,17 +317,22 @@ def generate(self, input_ids: list[int], max_new_tokens: int) -> list[dict[str, def __del__(self): if cudart is None: return - if not hasattr(self, "_d_token_id"): - return - bufs = [self._d_token_id, self._d_position_id, self._d_mask, self._d_logits, - self._d_mel, self._d_enc_out] - bufs.extend(self._d_cache_k) - bufs.extend(self._d_cache_v) - bufs.extend(self._d_present_k) - bufs.extend(self._d_present_v) - bufs.extend(self._d_cross_k) - bufs.extend(self._d_cross_v) - for d_ptr in bufs: - cudart.cudaFree(d_ptr) - if hasattr(self, "stream"): - cudart.cudaStreamDestroy(self.stream) + scalar_buffers = ( + "_d_token_id", "_d_position_id", "_d_mask", "_d_logits", "_d_mel", "_d_enc_out", + ) + list_buffers = ( + "_d_cache_k", "_d_cache_v", "_d_present_k", "_d_present_v", "_d_cross_k", "_d_cross_v", + ) + for name in scalar_buffers: + d_ptr = getattr(self, name, None) + if d_ptr is not None: + cudart.cudaFree(d_ptr) + setattr(self, name, None) + for name in list_buffers: + for d_ptr in getattr(self, name, ()): + cudart.cudaFree(d_ptr) + setattr(self, name, []) + stream = getattr(self, "stream", None) + if stream is not None: + cudart.cudaStreamDestroy(stream) + self.stream = None diff --git a/python/tensorrt_model_connect/families/whisper/decoder_tp_builder.py b/python/tensorrt_model_connect/families/whisper/decoder_tp_builder.py index 0932dd53e5..fa114324a2 100644 --- a/python/tensorrt_model_connect/families/whisper/decoder_tp_builder.py +++ b/python/tensorrt_model_connect/families/whisper/decoder_tp_builder.py @@ -14,18 +14,19 @@ from __future__ import annotations -from typing import TYPE_CHECKING +import importlib import sys +from typing import TYPE_CHECKING import numpy as np from tensorrt_model_connect import trt_compat -from . import graph_ops from ...parallel_config import ( add_all_reduce_sum, normalize_parallel_config, ) +graph_ops = importlib.import_module(f"{__package__}.graph_ops") trt = trt_compat.get_trt() if TYPE_CHECKING: diff --git a/src/runtime/models/whisper/MODEL.toml b/src/runtime/models/whisper/MODEL.toml index adb6830e35..06c883fc5a 100644 --- a/src/runtime/models/whisper/MODEL.toml +++ b/src/runtime/models/whisper/MODEL.toml @@ -9,5 +9,5 @@ runtime_tests = [ "test_whisper_mel_spectrogram|test_whisper_mel_spectrogram.cpp|_|whisper_mel_spectrogram.cpp|_", "test_whisper_host_plan|test_whisper_host_plan.cpp|_|_|_", "test_whisper_decode_policy|test_whisper_decode_policy.cpp|_|_|_", - "test_whisper_pipeline|test_whisper_pipeline.cpp|trtmc_model_whisper|_|REQUIRES_TRT", + "test_whisper_pipeline|test_whisper_pipeline.cpp|trtmc_model_whisper|_|REQUIRES_TRT,REQUIRES_GPU", ] diff --git a/tests/e2e/models/whisper/test_whisper_builder_engine.py b/tests/e2e/models/whisper/test_whisper_builder_engine.py index c376409bce..09e72e9dec 100644 --- a/tests/e2e/models/whisper/test_whisper_builder_engine.py +++ b/tests/e2e/models/whisper/test_whisper_builder_engine.py @@ -34,6 +34,8 @@ from __future__ import annotations +import importlib + import numpy as np import pytest @@ -118,9 +120,8 @@ def rand(*shape: int) -> np.ndarray: def _whisper_tp_builder_module(): - return pytest.importorskip( - "tensorrt_model_connect.families.whisper.decoder_tp_builder", - reason="TensorRT is required for Whisper TP builder tests", + return importlib.import_module( + "tensorrt_model_connect.families.whisper.decoder_tp_builder" )