Save the etrecord the lowered paths already generate

examples/models/llama/export_llama_lib.py

@@ -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,28 @@ def _to_edge_and_lower_llama_xnnpack(
     )
 
 
+def _save_etrecord_if_generated(builder) -> None:
+    """Write the etrecord the lowering attached, if there is one.
+
+    Called after the model is written, and never allowed to raise, because a debug artifact
+    must not cost a caller the export itself. A record twice the size of the model makes a
+    full disk the likely trigger.
+    """
+    try:
+        etrecord = builder.export_program.get_etrecord()
+    except RuntimeError:
+        # Not generated, which is the normal case.
+        return
+
+    try:
+        etrecord.save("etrecord.bin")
+    except OSError as error:
+        logging.warning("Could not write etrecord.bin: %s", error)
+        return
+
+    logging.info("Generated etrecord.bin")
+
+
 def _to_edge_and_lower_llama_openvino(
     builder_exported,
     modelname,
@@ -1693,6 +1714,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
 
@@ -1875,6 +1897,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,7 +5,10 @@
 # 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
 
 from executorch.devtools.backend_debug import get_delegation_info
 
@@ -40,7 +43,85 @@
 ]
 
 
+def _tiny_llm_builder():
+    """An exported LLMEdgeManager small enough to lower in a unit test.
+
+    Enough to exercise the lowering helpers without a checkpoint or a real llama model.
+    """
+    import torch
+    from executorch.extension.llm.export.builder import LLMEdgeManager
+
+    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),),
+    ).export()
+
+
 class ExportLlamaLibTest(unittest.TestCase):
+    def _run_tiny_export(self, generate_etrecord):
+        llm_config = LlmConfig()
+        llm_config.backend.xnnpack.enabled = True
+        llm_config.debug.generate_etrecord = generate_etrecord
+        # The recipe rejects dynamic shapes, and _validate_args runs before the lowering.
+        llm_config.model.enable_dynamic_shape = False
+        llm_config.export.output_name = "tiny.pte"
+        with patch(
+            "executorch.examples.models.llama.export_llama_lib._prepare_for_llama_export",
+            return_value=_tiny_llm_builder(),
+        ):
+            _export_llama(llm_config)
+
+    def test_export_saves_the_etrecord_it_generates(self):
+        """A lowering that generates an etrecord must also write it.
+
+        `to_edge_transform_and_lower` builds the record when asked and carries it to the ExecuTorch
+        program, but nothing saved it, so the flag produced no file and no warning while still
+        paying for the record.
+        """
+        with tempfile.TemporaryDirectory() as directory:
+            previous = os.getcwd()
+            os.chdir(directory)
+            try:
+                self._run_tiny_export(generate_etrecord=True)
+                self.assertEqual(sorted(os.listdir(".")), ["etrecord.bin", "tiny.pte"])
+            finally:
+                os.chdir(previous)
+
+    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:
+            previous = os.getcwd()
+            os.chdir(directory)
+            try:
+                self._run_tiny_export(generate_etrecord=False)
+                self.assertEqual(os.listdir("."), ["tiny.pte"])
+            finally:
+                os.chdir(previous)
+
+    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 and its failure is logged, so a full disk or an
+        unwritable directory loses the record and not the .pte.
+        """
+        with tempfile.TemporaryDirectory() as directory:
+            previous = os.getcwd()
+            os.chdir(directory)
+            try:
+                # A directory of that name makes the write fail without touching permissions.
+                os.mkdir("etrecord.bin")
+                self._run_tiny_export(generate_etrecord=True)
+                self.assertIn("tiny.pte", os.listdir("."))
+            finally:
+                os.chdir(previous)
+
     def test_has_expected_ops_and_op_counts(self):
         """
         Checks the presence of unwanted expensive ops.