diff --git a/python/tests/test_tensorrt.py b/python/tests/test_tensorrt.py index 18f919a..70b2e2e 100644 --- a/python/tests/test_tensorrt.py +++ b/python/tests/test_tensorrt.py @@ -164,7 +164,9 @@ def test_tensorrt_default_cache_dir_uses_model_root_dir(monkeypatch, tmp_path): assert session.model_root_dir == cfg.model_root_dir assert engine_paths == [ - cfg.model_root_dir / "models" / "en_PP-OCRv4_det_mobile_sm87_fp16.engine" + cfg.model_root_dir + / "models" + / "en_PP-OCRv4_det_mobile_sm87_fp16_tf32unset.engine" ] @@ -198,7 +200,7 @@ def test_tensorrt_explicit_model_path_skips_download(monkeypatch, tmp_path): assert builder_kwargs[0]["onnx_path"] == model_path assert builder_kwargs[0]["engine_path"] == ( - tmp_path / "trt_cache" / "custom_sm87_fp32.engine" + tmp_path / "trt_cache" / "custom_sm87_fp32_tf32unset.engine" ) @@ -206,7 +208,9 @@ def test_tensorrt_loads_cached_engine_without_rebuild(monkeypatch, tmp_path): tensorrt_main = import_tensorrt_main(monkeypatch) cfg = make_tensorrt_cfg(tmp_path, engine_cfg={"cache_dir": tmp_path / "trt_cache"}) - engine_path = tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16.engine" + engine_path = ( + tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16_tf32unset.engine" + ) engine_path.parent.mkdir() engine_path.write_bytes(b"cached engine") @@ -234,7 +238,9 @@ def test_tensorrt_force_rebuild_ignores_cached_engine(monkeypatch, tmp_path): tmp_path, engine_cfg={"cache_dir": tmp_path / "trt_cache", "force_rebuild": True}, ) - engine_path = tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16.engine" + engine_path = ( + tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16_tf32unset.engine" + ) engine_path.parent.mkdir() engine_path.write_bytes(b"cached engine") @@ -271,7 +277,9 @@ def test_tensorrt_falls_back_to_rebuild_when_cached_engine_load_fails( tensorrt_main = import_tensorrt_main(monkeypatch) cfg = make_tensorrt_cfg(tmp_path, engine_cfg={"cache_dir": tmp_path / "trt_cache"}) - engine_path = tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16.engine" + engine_path = ( + tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16_tf32unset.engine" + ) engine_path.parent.mkdir() engine_path.write_bytes(b"bad cached engine") @@ -302,6 +310,38 @@ def test_tensorrt_falls_back_to_rebuild_when_cached_engine_load_fails( assert build_count == 1 +def test_tensorrt_context_creation_failure_raises_tensorrt_error(monkeypatch, tmp_path): + tensorrt_main = import_tensorrt_main(monkeypatch) + + class FailedContextEngine: + def create_execution_context(self): + return None + + class FakeEngineBuilder: + def __init__(self, **kwargs): + pass + + def build(self): + return FailedContextEngine() + + monkeypatch.setattr( + tensorrt_main.DownloadFile, + "run", + lambda input_params: Path(input_params.save_path).write_bytes(b"fake onnx"), + ) + monkeypatch.setattr(tensorrt_main, "TRTEngineBuilder", FakeEngineBuilder) + + cfg = make_tensorrt_cfg(tmp_path, engine_cfg={"cache_dir": tmp_path / "trt_cache"}) + + try: + tensorrt_main.TRTInferSession(cfg) + except tensorrt_main.TensorRTError as exc: + assert "Failed to create TensorRT execution context" in str(exc) + assert "NVIDIA_TF32_OVERRIDE=unset" in str(exc) + else: + raise AssertionError("Expected TensorRTError") + + def test_tensorrt_get_dict_key_url_uses_fallback_engines(monkeypatch): tensorrt_main = import_tensorrt_main(monkeypatch)