From 01d4fcdc6332263c5bb4c606b808b86c53fc3473 Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Mon, 24 Aug 2026 11:05:44 -0700 Subject: [PATCH 01/11] refactor: lift engine runtime-settings marshalling out of the module Free _send_settings_to_engine(engine, rs) in _runtime_config.py; TorchTensorRTModule._send_to_engine becomes a one-line delegate. Pure refactor, zero behaviour change. Co-Authored-By: Claude Sonnet 4.6 --- .../dynamo/runtime/_TorchTensorRTModule.py | 20 ++--------------- py/torch_tensorrt/runtime/_runtime_config.py | 22 +++++++++++++++++++ 2 files changed, 24 insertions(+), 18 deletions(-) diff --git a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py index b52bb8d360..f97206dc9b 100644 --- a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py +++ b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py @@ -428,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: """ diff --git a/py/torch_tensorrt/runtime/_runtime_config.py b/py/torch_tensorrt/runtime/_runtime_config.py index 4fd16bc641..c16518759e 100644 --- a/py/torch_tensorrt/runtime/_runtime_config.py +++ b/py/torch_tensorrt/runtime/_runtime_config.py @@ -414,3 +414,25 @@ def set_dynamic_shapes_kernel_strategy( return runtime_config( target_or_targets, dynamic_shapes_kernel_specialization_strategy=strategy ) + + +def _send_settings_to_engine(engine: Any, rs: RuntimeSettings) -> None: + """Push ``rs`` to whichever TRT engine flavor is attached. + + Dispatches on engine flavor: Python ``TRTEngine`` uses the native + ``update_runtime_settings`` method; torchbind engines expect int-valued + strategies and a torchbind handle rather than the Python facade. + """ + from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine + from torch_tensorrt.runtime._runtime_cache import _to_torchbind_handle + + if isinstance(engine, TRTEngine): + engine.update_runtime_settings(rs) + else: + 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), + ) From eda39c7fd58be7c5ffd0d409ee10afcc591c1aa4 Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Mon, 24 Aug 2026 11:08:41 -0700 Subject: [PATCH 02/11] feat: apply runtime settings to engines that have no module Add apply_runtime_settings(target, settings) -> int that walks TRT engines via get_attr constants (AOT-loaded) and TorchTensorRTModule (in-process). Accepts nn.Module, ExportedProgram, or a sequence. Raises RuntimeError on zero engines; raises TypeError when runtime_cache is a path string and any engine is module-less (no owner to persist the handle). Returns engine count. Co-Authored-By: Claude Sonnet 4.6 --- docsrc/py_api/runtime.rst | 2 + .../runtime_performance/runtime_settings.rst | 57 +++++- py/torch_tensorrt/runtime/__init__.py | 1 + py/torch_tensorrt/runtime/_runtime_config.py | 184 +++++++++++++++++- 4 files changed, 240 insertions(+), 4 deletions(-) diff --git a/docsrc/py_api/runtime.rst b/docsrc/py_api/runtime.rst index f8c020262d..40bde9b4d5 100644 --- a/docsrc/py_api/runtime.rst +++ b/docsrc/py_api/runtime.rst @@ -27,6 +27,8 @@ Functions .. autofunction:: enable_output_allocator +.. autofunction:: apply_runtime_settings + Runtime backend --------------- diff --git a/docsrc/user_guide/runtime_performance/runtime_settings.rst b/docsrc/user_guide/runtime_performance/runtime_settings.rst index 94c3d3a303..5784b0bcc3 100644 --- a/docsrc/user_guide/runtime_performance/runtime_settings.rst +++ b/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,57 @@ 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`. + +: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 +482,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(...))`` diff --git a/py/torch_tensorrt/runtime/__init__.py b/py/torch_tensorrt/runtime/__init__.py index 3c9777c5f2..4964494a49 100644 --- a/py/torch_tensorrt/runtime/__init__.py +++ b/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, ) diff --git a/py/torch_tensorrt/runtime/_runtime_config.py b/py/torch_tensorrt/runtime/_runtime_config.py index c16518759e..9dc974e18e 100644 --- a/py/torch_tensorrt/runtime/_runtime_config.py +++ b/py/torch_tensorrt/runtime/_runtime_config.py @@ -1,6 +1,6 @@ """Runtime settings + the TRTRuntimeConfig shim + the ``runtime_config`` CM. -This module groups three closely related concepts together: +This module groups four closely related concepts together: * :class:`RuntimeSettings` -- the user-facing, frozen dataclass of runtime-only knobs sampled at IExecutionContext creation (cuda_graph_strategy, @@ -13,11 +13,15 @@ * :func:`runtime_config` -- the runtime-mode context manager that toggles settings on every TRT submodule under a target for the duration of a ``with`` block. +* :func:`apply_runtime_settings` -- permanent apply to every TRT engine under a + target, including engines loaded without a :class:`TorchTensorRTModule`. Three ways to use ``RuntimeSettings``: 1. **Runtime context manager** -- toggle settings inside a ``with`` block. 2. **Programmatic** -- assign ``module.runtime_settings = rs`` directly. +3. **AOT artifact** -- call :func:`apply_runtime_settings` on a loaded + :class:`ExportedProgram` or ``GraphModule``. ``RuntimeSettings`` is intentionally NOT part of ``CompilationSettings`` and is NOT serialized into the engine tuple. It's purely an in-memory initialization @@ -30,7 +34,16 @@ import logging import warnings from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Optional, + Sequence, + Set, + Tuple, + Union, +) import torch from torch_tensorrt._features import ENABLED_FEATURES @@ -57,6 +70,7 @@ "disabled": 0, "whole_graph_capture": 1, } +_TORCHBIND_ENGINE_FQN = "__torch__.torch.classes.tensorrt.Engine" @dataclass(frozen=True) @@ -75,6 +89,10 @@ class RuntimeSettings: handle. A handle is the shared-cache form, typically obtained from :func:`torch_tensorrt.runtime.runtime_cache` -- multiple engines attaching the same handle share one ``IRuntimeCache``. + For engines without a :class:`TorchTensorRTModule` (e.g. loaded + via :func:`torch_tensorrt.load`), only ``None`` or a + :class:`RuntimeCache` is accepted; a path string raises + ``TypeError`` at the :func:`apply_runtime_settings` call site. Equality compares all fields; for ``runtime_cache``, handle equality is by identity (same handle ⇒ same cache). @@ -436,3 +454,165 @@ def _send_settings_to_engine(engine: Any, rs: RuntimeSettings) -> None: _CUDA_GRAPH_STRATEGY_MAP[rs.cuda_graph_strategy], _to_torchbind_handle(rs.runtime_cache), ) + + +def _is_trt_engine(obj: Any) -> bool: + """True iff ``obj`` is a TRT engine on either runtime. + + ``isinstance(obj, torch.classes.tensorrt.Engine)`` is unusable -- it raises + ``TypeError`` on cpp rt and ``RuntimeError`` on python-only rt. Compare + ``_type().qualified_name()`` on the torchbind flavor, guarded against the + ``AttributeError`` the Python engine raises on that method. + """ + from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine + + if isinstance(obj, TRTEngine): + return True + if isinstance(obj, torch.ScriptObject): + try: + return bool(obj._type().qualified_name() == _TORCHBIND_ENGINE_FQN) + except AttributeError: + pass + return False + + +def _iter_trt_engines( + target_or_targets: Any, +) -> Any: + """Yield ``(owner_or_None, engine)`` for every TRT engine reachable from ``target_or_targets``. + + ``owner_or_None`` is the :class:`TorchTensorRTModule` that holds the + engine, or ``None`` for a bare engine constant (e.g. in an AOT-loaded + ``GraphModule``). + + Accepts an ``nn.Module``, a ``torch.export.ExportedProgram``, or a + sequence of those. Results are deduped by ``id(engine)``. + """ + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule + + seen: Set[int] = set() + + def _visit_ep(ep: Any) -> Any: + for obj in ep.constants.values(): + if _is_trt_engine(obj) and id(obj) not in seen: + seen.add(id(obj)) + yield (None, obj) + + def _visit_module(root: torch.nn.Module) -> Any: + for _, mod in root.named_modules(): + if isinstance(mod, TorchTensorRTModule) and mod.engine is not None: + if id(mod.engine) not in seen: + seen.add(id(mod.engine)) + yield (mod, mod.engine) + if hasattr(mod, "graph"): + for node in mod.graph.nodes: + if node.op == "get_attr": + parts = node.target.split(".") + obj: Any = mod + try: + for part in parts: + obj = getattr(obj, part) + except AttributeError: + continue + if _is_trt_engine(obj) and id(obj) not in seen: + seen.add(id(obj)) + yield (None, obj) + + if isinstance(target_or_targets, torch.export.ExportedProgram): + yield from _visit_ep(target_or_targets) + elif isinstance(target_or_targets, torch.nn.Module): + yield from _visit_module(target_or_targets) + elif hasattr(target_or_targets, "__iter__") and not isinstance( + target_or_targets, (str, bytes) + ): + for t in target_or_targets: + if isinstance(t, torch.export.ExportedProgram): + yield from _visit_ep(t) + elif isinstance(t, torch.nn.Module): + yield from _visit_module(t) + else: + raise TypeError( + f"_iter_trt_engines(): each target must be an nn.Module or " + f"ExportedProgram; got {type(t).__name__}. " + "For a torch_tensorrt.load() result, pass the ExportedProgram " + "directly or call .module() on it." + ) + else: + raise TypeError( + f"_iter_trt_engines(): target must be an nn.Module, an " + f"ExportedProgram, or a sequence of those; got " + f"{type(target_or_targets).__name__}. " + "For a torch_tensorrt.load() result, pass the ExportedProgram " + "directly or call .module() on it." + ) + + +def apply_runtime_settings( + target_or_targets: Any, + settings: "RuntimeSettings", +) -> int: + """Apply ``settings`` permanently to every TRT engine reachable from ``target_or_targets``. + + Returns the number of engines updated. Raises :exc:`RuntimeError` if no + TRT engines are found (the shape of the silent-no-op bug this function + removes) and :exc:`TypeError` if ``settings.runtime_cache`` is a path + string and any reachable engine has no :class:`TorchTensorRTModule` to own + the resulting handle. + + Accepted targets: + + * :class:`torch.nn.Module` -- compiled result of + :func:`torch_tensorrt.compile`. + * :class:`torch.export.ExportedProgram` -- loaded result of + :func:`torch_tensorrt.load`. + * A sequence (list / tuple) of the above. + + **Ownership rule for module-less engines** (e.g. an AOT-loaded artifact): + ``settings.runtime_cache`` must be ``None`` or a :class:`RuntimeCache` you + own. A path string is accepted only where a :class:`TorchTensorRTModule` + can own the result and save it on ``__del__``. If you pass + ``RuntimeSettings()`` (whose default ``runtime_cache`` is a path string), + you will get a :exc:`TypeError`. Pass + ``RuntimeSettings(runtime_cache=None)`` or supply a + :class:`RuntimeCache` explicitly. + + Settings are never serialized; they do not survive + :func:`torch_tensorrt.save`. Re-apply after each :func:`torch_tensorrt.load`. + """ + if not isinstance(settings, RuntimeSettings): + raise TypeError( + f"apply_runtime_settings(): 'settings' must be a RuntimeSettings; " + f"got {type(settings).__name__}." + ) + + # Drain traversal before mutating (validate-then-apply). + engines = list(_iter_trt_engines(target_or_targets)) + + if not engines: + raise RuntimeError( + "apply_runtime_settings(): no TRT engines found under the target(s). " + "If the model fell back entirely to PyTorch (no TRT subgraphs were " + "compiled), no engines exist to configure." + ) + + # A path string needs a module to own the resulting RuntimeCache and save + # it on __del__. Module-less engines have no such owner. + if isinstance(settings.runtime_cache, str): + module_less_count = sum(1 for owner, _ in engines if owner is None) + if module_less_count: + raise TypeError( + f"apply_runtime_settings(): settings.runtime_cache is a path " + f"string ({settings.runtime_cache!r}), but {module_less_count} " + "engine(s) in the target have no TorchTensorRTModule to own the " + "resulting RuntimeCache and persist it on __del__. " + "Pass runtime_cache=None (no JIT cache) or a RuntimeCache " + "you own and save explicitly." + ) + + for owner, engine in engines: + if owner is not None: + owner.runtime_settings = settings + else: + _send_settings_to_engine(engine, settings) + + return len(engines) From 03e1d545d3f8dfb438aaaea73d5826934b7130c1 Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Mon, 24 Aug 2026 11:09:31 -0700 Subject: [PATCH 03/11] fix: raise instead of silently skipping engines the runtime CMs cannot restore runtime_config and runtime_cache now use _iter_trt_engines for discovery. On encountering a module-less engine (e.g. from an AOT-loaded artifact) they raise TypeError naming apply_runtime_settings, before mutating any engine state. Prior behaviour was a silent no-op (runtime_config) or an unhelpful module-only error (runtime_cache). Co-Authored-By: Claude Sonnet 4.6 --- py/torch_tensorrt/runtime/_runtime_cache.py | 36 ++++++++++--------- py/torch_tensorrt/runtime/_runtime_config.py | 37 ++++++++++++-------- 2 files changed, 42 insertions(+), 31 deletions(-) diff --git a/py/torch_tensorrt/runtime/_runtime_cache.py b/py/torch_tensorrt/runtime/_runtime_cache.py index 5c5cbebe05..931d71b034 100644 --- a/py/torch_tensorrt/runtime/_runtime_cache.py +++ b/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 diff --git a/py/torch_tensorrt/runtime/_runtime_config.py b/py/torch_tensorrt/runtime/_runtime_config.py index 9dc974e18e..135add83fb 100644 --- a/py/torch_tensorrt/runtime/_runtime_config.py +++ b/py/torch_tensorrt/runtime/_runtime_config.py @@ -372,22 +372,29 @@ def __init__( self._saved: Dict[Any, RuntimeSettings] = {} def __enter__(self) -> Union["torch.nn.Module", Tuple["torch.nn.Module", ...]]: - # Deferred import to avoid a circular dependency at module-load time. - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( - TorchTensorRTModule, - ) + # Drain the traversal before mutating; raise on unsupported engines + # before any settings are changed. + 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_config() 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())." + ) + + for owner, _ in engines: + if owner in self._saved: + # The same TRTModule appears under multiple targets in the + # list (or the tree contains a cycle). Don't snapshot twice. + continue + current = owner.runtime_settings + self._saved[owner] = current + merged = current.merge(**self._overrides) + owner.runtime_settings = merged - for target in self._targets: - for _, mod in target.named_modules(): - if isinstance(mod, TorchTensorRTModule) and mod.engine is not None: - current = mod.runtime_settings - if mod in self._saved: - # The same TRTModule appears under multiple targets in the - # list (or the tree contains a cycle). Don't snapshot twice. - continue - self._saved[mod] = current - merged = current.merge(**self._overrides) - mod.runtime_settings = merged return self._targets if self._yield_tuple else self._targets[0] def __exit__(self, *args: Any) -> None: From 07a6f6d5079c0e99a82bc97c74385f3066c802ed Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Mon, 24 Aug 2026 11:10:29 -0700 Subject: [PATCH 04/11] test: cover applying runtime settings across module and module-less engines New test_005_apply_runtime_settings.py: module-less engine cache round-trip, path-string rejection, cuda_graph_strategy dispatch, ExportedProgram reaching shared engine objects, mixed-target atomic failure, and CM raise with apply_runtime_settings in the message. All AOT tests are shown failing on base by ImportError (function does not exist) plus assertions base cannot satisfy. Co-Authored-By: Claude Sonnet 4.6 --- .../test_005_apply_runtime_settings.py | 288 ++++++++++++++++++ 1 file changed, 288 insertions(+) create mode 100644 tests/py/dynamo/runtime/test_005_apply_runtime_settings.py diff --git a/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py b/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py new file mode 100644 index 0000000000..de66dac2f3 --- /dev/null +++ b/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py @@ -0,0 +1,288 @@ +# type: ignore +"""Tests for apply_runtime_settings and the updated CM raise behaviour. + +Tests that use save/load run both the module (in-process) and module-less +(AOT-loaded) paths. The AOT tests additionally verify that the CM: + +* raises TypeError on module-less engines (commit 3) +* names apply_runtime_settings in the error message +""" + +import os +import tempfile +import unittest + +import torch +import torch_tensorrt +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt._features import ENABLED_FEATURES +from torch_tensorrt.runtime import ( + RuntimeCache, + RuntimeSettings, + apply_runtime_settings, + runtime_cache, + runtime_config, +) + + +class SimpleModel(torch.nn.Module): + def forward(self, x): + return torch.relu(x) + 1.0 + + +def _compile_simple(): + model = SimpleModel().eval().cuda() + inputs = [torch.randn(2, 3).cuda()] + compiled = torch_tensorrt.compile( + model, + ir="dynamo", + inputs=inputs, + min_block_size=1, + ) + torch._dynamo.reset() + return compiled, inputs + + +def _save_load(compiled, inputs): + """Save ``compiled`` to a temp file and return the loaded ExportedProgram and GraphModule.""" + with tempfile.NamedTemporaryFile(suffix=".ep", delete=False) as f: + ep_path = f.name + try: + torch_tensorrt.save(compiled, ep_path, arg_inputs=inputs) + loaded_ep = torch_tensorrt.load(ep_path) + finally: + try: + os.unlink(ep_path) + except OSError: + pass + loaded_gm = loaded_ep.module() if hasattr(loaded_ep, "module") else loaded_ep + return loaded_ep, loaded_gm + + +# --------------------------------------------------------------------------- +# Tests that do NOT require an RTX build +# --------------------------------------------------------------------------- + + +class TestApplyRuntimeSettingsTypeErrors(TestCase): + """Rejection of bad arguments; no engine compile required.""" + + def test_settings_wrong_type_raises(self): + model = torch.nn.Linear(3, 3).cuda() + with self.assertRaises(TypeError) as cm: + apply_runtime_settings(model, {"cuda_graph_strategy": "disabled"}) + self.assertIn("RuntimeSettings", str(cm.exception)) + + def test_target_wrong_type_raises(self): + with self.assertRaises(TypeError): + apply_runtime_settings("not_a_module", RuntimeSettings(runtime_cache=None)) + + def test_zero_engines_raises(self): + # A plain nn.Module has no TRT engines. + model = torch.nn.Linear(3, 3).cuda() + with self.assertRaises(RuntimeError) as cm: + apply_runtime_settings(model, RuntimeSettings(runtime_cache=None)) + self.assertIn("no TRT engines", str(cm.exception)) + + +# --------------------------------------------------------------------------- +# Tests that require TRT-RTX +# --------------------------------------------------------------------------- + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "apply_runtime_settings dispatch requires TRT-RTX", +) +class TestApplyRuntimeSettingsModuleOwned(TestCase): + """Module-owned engines: string cache still accepted (module owns it).""" + + def test_module_path_string_accepted(self): + compiled, inputs = _compile_simple() + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + n = apply_runtime_settings( + compiled, RuntimeSettings(runtime_cache=cache_path) + ) + self.assertGreaterEqual(n, 1) + _ = compiled(*inputs) + self.assertTrue(os.path.exists(cache_path)) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + + def test_returns_engine_count(self): + compiled, _ = _compile_simple() + n = apply_runtime_settings(compiled, RuntimeSettings(runtime_cache=None)) + self.assertGreaterEqual(n, 1) + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "AOT load tests require TRT-RTX", +) +class TestApplyRuntimeSettingsModuleLess(TestCase): + """Module-less engines from save/load.""" + + def test_path_string_raises_for_module_less_engine(self): + compiled, inputs = _compile_simple() + _, loaded_gm = _save_load(compiled, inputs) + with self.assertRaises(TypeError) as cm: + apply_runtime_settings( + loaded_gm, + RuntimeSettings(runtime_cache="/tmp/should_not_be_created.bin"), + ) + msg = str(cm.exception) + self.assertIn("runtime_cache", msg) + self.assertIn("RuntimeCache", msg) + + def test_default_runtime_settings_raises_for_module_less_engine(self): + # RuntimeSettings() default runtime_cache is a path string — a common footgun. + compiled, inputs = _compile_simple() + _, loaded_gm = _save_load(compiled, inputs) + with self.assertRaises(TypeError) as cm: + apply_runtime_settings(loaded_gm, RuntimeSettings()) + self.assertIn("runtime_cache", str(cm.exception)) + + def test_none_cache_applies_and_forward_runs(self): + compiled, inputs = _compile_simple() + _, loaded_gm = _save_load(compiled, inputs) + n = apply_runtime_settings(loaded_gm, RuntimeSettings(runtime_cache=None)) + self.assertGreaterEqual(n, 1) + out = loaded_gm(*inputs) + self.assertEqual(out.shape, inputs[0].shape) + + def test_runtime_cache_applies_and_persists(self): + compiled, inputs = _compile_simple() + _, loaded_gm = _save_load(compiled, inputs) + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + cache = RuntimeCache(path=cache_path, autosave_on_del=False) + n = apply_runtime_settings( + loaded_gm, + RuntimeSettings(runtime_cache=cache), + ) + self.assertGreaterEqual(n, 1) + _ = loaded_gm(*inputs) + self.assertTrue(cache.has_cache()) + cache.save() + size = os.path.getsize(cache_path) + self.assertGreater(size, 0) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + + def test_cuda_graph_strategy_field_takes_effect(self): + """Non-cache field applied to a module-less engine; proves field-agnostic dispatch.""" + compiled, inputs = _compile_simple() + _, loaded_gm = _save_load(compiled, inputs) + n = apply_runtime_settings( + loaded_gm, + RuntimeSettings( + cuda_graph_strategy="whole_graph_capture", + runtime_cache=None, + ), + ) + self.assertGreaterEqual(n, 1) + # Forward must still run after strategy change. + out = loaded_gm(*inputs) + self.assertEqual(out.shape, inputs[0].shape) + + def test_exported_program_reaches_same_engines_as_module(self): + """apply_runtime_settings on ExportedProgram reaches the engines ep.module() uses.""" + compiled, inputs = _compile_simple() + loaded_ep, loaded_gm = _save_load(compiled, inputs) + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + cache = RuntimeCache(path=cache_path, autosave_on_del=False) + # Apply via ExportedProgram. + n_ep = apply_runtime_settings( + loaded_ep, RuntimeSettings(runtime_cache=cache) + ) + self.assertGreaterEqual(n_ep, 1) + # Running via ep.module() uses the same engine objects -> cache populated. + _ = loaded_gm(*inputs) + self.assertTrue(cache.has_cache()) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "Mixed target test requires TRT-RTX", +) +class TestApplyRuntimeSettingsMixedTarget(TestCase): + """Module + module-less engines in one call: string must fail whole-call.""" + + def test_mixed_target_string_fails_atomically(self): + compiled, inputs = _compile_simple() + _, loaded_gm = _save_load(compiled, inputs) + + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + TorchTensorRTModule, + ) + + # Snapshot prior settings on the module engine. + prior = { + mod: mod.runtime_settings + for _, mod in compiled.named_modules() + if isinstance(mod, TorchTensorRTModule) and mod.engine is not None + } + + with self.assertRaises(TypeError): + apply_runtime_settings( + [compiled, loaded_gm], + RuntimeSettings(runtime_cache="/tmp/should_not_apply.bin"), + ) + + # Module engine settings must be unchanged (validate-before-mutate). + for mod, saved in prior.items(): + self.assertEqual(mod.runtime_settings, saved) + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "CM raise tests require TRT-RTX", +) +class TestContextManagerRaisesOnModuleLess(TestCase): + """runtime_config and runtime_cache raise on module-less engines.""" + + def test_runtime_config_raises_on_loaded_gm(self): + compiled, inputs = _compile_simple() + _, loaded_gm = _save_load(compiled, inputs) + with self.assertRaises(TypeError) as cm: + with runtime_config(loaded_gm, runtime_cache=None): + pass + msg = str(cm.exception) + self.assertIn("apply_runtime_settings", msg) + + def test_runtime_cache_raises_on_loaded_gm(self): + compiled, inputs = _compile_simple() + _, loaded_gm = _save_load(compiled, inputs) + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + with self.assertRaises(TypeError) as cm: + with runtime_cache(loaded_gm, cache_path): + pass + msg = str(cm.exception) + self.assertIn("apply_runtime_settings", msg) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + + +if __name__ == "__main__": + run_tests() From 0ca4ac4960108701a0085e00cd1bae0c4f5675ec Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Mon, 24 Aug 2026 11:12:48 -0700 Subject: [PATCH 05/11] test: collapse the duplicated runtime-settings helpers onto the public API Replace the local _apply_runtime_settings helper in 7 test files (31 call sites) with the public apply_runtime_settings. The local bodies were semantically identical (walk named_modules, set mod.runtime_settings = rs); the public entry point is a strict superset. Last commit so reviewers can skip the mechanical diff. Co-Authored-By: Claude Sonnet 4.6 --- .../models/test_cuda_graph_strategy_models.py | 21 +++-------- ...t_dynamic_shapes_kernel_strategy_models.py | 17 ++------- .../models/test_runtime_cache_models.py | 37 +++++-------------- .../dynamo/runtime/test_000_runtime_cache.py | 27 ++++---------- .../runtime/test_001_cuda_graph_strategy.py | 15 ++------ ...test_001_dynamic_shapes_kernel_strategy.py | 15 ++------ .../runtime/test_004_runtime_settings.py | 14 ++----- 7 files changed, 35 insertions(+), 111 deletions(-) diff --git a/tests/py/dynamo/models/test_cuda_graph_strategy_models.py b/tests/py/dynamo/models/test_cuda_graph_strategy_models.py index beeff8cccd..b210678dc5 100644 --- a/tests/py/dynamo/models/test_cuda_graph_strategy_models.py +++ b/tests/py/dynamo/models/test_cuda_graph_strategy_models.py @@ -5,18 +5,7 @@ import torch_tensorrt as torchtrt from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt._features import ENABLED_FEATURES -from torch_tensorrt.runtime import RuntimeSettings - - -def _apply_runtime_settings(compiled, rs): - """Walk a compiled module and apply RuntimeSettings to every TRT submodule.""" - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( - TorchTensorRTModule, - ) - - for _, m in compiled.named_modules(): - if isinstance(m, TorchTensorRTModule): - m.runtime_settings = rs +from torch_tensorrt.runtime import RuntimeSettings, apply_runtime_settings class ConvModel(torch.nn.Module): @@ -71,7 +60,7 @@ def test_resnet18_whole_graph_capture(self): use_python_runtime=True, min_block_size=1, ) - _apply_runtime_settings( + apply_runtime_settings( compiled, RuntimeSettings(cuda_graph_strategy="whole_graph_capture") ) torch._dynamo.reset() @@ -105,7 +94,7 @@ def test_resnet18_disabled_strategy(self): use_python_runtime=True, min_block_size=1, ) - _apply_runtime_settings( + apply_runtime_settings( compiled, RuntimeSettings(cuda_graph_strategy="disabled") ) torch._dynamo.reset() @@ -145,7 +134,7 @@ def test_dynamic_batch_whole_graph_capture(self): use_python_runtime=True, min_block_size=1, ) - _apply_runtime_settings( + apply_runtime_settings( compiled, RuntimeSettings(cuda_graph_strategy="whole_graph_capture") ) torch._dynamo.reset() @@ -181,7 +170,7 @@ def test_dynamic_batch_with_subgraph_cudagraphs(self): use_python_runtime=True, min_block_size=1, ) - _apply_runtime_settings( + apply_runtime_settings( compiled, RuntimeSettings(cuda_graph_strategy="whole_graph_capture") ) torch._dynamo.reset() diff --git a/tests/py/dynamo/models/test_dynamic_shapes_kernel_strategy_models.py b/tests/py/dynamo/models/test_dynamic_shapes_kernel_strategy_models.py index 962f0e9955..ae01be351b 100644 --- a/tests/py/dynamo/models/test_dynamic_shapes_kernel_strategy_models.py +++ b/tests/py/dynamo/models/test_dynamic_shapes_kernel_strategy_models.py @@ -6,18 +6,7 @@ from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt._features import ENABLED_FEATURES from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity -from torch_tensorrt.runtime import RuntimeSettings - - -def _apply_runtime_settings(compiled, rs): - """Walk a compiled module and apply RuntimeSettings to every TRT submodule.""" - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( - TorchTensorRTModule, - ) - - for _, m in compiled.named_modules(): - if isinstance(m, TorchTensorRTModule): - m.runtime_settings = rs +from torch_tensorrt.runtime import RuntimeSettings, apply_runtime_settings @unittest.skipIf( @@ -50,7 +39,7 @@ def _compile_and_verify(self, model, strategy): use_python_runtime=True, min_block_size=1, ) - _apply_runtime_settings( + apply_runtime_settings( compiled, RuntimeSettings(dynamic_shapes_kernel_specialization_strategy=strategy), ) @@ -118,7 +107,7 @@ def forward(self, x): use_python_runtime=True, min_block_size=1, ) - _apply_runtime_settings( + apply_runtime_settings( compiled, RuntimeSettings(dynamic_shapes_kernel_specialization_strategy=strategy), ) diff --git a/tests/py/dynamo/models/test_runtime_cache_models.py b/tests/py/dynamo/models/test_runtime_cache_models.py index 61d7b3670b..9716ab5963 100644 --- a/tests/py/dynamo/models/test_runtime_cache_models.py +++ b/tests/py/dynamo/models/test_runtime_cache_models.py @@ -11,18 +11,7 @@ from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt._features import ENABLED_FEATURES from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity -from torch_tensorrt.runtime import RuntimeSettings - - -def _apply_runtime_settings(compiled, rs): - """Walk a compiled module and apply RuntimeSettings to every TRT submodule.""" - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( - TorchTensorRTModule, - ) - - for _, m in compiled.named_modules(): - if isinstance(m, TorchTensorRTModule): - m.runtime_settings = rs +from torch_tensorrt.runtime import RuntimeSettings, apply_runtime_settings @unittest.skipIf( @@ -57,9 +46,7 @@ def test_resnet18_with_runtime_cache(self): use_python_runtime=True, min_block_size=1, ) - _apply_runtime_settings( - compiled, RuntimeSettings(runtime_cache=self.cache_path) - ) + apply_runtime_settings(compiled, RuntimeSettings(runtime_cache=self.cache_path)) ref_output = model(input_tensor) trt_output = compiled(input_tensor) @@ -96,7 +83,7 @@ def test_resnet18_cache_reuse(self): # First compilation — cold cache compiled1 = torchtrt.compile(model, **compile_kwargs) - _apply_runtime_settings(compiled1, rs) + apply_runtime_settings(compiled1, rs) _ = compiled1(input_tensor) del compiled1 gc.collect() @@ -106,7 +93,7 @@ def test_resnet18_cache_reuse(self): # Second compilation — warm cache compiled2 = torchtrt.compile(model, **compile_kwargs) - _apply_runtime_settings(compiled2, rs) + apply_runtime_settings(compiled2, rs) output2 = compiled2(input_tensor) cos_sim = cosine_similarity(ref_output, output2) @@ -135,9 +122,7 @@ def test_mobilenet_v2_with_runtime_cache(self): use_python_runtime=True, min_block_size=1, ) - _apply_runtime_settings( - compiled, RuntimeSettings(runtime_cache=self.cache_path) - ) + apply_runtime_settings(compiled, RuntimeSettings(runtime_cache=self.cache_path)) ref_output = model(input_tensor) trt_output = compiled(input_tensor) @@ -194,9 +179,7 @@ def forward(self, x): use_python_runtime=True, min_block_size=1, ) - _apply_runtime_settings( - compiled, RuntimeSettings(runtime_cache=self.cache_path) - ) + apply_runtime_settings(compiled, RuntimeSettings(runtime_cache=self.cache_path)) # Test with batch size 1 input_bs1 = torch.randn(1, 3, 32, 32).cuda() @@ -253,7 +236,7 @@ def forward(self, x): # First run with batch=2 — saves cache compiled1 = torchtrt.compile(model, **compile_kwargs) - _apply_runtime_settings(compiled1, rs) + apply_runtime_settings(compiled1, rs) input_bs2 = torch.randn(2, 3, 16, 16).cuda() _ = compiled1(input_bs2) del compiled1 @@ -263,7 +246,7 @@ def forward(self, x): # Second run with batch=3 — loads same cache compiled2 = torchtrt.compile(model, **compile_kwargs) - _apply_runtime_settings(compiled2, rs) + apply_runtime_settings(compiled2, rs) input_bs3 = torch.randn(3, 3, 16, 16).cuda() ref_bs3 = model(input_bs3) out_bs3 = compiled2(input_bs3) @@ -316,7 +299,7 @@ def forward(self, x): # Cold cache compilation + inference compiled1 = torchtrt.compile(model, **compile_kwargs) - _apply_runtime_settings(compiled1, rs) + apply_runtime_settings(compiled1, rs) torch.cuda.synchronize() start = time.perf_counter() _ = compiled1(input_tensor) @@ -328,7 +311,7 @@ def forward(self, x): # Warm cache compilation + inference compiled2 = torchtrt.compile(model, **compile_kwargs) - _apply_runtime_settings(compiled2, rs) + apply_runtime_settings(compiled2, rs) torch.cuda.synchronize() start = time.perf_counter() _ = compiled2(input_tensor) diff --git a/tests/py/dynamo/runtime/test_000_runtime_cache.py b/tests/py/dynamo/runtime/test_000_runtime_cache.py index 1c43e63a9b..010b03be59 100644 --- a/tests/py/dynamo/runtime/test_000_runtime_cache.py +++ b/tests/py/dynamo/runtime/test_000_runtime_cache.py @@ -13,7 +13,11 @@ from torch_tensorrt._features import ENABLED_FEATURES from torch_tensorrt.dynamo._defaults import TIMING_CACHE_PATH from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity -from torch_tensorrt.runtime import RuntimeSettings, runtime_cache +from torch_tensorrt.runtime import ( + RuntimeSettings, + apply_runtime_settings, + runtime_cache, +) class SimpleModel(torch.nn.Module): @@ -35,21 +39,6 @@ def _fresh_conv_model_and_inputs(seed=0): return ConvModel().eval().cuda(), [torch.randn(2, 3, 16, 16).cuda()] -def _apply_runtime_settings(compiled, rs): - """Apply ``RuntimeSettings`` to every ``TorchTensorRTModule`` under ``compiled``. - - Mirrors what user code would do: ``mod.runtime_settings = rs`` after - compile. The compile-time hint (``torchtrt.compile(runtime_settings=...)``) - was dropped now that lazy ``IExecutionContext`` creation absorbs the - one-create benefit it used to provide. - """ - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule - - for _, mod in compiled.named_modules(): - if isinstance(mod, TorchTensorRTModule): - mod.runtime_settings = rs - - def _compile(model, inputs, *, runtime_cache_path=None): """Compile ``model`` through whichever runtime the build selects. @@ -65,7 +54,7 @@ def _compile(model, inputs, *, runtime_cache_path=None): ) torch._dynamo.reset() if runtime_cache_path is not None: - _apply_runtime_settings( + apply_runtime_settings( compiled, RuntimeSettings(runtime_cache=runtime_cache_path) ) return compiled @@ -865,7 +854,7 @@ def test_compiled_module_persists_its_implicit_cache(self): path = os.path.join(tmp, "rc.bin") model, inputs = _fresh_conv_model_and_inputs() compiled = _compile(model, inputs) - _apply_runtime_settings(compiled, RuntimeSettings(runtime_cache=path)) + apply_runtime_settings(compiled, RuntimeSettings(runtime_cache=path)) compiled(*inputs) del compiled gc.collect() @@ -880,7 +869,7 @@ def test_empty_path_string_is_normalized_to_none(self): model, inputs = _fresh_conv_model_and_inputs() compiled = _compile(model, inputs) - _apply_runtime_settings(compiled, RuntimeSettings(runtime_cache="")) + apply_runtime_settings(compiled, RuntimeSettings(runtime_cache="")) # Engine-flavor agnostic: the module's resolved settings are what gets # dispatched, on both the Python and cpp runtimes. diff --git a/tests/py/dynamo/runtime/test_001_cuda_graph_strategy.py b/tests/py/dynamo/runtime/test_001_cuda_graph_strategy.py index 4d5032ec68..154bde8f95 100644 --- a/tests/py/dynamo/runtime/test_001_cuda_graph_strategy.py +++ b/tests/py/dynamo/runtime/test_001_cuda_graph_strategy.py @@ -5,7 +5,7 @@ from parameterized import parameterized from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt._features import ENABLED_FEATURES -from torch_tensorrt.runtime import RuntimeSettings +from torch_tensorrt.runtime import RuntimeSettings, apply_runtime_settings class CudaGraphConvModel(torch.nn.Module): @@ -17,15 +17,6 @@ def forward(self, x): return torch.relu(self.conv(x)) -def _apply_runtime_settings(compiled, rs): - """Apply ``RuntimeSettings`` to every inner ``TorchTensorRTModule``.""" - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule - - for _, mod in compiled.named_modules(): - if isinstance(mod, TorchTensorRTModule): - mod.runtime_settings = rs - - def _compile_conv(strategy): """Compile CudaGraphConvModel + apply cuda_graph_strategy post-compile.""" model = CudaGraphConvModel().eval().cuda() @@ -38,7 +29,7 @@ def _compile_conv(strategy): min_block_size=1, ) torch._dynamo.reset() - _apply_runtime_settings(compiled, RuntimeSettings(cuda_graph_strategy=strategy)) + apply_runtime_settings(compiled, RuntimeSettings(cuda_graph_strategy=strategy)) return compiled, inputs @@ -67,7 +58,7 @@ def _compile_simple(*, runtime_settings=None): ) torch._dynamo.reset() if runtime_settings is not None: - _apply_runtime_settings(compiled, runtime_settings) + apply_runtime_settings(compiled, runtime_settings) return compiled diff --git a/tests/py/dynamo/runtime/test_001_dynamic_shapes_kernel_strategy.py b/tests/py/dynamo/runtime/test_001_dynamic_shapes_kernel_strategy.py index af606d4be1..0b945bff61 100644 --- a/tests/py/dynamo/runtime/test_001_dynamic_shapes_kernel_strategy.py +++ b/tests/py/dynamo/runtime/test_001_dynamic_shapes_kernel_strategy.py @@ -5,7 +5,7 @@ from parameterized import parameterized from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt._features import ENABLED_FEATURES -from torch_tensorrt.runtime import RuntimeSettings +from torch_tensorrt.runtime import RuntimeSettings, apply_runtime_settings _STRATEGIES = [("lazy",), ("eager",), ("none",)] @@ -42,22 +42,13 @@ def _compile_dynamic_conv(strategy): min_block_size=1, ) torch._dynamo.reset() - _apply_runtime_settings( + apply_runtime_settings( compiled, RuntimeSettings(dynamic_shapes_kernel_specialization_strategy=strategy), ) return compiled -def _apply_runtime_settings(compiled, rs): - """Apply ``RuntimeSettings`` to every inner ``TorchTensorRTModule``.""" - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule - - for _, mod in compiled.named_modules(): - if isinstance(mod, TorchTensorRTModule): - mod.runtime_settings = rs - - def _compile_simple(*, runtime_settings=None): """Compile SimpleModel with dynamic shapes through the build-selected runtime.""" model = SimpleModel().eval().cuda() @@ -77,7 +68,7 @@ def _compile_simple(*, runtime_settings=None): ) torch._dynamo.reset() if runtime_settings is not None: - _apply_runtime_settings(compiled, runtime_settings) + apply_runtime_settings(compiled, runtime_settings) return compiled diff --git a/tests/py/dynamo/runtime/test_004_runtime_settings.py b/tests/py/dynamo/runtime/test_004_runtime_settings.py index 97d36863c4..5facdec383 100644 --- a/tests/py/dynamo/runtime/test_004_runtime_settings.py +++ b/tests/py/dynamo/runtime/test_004_runtime_settings.py @@ -11,6 +11,7 @@ from torch_tensorrt.runtime import ( RuntimeCache, RuntimeSettings, + apply_runtime_settings, runtime_config, ) @@ -20,15 +21,6 @@ def forward(self, x): return torch.relu(x) + 1.0 -def _apply_runtime_settings(compiled, rs): - """Apply ``RuntimeSettings`` to every inner ``TorchTensorRTModule``.""" - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule - - for _, mod in compiled.named_modules(): - if isinstance(mod, TorchTensorRTModule): - mod.runtime_settings = rs - - def _compile_simple(*, runtime_settings=None): model = SimpleModel().eval().cuda() inputs = [ @@ -47,7 +39,7 @@ def _compile_simple(*, runtime_settings=None): ) torch._dynamo.reset() if runtime_settings is not None: - _apply_runtime_settings(compiled, runtime_settings) + apply_runtime_settings(compiled, runtime_settings) return compiled @@ -310,7 +302,7 @@ def test_setter_after_load_state_dict_does_not_raise(self): dst.load_state_dict(state) # routes through set_extra_state # B2 used to AttributeError on the next line because the slot # wasn't initialized after ``set_extra_state``. - _apply_runtime_settings( + apply_runtime_settings( dst, RuntimeSettings(cuda_graph_strategy="whole_graph_capture") ) From 7dd89af68a8f23e48aeee16a0d27986a290b5cb9 Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Mon, 24 Aug 2026 12:42:03 -0700 Subject: [PATCH 06/11] docs: document and test caller ownership of RuntimeCache.load on module-less engines The module path auto-calls RuntimeCache.load via _resolve_runtime_cache, warming the implicit handle from disk before first use. There is no equivalent hook on the module-less path -- the caller owns the handle and must warm it. This asymmetry was intentional but undocumented. - apply_runtime_settings docstring: state that caller owns .load() as well as .save() on the module-less path, and explain why the module path is different (_resolve_runtime_cache warm-load). - runtime_settings.rst: add ownership paragraph to the AOT subsection. - test_005_apply_runtime_settings.py: add test_warm_load_bytes_transferred to TestApplyRuntimeSettingsModuleLess; writes a cache file, warms a new RuntimeCache via load_from_stream, attaches via apply_runtime_settings, and asserts load_from_stream returned > 0 bytes (non-vacuous). Co-Authored-By: Claude Sonnet 4.6 --- .../runtime_performance/runtime_settings.rst | 7 +++ py/torch_tensorrt/runtime/_runtime_config.py | 9 ++++ .../test_005_apply_runtime_settings.py | 48 +++++++++++++++++++ 3 files changed, 64 insertions(+) diff --git a/docsrc/user_guide/runtime_performance/runtime_settings.rst b/docsrc/user_guide/runtime_performance/runtime_settings.rst index 5784b0bcc3..bf5cb8a679 100644 --- a/docsrc/user_guide/runtime_performance/runtime_settings.rst +++ b/docsrc/user_guide/runtime_performance/runtime_settings.rst @@ -174,6 +174,13 @@ there is no module to build and save the handle. If you call ``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()``: diff --git a/py/torch_tensorrt/runtime/_runtime_config.py b/py/torch_tensorrt/runtime/_runtime_config.py index 135add83fb..eb091bc728 100644 --- a/py/torch_tensorrt/runtime/_runtime_config.py +++ b/py/torch_tensorrt/runtime/_runtime_config.py @@ -583,6 +583,15 @@ def apply_runtime_settings( ``RuntimeSettings(runtime_cache=None)`` or supply a :class:`RuntimeCache` explicitly. + **Warm-load is also the caller's responsibility.** When a + :class:`TorchTensorRTModule` is present it calls :meth:`RuntimeCache.load` + automatically (via ``_resolve_runtime_cache``), so cached kernels are + available from the first execute without any caller action. For module-less + engines there is no equivalent hook — the caller must call + :meth:`RuntimeCache.load` (or :meth:`RuntimeCache.load_from_stream`) before + passing the handle to :func:`apply_runtime_settings`, or the engine will + start with an empty cache regardless of what is on disk. + Settings are never serialized; they do not survive :func:`torch_tensorrt.save`. Re-apply after each :func:`torch_tensorrt.load`. """ diff --git a/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py b/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py index de66dac2f3..67ef0d83e9 100644 --- a/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py +++ b/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py @@ -8,6 +8,7 @@ * names apply_runtime_settings in the error message """ +import io import os import tempfile import unittest @@ -216,6 +217,53 @@ def test_exported_program_reaches_same_engines_as_module(self): except OSError: pass + def test_warm_load_bytes_transferred(self): + """Caller must call .load() / .load_from_stream() to warm the cache before attach. + + The module path auto-calls load() via _resolve_runtime_cache; there is no + equivalent hook on the module-less path. Verify warm bytes actually land: + load_from_stream returns a byte count, so assertGreater(..., 0) is the + non-vacuous check. + """ + compiled, inputs = _compile_simple() + _, loaded_gm = _save_load(compiled, inputs) + + # Pass 1: populate a cache to get real bytes on disk. + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + cache_populate = RuntimeCache(path=cache_path, autosave_on_del=False) + apply_runtime_settings( + loaded_gm, RuntimeSettings(runtime_cache=cache_populate) + ) + _ = loaded_gm(*inputs) + cache_populate.save() + with open(cache_path, "rb") as fh: + saved_bytes = fh.read() + self.assertGreater(len(saved_bytes), 0, "first save produced empty file") + + # Pass 2: warm a new cache from those bytes via load_from_stream. + _, loaded_gm2 = _save_load(compiled, inputs) + cache_warm = RuntimeCache(autosave_on_del=False) + n_bytes = cache_warm.load_from_stream(io.BytesIO(saved_bytes)) + self.assertGreater( + n_bytes, + 0, + "load_from_stream transferred 0 bytes — warm load did nothing", + ) + + n_engines = apply_runtime_settings( + loaded_gm2, RuntimeSettings(runtime_cache=cache_warm) + ) + self.assertGreaterEqual(n_engines, 1) + _ = loaded_gm2(*inputs) + self.assertTrue(cache_warm.has_cache()) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + @unittest.skipIf( not ENABLED_FEATURES.tensorrt_rtx, From cac959a7d5b65e7c5677732387155d8a2e7e8baa Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Fri, 28 Aug 2026 11:03:26 -0700 Subject: [PATCH 07/11] =?UTF-8?q?refactor:=20inline=20=5Fsend=5Fto=5Fengin?= =?UTF-8?q?e=20=E2=80=94=20it=20was=20a=20one-line=20delegate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _send_to_engine called _send_settings_to_engine verbatim; the call site now imports and calls _send_settings_to_engine directly. Co-Authored-By: Claude Sonnet 4.6 --- .../dynamo/runtime/_TorchTensorRTModule.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py index f97206dc9b..84e2c36716 100644 --- a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py +++ b/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,12 +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.runtime._runtime_config import _send_settings_to_engine - - _send_settings_to_engine(self.engine, rs) - def setup_engine(self) -> None: """ Setup engine for a module which has deferred engine setup. From 1574143e367af9fcbdcf502e31f83aed45d508fa Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Fri, 28 Aug 2026 11:11:12 -0700 Subject: [PATCH 08/11] docs: promote apply_runtime_settings as the canonical permanent-assign API Direct assignment (mod.runtime_settings = rs) was listed as a separate "way" to apply settings; it is now an implementation detail that apply_runtime_settings calls internally for module-owned engines. - runtime_settings.rst: collapse four ways -> three ways; remove the Direct assignment subsection; broaden the apply_runtime_settings section to cover both compiled and AOT-loaded targets; update the Advanced, Best practices, and Quick reference sections to use apply_runtime_settings throughout; replace "Setter is per-module" footgun with a positive note about automatic subgraph walking. - _runtime_config.py module docstring: four concepts -> three; remove the Programmatic bullet from the usage list. Co-Authored-By: Claude Sonnet 4.6 --- .../runtime_performance/runtime_settings.rst | 152 +++++++----------- py/torch_tensorrt/runtime/_runtime_config.py | 9 +- 2 files changed, 63 insertions(+), 98 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/runtime_settings.rst b/docsrc/user_guide/runtime_performance/runtime_settings.rst index bf5cb8a679..23dd163bcf 100644 --- a/docsrc/user_guide/runtime_performance/runtime_settings.rst +++ b/docsrc/user_guide/runtime_performance/runtime_settings.rst @@ -29,21 +29,55 @@ emits a ``UserWarning``. ---- -The four ways to apply settings -------------------------------- +The three ways to apply settings +--------------------------------- -Direct assignment — permanent -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +``apply_runtime_settings(...)`` — permanent apply +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Use for any permanent assignment — whether the engine is wrapped in a +:class:`TorchTensorRTModule` (in-process compiled result) or bare (AOT-loaded +artifact with no module wrapper): .. code-block:: python import torch_tensorrt - from torch_tensorrt.runtime import RuntimeSettings + from torch_tensorrt.runtime import RuntimeSettings, apply_runtime_settings + # In-process compiled model — one or more TRT subgraphs, all updated: 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")) + + # AOT-loaded artifact — no TorchTensorRTModule, same call: + ep = torch_tensorrt.load("model.ep") + gm = ep.module() + from torch_tensorrt.runtime import RuntimeCache + cache = RuntimeCache(path="/var/cache/jit.bin") + cache.load() # warm from disk — caller's responsibility on the module-less path + apply_runtime_settings(gm, RuntimeSettings(runtime_cache=cache)) + out = gm(x) + cache.save() + +:func:`apply_runtime_settings` returns the number of engines updated and raises +:exc:`RuntimeError` if no TRT engines are found. For module-less engines, +``settings.runtime_cache`` must be ``None`` or a :class:`RuntimeCache` you own +(a path string raises :exc:`TypeError` — there is no module to build and save +the handle). See :func:`~torch_tensorrt.runtime.apply_runtime_settings` for +the full ownership rules. + +: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:: -Use when you want the setting to apply for the module's lifetime. + Runtime settings are never serialized. They do not survive + :func:`torch_tensorrt.save`; re-apply after each :func:`torch_tensorrt.load`. ``runtime_config(...)`` context manager — scoped override ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -136,64 +170,6 @@ 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 @@ -341,10 +317,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) @@ -354,7 +330,7 @@ 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() @@ -372,7 +348,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**: @@ -381,7 +357,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 @@ -391,8 +367,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 @@ -412,10 +388,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 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -434,22 +410,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 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -471,8 +439,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 @@ -489,5 +457,3 @@ 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(...))`` diff --git a/py/torch_tensorrt/runtime/_runtime_config.py b/py/torch_tensorrt/runtime/_runtime_config.py index eb091bc728..fed30d2d19 100644 --- a/py/torch_tensorrt/runtime/_runtime_config.py +++ b/py/torch_tensorrt/runtime/_runtime_config.py @@ -1,6 +1,6 @@ """Runtime settings + the TRTRuntimeConfig shim + the ``runtime_config`` CM. -This module groups four closely related concepts together: +This module groups three closely related concepts together: * :class:`RuntimeSettings` -- the user-facing, frozen dataclass of runtime-only knobs sampled at IExecutionContext creation (cuda_graph_strategy, @@ -16,12 +16,11 @@ * :func:`apply_runtime_settings` -- permanent apply to every TRT engine under a target, including engines loaded without a :class:`TorchTensorRTModule`. -Three ways to use ``RuntimeSettings``: +Two ways to use ``RuntimeSettings``: 1. **Runtime context manager** -- toggle settings inside a ``with`` block. -2. **Programmatic** -- assign ``module.runtime_settings = rs`` directly. -3. **AOT artifact** -- call :func:`apply_runtime_settings` on a loaded - :class:`ExportedProgram` or ``GraphModule``. +2. **Permanent apply** -- call :func:`apply_runtime_settings` on any compiled + module, :class:`ExportedProgram`, or AOT-loaded ``GraphModule``. ``RuntimeSettings`` is intentionally NOT part of ``CompilationSettings`` and is NOT serialized into the engine tuple. It's purely an in-memory initialization From 863e7d1a694f937b02bc82292c1647ecfaa22beb Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Fri, 28 Aug 2026 18:09:02 -0700 Subject: [PATCH 09/11] =?UTF-8?q?docs:=20reorder=20apply=5Fruntime=5Fsetti?= =?UTF-8?q?ngs=20section=20=E2=80=94=20modules=20first,=20AOT=20after=20ca?= =?UTF-8?q?ller-owned=20cache?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The section now flows: compiled module example -> context managers -> caller-owned RuntimeCache -> AOT/ExportedProgram. The ExportedProgram and module-less ownership rules are deferred until after the caller-owned RuntimeCache concept is introduced, which they build on. Co-Authored-By: Claude Sonnet 4.6 --- .../runtime_performance/runtime_settings.rst | 77 +++++++++++-------- 1 file changed, 43 insertions(+), 34 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/runtime_settings.rst b/docsrc/user_guide/runtime_performance/runtime_settings.rst index 23dd163bcf..b293fad395 100644 --- a/docsrc/user_guide/runtime_performance/runtime_settings.rst +++ b/docsrc/user_guide/runtime_performance/runtime_settings.rst @@ -35,49 +35,19 @@ The three ways to apply settings ``apply_runtime_settings(...)`` — permanent apply ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Use for any permanent assignment — whether the engine is wrapped in a -:class:`TorchTensorRTModule` (in-process compiled result) or bare (AOT-loaded -artifact with no module wrapper): +Use for any permanent assignment on in-process compiled models: .. code-block:: python import torch_tensorrt from torch_tensorrt.runtime import RuntimeSettings, apply_runtime_settings - # In-process compiled model — one or more TRT subgraphs, all updated: mod = torch_tensorrt.compile(model, inputs=inputs) apply_runtime_settings(mod, RuntimeSettings(runtime_cache="/var/cache/jit.bin")) - # AOT-loaded artifact — no TorchTensorRTModule, same call: - ep = torch_tensorrt.load("model.ep") - gm = ep.module() - from torch_tensorrt.runtime import RuntimeCache - cache = RuntimeCache(path="/var/cache/jit.bin") - cache.load() # warm from disk — caller's responsibility on the module-less path - apply_runtime_settings(gm, RuntimeSettings(runtime_cache=cache)) - out = gm(x) - cache.save() - -:func:`apply_runtime_settings` returns the number of engines updated and raises -:exc:`RuntimeError` if no TRT engines are found. For module-less engines, -``settings.runtime_cache`` must be ``None`` or a :class:`RuntimeCache` you own -(a path string raises :exc:`TypeError` — there is no module to build and save -the handle). See :func:`~torch_tensorrt.runtime.apply_runtime_settings` for -the full ownership rules. - -: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`. +: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 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -334,6 +304,45 @@ Or with explicit save/load: 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 From fab9a2f66eb1ef2e4097e9a20c1c2e28f6ff77b1 Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Fri, 28 Aug 2026 18:58:18 -0700 Subject: [PATCH 10/11] refactor: consolidate apply_runtime_settings into test_004, clean up _iter_trt_engines - test_005_apply_runtime_settings.py removed; all tests moved into test_004_runtime_settings.py where they belong alongside the rest of the RuntimeSettings test suite. Added _compile_with_inputs() and _save_load() helpers for the save/load path. - _iter_trt_engines: add _visit_one compositor to eliminate the duplicated EP/Module dispatch in the sequence branch; replace hasattr(__iter__) guard with isinstance(Iterable) + positive scalar-type check; remove redundant list() conversion on self._targets. - Import Iterable from typing. Co-Authored-By: Claude Sonnet 4.6 --- py/torch_tensorrt/runtime/_runtime_config.py | 42 ++- .../runtime/test_004_runtime_settings.py | 285 +++++++++++++++ .../test_005_apply_runtime_settings.py | 336 ------------------ 3 files changed, 305 insertions(+), 358 deletions(-) delete mode 100644 tests/py/dynamo/runtime/test_005_apply_runtime_settings.py diff --git a/py/torch_tensorrt/runtime/_runtime_config.py b/py/torch_tensorrt/runtime/_runtime_config.py index fed30d2d19..845a2ff7ae 100644 --- a/py/torch_tensorrt/runtime/_runtime_config.py +++ b/py/torch_tensorrt/runtime/_runtime_config.py @@ -37,6 +37,7 @@ TYPE_CHECKING, Any, Dict, + Iterable, Optional, Sequence, Set, @@ -373,7 +374,7 @@ def __init__( def __enter__(self) -> Union["torch.nn.Module", Tuple["torch.nn.Module", ...]]: # Drain the traversal before mutating; raise on unsupported engines # before any settings are changed. - engines = list(_iter_trt_engines(list(self._targets))) + engines = list(_iter_trt_engines(self._targets)) module_less = [(owner, eng) for owner, eng in engines if owner is None] if module_less: @@ -524,32 +525,29 @@ def _visit_module(root: torch.nn.Module) -> Any: seen.add(id(obj)) yield (None, obj) - if isinstance(target_or_targets, torch.export.ExportedProgram): - yield from _visit_ep(target_or_targets) - elif isinstance(target_or_targets, torch.nn.Module): - yield from _visit_module(target_or_targets) - elif hasattr(target_or_targets, "__iter__") and not isinstance( - target_or_targets, (str, bytes) - ): + def _visit_one(t: Any) -> Any: + if isinstance(t, torch.export.ExportedProgram): + yield from _visit_ep(t) + elif isinstance(t, torch.nn.Module): + yield from _visit_module(t) + else: + raise TypeError( + f"_iter_trt_engines(): each target must be an nn.Module or " + f"ExportedProgram; got {type(t).__name__}. " + "For a torch_tensorrt.load() result, pass the ExportedProgram " + "directly or call .module() on it." + ) + + if isinstance(target_or_targets, (torch.nn.Module, torch.export.ExportedProgram)): + yield from _visit_one(target_or_targets) + elif isinstance(target_or_targets, Iterable): for t in target_or_targets: - if isinstance(t, torch.export.ExportedProgram): - yield from _visit_ep(t) - elif isinstance(t, torch.nn.Module): - yield from _visit_module(t) - else: - raise TypeError( - f"_iter_trt_engines(): each target must be an nn.Module or " - f"ExportedProgram; got {type(t).__name__}. " - "For a torch_tensorrt.load() result, pass the ExportedProgram " - "directly or call .module() on it." - ) + yield from _visit_one(t) else: raise TypeError( f"_iter_trt_engines(): target must be an nn.Module, an " f"ExportedProgram, or a sequence of those; got " - f"{type(target_or_targets).__name__}. " - "For a torch_tensorrt.load() result, pass the ExportedProgram " - "directly or call .module() on it." + f"{type(target_or_targets).__name__}." ) diff --git a/tests/py/dynamo/runtime/test_004_runtime_settings.py b/tests/py/dynamo/runtime/test_004_runtime_settings.py index 5facdec383..0fb482a35b 100644 --- a/tests/py/dynamo/runtime/test_004_runtime_settings.py +++ b/tests/py/dynamo/runtime/test_004_runtime_settings.py @@ -1,6 +1,9 @@ """Whitebox tests for the RuntimeSettings data model + dispatch.""" import dataclasses +import io +import os +import tempfile import unittest import torch @@ -12,6 +15,7 @@ RuntimeCache, RuntimeSettings, apply_runtime_settings, + runtime_cache, runtime_config, ) @@ -393,5 +397,286 @@ def test_enable_cudagraphs_strategy_kwarg_rejected_on_non_rtx(self): self.assertIn("TRT-RTX-only", str(cm.exception)) +def _compile_with_inputs(): + """Compile SimpleModel and return (compiled, inputs) for save/load tests.""" + model = SimpleModel().eval().cuda() + inputs = [torch.randn(2, 3).cuda()] + compiled = torchtrt.compile( + model, + ir="dynamo", + inputs=inputs, + min_block_size=1, + ) + torch._dynamo.reset() + return compiled, inputs + + +def _save_load(compiled, inputs): + """Save ``compiled`` to a temp file and return the loaded ExportedProgram and GraphModule.""" + with tempfile.NamedTemporaryFile(suffix=".ep", delete=False) as f: + ep_path = f.name + try: + torchtrt.save(compiled, ep_path, arg_inputs=inputs) + loaded_ep = torchtrt.load(ep_path) + finally: + try: + os.unlink(ep_path) + except OSError: + pass + loaded_gm = loaded_ep.module() if hasattr(loaded_ep, "module") else loaded_ep + return loaded_ep, loaded_gm + + +class TestApplyRuntimeSettingsTypeErrors(TestCase): + """Rejection of bad arguments; no engine compile required.""" + + def test_settings_wrong_type_raises(self): + model = torch.nn.Linear(3, 3).cuda() + with self.assertRaises(TypeError) as cm: + apply_runtime_settings(model, {"cuda_graph_strategy": "disabled"}) + self.assertIn("RuntimeSettings", str(cm.exception)) + + def test_target_wrong_type_raises(self): + with self.assertRaises(TypeError): + apply_runtime_settings("not_a_module", RuntimeSettings(runtime_cache=None)) + + def test_zero_engines_raises(self): + model = torch.nn.Linear(3, 3).cuda() + with self.assertRaises(RuntimeError) as cm: + apply_runtime_settings(model, RuntimeSettings(runtime_cache=None)) + self.assertIn("no TRT engines", str(cm.exception)) + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "apply_runtime_settings dispatch requires TRT-RTX", +) +class TestApplyRuntimeSettingsModuleOwned(TestCase): + """Module-owned engines: string cache still accepted (module owns it).""" + + def test_module_path_string_accepted(self): + compiled, inputs = _compile_with_inputs() + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + n = apply_runtime_settings( + compiled, RuntimeSettings(runtime_cache=cache_path) + ) + self.assertGreaterEqual(n, 1) + _ = compiled(*inputs) + self.assertTrue(os.path.exists(cache_path)) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + + def test_returns_engine_count(self): + compiled, _ = _compile_with_inputs() + n = apply_runtime_settings(compiled, RuntimeSettings(runtime_cache=None)) + self.assertGreaterEqual(n, 1) + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "AOT load tests require TRT-RTX", +) +class TestApplyRuntimeSettingsModuleLess(TestCase): + """Module-less engines from save/load.""" + + def test_path_string_raises_for_module_less_engine(self): + compiled, inputs = _compile_with_inputs() + _, loaded_gm = _save_load(compiled, inputs) + with self.assertRaises(TypeError) as cm: + apply_runtime_settings( + loaded_gm, + RuntimeSettings(runtime_cache="/tmp/should_not_be_created.bin"), + ) + msg = str(cm.exception) + self.assertIn("runtime_cache", msg) + self.assertIn("RuntimeCache", msg) + + def test_default_runtime_settings_raises_for_module_less_engine(self): + compiled, inputs = _compile_with_inputs() + _, loaded_gm = _save_load(compiled, inputs) + with self.assertRaises(TypeError) as cm: + apply_runtime_settings(loaded_gm, RuntimeSettings()) + self.assertIn("runtime_cache", str(cm.exception)) + + def test_none_cache_applies_and_forward_runs(self): + compiled, inputs = _compile_with_inputs() + _, loaded_gm = _save_load(compiled, inputs) + n = apply_runtime_settings(loaded_gm, RuntimeSettings(runtime_cache=None)) + self.assertGreaterEqual(n, 1) + out = loaded_gm(*inputs) + self.assertEqual(out.shape, inputs[0].shape) + + def test_runtime_cache_applies_and_persists(self): + compiled, inputs = _compile_with_inputs() + _, loaded_gm = _save_load(compiled, inputs) + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + cache = RuntimeCache(path=cache_path, autosave_on_del=False) + n = apply_runtime_settings( + loaded_gm, + RuntimeSettings(runtime_cache=cache), + ) + self.assertGreaterEqual(n, 1) + _ = loaded_gm(*inputs) + self.assertTrue(cache.has_cache()) + cache.save() + size = os.path.getsize(cache_path) + self.assertGreater(size, 0) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + + def test_cuda_graph_strategy_field_takes_effect(self): + compiled, inputs = _compile_with_inputs() + _, loaded_gm = _save_load(compiled, inputs) + n = apply_runtime_settings( + loaded_gm, + RuntimeSettings( + cuda_graph_strategy="whole_graph_capture", + runtime_cache=None, + ), + ) + self.assertGreaterEqual(n, 1) + out = loaded_gm(*inputs) + self.assertEqual(out.shape, inputs[0].shape) + + def test_exported_program_reaches_same_engines_as_module(self): + compiled, inputs = _compile_with_inputs() + loaded_ep, loaded_gm = _save_load(compiled, inputs) + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + cache = RuntimeCache(path=cache_path, autosave_on_del=False) + n_ep = apply_runtime_settings( + loaded_ep, RuntimeSettings(runtime_cache=cache) + ) + self.assertGreaterEqual(n_ep, 1) + _ = loaded_gm(*inputs) + self.assertTrue(cache.has_cache()) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + + def test_warm_load_bytes_transferred(self): + """Caller must call .load() / .load_from_stream() to warm the cache before attach. + + The module path auto-calls load() via _resolve_runtime_cache; there is no + equivalent hook on the module-less path. Verify warm bytes actually land: + load_from_stream returns a byte count, so assertGreater(..., 0) is the + non-vacuous check. + """ + compiled, inputs = _compile_with_inputs() + _, loaded_gm = _save_load(compiled, inputs) + + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + cache_populate = RuntimeCache(path=cache_path, autosave_on_del=False) + apply_runtime_settings( + loaded_gm, RuntimeSettings(runtime_cache=cache_populate) + ) + _ = loaded_gm(*inputs) + cache_populate.save() + with open(cache_path, "rb") as fh: + saved_bytes = fh.read() + self.assertGreater(len(saved_bytes), 0, "first save produced empty file") + + _, loaded_gm2 = _save_load(compiled, inputs) + cache_warm = RuntimeCache(autosave_on_del=False) + n_bytes = cache_warm.load_from_stream(io.BytesIO(saved_bytes)) + self.assertGreater( + n_bytes, + 0, + "load_from_stream transferred 0 bytes — warm load did nothing", + ) + + n_engines = apply_runtime_settings( + loaded_gm2, RuntimeSettings(runtime_cache=cache_warm) + ) + self.assertGreaterEqual(n_engines, 1) + _ = loaded_gm2(*inputs) + self.assertTrue(cache_warm.has_cache()) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "Mixed target test requires TRT-RTX", +) +class TestApplyRuntimeSettingsMixedTarget(TestCase): + """Module + module-less engines in one call: string must fail whole-call.""" + + def test_mixed_target_string_fails_atomically(self): + compiled, inputs = _compile_with_inputs() + _, loaded_gm = _save_load(compiled, inputs) + + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + TorchTensorRTModule, + ) + + prior = { + mod: mod.runtime_settings + for _, mod in compiled.named_modules() + if isinstance(mod, TorchTensorRTModule) and mod.engine is not None + } + + with self.assertRaises(TypeError): + apply_runtime_settings( + [compiled, loaded_gm], + RuntimeSettings(runtime_cache="/tmp/should_not_apply.bin"), + ) + + for mod, saved in prior.items(): + self.assertEqual(mod.runtime_settings, saved) + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "CM raise tests require TRT-RTX", +) +class TestContextManagerRaisesOnModuleLess(TestCase): + """runtime_config and runtime_cache raise on module-less engines.""" + + def test_runtime_config_raises_on_loaded_gm(self): + compiled, inputs = _compile_with_inputs() + _, loaded_gm = _save_load(compiled, inputs) + with self.assertRaises(TypeError) as cm: + with runtime_config(loaded_gm, runtime_cache=None): + pass + msg = str(cm.exception) + self.assertIn("apply_runtime_settings", msg) + + def test_runtime_cache_raises_on_loaded_gm(self): + compiled, inputs = _compile_with_inputs() + _, loaded_gm = _save_load(compiled, inputs) + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + cache_path = f.name + try: + with self.assertRaises(TypeError) as cm: + with runtime_cache(loaded_gm, cache_path): + pass + msg = str(cm.exception) + self.assertIn("apply_runtime_settings", msg) + finally: + try: + os.unlink(cache_path) + except OSError: + pass + + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py b/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py deleted file mode 100644 index 67ef0d83e9..0000000000 --- a/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py +++ /dev/null @@ -1,336 +0,0 @@ -# type: ignore -"""Tests for apply_runtime_settings and the updated CM raise behaviour. - -Tests that use save/load run both the module (in-process) and module-less -(AOT-loaded) paths. The AOT tests additionally verify that the CM: - -* raises TypeError on module-less engines (commit 3) -* names apply_runtime_settings in the error message -""" - -import io -import os -import tempfile -import unittest - -import torch -import torch_tensorrt -from torch.testing._internal.common_utils import TestCase, run_tests -from torch_tensorrt._features import ENABLED_FEATURES -from torch_tensorrt.runtime import ( - RuntimeCache, - RuntimeSettings, - apply_runtime_settings, - runtime_cache, - runtime_config, -) - - -class SimpleModel(torch.nn.Module): - def forward(self, x): - return torch.relu(x) + 1.0 - - -def _compile_simple(): - model = SimpleModel().eval().cuda() - inputs = [torch.randn(2, 3).cuda()] - compiled = torch_tensorrt.compile( - model, - ir="dynamo", - inputs=inputs, - min_block_size=1, - ) - torch._dynamo.reset() - return compiled, inputs - - -def _save_load(compiled, inputs): - """Save ``compiled`` to a temp file and return the loaded ExportedProgram and GraphModule.""" - with tempfile.NamedTemporaryFile(suffix=".ep", delete=False) as f: - ep_path = f.name - try: - torch_tensorrt.save(compiled, ep_path, arg_inputs=inputs) - loaded_ep = torch_tensorrt.load(ep_path) - finally: - try: - os.unlink(ep_path) - except OSError: - pass - loaded_gm = loaded_ep.module() if hasattr(loaded_ep, "module") else loaded_ep - return loaded_ep, loaded_gm - - -# --------------------------------------------------------------------------- -# Tests that do NOT require an RTX build -# --------------------------------------------------------------------------- - - -class TestApplyRuntimeSettingsTypeErrors(TestCase): - """Rejection of bad arguments; no engine compile required.""" - - def test_settings_wrong_type_raises(self): - model = torch.nn.Linear(3, 3).cuda() - with self.assertRaises(TypeError) as cm: - apply_runtime_settings(model, {"cuda_graph_strategy": "disabled"}) - self.assertIn("RuntimeSettings", str(cm.exception)) - - def test_target_wrong_type_raises(self): - with self.assertRaises(TypeError): - apply_runtime_settings("not_a_module", RuntimeSettings(runtime_cache=None)) - - def test_zero_engines_raises(self): - # A plain nn.Module has no TRT engines. - model = torch.nn.Linear(3, 3).cuda() - with self.assertRaises(RuntimeError) as cm: - apply_runtime_settings(model, RuntimeSettings(runtime_cache=None)) - self.assertIn("no TRT engines", str(cm.exception)) - - -# --------------------------------------------------------------------------- -# Tests that require TRT-RTX -# --------------------------------------------------------------------------- - - -@unittest.skipIf( - not ENABLED_FEATURES.tensorrt_rtx, - "apply_runtime_settings dispatch requires TRT-RTX", -) -class TestApplyRuntimeSettingsModuleOwned(TestCase): - """Module-owned engines: string cache still accepted (module owns it).""" - - def test_module_path_string_accepted(self): - compiled, inputs = _compile_simple() - with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: - cache_path = f.name - try: - n = apply_runtime_settings( - compiled, RuntimeSettings(runtime_cache=cache_path) - ) - self.assertGreaterEqual(n, 1) - _ = compiled(*inputs) - self.assertTrue(os.path.exists(cache_path)) - finally: - try: - os.unlink(cache_path) - except OSError: - pass - - def test_returns_engine_count(self): - compiled, _ = _compile_simple() - n = apply_runtime_settings(compiled, RuntimeSettings(runtime_cache=None)) - self.assertGreaterEqual(n, 1) - - -@unittest.skipIf( - not ENABLED_FEATURES.tensorrt_rtx, - "AOT load tests require TRT-RTX", -) -class TestApplyRuntimeSettingsModuleLess(TestCase): - """Module-less engines from save/load.""" - - def test_path_string_raises_for_module_less_engine(self): - compiled, inputs = _compile_simple() - _, loaded_gm = _save_load(compiled, inputs) - with self.assertRaises(TypeError) as cm: - apply_runtime_settings( - loaded_gm, - RuntimeSettings(runtime_cache="/tmp/should_not_be_created.bin"), - ) - msg = str(cm.exception) - self.assertIn("runtime_cache", msg) - self.assertIn("RuntimeCache", msg) - - def test_default_runtime_settings_raises_for_module_less_engine(self): - # RuntimeSettings() default runtime_cache is a path string — a common footgun. - compiled, inputs = _compile_simple() - _, loaded_gm = _save_load(compiled, inputs) - with self.assertRaises(TypeError) as cm: - apply_runtime_settings(loaded_gm, RuntimeSettings()) - self.assertIn("runtime_cache", str(cm.exception)) - - def test_none_cache_applies_and_forward_runs(self): - compiled, inputs = _compile_simple() - _, loaded_gm = _save_load(compiled, inputs) - n = apply_runtime_settings(loaded_gm, RuntimeSettings(runtime_cache=None)) - self.assertGreaterEqual(n, 1) - out = loaded_gm(*inputs) - self.assertEqual(out.shape, inputs[0].shape) - - def test_runtime_cache_applies_and_persists(self): - compiled, inputs = _compile_simple() - _, loaded_gm = _save_load(compiled, inputs) - with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: - cache_path = f.name - try: - cache = RuntimeCache(path=cache_path, autosave_on_del=False) - n = apply_runtime_settings( - loaded_gm, - RuntimeSettings(runtime_cache=cache), - ) - self.assertGreaterEqual(n, 1) - _ = loaded_gm(*inputs) - self.assertTrue(cache.has_cache()) - cache.save() - size = os.path.getsize(cache_path) - self.assertGreater(size, 0) - finally: - try: - os.unlink(cache_path) - except OSError: - pass - - def test_cuda_graph_strategy_field_takes_effect(self): - """Non-cache field applied to a module-less engine; proves field-agnostic dispatch.""" - compiled, inputs = _compile_simple() - _, loaded_gm = _save_load(compiled, inputs) - n = apply_runtime_settings( - loaded_gm, - RuntimeSettings( - cuda_graph_strategy="whole_graph_capture", - runtime_cache=None, - ), - ) - self.assertGreaterEqual(n, 1) - # Forward must still run after strategy change. - out = loaded_gm(*inputs) - self.assertEqual(out.shape, inputs[0].shape) - - def test_exported_program_reaches_same_engines_as_module(self): - """apply_runtime_settings on ExportedProgram reaches the engines ep.module() uses.""" - compiled, inputs = _compile_simple() - loaded_ep, loaded_gm = _save_load(compiled, inputs) - with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: - cache_path = f.name - try: - cache = RuntimeCache(path=cache_path, autosave_on_del=False) - # Apply via ExportedProgram. - n_ep = apply_runtime_settings( - loaded_ep, RuntimeSettings(runtime_cache=cache) - ) - self.assertGreaterEqual(n_ep, 1) - # Running via ep.module() uses the same engine objects -> cache populated. - _ = loaded_gm(*inputs) - self.assertTrue(cache.has_cache()) - finally: - try: - os.unlink(cache_path) - except OSError: - pass - - def test_warm_load_bytes_transferred(self): - """Caller must call .load() / .load_from_stream() to warm the cache before attach. - - The module path auto-calls load() via _resolve_runtime_cache; there is no - equivalent hook on the module-less path. Verify warm bytes actually land: - load_from_stream returns a byte count, so assertGreater(..., 0) is the - non-vacuous check. - """ - compiled, inputs = _compile_simple() - _, loaded_gm = _save_load(compiled, inputs) - - # Pass 1: populate a cache to get real bytes on disk. - with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: - cache_path = f.name - try: - cache_populate = RuntimeCache(path=cache_path, autosave_on_del=False) - apply_runtime_settings( - loaded_gm, RuntimeSettings(runtime_cache=cache_populate) - ) - _ = loaded_gm(*inputs) - cache_populate.save() - with open(cache_path, "rb") as fh: - saved_bytes = fh.read() - self.assertGreater(len(saved_bytes), 0, "first save produced empty file") - - # Pass 2: warm a new cache from those bytes via load_from_stream. - _, loaded_gm2 = _save_load(compiled, inputs) - cache_warm = RuntimeCache(autosave_on_del=False) - n_bytes = cache_warm.load_from_stream(io.BytesIO(saved_bytes)) - self.assertGreater( - n_bytes, - 0, - "load_from_stream transferred 0 bytes — warm load did nothing", - ) - - n_engines = apply_runtime_settings( - loaded_gm2, RuntimeSettings(runtime_cache=cache_warm) - ) - self.assertGreaterEqual(n_engines, 1) - _ = loaded_gm2(*inputs) - self.assertTrue(cache_warm.has_cache()) - finally: - try: - os.unlink(cache_path) - except OSError: - pass - - -@unittest.skipIf( - not ENABLED_FEATURES.tensorrt_rtx, - "Mixed target test requires TRT-RTX", -) -class TestApplyRuntimeSettingsMixedTarget(TestCase): - """Module + module-less engines in one call: string must fail whole-call.""" - - def test_mixed_target_string_fails_atomically(self): - compiled, inputs = _compile_simple() - _, loaded_gm = _save_load(compiled, inputs) - - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( - TorchTensorRTModule, - ) - - # Snapshot prior settings on the module engine. - prior = { - mod: mod.runtime_settings - for _, mod in compiled.named_modules() - if isinstance(mod, TorchTensorRTModule) and mod.engine is not None - } - - with self.assertRaises(TypeError): - apply_runtime_settings( - [compiled, loaded_gm], - RuntimeSettings(runtime_cache="/tmp/should_not_apply.bin"), - ) - - # Module engine settings must be unchanged (validate-before-mutate). - for mod, saved in prior.items(): - self.assertEqual(mod.runtime_settings, saved) - - -@unittest.skipIf( - not ENABLED_FEATURES.tensorrt_rtx, - "CM raise tests require TRT-RTX", -) -class TestContextManagerRaisesOnModuleLess(TestCase): - """runtime_config and runtime_cache raise on module-less engines.""" - - def test_runtime_config_raises_on_loaded_gm(self): - compiled, inputs = _compile_simple() - _, loaded_gm = _save_load(compiled, inputs) - with self.assertRaises(TypeError) as cm: - with runtime_config(loaded_gm, runtime_cache=None): - pass - msg = str(cm.exception) - self.assertIn("apply_runtime_settings", msg) - - def test_runtime_cache_raises_on_loaded_gm(self): - compiled, inputs = _compile_simple() - _, loaded_gm = _save_load(compiled, inputs) - with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: - cache_path = f.name - try: - with self.assertRaises(TypeError) as cm: - with runtime_cache(loaded_gm, cache_path): - pass - msg = str(cm.exception) - self.assertIn("apply_runtime_settings", msg) - finally: - try: - os.unlink(cache_path) - except OSError: - pass - - -if __name__ == "__main__": - run_tests() From 1e69635d72cb6364b9bfbdb7ad005f17684f6760 Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Fri, 28 Aug 2026 19:07:10 -0700 Subject: [PATCH 11/11] refactor: replace module_less list with any() in CM enter checks The list was materialized only to test truthiness and format a count into the error message. any() short-circuits on the first None owner without a second pass over the already-drained engines list. Also removes the redundant list() wrapping of self._targets in _RuntimeCacheContextManager.__enter__. Co-Authored-By: Claude Sonnet 4.6 --- py/torch_tensorrt/runtime/_runtime_cache.py | 9 ++++----- py/torch_tensorrt/runtime/_runtime_config.py | 7 +++---- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/py/torch_tensorrt/runtime/_runtime_cache.py b/py/torch_tensorrt/runtime/_runtime_cache.py index 931d71b034..49459dbb4d 100644 --- a/py/torch_tensorrt/runtime/_runtime_cache.py +++ b/py/torch_tensorrt/runtime/_runtime_cache.py @@ -530,13 +530,12 @@ def __enter__(self) -> RuntimeCache: ) # 1. Discover all TRT engines under the targets, validate before mutating. - engines = list(_iter_trt_engines(list(self._targets))) + engines = list(_iter_trt_engines(self._targets)) - module_less = [(owner, eng) for owner, eng in engines if owner is None] - if module_less: + if any(owner is None for owner, _ in engines): raise TypeError( - f"runtime_cache() encountered {len(module_less)} module-less " - "TRT engine(s) that it cannot snapshot and restore on exit. " + "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())." ) diff --git a/py/torch_tensorrt/runtime/_runtime_config.py b/py/torch_tensorrt/runtime/_runtime_config.py index 845a2ff7ae..587c245e00 100644 --- a/py/torch_tensorrt/runtime/_runtime_config.py +++ b/py/torch_tensorrt/runtime/_runtime_config.py @@ -376,11 +376,10 @@ def __enter__(self) -> Union["torch.nn.Module", Tuple["torch.nn.Module", ...]]: # before any settings are changed. engines = list(_iter_trt_engines(self._targets)) - module_less = [(owner, eng) for owner, eng in engines if owner is None] - if module_less: + if any(owner is None for owner, _ in engines): raise TypeError( - f"runtime_config() encountered {len(module_less)} module-less " - "TRT engine(s) that it cannot snapshot and restore on exit. " + "runtime_config() 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())." )