diff --git a/py/torch_tensorrt/runtime/_runtime_cache.py b/py/torch_tensorrt/runtime/_runtime_cache.py index 5c5cbebe05..066a7ab3aa 100644 --- a/py/torch_tensorrt/runtime/_runtime_cache.py +++ b/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) ) diff --git a/py/torch_tensorrt/runtime/_runtime_config.py b/py/torch_tensorrt/runtime/_runtime_config.py index 1b2cbee643..16e6c0ba6c 100644 --- a/py/torch_tensorrt/runtime/_runtime_config.py +++ b/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 "