feat: apply_runtime_settings — permanent runtime config for AOT-loaded engines - #4570
Draft
tp5uiuc wants to merge 9 commits into
Draft
feat: apply_runtime_settings — permanent runtime config for AOT-loaded engines#4570tp5uiuc wants to merge 9 commits into
tp5uiuc wants to merge 9 commits into
Conversation
A TRTEngine built without a module came up with the default RuntimeSettings, whose runtime_cache is a path string. Nothing owned the wrapper that string implies, so attaching it handed the live IRuntimeCache to the engine's IRuntimeConfig and then let it be collected -- a use-after-free at the next createExecutionContext. Engines now default to runtime_cache=None, matching the cpp side, where RuntimeSettings::runtime_cache is an intrusive_ptr defaulting to nullptr and no string form exists. The implicit cache belongs to the module, which resolves its path string to a RuntimeCache and pushes it down in setup_engine; that path is unchanged. TorchTensorRTModule resets its own RuntimeSettings on the post-load paths (set_extra_state, __setstate__) and those resets move to runtime_cache=None too, so the module and the engine it rebuilds agree. Leaving them at the string default would let a later runtime_config(...) block resolve the stale path and install an autosaving handle at the shared default location on exit -- switching caching on via a call that never mentioned it. An engine reached without a module -- built from packed engine info, or loaded as a graph constant from a saved ExportedProgram -- now runs with no cache instead of a dangling one, and a caller can attach a RuntimeCache explicitly. No configuration loses working behaviour: on the Python runtime this path raised, on the cpp runtime it already attached nothing, and on standard TensorRT the runtime config is never initialized.
_apply_settings had three arms, and the str one built a RuntimeCache it did not outlive. Its own docstring already claimed raw strings were not accepted here; the code twenty lines below accepted them. Engines now take only something that owns what it points at -- the Python equivalent of the cpp intrusive_ptr<RuntimeCacheHandle>. A str raises TypeError naming the module as the place path strings are resolved. The class's own default follows: a TRTRuntimeConfig built with no settings would otherwise start from the string form this commit exists to abolish. TorchTensorRTModule._resolve_runtime_cache normalizes an empty-string runtime_cache to None rather than passing it through, so no str can reach an engine from the module. Also corrects the RuntimeSettings.runtime_cache docstring, which promised the engine owned the implicit handle and saved it on __del__, and a reference to a method renamed some time ago.
TestEngineOwnsNoCache pins the contract at the engine: a module-less engine defaults to no cache, executes without one, accepts an explicitly attached RuntimeCache, and raises TypeError on a path string. TestModuleStillOwnsImplicitCache guards the other direction -- the compile path must keep building, attaching and persisting its implicit cache -- and covers the empty-string normalization. TestPostLoadOwnsNoCache pins the same contract on the reset paths, where module and engine could otherwise disagree: the config's own default, and what torch.load / load_state_dict leave behind. Its last test is the one that matters -- a cuda-graph-only context manager over a loaded module must not install a cache on enter or leave one installed on exit, because re-applying a path string through the setter creates a handle rather than restoring one.
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…t 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 <noreply@anthropic.com>
…ngines 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 <noreply@anthropic.com>
…c 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 <noreply@anthropic.com>
…le-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 <noreply@anthropic.com>
github-actions
Bot
added
documentation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds
torch_tensorrt.runtime.apply_runtime_settings(target, settings), apublic API to permanently configure runtime settings on engines loaded via
torch_tensorrt.load()(AOT-loaded artifacts with noTorchTensorRTModulewrapper).
What — A new entry point that applies a
RuntimeSettingsobject to everyTRT engine reachable from an
nn.Module, a list of modules, or anExportedProgram. Returns the number of engines updated.Why — The existing
runtime_configandruntime_cachecontext managersauto-restore on exit and anchor to
TorchTensorRTModule; they cannot reachengines produced by the AOT save/load path. There was no supported way to
attach a
RuntimeCache(or any other setting) to a bare loaded artifact.How
py/torch_tensorrt/runtime/_runtime_config.py: add_is_trt_engine,_iter_trt_engines,_send_settings_to_engine, andapply_runtime_settings;update both context managers to use
_iter_trt_enginesand raiseTypeErroron module-less engines (naming
apply_runtime_settingsin the message).py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py: reduce_send_to_engineto a one-line delegate to_send_settings_to_engine.py/torch_tensorrt/runtime/__init__.py: exportapply_runtime_settings.docsrc/py_api/runtime.rst+docsrc/user_guide/runtime_performance/runtime_settings.rst:add API entry and AOT section with example and caller-ownership rules.
tests/py/dynamo/runtime/test_005_apply_runtime_settings.py: new test filecovering type errors, module-owned engines, module-less engines, mixed targets,
CM raises, and warm-load byte-count verification.
tests/py/dynamo/runtime/test_{000,001,004}_*.py,tests/py/dynamo/models/test_*_models.py: replace local_apply_runtime_settingshelpers with the public API.
Key constraints applied: string
runtime_cachepaths are rejected formodule-less engines (caller owns the handle and must call
cache.load()beforeapplying);
ExportedProgramsupported viaep.constants; apply isvalidate-before-mutate (all-or-nothing across a mixed target list).
Testing
Full
tests/py/dynamo/runtime/suite and the three model test files pass onboth the Python-only and C++ runtime builds. New test file adds 11 tests; all
green.
Note: this PR stacks on #4482 (
fix/python-runtime-cache-on-cpp-build);those commits appear in the diff until #4482 merges.
Type of change
Checklist: