fix(executorch): move compile_and_save's engines into the program instead of copying them

docsrc/user_guide/runtime_performance/saving_models.rst

@@ -281,7 +281,7 @@ warning and are ignored. With this format, an unknown keyword raises a ``TypeErr
   :ref:`Zero-copy KV cache <executorch_zero_copy_kv>`.
 * ``generate_etrecord``: also write an ETRecord, a debug record for the ExecuTorch Developer Tools
   ``Inspector``, as ``<name>_etrecord.bin``. It does not work for a program with TensorRT
-  engines: writing the record fails with a ``SerializeError``.
+  engines: writing the record fails with a ``SerializeError``. ``compile_and_save`` can write one.
 
 :ref:`executorch_deployment` explains the run time side of ``use_cuda_graphs``,
 ``weight_streaming_budget_per_engine`` and ``external_engine_data``.

py/torch_tensorrt/_compile.py

@@ -77,6 +77,7 @@
     from torch_tensorrt.dynamo._tracer import trace as dynamo_trace
     from torch_tensorrt.dynamo.runtime import TorchTensorRTModule
     from torch_tensorrt.dynamo.runtime._offloaded_engines import (
+        defer_engine_setup_to_load,
         offload_engines_to_disk,
     )
     from torch_tensorrt.dynamo.utils import get_torch_inputs
