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..bf5cb8a679 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,64 @@ Stream-mode behavior: * On exit: cache serialized, ``stream.write(bytes)`` once. * ``rc.path`` reports ``""`` in stream-mode. +``apply_runtime_settings(...)`` — permanent apply for AOT artifacts +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Engines loaded with :func:`torch_tensorrt.load` have no +:class:`TorchTensorRTModule`. :func:`~torch_tensorrt.runtime.runtime_config` +and :func:`~torch_tensorrt.runtime.runtime_cache` cannot restore settings on +exit for such engines (there is no getter on the torchbind engine) and will +raise if they encounter one. Use the permanent-apply entry point instead: + +.. code-block:: python + + import torch_tensorrt + from torch_tensorrt.runtime import RuntimeCache, RuntimeSettings, apply_runtime_settings + + ep = torch_tensorrt.load("model.ep") + gm = ep.module() + + cache = RuntimeCache(path="/var/cache/jit.bin") + cache.load() # warm from disk if it exists + + n = apply_runtime_settings( + gm, + RuntimeSettings( + cuda_graph_strategy="whole_graph_capture", + runtime_cache=cache, + ), + ) + print(f"Applied to {n} engine(s)") + out = gm(x) + cache.save() # persist newly JIT'd kernels + +**Ownership rule:** ``settings.runtime_cache`` must be ``None`` or a +:class:`RuntimeCache` you own -- a path string raises ``TypeError`` because +there is no module to build and save the handle. If you call +``apply_runtime_settings(gm, RuntimeSettings())`` (the default +``runtime_cache`` is a path string), you will hit this error. +Pass ``runtime_cache=None`` or a :class:`RuntimeCache`. + +**You own** ``.load()`` **as well as** ``.save()``. A +:class:`TorchTensorRTModule` calls :meth:`RuntimeCache.load` automatically when +it resolves a path string (via ``_resolve_runtime_cache``), so in-process +compiled models warm the cache implicitly. There is no equivalent hook on the +module-less path -- call ``cache.load()`` (shown above) before passing the +handle, or the engine starts with an empty cache regardless of what is on disk. + +:func:`apply_runtime_settings` also accepts a :class:`~torch.export.ExportedProgram` +directly (the :func:`torch_tensorrt.load` return value), which is equivalent to +passing ``ep.module()``: + +.. code-block:: python + + apply_runtime_settings(ep, RuntimeSettings(runtime_cache=cache)) + +.. note:: + + Runtime settings are never serialized. They do not survive + :func:`torch_tensorrt.save`; re-apply after each :func:`torch_tensorrt.load`. + ---- Composing the context managers @@ -431,3 +489,5 @@ Quick reference - ``RuntimeSettings(runtime_cache=None)`` or ``runtime_cache(mod, "")`` * - Non-cuda-graph settings alongside cudagraphs capture - nest ``runtime_config(...)`` *outside* ``enable_cudagraphs(...)`` + * - Set a runtime knob on a loaded artifact (no module) + - ``apply_runtime_settings(gm_or_ep, RuntimeSettings(...))`` diff --git a/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py b/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py index 3b28641c35..509846ca5e 100644 --- a/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py +++ b/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py @@ -282,11 +282,18 @@ def __init__( # engines compiled with native multi-device collective layers. self._nccl_comm: Optional[Any] = None - # Owns RuntimeSettings + the live trt.IRuntimeConfig + the - # engine-implicit RuntimeCache. Hides all RTX feature gates. - # ``RuntimeSettings`` default; callers wanting non-defaults assign via - # the module's ``runtime_settings`` setter after compile. - self._trt_runtime_config: TRTRuntimeConfig = TRTRuntimeConfig(RuntimeSettings()) + # Owns RuntimeSettings + the live trt.IRuntimeConfig. Hides all RTX + # feature gates. + # + # ``runtime_cache=None``: an engine never owns a runtime cache. The + # module owns the implicit one and pushes it down via + # ``setup_engine``; an engine with no module (a packed-engine-info + # build, or a constant in an AOT-loaded ExportedProgram) runs without + # one until a caller attaches a ``RuntimeCache`` explicitly. Mirrors + # the cpp default (``RuntimeSettings::runtime_cache = nullptr``). + self._trt_runtime_config: TRTRuntimeConfig = TRTRuntimeConfig( + RuntimeSettings(runtime_cache=None) + ) # Multiple optimization profiles. Manual selection by default: # ``_active_profile_index`` is the profile currently loaded in the TRT # context (default 0, reused across calls). ``_auto_select_profiles`` @@ -399,10 +406,11 @@ def __setstate__(self, state: Any) -> None: # NCCL communicators cannot be pickled; rebind lazily on the next # forward pass via setup_nccl_comm(). self._nccl_comm = None - # RuntimeSettings are NOT serialized -- restore defaults. Callers - # who want runtime-mode overrides must reapply them post-load via - # ``mod.runtime_settings = ...`` (per ``TorchTensorRTModule``) or a runtime CM. - self._trt_runtime_config = TRTRuntimeConfig(RuntimeSettings()) + # RuntimeSettings are NOT serialized -- restore defaults, runtime cache + # included (see ``__init__``). Callers who want runtime-mode overrides + # must reapply them post-load via ``mod.runtime_settings = ...`` (per + # ``TorchTensorRTModule``) or a runtime CM. + self._trt_runtime_config = TRTRuntimeConfig(RuntimeSettings(runtime_cache=None)) self._active_profile_index = 0 self._auto_select_profiles = False diff --git a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py index 6028359130..f97206dc9b 100644 --- a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py +++ b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py @@ -357,6 +357,10 @@ def runtime_settings(self, rs: RuntimeSettings) -> None: def _resolve_runtime_cache(self, rs: RuntimeSettings) -> RuntimeSettings: """Normalize ``rs.runtime_cache`` to ``None`` | ``RuntimeCache`` (never a path str). + This module is the only place a path string is resolved; engines reject + one (see ``TRTRuntimeConfig._apply_settings``), so the returned settings + must never carry a ``str``. + Manages the ``_implicit_cache_handle`` slot as a side effect: builds a fresh wrapper for a new path, reuses the existing one for the same path, releases it (with save-on-swap) for non-path inputs. @@ -372,7 +376,9 @@ def _resolve_runtime_cache(self, rs: RuntimeSettings) -> RuntimeSettings: if not (isinstance(rc, str) and rc): if rc is not self._implicit_cache_handle: self._set_managed_handle(None) - return rs + # An empty string means "no cache", but engines take only None or a + # RuntimeCache -- normalize so no str ever reaches one. + return rs.merge(runtime_cache=None) if isinstance(rc, str) else rs # Branch 2: same path + wrapper still usable -> reuse. Keeps the CM # enter/exit cycle cheap (no teardown/rebuild loses in-memory kernels). @@ -422,25 +428,9 @@ def _wrapper_still_attached(self, w: Any) -> bool: def _send_to_engine(self, rs: RuntimeSettings) -> None: """Push ``rs`` to whichever engine flavor is attached.""" - from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine - from torch_tensorrt.runtime._runtime_cache import _to_torchbind_handle - from torch_tensorrt.runtime._runtime_config import ( - _CUDA_GRAPH_STRATEGY_MAP, - _DYNAMIC_SHAPES_KERNEL_STRATEGY_MAP, - ) + from torch_tensorrt.runtime._runtime_config import _send_settings_to_engine - if isinstance(self.engine, TRTEngine): - self.engine.update_runtime_settings(rs) - else: - # Strategies cross the boundary as ints (TorchBind ``int64_t``, - # mirroring the nvinfer1 enum integers on the cpp side). - self.get_engine().update_runtime_settings( - _DYNAMIC_SHAPES_KERNEL_STRATEGY_MAP[ - rs.dynamic_shapes_kernel_specialization_strategy - ], - _CUDA_GRAPH_STRATEGY_MAP[rs.cuda_graph_strategy], - _to_torchbind_handle(rs.runtime_cache), - ) + _send_settings_to_engine(self.engine, rs) def setup_engine(self) -> None: """ @@ -566,9 +556,11 @@ def set_extra_state(self, state: SerializedTorchTensorRTModuleFmt) -> None: self.settings = metadata["settings"] self.symbolic_shape_expressions = metadata["inout_symexprs"] - # RuntimeSettings are NOT serialized; restore defaults. Caller can - # reapply via ``mod.runtime_settings = ...`` (per submodule) or a CM after load. - self._runtime_settings = RuntimeSettings() + # RuntimeSettings are NOT serialized; the reset leaves no runtime + # cache, matching the freshly-built engine below. A caller who wants + # one reapplies via ``mod.runtime_settings = ...`` (per submodule) or + # a CM after load. + self._runtime_settings = RuntimeSettings(runtime_cache=None) # Mirror the settings reset on the implicit cache handle so a # stale wrapper from prior use doesn't survive load_state_dict and # silently write the fresh engine's cache bytes to the old path. @@ -682,7 +674,7 @@ def __getstate__(self) -> dict[str, Any]: return state def __setstate__(self, state: dict[str, Any]) -> None: - state.setdefault("_runtime_settings", RuntimeSettings()) + state.setdefault("_runtime_settings", RuntimeSettings(runtime_cache=None)) state.setdefault("_implicit_cache_handle", None) set_state = getattr(super(), "__setstate__", None) if set_state is not None: 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_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 1b2cbee643..ab82fecec1 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) @@ -68,13 +82,17 @@ class RuntimeSettings: TRT-RTX-only; no-op on standard TensorRT. cuda_graph_strategy: ``"disabled" | "whole_graph_capture"``. TRT-RTX-only. runtime_cache: ``None``, a disk path string, or a - :class:`RuntimeCache`. ``None`` ⇒ no cache attached. A - string is honored at engine construction time and primes a - per-engine disk-backed cache (engine owns the implicit handle and - it saves on ``__del__``). 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``. + :class:`RuntimeCache`. ``None`` ⇒ no cache attached. A string is + resolved by :class:`TorchTensorRTModule`, which builds the + engine-implicit :class:`RuntimeCache`, owns it, and saves it on + ``__del__``; engines themselves only ever receive ``None`` or a + 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). @@ -149,7 +167,9 @@ class TRTRuntimeConfig: """ def __init__(self, settings: Optional[RuntimeSettings] = None) -> None: - self._settings: RuntimeSettings = settings or RuntimeSettings() + self._settings: RuntimeSettings = settings or RuntimeSettings( + runtime_cache=None + ) # Live trt.IRuntimeConfig (RTX) or None (non-RTX / pre-init). self._live: Any = None @@ -164,8 +184,8 @@ def set_settings(self, new: RuntimeSettings) -> bool: On change, invalidates the live ``IRuntimeConfig`` and signals callers to recreate the ``IExecutionContext``. Disk persistence of any prior implicit cache handle is the module's responsibility (see - ``TorchTensorRTModule._materialize_implicit_handle``); this method is - a pure-execution swap. + ``TorchTensorRTModule._set_managed_handle``); this method is a + pure-execution swap. """ if new == self._settings: return False @@ -248,15 +268,18 @@ def is_monolithic_capturable( def _apply_settings(self) -> None: """Apply ``self._settings`` to the live ``trt.IRuntimeConfig``. - Resolves ``runtime_cache``: - - ``None`` ⇒ no cache attached. - - ``RuntimeCache`` ⇒ caller owns lifecycle. ``ensure_cache`` - materializes the inner ``IRuntimeCache`` on first use and drains - any pending warm bytes loaded into the handle's pending buffer at - construction time (by ``_TorchTensorRTModule._resolve_runtime_cache`` - for engine-implicit handles, or by the ``runtime_cache`` CM for - shared ones). String paths are pre-wrapped into handles upstream; - raw strings are not accepted here. + Resolves ``runtime_cache``, which must be ``None`` or a + :class:`RuntimeCache` -- something that *owns* what it points at. A + ``RuntimeCache``'s ``ensure_cache`` materializes the inner + ``IRuntimeCache`` on first use and drains any pending warm bytes loaded + into the handle's pending buffer at construction time (by + ``_TorchTensorRTModule._resolve_runtime_cache`` for engine-implicit + handles, or by the ``runtime_cache`` CM for shared ones). + + Path strings are resolved upstream by ``TorchTensorRTModule`` and are + rejected here: a string owns nothing, so honoring one would mean + building a handle this method does not outlive, handing its + ``IRuntimeCache`` to ``_live``, and letting it be collected on return. """ # Deferred imports: trt is import-aliased to tensorrt_rtx on RTX builds, # and _runtime_cache imports this module's RuntimeSettings. @@ -275,36 +298,15 @@ def _apply_settings(self) -> None: rc = self._settings.runtime_cache if rc is None: - logger.debug("Runtime cache disabled (no RuntimeCache / path provided).") + logger.debug("Runtime cache disabled (no RuntimeCache provided).") elif isinstance(rc, RuntimeCache): - cache = rc.ensure_cache(self._live) - self._live.set_runtime_cache(cache) - elif isinstance(rc, str): - # ``TorchTensorRTModule._resolve_runtime_cache`` pre-wraps path - # strings on the compile / configure path, but engines created - # directly (e.g. the Python ``TRTEngine`` constructed from a - # cross-runtime ``.pt2`` load — see - # ``test_cross_runtime_serde::test_save_python_load_python``) - # get a default ``RuntimeSettings(runtime_cache=RUNTIME_CACHE_PATH)`` - # that's never seen by the module's resolver. Wrap defensively - # here so the load path doesn't crash; this also keeps the - # documented contract that callers MAY pass a path string. - # - # ``RuntimeSettings`` is a frozen dataclass, so we can't store the - # wrapper back onto ``self._settings``; just use it locally. The - # wrapper is GC'd after this call, which is fine: ensure_cache has - # already materialized the underlying IRuntimeCache on ``_live``. - wrapped = RuntimeCache(path=rc, autosave_on_del=True) - try: - wrapped.load() - except Exception as e: - logger.warning(f"Failed to warm-load runtime cache from {rc!r}: {e}") - cache = wrapped.ensure_cache(self._live) - self._live.set_runtime_cache(cache) + self._live.set_runtime_cache(rc.ensure_cache(self._live)) else: raise TypeError( - f"runtime_cache must be None, str, or RuntimeCache by the " - f"time it reaches TRTRuntimeConfig; got {type(rc).__name__}." + f"runtime_cache must be None or a RuntimeCache by the time it " + f"reaches TRTRuntimeConfig; got {type(rc).__name__}. Path " + f"strings are resolved by TorchTensorRTModule -- an engine " + f"used without one must be given a RuntimeCache explicitly." ) logger.info("TensorRT-RTX runtime config configured") @@ -364,22 +366,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: @@ -424,3 +433,196 @@ 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), + ) + + +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. + + **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`. + """ + 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) 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 5616a234c8..254a59d2e2 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 @@ -757,5 +746,221 @@ def test_pending_warm_bytes_populated_at_construction(self): self.assertTrue(found, "No TorchTensorRTModule with implicit handle found") +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "Runtime cache is only available with TensorRT-RTX", +) +@unittest.skipIf( + ENABLED_FEATURES.torch_tensorrt_runtime, + "Module-less TRTEngine construction requires the Python TRTEngine path", +) +class TestEngineOwnsNoCache(TestCase): + """An engine never owns a runtime cache; the module does. + + A module-less engine -- built from packed engine info, or loaded as a + graph constant from a saved ``ExportedProgram`` -- comes up with + ``runtime_cache=None`` and must run without one, rather than attaching a + handle nothing outlives. + """ + + def _bare_engine(self, compiled): + from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine + + mod = _find_python_trt_module(compiled) + self.assertIsNotNone(mod, "expected a TorchTensorRTModule after compile") + return TRTEngine(mod._pack_engine_info()) + + @staticmethod + def _first(out): + """``execute`` returns a bare Tensor for single-output engines.""" + return out if isinstance(out, torch.Tensor) else out[0] + + def test_bare_engine_defaults_to_no_cache(self): + model, inputs = _fresh_conv_model_and_inputs() + engine = self._bare_engine(_compile(model, inputs)) + self.assertIsNone(engine.runtime_settings.runtime_cache) + + def test_bare_engine_executes_without_a_cache(self): + model, inputs = _fresh_conv_model_and_inputs() + compiled = _compile(model, inputs) + ref = compiled(*inputs) + engine = self._bare_engine(compiled) + out = self._first(engine.execute(list(inputs))) + self.assertGreater(cosine_similarity(ref, out), COSINE_THRESHOLD) + + def test_bare_engine_takes_an_explicit_cache(self): + """The opt-in an AOT deployment is expected to make.""" + model, inputs = _fresh_conv_model_and_inputs() + compiled = _compile(model, inputs) + ref = compiled(*inputs) + engine = self._bare_engine(compiled) + + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "rc.bin") + handle = torchtrt.runtime.RuntimeCache(path=path) + engine.update_runtime_settings(RuntimeSettings(runtime_cache=handle)) + out = self._first(engine.execute(list(inputs))) + self.assertGreater(cosine_similarity(ref, out), COSINE_THRESHOLD) + handle.save() + self.assertTrue(os.path.exists(path)) + self.assertGreater(os.path.getsize(path), 0) + + def test_path_string_on_an_engine_raises(self): + """Strings are the module's business; an engine rejects them.""" + model, inputs = _fresh_conv_model_and_inputs() + engine = self._bare_engine(_compile(model, inputs)) + with tempfile.TemporaryDirectory() as tmp: + engine.update_runtime_settings( + RuntimeSettings(runtime_cache=os.path.join(tmp, "rc.bin")) + ) + with self.assertRaisesRegex(TypeError, "must be None or a RuntimeCache"): + engine.execute(list(inputs)) + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "Runtime cache is only available with TensorRT-RTX", +) +class TestModuleStillOwnsImplicitCache(TestCase): + """Guard against over-correction: the compile path must keep caching.""" + + def test_compiled_module_persists_its_implicit_cache(self): + with tempfile.TemporaryDirectory() as tmp: + 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)) + compiled(*inputs) + del compiled + gc.collect() + self.assertTrue(os.path.exists(path), "implicit cache was not saved") + self.assertGreater(os.path.getsize(path), 0) + + def test_empty_path_string_is_normalized_to_none(self): + """An empty string means "no cache" and must not reach the engine.""" + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + TorchTensorRTModule, + ) + + model, inputs = _fresh_conv_model_and_inputs() + compiled = _compile(model, inputs) + 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. + mods = [ + m for _, m in compiled.named_modules() if isinstance(m, TorchTensorRTModule) + ] + self.assertTrue(mods, "expected a TorchTensorRTModule after compile") + for m in mods: + self.assertIsNone(m.runtime_settings.runtime_cache) + compiled(*inputs) + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "Runtime cache is only available with TensorRT-RTX", +) +class TestPostLoadOwnsNoCache(TestCase): + """The reset paths must leave no cache, matching the engine they rebuild. + + ``set_extra_state`` / ``__setstate__`` restore ``RuntimeSettings`` defaults + after a load. If that reset kept the default path *string*, the module would + disagree with the engine it just built (which comes up ``None``), and a later + ``runtime_config(...)`` block -- a call that need not mention caching at all -- + would resolve the string and leave an autosaving handle installed on exit. + """ + + def _round_tripped(self, tmp): + model, inputs = _fresh_conv_model_and_inputs() + compiled = _compile(model, inputs) + path = os.path.join(tmp, "mod.pt") + torch.save(compiled, path) + return torch.load(path, weights_only=False), inputs + + def test_config_default_has_no_cache(self): + """A default-constructed config must not start from a path string.""" + from torch_tensorrt.runtime._runtime_config import TRTRuntimeConfig + + self.assertIsNone(TRTRuntimeConfig().settings.runtime_cache) + + def test_torch_load_leaves_no_cache(self): + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + TorchTensorRTModule, + ) + + with tempfile.TemporaryDirectory() as tmp: + loaded, _ = self._round_tripped(tmp) + mods = [ + m + for _, m in loaded.named_modules() + if isinstance(m, TorchTensorRTModule) + ] + self.assertTrue(mods, "expected a TorchTensorRTModule after load") + for m in mods: + self.assertIsNone(m.runtime_settings.runtime_cache) + + def test_load_state_dict_leaves_no_cache(self): + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + TorchTensorRTModule, + ) + + model, inputs = _fresh_conv_model_and_inputs() + compiled = _compile(model, inputs) + fresh, _ = _fresh_conv_model_and_inputs() + target = _compile(fresh, inputs) + target.load_state_dict(compiled.state_dict()) + + mods = [ + m for _, m in target.named_modules() if isinstance(m, TorchTensorRTModule) + ] + self.assertTrue(mods, "expected a TorchTensorRTModule after load_state_dict") + for m in mods: + self.assertIsNone(m.runtime_settings.runtime_cache) + + def test_unrelated_context_manager_does_not_install_a_cache(self): + """A cuda-graph-only CM must not switch caching on, on enter or on exit. + + The CM snapshots the module's pre-resolution view; re-applying a path + string through the setter *creates* a handle rather than restoring one, + so a stale string here would survive ``__exit__`` pointed at the shared + default path with ``autosave_on_del=True``. + """ + from torch_tensorrt.dynamo._defaults import RUNTIME_CACHE_PATH + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + TorchTensorRTModule, + ) + from torch_tensorrt.runtime import runtime_config + + default_existed = os.path.exists(RUNTIME_CACHE_PATH) + with tempfile.TemporaryDirectory() as tmp: + loaded, inputs = self._round_tripped(tmp) + mods = [ + m + for _, m in loaded.named_modules() + if isinstance(m, TorchTensorRTModule) + ] + self.assertTrue(mods, "expected a TorchTensorRTModule after load") + + with runtime_config(loaded, cuda_graph_strategy="whole_graph_capture"): + for m in mods: + self.assertIsNone( + m._implicit_cache_handle, "CM installed a cache on enter" + ) + loaded(*inputs) + + for m in mods: + self.assertIsNone( + m._implicit_cache_handle, "CM left a cache installed on exit" + ) + self.assertIsNone(m.runtime_settings.runtime_cache) + + if not default_existed: + self.assertFalse( + os.path.exists(RUNTIME_CACHE_PATH), + "an unrelated CM wrote the shared default cache file", + ) + + if __name__ == "__main__": run_tests() 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") ) 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..67ef0d83e9 --- /dev/null +++ b/tests/py/dynamo/runtime/test_005_apply_runtime_settings.py @@ -0,0 +1,336 @@ +# 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()