Changed the way to deserialize and move it to settings.py

Author: cehongwang

py/torch_tensorrt/dynamo/_settings.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
-from typing import Collection, Optional, Set, Tuple, Union
+from typing import Any, Collection, Optional, Set, Tuple, Union
 
+import torch
 from torch.fx.node import Target
 from torch_tensorrt._Device import Device
 from torch_tensorrt._enums import EngineCapability, dtype
@@ -143,6 +144,16 @@ class CompilationSettings:
     use_distributed_mode_trace: bool = USE_DISTRIBUTED_MODE_TRACE
     offload_module_to_cpu: bool = OFFLOAD_MODULE_TO_CPU
 
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state["torch_executed_ops"] = {str(op) for op in state["torch_executed_ops"]}
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        ops_str = self.torch_executed_ops
+        self.torch_executed_ops = {getattr(torch.ops, op) for op in ops_str}
+
 
 _SETTINGS_TO_BE_ENGINE_INVARIANT = (
     "enabled_precisions",

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -4,7 +4,7 @@
 import copy
 import logging
 import pickle
-from typing import Any, List, Optional, Set, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 from torch_tensorrt._Device import Device
@@ -227,11 +227,6 @@ def setup_engine(self) -> None:
 
     def encode_metadata(self, metadata: Any) -> str:
         metadata = copy.deepcopy(metadata)
-        metadata["settings"].torch_executed_ops = (
-            TorchTensorRTModule.serialize_aten_ops(
-                metadata["settings"].torch_executed_ops
-            )
-        )
         dumped_metadata = pickle.dumps(metadata)
         encoded_metadata = base64.b64encode(dumped_metadata).decode("utf-8")
         return encoded_metadata
@@ -240,21 +235,8 @@ def encode_metadata(self, metadata: Any) -> str:
     def decode_metadata(encoded_metadata: bytes) -> Any:
         dumped_metadata = base64.b64decode(encoded_metadata.encode("utf-8"))
         metadata = pickle.loads(dumped_metadata)
-        metadata["settings"].torch_executed_ops = (
-            TorchTensorRTModule.deserialize_aten_ops(
-                metadata["settings"].torch_executed_ops
-            )
-        )
         return metadata
 
-    @staticmethod
-    def serialize_aten_ops(aten_ops: Set[torch._ops.OpOverload]) -> Set[str]:
-        return {str(op) for op in aten_ops}
-
-    @staticmethod
-    def deserialize_aten_ops(aten_ops: Set[str]) -> Set[torch._ops.OpOverload]:
-        return {eval("torch.ops." + str(v)) for v in aten_ops}
-
     def get_extra_state(self) -> SerializedTorchTensorRTModuleFmt:
         if self.engine:
             return (