@@ -873,6 +874,7 @@ def save(
                 ``<model>_etrecord.bin`` next to the ``.pte`` for use with the
                 ExecuTorch Inspector. It does not work for a program with
                 TensorRT engines: writing the record raises a ``SerializeError``.
+                ``compile_and_save`` can write one.
                 ``weight_streaming_budget_per_engine=`` takes an ``Optional[int]``
                 number of bytes of engine weights that may stay resident in GPU
                 memory, with the rest streamed from host memory. It applies to
@@ -1550,6 +1552,11 @@ def compile_and_save(
     but each engine is written to a file as soon as it is built, so no built engine is held on the GPU
     or in host memory while the next ones build. With ``output_format="executorch"`` no engine is ever
     loaded: the save streams each file into the ``.pte`` (or the ``.ptd`` with ``external_engine_data``).
+    Without ``generate_etrecord``, the save moves each file there instead, so no engine is on disk twice,
+    when ``file_path`` ends with ``.pte`` and the temporary directory (``TMPDIR``) is on the same file
+    system. With
+    ``generate_etrecord``, the save copies the engines and writes the ETRecord after the program, so the
+    program is kept if the ETRecord can't be written.
     The other formats save each engine from a loaded one, so their engines are loaded once all of them
     are built, and their save needs as much memory as after ``compile``.
 
@@ -1600,7 +1607,9 @@ def compile_and_save(
         save(trt_gm, os.fspath(file_path), output_format=output_format, **save_kwargs)
 
 
-def _write_external_tensor_data(executorch_program: Any, file_path: str) -> None:
+def _write_external_tensor_data(
+    executorch_program: Any, file_path: str, consume: bool = False
+) -> None:
     """Write an ExecuTorch program's external named tensor data (``.ptd``) next to the ``.pte``.
 
     The CUDA (AOTInductor) backend emits its weights as external named data
@@ -1616,7 +1625,11 @@ def _write_external_tensor_data(executorch_program: Any, file_path: str) -> None
     # Direct attribute access (no getattr default) so a future rename fails loud.
     if executorch_program._tensor_data:
         out_dir = os.path.dirname(os.path.abspath(file_path))
-        executorch_program.write_tensor_data_to_file(out_dir)
+        # ExecuTorch releases before the move take no consume argument.
+        if consume:
+            executorch_program.write_tensor_data_to_file(out_dir, consume=True)
+        else:
+            executorch_program.write_tensor_data_to_file(out_dir)
         # write_tensor_data_to_file adds .ptd unless the name already ends with it.
         paths = [
             os.path.join(out_dir, name if name.endswith(".ptd") else f"{name}.ptd")
@@ -1752,9 +1765,21 @@ def _lower_and_write_executorch(
         # the same call is the caller's to make, because finalization happens on
         # ExecuTorch's own manager and out of this library's sight.
         check_zero_copy_kv(executorch_program)
-    with open(file_path, "wb") as f:
-        executorch_program.write_to_file(f)
-    _write_external_tensor_data(executorch_program, file_path)
+    # Inside compile_and_save the engine files are deleted after this save, so it moves them into
+    # the program. The ETRecord reads the engines after the program is written, so with one the save
+    # copies them, which also keeps the program when the ETRecord can't be saved. ExecuTorch moves
+    # data only in its own save, which takes a .pte path.
+    consume = (
+        defer_engine_setup_to_load()
+        and not generate_etrecord
+        and file_path.endswith(".pte")
+    )
+    if consume:
+        executorch_program.save(file_path, consume=True)
+    else:
+        with open(file_path, "wb") as f:
+            executorch_program.write_to_file(f)
+    _write_external_tensor_data(executorch_program, file_path, consume)
     if generate_etrecord:
         # Follows ExecuTorch's example convention (e.g. examples/cuda/scripts/export.py):
         # persist the ETRecord as "<pte_base>_etrecord.bin" next to the .pte.

py/torch_tensorrt/dynamo/runtime/_offloaded_engines.py

@@ -8,6 +8,7 @@
 engines built so far are then held neither on the GPU nor in host memory while the next ones build.
 ``compile_module`` leaves them in their files: an ExecuTorch save streams each file into the program, and
 ``compile_and_save`` loads them before any other save. The files are deleted when the block ends.
+A file the save moved stays mapped, empty, until then, so nothing may read its engine after the save.
 """
 
 import contextlib

py/torch_tensorrt/executorch/_export.py

@@ -599,6 +599,7 @@ def export(
         generate_etrecord (bool): Ask ExecuTorch for an ETRecord for later debugging.
             This copies the whole program, engines included. Saving the record does not
             work for a program with TensorRT engines: it raises a ``SerializeError``.
+            ``torch_tensorrt.compile_and_save`` can write one.
         weight_streaming_budget_per_engine (Optional[int]): Bytes of engine weights that
             may stay resident in GPU memory, with the rest streamed from host memory.
             It applies to **each** TensorRT engine separately, not as a total for the

py/torch_tensorrt/executorch/backend.py

@@ -699,7 +699,9 @@ def preprocess(
 
         engine_data: Any
         if engine_file is not None:
-            engine_data = FileBackedData(engine_file)
+            # compile_and_save deletes the file after the save, so the data can own it, which lets
+            # a consuming write move the engine into the program instead of copying it.
+            engine_data = FileBackedData(engine_file, cleanup=True)
             engine_data_sha256 = engine_data.sha256()
         else:
             assert isinstance(engine_bytes, bytes)

tests/py/dynamo/executorch/test_api.py

@@ -2824,28 +2824,56 @@ def test_lift_keeps_parameter_for_any_dtype(dtype, requires_grad):
 # next to the .pte.
 
 
-def _patch_executorch_lowering(monkeypatch, captured):
+def _patch_executorch_lowering(monkeypatch, captured, log=None, etrecord_error=None):
     """Spy the seams _save_as_executorch delegates to after #4440, without running a
     real TRT lowering. Records the kwargs forwarded into
     torch_tensorrt.executorch.export() (captured["export_kwargs"]); lets the real
     export() run against an engine-free program while stubbing only the innermost
     ExecuTorch lowering, so the kwargs export() finally hands that lowering --
     including the get_edge_compile_config() it substitutes for an omitted
     compile_config -- land in `captured`; and records backend_config from
-    to_executorch(). Fills `captured`; returns nothing."""
+    to_executorch(). Fills `captured`; returns nothing. With `log`, records the order of
+    the ETRecord, program and .ptd writes, and whether each write consumes;
+    `etrecord_error` makes the ETRecord save raise."""
     import executorch.exir as exir
     import torch_tensorrt._compile as tc
     import torch_tensorrt.executorch as tte
 
     class _FakeETRecord:
         def save(self, path):
+            if log is not None:
+                log.append("etrecord")
+            if etrecord_error is not None:
+                raise etrecord_error
             with open(path, "wb") as fh:
                 fh.write(b"etrecord")
 
+    def _consumes(kwargs):
+        # ExecuTorch releases before the move reject consume, so only a move may pass it.
+        if kwargs not in ({}, {"consume": True}):
+            raise TypeError(f"unexpected keyword arguments {kwargs}")
+        return bool(kwargs)
+
     class _FakeExec:
+        _tensor_data = {"model": b""}
+
+        def save(self, path, **kwargs):
+            consume = _consumes(kwargs)
+            if log is not None:
+                log.append(("program", consume))
+            with open(path, "wb") as fh:
+                fh.write(b"")
+
         def write_to_file(self, f):
+            if log is not None:
+                log.append(("program", False))
             f.write(b"")
 
+        def write_tensor_data_to_file(self, outdir, **kwargs):
+            consume = _consumes(kwargs)
+            if log is not None:
+                log.append(("ptd", consume))
+
         def get_etrecord(self):
             return _FakeETRecord()
 
@@ -2875,7 +2903,6 @@ def _spy_export(source, **kw):
         return real_export(source, **kw)
 
     monkeypatch.setattr(tte, "export", _spy_export)
-    monkeypatch.setattr(tc, "_write_external_tensor_data", lambda prog, path: None)
     # ENABLED_FEATURES is an immutable namedtuple; swap the whole module attribute.
     monkeypatch.setattr(
         tc, "ENABLED_FEATURES", types.SimpleNamespace(torch_tensorrt_runtime=True)
@@ -2933,6 +2960,101 @@ def test_save_executorch_forwards_lowering_kwargs(monkeypatch, tmp_path):
     assert (tmp_path / "model_etrecord.bin").exists()
 
 
+@pytest.mark.unit
+def test_save_executorch_keeps_the_program_when_the_etrecord_fails(
+    monkeypatch, tmp_path
+):
+    pytest.importorskip("executorch.exir")
+    import torch_tensorrt._compile as tc
+
+    log = []
+    _patch_executorch_lowering(
+        monkeypatch, {}, log=log, etrecord_error=RuntimeError("cannot serialize")
+    )
+    out = tmp_path / "model.pte"
+    ep = torch.export.export(_AddOne(), (torch.randn(2, 2),))
+    with pytest.raises(RuntimeError, match="cannot serialize"):
+        tc._save_as_executorch(ep, str(out), generate_etrecord=True)
+
+    assert log == [("program", False), ("ptd", False), "etrecord"]
+    assert out.exists()
+
+
+@pytest.mark.unit
+def test_save_executorch_does_not_consume(monkeypatch, tmp_path):
+    pytest.importorskip("executorch.exir")
+    import torch_tensorrt._compile as tc
+
+    log = []
+    _patch_executorch_lowering(monkeypatch, {}, log=log)
+    ep = torch.export.export(_AddOne(), (torch.randn(2, 2),))
+    tc._save_as_executorch(ep, str(tmp_path / "model.pte"))
+
+    assert log == [("program", False), ("ptd", False)]
+
+
+@pytest.mark.unit
+def test_compile_and_save_moves_the_engines(monkeypatch, tmp_path):
+    pytest.importorskip("executorch.exir")
+    import torch_tensorrt._compile as tc
+    from torch_tensorrt.dynamo.runtime._offloaded_engines import (
+        offload_engines_to_disk,
+    )
+
+    log = []
+    _patch_executorch_lowering(monkeypatch, {}, log=log)
+    ep = torch.export.export(_AddOne(), (torch.randn(2, 2),))
+    # The block compile_and_save saves in, where the engine files are deleted after the save.
+    with offload_engines_to_disk():
+        tc._save_as_executorch(ep, str(tmp_path / "model.pte"))
+
+    assert log == [("program", True), ("ptd", True)]
+
+
+@pytest.mark.unit
+def test_compile_and_save_copies_the_engines_into_a_path_without_pte(
+    monkeypatch, tmp_path
+):
+    pytest.importorskip("executorch.exir")
+    import torch_tensorrt._compile as tc
+    from torch_tensorrt.dynamo.runtime._offloaded_engines import (
+        offload_engines_to_disk,
+    )
+
+    log = []
+    _patch_executorch_lowering(monkeypatch, {}, log=log)
+    ep = torch.export.export(_AddOne(), (torch.randn(2, 2),))
+    # ExecuTorch moves data only in its save, which needs a .pte path.
+    with offload_engines_to_disk():
+        tc._save_as_executorch(ep, str(tmp_path / "model.bin"))
+
+    assert log == [("program", False), ("ptd", False)]
+
+
+@pytest.mark.unit
+def test_compile_and_save_keeps_the_program_when_the_etrecord_fails(
+    monkeypatch, tmp_path
+):
+    pytest.importorskip("executorch.exir")
+    import torch_tensorrt._compile as tc
+    from torch_tensorrt.dynamo.runtime._offloaded_engines import (
+        offload_engines_to_disk,
+    )
+
+    log = []
+    _patch_executorch_lowering(
+        monkeypatch, {}, log=log, etrecord_error=OSError("No space left on device")
+    )
+    out = tmp_path / "model.pte"
+    ep = torch.export.export(_AddOne(), (torch.randn(2, 2),))
+    with offload_engines_to_disk():
+        with pytest.raises(OSError, match="No space left"):
+            tc._save_as_executorch(ep, str(out), generate_etrecord=True)
+
+    assert log == [("program", False), ("ptd", False), "etrecord"]
+    assert out.exists()
+
+
 @pytest.mark.unit
 def test_save_executorch_defaults_when_lowering_kwargs_omitted(monkeypatch, tmp_path):
     pytest.importorskip("executorch.exir")

tests/py/dynamo/executorch/test_compile_and_save.py

@@ -31,7 +31,7 @@ def forward(self, x):
 @pytest.mark.unit
 def test_engines_stream_from_their_files_into_the_program(tmp_path, monkeypatch):
     monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
-    setups, streamed = [], {}
+    setups, streamed, owned = [], {}, []
     real_setup = TorchTensorRTModule.setup_engine
     real_add = NamedDataStore.add_named_data
 
@@ -43,6 +43,8 @@ def recording_add(self, key, data, *args, **kwargs):
         if isinstance(data, FileBackedData):
             with open(data._path, "rb") as f:
                 streamed[data._path] = f.read()
+            # Data that owns its file is what a consuming save may move.
+            owned.append(data._finalizer is not None)
         else:
             streamed[key] = data
         return real_add(self, key, data, *args, **kwargs)
@@ -65,6 +67,7 @@ def recording_add(self, key, data, *args, **kwargs):
     # No engine was ever set up, and each one reached the program as its own file, from the cache directory.
     assert setups == []
     assert len(streamed) == 2
+    assert owned and all(owned)
     offloaded = str(_cache_root() / "offloaded_engines")
     assert all(
         path.endswith(".engine") and os.path.dirname(path) == offloaded
@@ -74,3 +77,69 @@ def recording_add(self, key, data, *args, **kwargs):
     assert all(pte.find(engine) >= 0 for engine in streamed.values())
     # The files are gone once the program is saved.
     assert not glob.glob(str(tmp_path / "**" / "*.engine"), recursive=True)
+
+
+@pytest.mark.unit
+def test_engines_move_into_the_program_without_a_second_copy(tmp_path, monkeypatch):
+    monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
+    engines, copied = {}, []
+    real_add = NamedDataStore.add_named_data
+    real_write = FileBackedData.write_to_file
+
+    def recording_add(self, key, data, *args, **kwargs):
+        if isinstance(data, FileBackedData):
+            with open(data._path, "rb") as f:
+                engines[data._path] = f.read()
+        return real_add(self, key, data, *args, **kwargs)
+
+    def recording_write(self, outfile):
+        copied.append(self._path)
+        real_write(self, outfile)
+
+    monkeypatch.setattr(NamedDataStore, "add_named_data", recording_add)
+    monkeypatch.setattr(FileBackedData, "write_to_file", recording_write)
+    torch.manual_seed(0)
+    x = torch.randn(4, 64, device="cuda")
+    pte_path = tmp_path / "model.pte"
+    torch_tensorrt.compile_and_save(
+        torch.export.export(_TwoEngines().eval().cuda(), (x,)),
+        pte_path,
+        output_format="executorch",
+        arg_inputs=[x],
+        min_block_size=1,
+        torch_executed_ops={torch.ops.aten.sin.default},
+    )
+    # ExecuTorch copies data it can't move through write_to_file, so no call there means each
+    # engine was moved, and the program holds them.
+    assert len(engines) == 2 and copied == []
+    pte = pte_path.read_bytes()
+    assert all(pte.find(engine) >= 0 for engine in engines.values())
+
+
+@pytest.mark.unit
+def test_an_etrecord_makes_the_save_copy_the_engines(tmp_path, monkeypatch):
+    monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
+    left = []
+    real_write = FileBackedData.write_to_file
+
+    def recording_write(self, outfile):
+        real_write(self, outfile)
+        left.append(os.path.getsize(self._path))
+
+    monkeypatch.setattr(FileBackedData, "write_to_file", recording_write)
+    torch.manual_seed(0)
+    x = torch.randn(4, 64, device="cuda")
+    pte_path = tmp_path / "model.pte"
+    torch_tensorrt.compile_and_save(
+        torch.export.export(_TwoEngines().eval().cuda(), (x,)),
+        pte_path,
+        output_format="executorch",
+        arg_inputs=[x],
+        min_block_size=1,
+        torch_executed_ops={torch.ops.aten.sin.default},
+        save_options={"generate_etrecord": True},
+    )
+    assert pte_path.stat().st_size > 0
+    assert (tmp_path / "model_etrecord.bin").stat().st_size > 0
+    # The ETRecord reads the engines after the program is written, so they stayed in their files.
+    assert left and all(size > 0 for size in left)