feat: apply_runtime_settings — permanent runtime config for AOT-loaded engines

docsrc/py_api/runtime.rst

@@ -27,6 +27,8 @@ Functions
 
 .. autofunction:: enable_output_allocator
 
+.. autofunction:: apply_runtime_settings
+
 Runtime backend
 ---------------
 

docsrc/user_guide/runtime_performance/runtime_settings.rst

@@ -29,8 +29,8 @@ emits a ``UserWarning``.
 
 ----
 
-The three ways to apply settings
---------------------------------
+The four ways to apply settings
+-------------------------------
 
 Direct assignment — permanent
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -136,6 +136,64 @@ Stream-mode behavior:
 * On exit: cache serialized, ``stream.write(bytes)`` once.
 * ``rc.path`` reports ``""`` in stream-mode.
 
+``apply_runtime_settings(...)`` — permanent apply for AOT artifacts
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Engines loaded with :func:`torch_tensorrt.load` have no
+:class:`TorchTensorRTModule`.  :func:`~torch_tensorrt.runtime.runtime_config`
+and :func:`~torch_tensorrt.runtime.runtime_cache` cannot restore settings on
+exit for such engines (there is no getter on the torchbind engine) and will
+raise if they encounter one.  Use the permanent-apply entry point instead:
+
+.. code-block:: python
+
+    import torch_tensorrt
+    from torch_tensorrt.runtime import RuntimeCache, RuntimeSettings, apply_runtime_settings
+
+    ep = torch_tensorrt.load("model.ep")
+    gm = ep.module()
+
+    cache = RuntimeCache(path="/var/cache/jit.bin")
+    cache.load()  # warm from disk if it exists
+
+    n = apply_runtime_settings(
+        gm,
+        RuntimeSettings(
+            cuda_graph_strategy="whole_graph_capture",
+            runtime_cache=cache,
+        ),
+    )
+    print(f"Applied to {n} engine(s)")
+    out = gm(x)
+    cache.save()  # persist newly JIT'd kernels
+
+**Ownership rule:** ``settings.runtime_cache`` must be ``None`` or a
+:class:`RuntimeCache` you own -- a path string raises ``TypeError`` because
+there is no module to build and save the handle.  If you call
+``apply_runtime_settings(gm, RuntimeSettings())`` (the default
+``runtime_cache`` is a path string), you will hit this error.
+Pass ``runtime_cache=None`` or a :class:`RuntimeCache`.
+
+**You own** ``.load()`` **as well as** ``.save()``.  A
+:class:`TorchTensorRTModule` calls :meth:`RuntimeCache.load` automatically when
+it resolves a path string (via ``_resolve_runtime_cache``), so in-process
+compiled models warm the cache implicitly.  There is no equivalent hook on the
+module-less path -- call ``cache.load()`` (shown above) before passing the
+handle, or the engine starts with an empty cache regardless of what is on disk.
+
+:func:`apply_runtime_settings` also accepts a :class:`~torch.export.ExportedProgram`
+directly (the :func:`torch_tensorrt.load` return value), which is equivalent to
+passing ``ep.module()``:
+
+.. code-block:: python
+
+    apply_runtime_settings(ep, RuntimeSettings(runtime_cache=cache))
+
+.. note::
+
+   Runtime settings are never serialized.  They do not survive
+   :func:`torch_tensorrt.save`; re-apply after each :func:`torch_tensorrt.load`.
+
 ----
 
 Composing the context managers
@@ -431,3 +489,5 @@ Quick reference
      - ``RuntimeSettings(runtime_cache=None)`` or ``runtime_cache(mod, "")``
    * - Non-cuda-graph settings alongside cudagraphs capture
      - nest ``runtime_config(...)`` *outside* ``enable_cudagraphs(...)``
+   * - Set a runtime knob on a loaded artifact (no module)
+     - ``apply_runtime_settings(gm_or_ep, RuntimeSettings(...))``

py/torch_tensorrt/dynamo/runtime/_TRTEngine.py

