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

@@ -30,20 +30,24 @@ emits a ``UserWarning``.
 ----
 
 The three ways to apply settings
---------------------------------
+---------------------------------
 
-Direct assignment — permanent
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+``apply_runtime_settings(...)`` — permanent apply
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use for any permanent assignment on in-process compiled models:
 
 .. code-block:: python
 
     import torch_tensorrt
-    from torch_tensorrt.runtime import RuntimeSettings
+    from torch_tensorrt.runtime import RuntimeSettings, apply_runtime_settings
 
     mod = torch_tensorrt.compile(model, inputs=inputs)
-    mod.runtime_settings = RuntimeSettings(runtime_cache="/var/cache/jit.bin")
+    apply_runtime_settings(mod, RuntimeSettings(runtime_cache="/var/cache/jit.bin"))
 
-Use when you want the setting to apply for the module's lifetime.
+:func:`apply_runtime_settings` walks all TRT subgraphs under ``mod`` and
+applies ``settings`` to each one. It returns the number of engines updated and
+raises :exc:`RuntimeError` if no TRT engines are found.
 
 ``runtime_config(...)`` context manager — scoped override
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -283,10 +287,10 @@ Construct your own handle if you want full lifetime control:
 
 .. code-block:: python
 
-    from torch_tensorrt.runtime import RuntimeCache, RuntimeSettings
+    from torch_tensorrt.runtime import RuntimeCache, RuntimeSettings, apply_runtime_settings
 
     handle = RuntimeCache(path="/var/cache/jit.bin", autosave_on_del=True)
-    mod.runtime_settings = RuntimeSettings(runtime_cache=handle)
+    apply_runtime_settings(mod, RuntimeSettings(runtime_cache=handle))
     out = mod(x)
     # handle.save() will fire when handle goes out of scope (autosave_on_del=True)
 
@@ -296,10 +300,49 @@ Or with explicit save/load:
 
     handle = RuntimeCache(path="/var/cache/jit.bin")  # autosave_on_del=False default
     handle.load()
-    mod.runtime_settings = RuntimeSettings(runtime_cache=handle)
+    apply_runtime_settings(mod, RuntimeSettings(runtime_cache=handle))
     out = mod(x)
     handle.save()
 
