diff --git a/docsrc/user_guide/runtime_performance/executorch.rst b/docsrc/user_guide/runtime_performance/executorch.rst index 6f00e262f9..5245fbeed4 100644 --- a/docsrc/user_guide/runtime_performance/executorch.rst +++ b/docsrc/user_guide/runtime_performance/executorch.rst @@ -21,8 +21,8 @@ program, so the whole model stays on the GPU. .. note:: - The ExecuTorch delegate runs on Linux, x86_64 and aarch64, and on Windows x86_64. It - requires CUDA 13. Ordinary Torch-TensorRT builds keep their separate CUDA 12 support. + The prebuilt ExecuTorch delegate runs on Linux, x86_64 and aarch64, and on Windows x86_64. + It requires CUDA 13. Ordinary Torch-TensorRT builds keep their separate CUDA 12 support. ---- @@ -68,20 +68,18 @@ README shows. The delegate is built for the standard Linux wheels (x86-64 and aarch64), and for Windows x86-64. On aarch64 it needs glibc 2.35 or newer, the same as TensorRT itself. -The TensorRT-RTX and Python-only builds have no delegate. To export a coalesced program you also -need a CUDA toolkit, because ExecuTorch's CUDA backend compiles the leftover operators with -``nvcc``. +The JetPack, TensorRT-RTX and Python-only builds have no delegate. To export a coalesced program +you also need a CUDA toolkit, because ExecuTorch's CUDA backend compiles the leftover operators +with ``nvcc``. ---- Compile and Save ----------------- -The usual ``ir="dynamo"`` path, with two extra arguments to ``torch_tensorrt.save``: - -* ``output_format="executorch"`` selects the ``.pte`` writer. -* ``retrace=False`` keeps the compiled graph as it is. Re-exporting would drop the TensorRT - engines before ExecuTorch's partitioner sees them. +The usual ``ir="dynamo"`` path, with one extra argument to ``torch_tensorrt.save``: +``output_format="executorch"`` selects the ``.pte`` writer. For this format ``save`` keeps +the compiled graph as it is rather than re-exporting it, so it needs no example inputs. .. code-block:: python @@ -98,24 +96,19 @@ The usual ``ir="dynamo"`` path, with two extra arguments to ``torch_tensorrt.sav min_block_size=1, ) - torch_tensorrt.save( - trt_gm, - "model.pte", - output_format="executorch", - arg_inputs=example_inputs, - retrace=False, - ) + torch_tensorrt.save(trt_gm, "model.pte", output_format="executorch") -:ref:`executorch_save` covers the rest: dynamic shapes, several methods in one file, a -zero-copy KV cache, and the two-step ``torch_tensorrt.executorch.export()`` path for -programs that need work before they are written. +:ref:`executorch_save` covers the rest: several methods in one file, a zero-copy key-value +(KV) cache, and the two-step ``torch_tensorrt.executorch.export()`` path for programs that +need work before they are written. Coalescing the leftover operators onto the CUDA backend ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -TensorRT has no converter for every ATen operator, and by default the ones it rejects run on -the CPU, costing a copy each way on every call. Pass a ``CudaPartitioner`` and ExecuTorch's -CUDA backend compiles them instead, so nothing leaves the GPU: +TensorRT has no converter for every ATen operator (PyTorch's core operator set), and by +default the ones it rejects run on the CPU, costing a copy each way on every call. Pass a +``CudaPartitioner`` and ExecuTorch's CUDA backend compiles them instead, so nothing leaves +the GPU: .. code-block:: python @@ -126,8 +119,6 @@ CUDA backend compiles them instead, so nothing leaves the GPU: trt_gm, "coalesced.pte", output_format="executorch", - arg_inputs=example_inputs, - retrace=False, partitioners=[ CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")]) ], @@ -140,9 +131,9 @@ where TensorRT cannot take ``erfinv``, the saved program's delegate list reads .. warning:: The CUDA backend names its external weight file per device, not per model, so saving two - coalesced programs into one directory overwrites the first one's weights. That program - still loads, still reports finding its weights, and returns a wrong answer with no error. - Give each export its own directory. + coalesced programs into one directory overwrites the first one's weights. The first program + can then return wrong answers with no error, or fail to load. Give each export its own + directory. .. _executorch_external_engine_data: @@ -211,8 +202,6 @@ copies are pure overhead: trt_gm, "device_resident.pte", output_format="executorch", - arg_inputs=example_inputs, - retrace=False, backend_config=ExecutorchBackendConfig( propagate_device_config=PropagateDeviceConfig( skip_h2d_for_method_inputs=True, @@ -223,20 +212,21 @@ copies are pure overhead: ), ) -Three things are easy to miss: +Three things to know: * **Both skip flags need** ``enable_non_cpu_memory_planning=True``, because the copies are inserted during device-aware memory planning. Asking for a skip with it off raises. * **Inputs must be unplanned**, through ``MemoryPlanningPass(alloc_graph_input=False)``. Otherwise the program reserves its own input buffer and the runtime fills it from the caller, which puts the copy straight back. -* **Leave outputs planned if Python will run the program**, so the program's device arena - owns the output. Add ``alloc_graph_output=False`` only for a C++ caller that supplies the - address with ``Module::set_output``. Python cannot. +* **The outputs can stay in the program's own GPU memory**, its device arena. To give each output + its own buffer instead, also pass ``alloc_graph_output=False``. Then each output needs a GPU + buffer before the call: in C++ you set it with ``Module::set_output``, and the Python runtime + allocates one for you. -The choice is baked into the file, so feed such a program CUDA tensors. A host tensor still -gives the right answer, but it copies on every call, which is the cost you exported to -avoid, and nothing warns you. +The choice is baked into the file, so feed such a program CUDA tensors. ExecuTorch's Python +runtime refuses a CPU tensor. From C++, a CPU tensor costs a copy on every call when the first +delegate to use that input is TensorRT, and fails when it is the CUDA backend. ---- @@ -261,14 +251,16 @@ registers the backend. Everything after it is ExecuTorch's own API: If the delegate cannot load, that import raises at once, instead of surfacing later as a program that will not load. -A coalesced program needs nothing extra: both backends are registered and the program records -which parts go where. Torch-TensorRT is only needed to export, not to run. +A coalesced program needs no extra import: both backends are registered and the program +records which parts go where. If its export wrote a data file such as ``aoti_cuda_blob.ptd``, +pass it with ``load_program(..., data_path=Path("aoti_cuda_blob.ptd"))``. Torch-TensorRT is only +needed to export, not to run. .. note:: ``torch_tensorrt.load(path, format="executorch")`` still works but is deprecated. It copies CUDA inputs to the CPU and supports embedded weights only. Use the Runtime API - above instead. A device-resident program requires it. + above instead. A device-resident program runs only through the Runtime API. ---- @@ -417,12 +409,12 @@ To put every delegate on a stream of your own, scope a guard over the whole exec :ref:`Running a coalesced .pte ` has the same rule for a decode loop. -That stream can be a green context, which holds a fixed number of streaming multiprocessors, +That stream can be a green context, which holds a fixed number of streaming multiprocessors (SMs), so the model stays inside that partition and the rest of the GPU is free. Create it with ``cuGreenCtxStreamCreate`` and scope the same guard over it. The limit rides the stream, so the green context does not have to be made current. -The reference runner has this built in. Pass the SM count: +The reference runner has this built in. Pass the number of SMs: .. code-block:: bash @@ -444,8 +436,6 @@ launch many short kernels, where CPU launch work dominates each call. trt_gm, "model.pte", output_format="executorch", - arg_inputs=example_inputs, - retrace=False, use_cuda_graphs=True, ) @@ -457,17 +447,23 @@ It is not free. Each engine keeps a device buffer per input and output, and ever call copies each one in and out. With few kernels, or large inputs and outputs, those copies can cost more than the launches they save. -Some engines never replay and still pay the copies: changing shapes, aliased outputs, the -shared activation scratch below, a green-context stream, GPUs without stream-ordered memory, -and drivers older than CUDA 12.5. The delegate logs the reason once. +Some engines never replay: changing shapes, aliased outputs, a green-context stream, GPUs +without stream-ordered memory, and drivers older than CUDA 12.5. Only changing shapes still pay +the copies; the others take the ordinary path without the extra buffers. The delegate logs the +reason once for aliased outputs, missing stream-ordered memory and an old driver. + +Engines on the shared activation scratch below do replay, each with its own buffers for its +inputs and outputs. When the pool moves to a larger buffer, an engine records again, as it does +after a shape change. .. warning:: While a recording runs, no other thread may create or destroy a TensorRT engine or - execution context, or call a whole-device sync such as ``cudaDeviceSynchronize`` or - ``torch.cuda.synchronize()``. Every input shape change starts another recording, so - loading everything up front does not end the risk. Leave replay off when other threads - may do any of that. + execution context outside this delegate, or call a whole-device sync such as + ``cudaDeviceSynchronize`` or ``torch.cuda.synchronize()``. Every input shape change, and + every growth of the shared activation scratch, starts another recording, so loading + everything up front does not end the risk. Leave replay off when other threads may do any + of that. A C++ host can refuse a program's saved request for one load: @@ -525,10 +521,10 @@ and a device that has run a pooled engine must not be reset with ``cudaDeviceRes Shared engines ^^^^^^^^^^^^^^^ -Loading the same program twice in one process, say once per robot arm, used to hold its -weights in device memory twice. Handles built from the same engine bytes, on the same device -and with the same weight streaming request, now share one engine. Each still gets its own -execution context, buffers and lock. +Loading the same program twice in one process, say once per robot arm, keeps one copy of its +weights on the GPU. Handles built from the same engine bytes, on the same device and with the +same weight streaming request, share one engine. Each still gets its own execution context, +buffers and lock. This is on by default. On an 8 GB Jetson Orin Nano, loading one robot policy twice grew reported GPU memory by 806 MiB without sharing and by 122 MiB with it. The match is found by @@ -557,12 +553,12 @@ Comparison: .pte vs .pt2 * - libtorch at runtime - Not required - Required - * - Non-TRT ops + * - Operators TensorRT cannot run - ExecuTorch CUDA backend, or CPU fallback - Compiled by AOTInductor * - Platform - - Linux + Windows x86_64, CUDA 13 - - Linux + Windows + - Linux (x86_64, aarch64) and Windows x86_64, CUDA 13 + - Linux ---- @@ -570,6 +566,10 @@ Examples --------- * ``examples/torchtrt_executorch_example/export_static_shape.py`` writes a minimal ``.pte``. +* ``examples/torchtrt_executorch_example/export_dynamic_shape.py`` exports a model whose batch + size can change. +* ``examples/torchtrt_executorch_example/export_kv_cache_decode.py`` exports a decoder step + whose key-value cache stays in memory you own. * ``examples/torchtrt_executorch_example/export_coalesced.py`` splits one graph across the TensorRT delegate and ExecuTorch's CUDA backend. * ``examples/torchtrt_executorch_example/export_device_resident.py`` does the same with no diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index c39faec4e9..800a30f58b 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -22,7 +22,7 @@ specifying the `output_format` flag. Here are the options `output_format` will a * `exported_program` : This is the default. We perform transformations on the graphmodule first and use `torch.export.save` to save the module. * `torchscript` : We trace the graphmodule via `torch.jit.trace` and save it via `torch.jit.save`. * `PT2 Format` : This is a next generation runtime for PyTorch models, allowing them to run in Python and in C++ -* `executorch` : We lower the graphmodule to an ExecuTorch ``.pte`` program, delegating the TensorRT engines to the ExecuTorch backend. Linux-only; requires the ``executorch`` package. +* `executorch` : We lower the graphmodule to an ExecuTorch ``.pte`` program, delegating the TensorRT engines to the ExecuTorch backend. Runs on Linux (x86-64 and aarch64) and Windows x86-64; see :ref:`executorch_deployment` to install it. a) ExportedProgram ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -233,24 +233,7 @@ c) ExecuTorch (.pte) The ``executorch`` output format lowers the compiled module to an ExecuTorch ``.pte`` program, delegating the TensorRT engines to the Torch-TensorRT ExecuTorch -backend. This CUDA integration links CUDA 13 libraries, so it supports Linux with -any CUDA 13 PyTorch build and needs an ExecuTorch nightly wheel. Install from the -channel matching your CUDA, for example a CUDA 13.2 environment: - -.. code-block:: bash - - pip install --pre "torch_tensorrt[executorch]" \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 - -Substitute the channel for your CUDA, such as ``cu134`` for CUDA 13.4. ExecuTorch -publishes a channel later than PyTorch does, so a very new CUDA minor may not have -ExecuTorch wheels yet. Keep PyTorch, ExecuTorch and Torch-TensorRT on the -same CUDA channel. ``--pre`` permits prereleases; it does not request an upgrade -of an installed package. An existing version that declares the extra can gain -its missing dependencies, and dependency resolution may change other packages. -If an older version does not declare the extra, deliberately install the intended -compatible Torch-TensorRT wheel first. A fresh virtual environment avoids changing -an existing working stack. +backend. To install it, follow the Installation section of :ref:`executorch_deployment`. There are two ways to produce a ``.pte``, and they suit different needs: @@ -499,16 +482,10 @@ the third is refused where it is visible in the ``.pte``: surface as wrong results *or* as an illegal memory access. See :ref:`Running a coalesced .pte `. -* **Synchronizing that stream before reading a cache on the host**, for any - zero-copy ``.pte`` a runner drives on a caller stream, coalesced or not. A - delegate whose aliased outputs are threaded through it reflects each one into - its delegate output and so waits for the engine before returning; zero-copy - elides those outputs, so there is nothing to reflect and ``execute()`` returns - with the engine still running. A single-delegate decode loop owes the - synchronization as much as a coalesced program does. It is only the caller - stream that brings the duty: with none installed ``execute()`` synchronizes - before it returns, as it does whenever a delegate input or output is staged - through the host. +* **Synchronizing that stream before reading a cache on the host**, if the ``.pte`` is + coalesced and ExecuTorch's CUDA backend runs after the engine. The TensorRT delegate waits + for its engine before it returns, but the CUDA backend may return early. See + :ref:`Running a coalesced .pte `. * **Sharing one cache between methods.** Zero-copy is per method: it makes each method's engine write that method's buffer. Giving a prefill and a decode @@ -556,43 +533,29 @@ must be pointed at those data files to load them. .. warning:: - The CUDA backend names its external weight blob per-device (e.g. - ``aoti_cuda_blob.ptd``), not per-model, so saving two different coalesced - ``.pte`` into the same directory overwrites the blob and the first ``.pte`` - will fail to load. Save each coalesced model into its own directory. + The CUDA backend names its external weight file per device (for example + ``aoti_cuda_blob.ptd``), not per model, so saving two coalesced programs into + one directory overwrites the first one's weights. The first program can then + return wrong answers with no error, or fail to load. Save each coalesced model + into its own directory. .. _executorch_single_stream: -**Running a coalesced .pte: use a single CUDA stream** - -A coalesced ``.pte`` runs on more than one backend delegate (the TensorRT delegate -plus the CUDA backend). By default each backend enqueues its GPU work on its *own* -CUDA stream, and delegate execution is asynchronous -- a delegate returns after -*enqueuing* its work, not after it completes. Because separate CUDA streams are not -ordered relative to one another, at a delegate boundary the consuming delegate can -begin before the producing delegate's writes have finished, reading incomplete -data. This is a race: it is intermittent and can surface as wrong results or an -illegal memory access. - -The runtime does not impose a shared stream across delegates, so it is the -**runner's responsibility** to run all delegates on one CUDA stream. Create a -single stream and scope ``executorch::extension::cuda::CallerStreamGuard`` over it -for the duration of execution. That one guard reaches every CUDA-capable -delegate: they resolve a single shared ``libextension_cuda``, so the TensorRT -backend and the CUDA backend read the same caller-stream storage. All GPU work is -then enqueued in order and every cross-boundary dependency is satisfied, while -execution stays asynchronous. - -If the runner reads a delegate's outputs between calls (for example, an -autoregressive decode loop), synchronize the shared stream before reading: the -work may still be in flight when ``execute()`` returns, and a host-side copy on -the default stream will not wait for a non-blocking stream. A model that threads -its aliased outputs through the delegate is insulated from this in practice: -reflecting each aliased output into its delegate output makes the delegate wait -for the engine before it returns. Under -:ref:`zero-copy KV ` those outputs are elided, so there -is nothing to reflect and the delegate returns with the engine still running -- -the synchronization is then the only thing making a host read see the new values. +**Running a coalesced .pte on one CUDA stream** + +A coalesced ``.pte`` runs on two delegates: the TensorRT delegate and ExecuTorch's CUDA +backend. With no stream chosen, both use ``cudaStreamPerThread``, the calling thread's default +stream, so a program run from one thread is already in the right order. + +To use a stream of your own, put one ``executorch::extension::cuda::CallerStreamGuard`` around +the whole call. It reaches both delegates, because they share one ``libextension_cuda``. Do not +give the two delegates different streams: work on separate streams is not ordered, and the +result is a race that can show up as wrong values or an illegal memory access. + +The TensorRT delegate waits for its own work before it returns. ExecuTorch's CUDA backend can +return while its work is still queued, so when it runs last, ``execute()`` can return before the +outputs are ready. Synchronize the stream before you read an output on the host, for example +between the steps of a decode loop. **ExecuTorch lowering options** diff --git a/py/torch-tensorrt-executorch-runtime/README.md b/py/torch-tensorrt-executorch-runtime/README.md index 6d2b112cb4..d5235d3fce 100644 --- a/py/torch-tensorrt-executorch-runtime/README.md +++ b/py/torch-tensorrt-executorch-runtime/README.md @@ -46,16 +46,15 @@ wheel's directory on its own when you link an imported library. Then ship the delegate yourself, or point a library path at it at run time. For device-resident exports with `alloc_graph_output=False`, C++ Module callers must provide a CUDA output tensor with `Module::set_output` before execution. -Python callers cannot, so that arrangement is C++ only. +ExecuTorch's Python runtime allocates one for each output itself. -If you want one program that runs from both, turn off allocation for the inputs -and leave it on for the outputs. The caller then supplies the input buffer, which -is the copy worth avoiding, and the program's own device arena owns the output. +The simplest arrangement turns off allocation for the inputs and leaves it on for the outputs. +The caller then supplies the input buffer, which is the copy worth avoiding, and the program's own +device arena owns the output. Measured on two architectures, that arrangement runs from Python and from C++ with -no boundary copies in the program and results identical to eager. Turning -allocation off for both works only from C++. Leaving it on for both is refused by the -export step, because a program that allocates its own inputs has nothing for a caller -to hand in. +no boundary copies in the program and results identical to eager. Leaving it on for both is +refused by the export step, because a program that allocates its own inputs has nothing for a +caller to hand in. Point CMake at both wheels. The example above calls `find_package(executorch)` as well, and that package lives in its own distribution. ExecuTorch is a @@ -164,8 +163,9 @@ If the delegate cannot be loaded, that import raises, rather than letting the failure surface later as a program that will not load. A program whose graph is split between TensorRT and ExecuTorch's own CUDA backend -needs nothing extra here: both backends are registered, and the program says which -parts go where. +needs no extra import: both backends are registered, and the program says which +parts go where. If its export wrote a data file such as `aoti_cuda_blob.ptd`, pass it +with `load_program(..., data_path=Path("aoti_cuda_blob.ptd"))`. ## Runtime libraries @@ -230,7 +230,7 @@ render it for your machine before building. The template takes seven values: export CUDA_HOME=/usr/local/cuda-13.2 export CU_VERSION=cu132 export CHANNEL=nightly -export BUILD_VERSION=0.2.0.dev0 +export BUILD_VERSION=0.2.1.dev0 export TORCH_INSTALL_PATH="$(python -c 'import torch; print(torch.__path__[0])')" export TENSORRT_CUDA_VERSION_UPPER_BOUND=14.0 export TENSORRT_RTX_CUDA_VERSION_UPPER_BOUND=14.0 @@ -248,7 +248,7 @@ python -m pip install pyyaml patchelf tensorrt-cu13 \ --extra-index-url "https://download.pytorch.org/whl/test/${CU_VERSION}" \ --extra-index-url https://pypi.nvidia.com/ \ "executorch==1.6.0.dev20261009" -export TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION="0.2.0.dev0+${CU_VERSION}" +export TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION="0.2.1.dev0+${CU_VERSION}" python -m pip wheel --no-build-isolation --no-deps \ --wheel-dir dist py/torch-tensorrt-executorch-runtime ``` @@ -345,11 +345,13 @@ backend it cannot find. Loading and running a program is ExecuTorch's API, not this package's. Tensor placement, method lookup and output devices are all documented by ExecuTorch. A program exported with `skip_h2d_for_method_inputs` keeps its inputs on the -device, because nothing here copies them. +device, because nothing here moves them off the GPU. -Such a program does not check that promise, and it will not tell you when you -break it. Hand it a host tensor and it runs, returns the right answer, and -stages the copy for you, once per call. So a program you believe is copying +ExecuTorch's Python runtime refuses a host tensor for such an input. From C++, when TensorRT is +the first delegate to read the input, the program does not check that promise, and it will not +tell you when you break it. +Hand it a host tensor and it runs, returns the right answer, and stages the copy +for you, once per call. So a program you believe is copying nothing can be copying on every run with nothing to show it. The other direction is refused, so the leniency goes one way only. If the copies are why you exported this way, measure them rather than trusting the flag. @@ -360,10 +362,9 @@ extension: it is private and can change without notice. The two were measured against each other on every arrangement here, host and device, including a graph split with ExecuTorch's own CUDA backend, and they agree to the digit. -The choice that does matter is who owns the output buffer. Leave allocation on for -the graph outputs if Python will run the program, because a Python caller has no -way to hand one in. Turn it off only for a C++ consumer that supplies the address -itself. +The choice that does matter is who owns the output buffer. With allocation on for +the graph outputs, the program's device arena owns them. With it off, a C++ caller +supplies each address, and ExecuTorch's Python runtime allocates one for each output. ## Use diff --git a/py/torch-tensorrt-executorch-runtime/version.txt b/py/torch-tensorrt-executorch-runtime/version.txt index 341cf11faf..7dff5b8921 100644 --- a/py/torch-tensorrt-executorch-runtime/version.txt +++ b/py/torch-tensorrt-executorch-runtime/version.txt @@ -1 +1 @@ -0.2.0 \ No newline at end of file +0.2.1 \ No newline at end of file diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index f7330d43f2..3de2e759fb 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -945,7 +945,7 @@ def save( if output_format == "executorch" and not _has_executorch_exir(): raise ImportError( "Saving with output_format='executorch' requires executorch.exir. " - "This CUDA integration supports Linux. Setup: " + "This CUDA integration supports Linux and Windows x86-64. Setup: " + executorch_install_command() ) if output_format == "executorch": @@ -1586,7 +1586,8 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None except ImportError: raise ImportError( "Could not import the ExecuTorch export integration for " - "output_format='executorch'. This CUDA integration supports Linux. " + "output_format='executorch'. This CUDA integration supports Linux and " + "Windows x86-64. " "Setup: " + executorch_install_command() ) # save() only ever lowers a single method, so keep rejecting the per-method diff --git a/tests/py/dynamo/executorch/test_executorch_pin.py b/tests/py/dynamo/executorch/test_executorch_pin.py index 62f67b96db..97de4fe695 100644 --- a/tests/py/dynamo/executorch/test_executorch_pin.py +++ b/tests/py/dynamo/executorch/test_executorch_pin.py @@ -936,15 +936,29 @@ def test_install_channel_guard_rejects_unsupported_nightly_recipes( @pytest.mark.unit @pytest.mark.parametrize( - ("path", "command", "allowed"), + ("path", "command", "error"), [ - ("docsrc/install.rst", 'pip install "torch-tensorrt[executorch]"', True), - ("README.md", 'pip install "torch-tensorrt[executorch]"', False), - ("docsrc/install.rst", 'pip install --pre "torch-tensorrt[executorch]"', False), + ("docsrc/install.rst", 'pip install "torch-tensorrt[executorch]"', None), + ( + "README.md", + 'pip install "torch-tensorrt[executorch]"', + "names no nightly channel", + ), + ( + "docsrc/install.rst", + 'pip install --pre "torch-tensorrt[executorch]"', + "names no nightly channel", + ), + ( + "docsrc/install.rst", + 'pip install "torch-tensorrt[executorch]" ' + "--index-url https://download.pytorch.org/whl/nightly/cu132", + "without --pre", + ), ], ) def test_only_user_docs_may_print_the_release_install( - tmp_path, monkeypatch, path, command, allowed + tmp_path, monkeypatch, path, command, error ): recipe = tmp_path / path recipe.parent.mkdir(parents=True, exist_ok=True) @@ -952,13 +966,32 @@ def test_only_user_docs_may_print_the_release_install( subprocess.run(["git", "init", "-q", str(tmp_path)], check=True) subprocess.run(["git", "-C", str(tmp_path), "add", path], check=True) monkeypatch.setattr(sys.modules[__name__], "REPO_ROOT", tmp_path) - if allowed: + if error is None: test_every_printed_install_instruction_names_the_nightly_channel() else: - with pytest.raises(AssertionError, match="names no nightly channel"): + with pytest.raises(AssertionError, match=error): test_every_printed_install_instruction_names_the_nightly_channel() +@pytest.mark.unit +def test_the_release_install_is_one_command() -> None: + """The guide installs this release from PyPI with the extra alone. + + The release's PyPI wheel is built for the CUDA version of PyTorch's PyPI wheel, and once the + pin is an ExecuTorch release its extra brings ExecuTorch's CUDA build and the delegate built + with it, so naming either would only let them drift apart. + """ + guide = REPO_ROOT / "docsrc/user_guide/runtime_performance/executorch.rst" + commands = [ + shlex.split(block.replace("\\\n", "")) + for block in re.findall( + r"^ *pip install (?:.*\\\n)*.*$", guide.read_text(encoding="utf-8"), re.M + ) + ] + releases = [argv for argv in commands if "--pre" not in argv] + assert releases == [["pip", "install", "torch_tensorrt[executorch]"]], commands + + @pytest.mark.unit @pytest.mark.parametrize("variable", [False, True]) def test_install_channel_check_detects_removed_suffix_fix(