Revised according to comments

Author: cehongwang

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -65,7 +65,7 @@ class UnsupportedOperatorException(RuntimeError):
 
 
 class TRTInterpreterResult(NamedTuple):
-    engine: trt.ICudaEngine | bytes
+    engine: trt.ICudaEngine
     input_names: Sequence[str]
     output_names: Sequence[str]
     weight_name_map: Optional[dict[Any, Any]]
@@ -770,29 +770,13 @@ def run(
         ):
             self._insert_engine_to_cache(hash_val, cuda_engine)
 
-        if self.compilation_settings.use_python_runtime:
-            return TRTInterpreterResult(
-                cuda_engine,
-                self._input_names,
-                self._output_names,
-                self.weight_name_map,
-                self.ctx.requires_output_allocator,
-            )
-        else:
-            serialized_engine = cuda_engine.serialize()
-            _LOGGER.info(f"TRT Engine uses: {serialized_engine.nbytes} bytes of Memory")
-
-            with io.BytesIO() as engine_bytes:
-                engine_bytes.write(serialized_engine)
-                engine_str = engine_bytes.getvalue()
-
-            return TRTInterpreterResult(
-                engine_str,
-                self._input_names,
-                self._output_names,
-                self.weight_name_map,
-                self.ctx.requires_output_allocator,
-            )
+        return TRTInterpreterResult(
+            cuda_engine,
+            self._input_names,
+            self._output_names,
+            self.weight_name_map,
+            self.ctx.requires_output_allocator,
+        )
 
     def run_node(self, n: torch.fx.Node) -> torch.fx.Node:
         self._cur_node_name = get_node_name(n)

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -89,18 +89,12 @@ def convert_module(
         module, inputs, settings, engine_cache=engine_cache
     )
 
+    rt_cls = PythonTorchTensorRTModule
+
     if ENABLED_FEATURES.torch_tensorrt_runtime and not settings.use_python_runtime:
         from torch_tensorrt.dynamo.runtime import TorchTensorRTModule
 
-        return TorchTensorRTModule(
-            serialized_engine=interpreter_result.engine,
-            input_binding_names=list(interpreter_result.input_names),
-            output_binding_names=list(interpreter_result.output_names),
-            name=name,
-            settings=settings,
-            weight_name_map=interpreter_result.weight_name_map,
-            requires_output_allocator=interpreter_result.requires_output_allocator,
-        )
+        rt_cls = TorchTensorRTModule
 
     elif (
         not ENABLED_FEATURES.torch_tensorrt_runtime and not settings.use_python_runtime
@@ -109,7 +103,7 @@ def convert_module(
             "Since Torch-TensorRT runtime is not available, using Python Runtime, some features may not be available"
         )
 
-    return PythonTorchTensorRTModule(
+    return rt_cls(
         cuda_engine=interpreter_result.engine,
         input_binding_names=list(interpreter_result.input_names),
         output_binding_names=list(interpreter_result.output_names),

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -15,6 +15,7 @@
 from torch_tensorrt.dynamo.debug._DebuggerConfig import DebuggerConfig
 from torch_tensorrt.dynamo.debug._supports_debugger import cls_supports_debugger
 from torch_tensorrt.dynamo.utils import DYNAMIC_DIM
+from torch_tensorrt.logging import TRT_LOGGER
 from torch_tensorrt.runtime._utils import (
     _is_switch_required,
     _select_rt_device,
@@ -123,6 +124,7 @@ class PythonTorchTensorRTModule(Module):  # type: ignore[misc]
     def __init__(
         self,
         cuda_engine: trt.ICudaEngine = None,
+        serialized_engine: Optional[bytes] = None,
         input_binding_names: Optional[List[str]] = None,
         output_binding_names: Optional[List[str]] = None,
         *,
@@ -181,7 +183,19 @@ def __init__(
         # Unused currently - to be used by Dynamic Shape support implementation
         self.memory_pool = None
 
-        self.engine = cuda_engine
+        if cuda_engine:
+            assert isinstance(
+                cuda_engine, trt.ICudaEngine
+            ), "Cuda engine must be a trt.ICudaEngine object"
+            self.engine = cuda_engine
+        elif serialized_engine:
+            assert isinstance(
+                serialized_engine, bytes
+            ), "Serialized engine must be a bytes object"
+            self.engine = serialized_engine
+        else:
+            raise ValueError("Serialized engine or cuda engine must be provided")
+
         self.input_names = (
             input_binding_names if input_binding_names is not None else []
         )
@@ -217,7 +231,7 @@ def __init__(
         self.output_allocator: Optional[DynamicOutputAllocator] = None
         self.use_output_allocator_outputs = False
 
-        if self.engine is not None and not self.settings.lazy_engine_init:
+        if self.engine and not self.settings.lazy_engine_init:
             self.setup_engine()
 
     def get_streamable_device_memory_budget(self) -> Any:
@@ -258,6 +272,17 @@ def set_default_device_memory_budget(self) -> int:
         return self._set_device_memory_budget(budget_bytes)
 
     def setup_engine(self) -> None:
+
+        if isinstance(self.engine, trt.ICudaEngine):
+            pass
+        elif isinstance(self.engine, bytes):
+            runtime = trt.Runtime(TRT_LOGGER)
+            self.engine = runtime.deserialize_cuda_engine(self.engine)
+        else:
+            raise ValueError(
+                "Expected engine as trt.ICudaEngine or serialized engine as bytes"
+            )
+
         assert (
             self.target_platform == Platform.current_platform()
         ), f"TensorRT engine was not built to target current platform (target: {self.target_platform}, current: {Platform.current_platform()})"
@@ -298,7 +323,7 @@ def _check_initialized(self) -> None:
             raise RuntimeError("PythonTorchTensorRTModule is not initialized.")
 
     def _on_state_dict(self, state_dict: Dict[str, Any], prefix: str, _: Any) -> None:
-        state_dict[prefix + "engine"] = self.serialized_engine
+        state_dict[prefix + "engine"] = self.engine
         state_dict[prefix + "input_names"] = self.input_names
         state_dict[prefix + "output_names"] = self.output_names
         state_dict[prefix + "platform"] = self.target_platform
@@ -313,7 +338,7 @@ def _load_from_state_dict(
         unexpected_keys: Any,
         error_msgs: Any,
     ) -> None:
-        self.serialized_engine = state_dict[prefix + "engine"]
+        self.engine = state_dict[prefix + "engine"]
         self.input_names = state_dict[prefix + "input_names"]
         self.output_names = state_dict[prefix + "output_names"]
         self.target_platform = state_dict[prefix + "platform"]

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -2,10 +2,12 @@
 
 import base64
 import copy
+import io
 import logging
 import pickle
 from typing import Any, List, Optional, Tuple, Union
 
+import tensorrt as trt
 import torch
 from torch_tensorrt._Device import Device
 from torch_tensorrt._enums import Platform
@@ -76,6 +78,7 @@ class TorchTensorRTModule(torch.nn.Module):  # type: ignore[misc]
 
     def __init__(
         self,
+        cuda_engine: Optional[trt.ICudaEngine | bytes] = None,
         serialized_engine: Optional[bytes] = None,
         input_binding_names: Optional[List[str]] = None,
         output_binding_names: Optional[List[str]] = None,
@@ -123,8 +126,22 @@ def __init__(
         """
         super(TorchTensorRTModule, self).__init__()
 
-        if not isinstance(serialized_engine, bytearray):
-            ValueError("Expected serialized engine as bytearray")
+        if serialized_engine:
+            assert isinstance(
+                serialized_engine, bytes
+            ), "Serialized engine must be a bytes object"
+            self.serialized_engine = serialized_engine
+
+        elif cuda_engine:
+            assert isinstance(
+                cuda_engine, trt.ICudaEngine
+            ), "Cuda engine must be a trt.ICudaEngine object"
+            serialized_engine = cuda_engine.serialize()
+            with io.BytesIO() as engine_bytes:
+                engine_bytes.write(serialized_engine)  # type: ignore
+                self.serialized_engine = engine_bytes.getvalue()
+        else:
+            raise ValueError("Serialized engine or cuda engine must be provided")
 
         self.input_binding_names = (
             input_binding_names if input_binding_names is not None else []
@@ -136,12 +153,11 @@ def __init__(
         self.hardware_compatible = settings.hardware_compatible
         self.settings = copy.deepcopy(settings)
         self.weight_name_map = weight_name_map
-        self.serialized_engine = serialized_engine
         self.engine = None
         self.requires_output_allocator = requires_output_allocator
 
         if (
-            serialized_engine
+            self.serialized_engine
             and not self.settings.lazy_engine_init
             and not self.settings.enable_cross_compile_for_windows
         ):