fix(runtime): back a python engine runtime cache with the python handle

py/torch_tensorrt/runtime/_runtime_cache.py

@@ -222,6 +222,9 @@ class RuntimeCache:
     ``ENABLED_FEATURES.torch_tensorrt_runtime`` is set, else the
     pure-Python :class:`_RuntimeCacheHandle`; both satisfy
     :class:`_RuntimeCacheHandleProtocol` interface which this class uses.
+    A caller that knows the cache will be consumed by a python engine passes
+    ``python_runtime=True`` to force the pure-Python backing, because only
+    that one can hand a live ``IRuntimeCache`` back to python.
 
     **Identity-based equality.** Two handles wrap distinct underlying
     ``IRuntimeCache`` instances even when they share a path, and the
@@ -232,21 +235,31 @@ def __init__(
         self,
         path: str = "",
         autosave_on_del: bool = False,
+        python_runtime: bool = False,
     ) -> None:
         # Set the atexit-token slot first so ``__del__`` can safely read it
         # even if a later step in ``__init__`` raises and leaves the object
         # partially constructed.
         self._atexit_token: Optional[Callable[..., None]] = None
 
-        # Pick the backing that matches the active runtime. The torchbind
-        # class ``torch.classes.tensorrt.RuntimeCacheHandle`` is registered by
-        # the C++ shared library; if the .so isn't loaded
+        # Pick the backing that matches the runtime that will consume the
+        # cache. The torchbind class ``torch.classes.tensorrt.RuntimeCacheHandle``
+        # is registered by the C++ shared library; if the .so isn't loaded
         # (``ENABLED_FEATURES.torch_tensorrt_runtime is False``) it doesn't
         # exist as an attribute, so we fall back to the pure-Python
         # ``_RuntimeCacheHandle``. Both satisfy
         # :class:`_RuntimeCacheHandleProtocol`, so the facade methods forward
         # without branching on which backing won.
-        if torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime:
+        #
+        # A loaded .so does not mean the consumer is a cpp engine: a python
+        # ``TRTEngine`` can be built directly from a packed engine tuple while
+        # the .so is loaded. Such a caller passes ``python_runtime=True``,
+        # because ``ensure_cache`` can only materialize through the python
+        # backing (torchbind cannot carry an ``IRuntimeCache`` across).
+        if (
+            torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime
+            and not python_runtime
+        ):
             self._handle: _RuntimeCacheHandleProtocol = (
                 torch.classes.tensorrt.RuntimeCacheHandle(path)
             )

py/torch_tensorrt/runtime/_runtime_config.py

@@ -152,6 +152,10 @@ def __init__(self, settings: Optional[RuntimeSettings] = None) -> None:
         self._settings: RuntimeSettings = settings or RuntimeSettings()
         # Live trt.IRuntimeConfig (RTX) or None (non-RTX / pre-init).
         self._live: Any = None
+        # Wrapper built by ``_apply_settings`` for a path-string setting. The
+        # IRuntimeConfig does not own the IRuntimeCache attached to it, so the
+        # wrapper has to stay alive as long as ``_live`` does.
+        self._implicit_cache: Optional["RuntimeCache"] = None
 
     @property
     def settings(self) -> RuntimeSettings:
@@ -189,6 +193,9 @@ def ensure_initialized(self, cuda_engine: Any) -> None:
     def reset(self) -> None:
         """Drop the live ``IRuntimeConfig``; the next ``ensure_initialized`` rebuilds."""
         self._live = None
+        # Release after ``_live``, never before: the config points at the cache.
+        # Dropping the last reference is what saves an implicit cache to disk.
+        self._implicit_cache = None
 
     def create_execution_context(
         self,
@@ -291,16 +298,22 @@ def _apply_settings(self) -> None:
             # documented contract that callers MAY pass a path string.
             #
             # ``RuntimeSettings`` is a frozen dataclass, so we can't store the
-            # wrapper back onto ``self._settings``; just use it locally. The
-            # wrapper is GC'd after this call, which is fine: ensure_cache has
-            # already materialized the underlying IRuntimeCache on ``_live``.
-            wrapped = RuntimeCache(path=rc, autosave_on_del=True)
+            # wrapper back onto ``self._settings``; it goes on
+            # ``self._implicit_cache`` instead, which keeps the IRuntimeCache
+            # alive for as long as the IRuntimeConfig points at it.
+            #
+            # ``python_runtime=True`` because this method only ever runs for a
+            # python engine. The default backing follows the loaded C++ library
+            # instead, and a cpp-backed handle materializes on the cpp side, so
+            # ``ensure_cache`` would hand back ``None``.
+            wrapped = RuntimeCache(path=rc, autosave_on_del=True, python_runtime=True)
             try:
                 wrapped.load()
             except Exception as e:
                 logger.warning(f"Failed to warm-load runtime cache from {rc!r}: {e}")
             cache = wrapped.ensure_cache(self._live)
             self._live.set_runtime_cache(cache)
+            self._implicit_cache = wrapped
         else:
             raise TypeError(
                 f"runtime_cache must be None, str, or RuntimeCache by the "