Save the etrecord the lowered paths already generate

examples/models/llama/export_llama_lib.py

@@ -14,8 +14,10 @@
 import json
 import logging
 import math
+import os
 import re
 import shlex
+import tempfile
 from functools import partial
 from importlib import resources as _resources
 from json import JSONDecodeError
@@ -24,7 +26,6 @@
 
 import torch
 from executorch.devtools.backend_debug import print_delegation_info
-from executorch.devtools.etrecord import generate_etrecord as generate_etrecord_func
 from executorch.examples.models.llama.hf_download import (
     download_and_convert_hf_checkpoint,
 )
@@ -762,7 +763,6 @@ def _ensure_static_kvq_out_variants(llm_config: LlmConfig) -> None:
         libraries = [llm_config.export.so_library]
     else:
         import glob
-        import os
 
         import executorch
         from executorch.extension.pybindings import portable_lib  # noqa # usort: skip
@@ -1186,7 +1186,6 @@ def _to_edge_and_lower_llama_xnnpack(
     for partitioner in partitioners:
         logging.info(f"--> {partitioner.__class__.__name__}")
 
-    # TODO: Enable generating ETRecord with XNNPack and to_edge_transform_and_lower().
     if generate_etrecord:
         builder_exported.generate_etrecord = True
 
@@ -1214,6 +1213,53 @@ def _to_edge_and_lower_llama_xnnpack(
     )
 
 
+def _etrecord_path_beside_model(builder) -> str:
+    """Where to write the etrecord so it lands beside the model.
+
+    `save_pte_program` uses an output name ending in `.pte` as the path verbatim, ignoring
+    `output_dir`, so the model's own directory is the only reliable anchor. Falls back to
+    `output_dir` before the model is saved.
+    """
+    saved = builder.get_saved_pte_filename()
+    directory = os.path.dirname(saved) if saved else builder.output_dir
+    return os.path.join(directory, "etrecord.bin")
+
+
+def _save_etrecord_if_generated(builder) -> None:
+    """Write the etrecord the lowering produced, if there is one.
+
+    The record goes beside the model rather than under `output_dir`, because an output name
+    ending in `.pte` is used as the path verbatim and so bypasses `output_dir`.
+
+    A failed write costs the record and not the export, because a debug artifact must not
+    cost a caller the model. The record is staged under a temporary name and moved into
+    place, so a failure part way cannot leave a half-written record under the real name.
+    """
+    try:
+        etrecord = builder.export_program.get_etrecord()
+    except RuntimeError:
+        # Not generated, which is the normal case.
+        return
+
+    path = _etrecord_path_beside_model(builder)
+    try:
+        with tempfile.NamedTemporaryFile(
+            dir=os.path.dirname(path) or ".", suffix=".bin", delete=False
+        ) as temporary:
+            staged = temporary.name
+        try:
+            etrecord.save(staged)
+            os.replace(staged, path)
+        except BaseException:
+            os.unlink(staged)
+            raise
+    except Exception as error:
+        logging.warning("Could not write %s: %s", path, error)
+        return
+
+    logging.info("Generated %s", path)
+
+
 def _to_edge_and_lower_llama_openvino(
     builder_exported,
     modelname,
@@ -1420,6 +1466,8 @@ def _to_edge_and_lower_llama(  # noqa: C901
     generate_etrecord: bool = False,
     verbose: bool = False,
 ):
+    # Set before the edge export, because that is where the record is created.
+    builder_exported.generate_etrecord = generate_etrecord
     builder_exported_to_edge = builder_exported.pt2e_quantize(
         quantizers
     ).export_to_edge()
@@ -1528,9 +1576,6 @@ def _to_edge_and_lower_llama(  # noqa: C901
         if not builder_exported_to_edge.edge_manager:
             raise ValueError("Unable to generate etrecord due to missing edge manager.")
 
-        logging.info("Generating etrecord")
-        # Copy the edge manager which will be serialized into etrecord. This is memory-wise expensive.
-        edge_manager_copy = copy.deepcopy(builder_exported_to_edge.edge_manager)
         builder = builder_exported_to_edge.to_backend(partitioners)
         if verbose:
             print_delegation_info(builder.edge_manager.exported_program().graph_module)
@@ -1543,15 +1588,8 @@ def _to_edge_and_lower_llama(  # noqa: C901
         builder = builder.to_executorch(
             passes=additional_passes,
         )
-
-        # Generate ETRecord
-        if edge_manager_copy:
-            generate_etrecord_func(
-                et_record="etrecord.bin",
-                edge_dialect_program=edge_manager_copy,
-                executorch_program=builder.export_program,
-            )
-            logging.info("Generated etrecord.bin")
+        # The record rides along with the program, so _save_etrecord_if_generated writes it
+        # after the model like every other path.
     else:
         builder = builder_exported_to_edge.to_backend(partitioners)
         if verbose:
@@ -1693,6 +1731,7 @@ def _export_llama_multimethod(llm_config: LlmConfig) -> LLMEdgeManager:
         first_builder.dtype,
     )
     first_builder.save_to_pte(output_file)
+    _save_etrecord_if_generated(first_builder)
 
     return first_builder
 
@@ -1873,6 +1912,7 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager:  # noqa: C901
         builder.dtype,
     )
     builder.save_to_pte(output_file)
+    _save_etrecord_if_generated(builder)
     return builder
 
 

examples/models/llama/tests/test_export_llama_lib.py

@@ -5,10 +5,15 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
+import tempfile
 import unittest
 from unittest.mock import patch
 
+import torch
 from executorch.devtools.backend_debug import get_delegation_info
+from executorch.devtools.etrecord import parse_etrecord
+from executorch.devtools.etrecord._etrecord import ETRecord
 
 try:
     from executorch.backends.arm.quantizer.arm_quantizer import (
@@ -30,8 +35,10 @@
     build_args_parser,
     get_quantizer_and_quant_params,
 )
+from executorch.extension.llm.export.builder import LLMEdgeManager
 from executorch.extension.llm.export.config.llm_config import (
     LlmConfig,
+    MethodConfig,
     Pt2eQuantize,
     VgfQuantizeScope,
 )
@@ -42,6 +49,25 @@
 ]
 
 
+def _tiny_llm_builder():
+    """An LLMEdgeManager small enough to lower in a unit test.
+
+    `_export_llama` exports it, so this returns it unexported.
+    """
+
+    class Tiny(torch.nn.Module):
+        def forward(self, tokens):
+            return tokens.to(torch.float32) * 2.0 + 1.0
+
+    return LLMEdgeManager(
+        model=Tiny(),
+        modelname="tiny",
+        max_seq_len=4,
+        use_kv_cache=False,
+        example_inputs=(torch.ones(1, 4, dtype=torch.long),),
+    )
+
+
 class ExportLlamaLibTest(unittest.TestCase):
     def _assert_routes_to(self, lowering, coreml=False, vulkan=False, qnn=False):
         """Assert which lowering an export routes to, without running one."""
@@ -96,6 +122,148 @@ def test_core_ml_with_vulkan_routes_to_the_combined_lowering(self):
         self.assertTrue(target.call_args.kwargs["coreml"])
         self.assertTrue(target.call_args.kwargs["vulkan"])
 
+    def _run_tiny_export(
+        self,
+        generate_etrecord,
+        directory,
+        output_name="tiny.pte",
+        output_dir=None,
+        xnnpack=True,
+    ):
+        """Export the tiny model, writing the .pte into `directory`.
+
+        The builder is pointed at a directory rather than chdir'ing the process, because these
+        tests run under pytest-xdist where the cwd is shared between tests in a worker.
+
+        `output_dir` defaults to `directory`, and is set separately only by the test that
+        checks a `.pte` output name does not separate the model from its record.
+
+        `xnnpack=False` selects the combined lowering, which reaches the edge export by a
+        different route than the backend-specific helpers.
+        """
+        llm_config = LlmConfig()
+        llm_config.backend.xnnpack.enabled = xnnpack
+        llm_config.debug.generate_etrecord = generate_etrecord
+        llm_config.model.enable_dynamic_shape = False
+        llm_config.export.output_dir = output_dir or directory
+        llm_config.export.output_name = os.path.join(directory, output_name)
+        builder = _tiny_llm_builder().set_output_dir(output_dir or directory)
+        with patch(
+            "executorch.examples.models.llama.export_llama_lib._prepare_for_llama_export",
+            return_value=builder,
+        ):
+            _export_llama(llm_config)
+        return sorted(os.listdir(directory))
+
+    def test_combined_lowering_saves_the_etrecord_beside_the_model(self):
+        """The combined lowering must write a record too.
+
+        It reaches the edge export by a different route than the backend helpers, so the flag
+        has to be set on the builder before that export rather than at the lowering.
+        """
+        with tempfile.TemporaryDirectory() as directory:
+            self.assertEqual(
+                self._run_tiny_export(
+                    generate_etrecord=True, directory=directory, xnnpack=False
+                ),
+                ["etrecord.bin", "tiny.pte"],
+            )
+            self.assertIsNotNone(
+                parse_etrecord(
+                    os.path.join(directory, "etrecord.bin")
+                ).edge_dialect_program
+            )
+
+    def test_export_saves_the_etrecord_beside_the_model(self):
+        """The record must be written, and land beside the model.
+
+        A `.pte` output name is used as the path verbatim, so `output_dir` points somewhere
+        else here: a record keyed off `output_dir` would separate it from the model.
+        """
+        with tempfile.TemporaryDirectory() as directory, tempfile.TemporaryDirectory() as elsewhere:
+            self.assertEqual(
+                self._run_tiny_export(
+                    generate_etrecord=True, directory=directory, output_dir=elsewhere
+                ),
+                ["etrecord.bin", "tiny.pte"],
+            )
+            # An empty file would satisfy the listing, so load it back.
+            self.assertIsNotNone(
+                parse_etrecord(
+                    os.path.join(directory, "etrecord.bin")
+                ).edge_dialect_program
+            )
+
+    def test_export_writes_no_etrecord_when_not_asked(self):
+        """The common case must stay silent rather than raise or leave a stray file."""
+        with tempfile.TemporaryDirectory() as directory:
+            with self.assertNoLogs(level="WARNING"):
+                landed = self._run_tiny_export(
+                    generate_etrecord=False, directory=directory
+                )
+            self.assertEqual(landed, ["tiny.pte"])
+
+    def test_export_keeps_the_model_when_the_etrecord_cannot_be_written(self):
+        """A debug artifact must not cost the caller the export.
+
+        The record is written after the model, so a failed record write costs the record and
+        not the .pte. A full disk is the likely trigger, since the record is the larger.
+        """
+        with tempfile.TemporaryDirectory() as directory:
+            # A directory of that name makes the rename fail without touching permissions.
+            os.mkdir(os.path.join(directory, "etrecord.bin"))
+            with self.assertLogs(level="WARNING") as logs:
+                landed = self._run_tiny_export(
+                    generate_etrecord=True, directory=directory
+                )
+            self.assertIn("tiny.pte", landed)
+            self.assertTrue(any("Could not write" in line for line in logs.output))
+            # The staged record must not be left behind next to the model.
+            self.assertEqual(landed, ["etrecord.bin", "tiny.pte"])
+
+    def test_export_keeps_the_previous_etrecord_when_a_rewrite_fails(self):
+        """A failed rewrite must not destroy the record that was already there.
+
+        The record format truncates its target on open, so writing in place would leave a
+        short file under the real name and report only a warning.
+        """
+        with tempfile.TemporaryDirectory() as directory:
+            self._run_tiny_export(generate_etrecord=True, directory=directory)
+            record = os.path.join(directory, "etrecord.bin")
+            good = os.path.getsize(record)
+
+            with patch.object(
+                ETRecord, "_save_graph_map", side_effect=RuntimeError("disk full")
+            ):
+                with self.assertLogs(level="WARNING"):
+                    self._run_tiny_export(generate_etrecord=True, directory=directory)
+
+            self.assertEqual(os.path.getsize(record), good)
+            self.assertIsNotNone(parse_etrecord(record).edge_dialect_program)
+
+    def test_multimethod_export_saves_the_etrecord_beside_the_model(self):
+        """The multimethod path must write the record too.
+
+        It builds a record the same way the single-method paths do, so leaving out the save
+        would reproduce the silent-flag bug on that path alone.
+        """
+        with tempfile.TemporaryDirectory() as directory:
+            llm_config = LlmConfig()
+            llm_config.backend.xnnpack.enabled = True
+            llm_config.debug.generate_etrecord = True
+            llm_config.multimethod.methods = [MethodConfig(method_name="forward")]
+            llm_config.export.output_dir = directory
+            llm_config.export.output_name = os.path.join(directory, "tiny.pte")
+            with patch.object(
+                export_llama_lib,
+                "_prepare_for_llama_export",
+                side_effect=lambda _: _tiny_llm_builder().set_output_dir(directory),
+            ):
+                _export_llama(llm_config)
+            self.assertEqual(
+                sorted(os.listdir(directory)), ["etrecord.bin", "tiny.pte"]
+            )
+
     def test_has_expected_ops_and_op_counts(self):
         """
         Checks the presence of unwanted expensive ops.

extension/export_util/utils.py

@@ -56,6 +56,7 @@ def _core_aten_to_edge(
     edge_constant_methods: Optional[Dict[str, Any]] = None,
     edge_compile_config=None,
     verbose=True,
+    generate_etrecord: bool = False,
 ) -> EdgeProgramManager:
     if not edge_compile_config:
         edge_compile_config = exir.EdgeCompileConfig(
@@ -65,6 +66,7 @@ def _core_aten_to_edge(
         core_aten_exir_ep,
         constant_methods=edge_constant_methods,
         compile_config=edge_compile_config,
+        generate_etrecord=generate_etrecord,
     )
     if verbose:
         logging.info(f"Exported graph:\n{edge_manager.exported_program()}")
@@ -81,6 +83,7 @@ def export_to_edge(
     edge_compile_config=_EDGE_COMPILE_CONFIG,
     strict=True,
     verbose=True,
+    generate_etrecord: bool = False,
 ) -> EdgeProgramManager:
     core_aten_ep = _to_core_aten(
         model,
@@ -91,7 +94,11 @@ def export_to_edge(
         verbose=verbose,
     )
     return _core_aten_to_edge(
-        core_aten_ep, edge_constant_methods, edge_compile_config, verbose=verbose
+        core_aten_ep,
+        edge_constant_methods,
+        edge_compile_config,
+        verbose=verbose,
+        generate_etrecord=generate_etrecord,
     )
 
 

extension/llm/export/builder.py

@@ -471,6 +471,7 @@ def export_to_edge(self) -> "LLMEdgeManager":
                     edge_constant_methods=self.metadata,
                     edge_compile_config=edge_config,
                     verbose=self.verbose,
+                    generate_etrecord=self.generate_etrecord,
                 )
         return self