@@ -282,11 +282,18 @@ def __init__(
         # engines compiled with native multi-device collective layers.
         self._nccl_comm: Optional[Any] = None
 
-        # Owns RuntimeSettings + the live trt.IRuntimeConfig + the
-        # engine-implicit RuntimeCache. Hides all RTX feature gates.
-        # ``RuntimeSettings`` default; callers wanting non-defaults assign via
-        # the module's ``runtime_settings`` setter after compile.
-        self._trt_runtime_config: TRTRuntimeConfig = TRTRuntimeConfig(RuntimeSettings())
+        # Owns RuntimeSettings + the live trt.IRuntimeConfig. Hides all RTX
+        # feature gates.
+        #
+        # ``runtime_cache=None``: an engine never owns a runtime cache. The
+        # module owns the implicit one and pushes it down via
+        # ``setup_engine``; an engine with no module (a packed-engine-info
+        # build, or a constant in an AOT-loaded ExportedProgram) runs without
+        # one until a caller attaches a ``RuntimeCache`` explicitly. Mirrors
+        # the cpp default (``RuntimeSettings::runtime_cache = nullptr``).
+        self._trt_runtime_config: TRTRuntimeConfig = TRTRuntimeConfig(
+            RuntimeSettings(runtime_cache=None)
+        )
         # Multiple optimization profiles. Manual selection by default:
         # ``_active_profile_index`` is the profile currently loaded in the TRT
         # context (default 0, reused across calls). ``_auto_select_profiles``
@@ -399,10 +406,11 @@ def __setstate__(self, state: Any) -> None:
         # NCCL communicators cannot be pickled; rebind lazily on the next
         # forward pass via setup_nccl_comm().
         self._nccl_comm = None
-        # RuntimeSettings are NOT serialized -- restore defaults. Callers
-        # who want runtime-mode overrides must reapply them post-load via
-        # ``mod.runtime_settings = ...`` (per ``TorchTensorRTModule``) or a runtime CM.
-        self._trt_runtime_config = TRTRuntimeConfig(RuntimeSettings())
+        # RuntimeSettings are NOT serialized -- restore defaults, runtime cache
+        # included (see ``__init__``). Callers who want runtime-mode overrides
+        # must reapply them post-load via ``mod.runtime_settings = ...`` (per
+        # ``TorchTensorRTModule``) or a runtime CM.
+        self._trt_runtime_config = TRTRuntimeConfig(RuntimeSettings(runtime_cache=None))
 
         self._active_profile_index = 0
         self._auto_select_profiles = False

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -357,6 +357,10 @@ def runtime_settings(self, rs: RuntimeSettings) -> None:
     def _resolve_runtime_cache(self, rs: RuntimeSettings) -> RuntimeSettings:
         """Normalize ``rs.runtime_cache`` to ``None`` | ``RuntimeCache`` (never a path str).
 
+        This module is the only place a path string is resolved; engines reject
+        one (see ``TRTRuntimeConfig._apply_settings``), so the returned settings
+        must never carry a ``str``.
+
         Manages the ``_implicit_cache_handle`` slot as a side effect: builds a
         fresh wrapper for a new path, reuses the existing one for the same
         path, releases it (with save-on-swap) for non-path inputs.
@@ -372,7 +376,9 @@ def _resolve_runtime_cache(self, rs: RuntimeSettings) -> RuntimeSettings:
         if not (isinstance(rc, str) and rc):
             if rc is not self._implicit_cache_handle:
                 self._set_managed_handle(None)
-            return rs
+            # An empty string means "no cache", but engines take only None or a
+            # RuntimeCache -- normalize so no str ever reaches one.
+            return rs.merge(runtime_cache=None) if isinstance(rc, str) else rs
 
         # Branch 2: same path + wrapper still usable -> reuse. Keeps the CM
         # enter/exit cycle cheap (no teardown/rebuild loses in-memory kernels).
@@ -422,25 +428,9 @@ def _wrapper_still_attached(self, w: Any) -> bool:
 
     def _send_to_engine(self, rs: RuntimeSettings) -> None:
         """Push ``rs`` to whichever engine flavor is attached."""
-        from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine
-        from torch_tensorrt.runtime._runtime_cache import _to_torchbind_handle
-        from torch_tensorrt.runtime._runtime_config import (
-            _CUDA_GRAPH_STRATEGY_MAP,
-            _DYNAMIC_SHAPES_KERNEL_STRATEGY_MAP,
-        )
+        from torch_tensorrt.runtime._runtime_config import _send_settings_to_engine
 
-        if isinstance(self.engine, TRTEngine):
-            self.engine.update_runtime_settings(rs)
-        else:
-            # Strategies cross the boundary as ints (TorchBind ``int64_t``,
-            # mirroring the nvinfer1 enum integers on the cpp side).
-            self.get_engine().update_runtime_settings(
-                _DYNAMIC_SHAPES_KERNEL_STRATEGY_MAP[
-                    rs.dynamic_shapes_kernel_specialization_strategy
-                ],
-                _CUDA_GRAPH_STRATEGY_MAP[rs.cuda_graph_strategy],
-                _to_torchbind_handle(rs.runtime_cache),
-            )
+        _send_settings_to_engine(self.engine, rs)
 
     def setup_engine(self) -> None:
         """
@@ -566,9 +556,11 @@ def set_extra_state(self, state: SerializedTorchTensorRTModuleFmt) -> None:
             self.settings = metadata["settings"]
             self.symbolic_shape_expressions = metadata["inout_symexprs"]
 
-            # RuntimeSettings are NOT serialized; restore defaults. Caller can
-            # reapply via ``mod.runtime_settings = ...`` (per submodule) or a CM after load.
-            self._runtime_settings = RuntimeSettings()
+            # RuntimeSettings are NOT serialized; the reset leaves no runtime
+            # cache, matching the freshly-built engine below. A caller who wants
+            # one reapplies via ``mod.runtime_settings = ...`` (per submodule) or
+            # a CM after load.
+            self._runtime_settings = RuntimeSettings(runtime_cache=None)
             # Mirror the settings reset on the implicit cache handle so a
             # stale wrapper from prior use doesn't survive load_state_dict and
             # silently write the fresh engine's cache bytes to the old path.
@@ -682,7 +674,7 @@ def __getstate__(self) -> dict[str, Any]:
         return state
 
     def __setstate__(self, state: dict[str, Any]) -> None:
-        state.setdefault("_runtime_settings", RuntimeSettings())
+        state.setdefault("_runtime_settings", RuntimeSettings(runtime_cache=None))
         state.setdefault("_implicit_cache_handle", None)
         set_state = getattr(super(), "__setstate__", None)
         if set_state is not None:

py/torch_tensorrt/runtime/__init__.py

@@ -14,6 +14,7 @@
 from torch_tensorrt.runtime._runtime_cache import RuntimeCache, runtime_cache
 from torch_tensorrt.runtime._runtime_config import (
     RuntimeSettings,
+    apply_runtime_settings,
     runtime_config,
     set_dynamic_shapes_kernel_strategy,
 )

py/torch_tensorrt/runtime/_runtime_cache.py

@@ -524,24 +524,28 @@ def _save_from(self, handle: "RuntimeCache") -> None:
     def __enter__(self) -> RuntimeCache:
         # Defer imports to avoid a circular dependency:
         # _runtime_cache -> _runtime_config -> _TorchTensorRTModule -> (indirect) _runtime_cache.
-        from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import (
-            TorchTensorRTModule,
+        from torch_tensorrt.runtime._runtime_config import (
+            _iter_trt_engines,
+            runtime_config,
         )
-        from torch_tensorrt.runtime._runtime_config import runtime_config
-
-        # 1. Find any TorchTensorRTModule under the targets; first one wins.
-        bootstrap_module = None
-        for target in self._targets:
-            for _, mod in target.named_modules():
-                if isinstance(mod, TorchTensorRTModule):
-                    bootstrap_module = mod
-                    break
-            if bootstrap_module is not None:
-                break
-        if bootstrap_module is None:
+
+        # 1. Discover all TRT engines under the targets, validate before mutating.
+        engines = list(_iter_trt_engines(list(self._targets)))
+
+        module_less = [(owner, eng) for owner, eng in engines if owner is None]
+        if module_less:
+            raise TypeError(
+                f"runtime_cache() encountered {len(module_less)} module-less "
+                "TRT engine(s) that it cannot snapshot and restore on exit. "
+                "Use apply_runtime_settings() for engines loaded without a "
+                "TorchTensorRTModule (e.g. via torch_tensorrt.load())."
+            )
+
+        if not engines:
             raise RuntimeError(
-                "runtime_cache() requires at least one TorchTensorRTModule "
-                "under the target(s)."
+                "runtime_cache() requires at least one TRT engine under the "
+                "target(s). The target may have fallen back entirely to PyTorch "
+                "or may not contain any compiled TRT subgraphs."
             )
 
         # 2. Build the handle in its pending state on both runtimes. The