🐛 [Bug] Serialization fails in Runtime Cache #4359

from pathlib import Path

import torch
import torch_tensorrt as torchtrt
import torchvision.models as models

torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

artifact = Path(__file__).resolve().parent / "trt_mod.pt2"


def main():
    model = models.resnet50(pretrained=False).eval().to("cuda")

    inputs = torch.randn(1, 3, 224, 224, device="cuda")

    exported = torch.export.export(model, (inputs,))

    trt_mod = torchtrt.dynamo.compile(
        exported,
        inputs=[inputs],
        min_block_size=1,
        immutable_weights=False,
        reuse_cached_engines=False,
        enabled_precisions={torch.float32},
        torch_executed_ops={torch.ops.aten.relu.default},
    )

    with torch.no_grad():
        results1 = trt_mod(inputs)

    torchtrt.save(
        trt_mod,
        str(artifact),
        output_format="exported_program",
        inputs=[inputs],
        retrace=False,
    )

    trt_ep = torch.export.load(str(artifact))

    # Inspect loaded engine types
    gm = trt_ep.module()

    with torch.no_grad():
        results2 = gm(inputs)

    print("max diff:", (results1 - results2).abs().max().item())
    # for name, val in gm._modules.items():
    #     print(f"  submodule {name}: {type(val).__name__}")
    # for attr_name in dir(gm):
    #     obj = getattr(gm, attr_name, None)
    #     if obj is not None and "engine" in attr_name.lower():
    #         print(f"  attr {attr_name}: {type(obj)} is_script_obj={isinstance(obj, torch._C.ScriptObject)}")
    # for fqn, const in trt_ep.constants.items():
    #     print(f"  constant {fqn}: {type(const)} is_script_obj={isinstance(const, torch._C.ScriptObject)}")
    # for node in gm.graph.nodes:
    #     if node.op == "get_attr":
    #         val = getattr(gm, node.target, None)
    #         print(f"  graph get_attr '{node.target}': {type(val)}")


if __name__ == "__main__":
    main()

WARNING: [Torch-TensorRT - Debug Build] - RuntimeCacheHandle::serialize() invoked on a non-RTX build; returning empty bytes.
WARNING: [Torch-TensorRT - Debug Build] - RuntimeCacheHandle::serialize() invoked on a non-RTX build; returning empty bytes.
WARNING: [Torch-TensorRT - Debug Build] - RuntimeCacheHandle::serialize() invoked on a non-RTX build; returning empty bytes.
WARNING: [Torch-TensorRT - Debug Build] - RuntimeCacheHandle::serialize() invoked on a non-RTX build; returning empty bytes.
Traceback (most recent call last):
  File "/home/other/test_two_runtime.py", line 65, in <module>
    main()
  File "/home/other/test_two_runtime.py", line 33, in main
    torchtrt.save(
  File "/home/TensorRT/py/torch_tensorrt/_compile.py", line 1001, in save
    exp_program = export(
                  ^^^^^^^
  File "/home/TensorRT/py/torch_tensorrt/dynamo/_exporter.py", line 42, in export
    patched_module = transform(gm, cross_compile_module)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/TensorRT/py/torch_tensorrt/dynamo/_exporter.py", line 92, in transform
    gm.graph.eliminate_dead_code()
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph.py", line 2652, in eliminate_dead_code
    self.lint()
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph.py", line 2559, in lint
    check_arg(arg, node)
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph.py", line 2544, in check_arg
    raise RuntimeError(
RuntimeError: Argument 'getitem' of Node 'relu' was used before it has been defined! Please check that Nodes in the graph are topologically ordered
...
    return (getitem_53,)
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down
Exception ignored in: <function RuntimeCache.__del__ at 0x7fe862091b20>
Traceback (most recent call last):
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 309, in __del__
  File "/home/TensorRT/py/torch_tensorrt/runtime/_runtime_cache.py", line 204, in path
ImportError: sys.meta_path is None, Python is likely shutting down