+AOT-loaded artifacts (``ExportedProgram`` / ``GraphModule``)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Engines loaded with :func:`torch_tensorrt.load` have no
+:class:`TorchTensorRTModule`, so the context managers cannot restore settings
+on exit and will raise. Use :func:`apply_runtime_settings` with a caller-owned
+:class:`RuntimeCache` 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()  # caller's responsibility — no module to auto-warm the cache
+
+    apply_runtime_settings(gm, RuntimeSettings(runtime_cache=cache))
+    out = gm(x)
+    cache.save()
+
+``settings.runtime_cache`` must be ``None`` or a :class:`RuntimeCache` you own
+for module-less engines — a path string raises :exc:`TypeError` because there
+is no module to build and save the handle.
+
+:func:`apply_runtime_settings` also accepts the :class:`~torch.export.ExportedProgram`
+directly, 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`.
+
 ----
 
 Best practices
@@ -314,7 +357,7 @@ before that and you get **one** context create:
 .. code-block:: python
 
     mod = torch_tensorrt.compile(...)
-    mod.runtime_settings = RuntimeSettings(cuda_graph_strategy="whole_graph_capture")
+    apply_runtime_settings(mod, RuntimeSettings(cuda_graph_strategy="whole_graph_capture"))
     out = mod(x)  # single createExecutionContext call here
 
 Apply settings *after* first execute and you get **two**:
@@ -323,7 +366,7 @@ Apply settings *after* first execute and you get **two**:
 
     mod = torch_tensorrt.compile(...)
     out = mod(x)  # context created with defaults
-    mod.runtime_settings = RuntimeSettings(cuda_graph_strategy="whole_graph_capture")
+    apply_runtime_settings(mod, RuntimeSettings(cuda_graph_strategy="whole_graph_capture"))
     out = mod(x)  # context invalidated + recreated
 
 On RTX, each ``createExecutionContext`` JIT-compiles the specialized kernel
@@ -333,8 +376,8 @@ NCCL engines pay the extra create
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 NCCL-collective engines eagerly materialize the context at setup (cross-rank
-barrier ordering). Any subsequent ``mod.runtime_settings = ...`` triggers a
-second create. This is a documented trade-off — apply settings before any
+barrier ordering). Any subsequent :func:`apply_runtime_settings` call triggers
+a second create. This is a documented trade-off — apply settings before any
 inference if you can, but the eager bind is non-negotiable for NCCL safety.
 
 Default ``runtime_cache`` is shared per-user — concurrent processes can lose kernels
@@ -354,10 +397,10 @@ matter:
 .. code-block:: python
 
     # Option 1: per-worker path
-    mod.runtime_settings = RuntimeSettings(runtime_cache=f"/var/cache/jit-worker-{worker_id}.bin")
+    apply_runtime_settings(mod, RuntimeSettings(runtime_cache=f"/var/cache/jit-worker-{worker_id}.bin"))
 
     # Option 2: opt out
-    mod.runtime_settings = RuntimeSettings(runtime_cache=None)
+    apply_runtime_settings(mod, RuntimeSettings(runtime_cache=None))
 
 Don't nest ``runtime_cache(...)`` CMs with the same path
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -376,22 +419,14 @@ re-attached (different ``IRuntimeCache`` from ``rc2``), and ``rc1.save()``
 overwrites ``/p`` with the now-stale ``rc1`` state. **Last writer wins;
 mid-block kernels are silently lost.**
 
-Setter is per-``TorchTensorRTModule``
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-``mod.runtime_settings = rs`` only affects ``self``. If you compile a model
-with multiple TRT subgraphs, walk the submodules:
-
-.. code-block:: python
-
-    from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule
-
-    for _, sub in compiled.named_modules():
-        if isinstance(sub, TorchTensorRTModule):
-            sub.runtime_settings = RuntimeSettings(...)
+``apply_runtime_settings`` reaches all subgraphs automatically
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-``runtime_config(...)`` and ``runtime_cache(...)`` do this walk automatically
-— that is the easier API for compound models.
+A compiled model with multiple TRT subgraphs has one
+:class:`TorchTensorRTModule` per subgraph. Calling
+:func:`apply_runtime_settings` on the top-level module (or an
+:class:`~torch.export.ExportedProgram`) walks all of them in one call — you do
+not need to iterate submodules manually. The context managers do the same walk.
 
 Non-TensorRT-RTX builds emit a warning, do nothing
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -413,8 +448,8 @@ Quick reference
 
    * - Goal
      - API
-   * - Set a runtime knob permanently on one module
-     - ``mod.runtime_settings = RuntimeSettings(...)``
+   * - Set a runtime knob permanently (compiled or AOT-loaded)
+     - ``apply_runtime_settings(mod_or_ep, RuntimeSettings(...))``
    * - Temporary override for one call site
      - ``with runtime_config(mod, **overrides):``
    * - Just the dynamic-shapes kernel strategy

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -350,7 +350,9 @@ def runtime_settings(self, rs: RuntimeSettings) -> None:
         rs_resolved = self._resolve_runtime_cache(rs)
         # 2. Push to the engine if it exists; if not we stash for later.
         if self.engine is not None:
-            self._send_to_engine(rs_resolved)
+            from torch_tensorrt.runtime._runtime_config import _send_settings_to_engine
+
+            _send_settings_to_engine(self.engine, rs_resolved)
         # 3. Store the resolved form so reads agree with what the engine sees.
         self._runtime_settings = rs_resolved
 
@@ -426,28 +428,6 @@ def _wrapper_still_attached(self, w: Any) -> bool:
         """
         return not ENABLED_FEATURES.torch_tensorrt_runtime or w.is_cpp_runtime()
 
-    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,
-        )
-
-        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),
-            )
-
     def setup_engine(self) -> None:
         """
         Setup engine for a module which has deferred engine setup.

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,27 @@ 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(self._targets))
+
+        if any(owner is None for owner, _ in engines):
+            raise TypeError(
+                "runtime_cache() encountered 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