Save flatc input files when serialization fails

Summary:
It's been tough to debug `flatc` serialization problems because the files are deleted after the call.

Now, if the serialization fails, copy the input files to a directory where we can look at them, if the environment var `ET_EXIR_SAVE_FLATC_INPUTS_ON_FAILURE=1` is set.

Reviewed By: lucylq

Differential Revision: D53736362

fbshipit-source-id: 4ecd57f13ec46ba6ee0d31528f59300bc3947fd0

Author: dbort

exir/_serialize/_flatbuffer.py

@@ -9,13 +9,19 @@
 import importlib.resources
 import os
 import re
+import shutil
 import subprocess
 import tempfile
 
 from dataclasses import dataclass
 from typing import Callable, Dict, List, Optional, Sequence
 
 
+# If this environment variable is set to true, save the flatc input files when
+# serialization fails.
+_SAVE_FLATC_ENV: str = "ET_EXIR_SAVE_FLATC_INPUTS_ON_FAILURE"
+
+
 def _is_valid_alignment(alignment: int) -> bool:
     """Returns True if the alignment is valid, or is None."""
     if alignment is None:
@@ -278,7 +284,33 @@ def _program_json_to_flatbuffer(
         with open(json_path, "wb") as json_file:
             json_file.write(program_json.encode("ascii"))
 
-        _flatc_compile(temp_dir, schema_info.root_path, json_path)
+        try:
+            _flatc_compile(temp_dir, schema_info.root_path, json_path)
+        except Exception as err:
+            # It's helpful to save the breaking files for debugging. Optionally
+            # move them out of the auto-deleting temporary directory. Don't do
+            # this by default because some input files can be many GB in size,
+            # and these copies won't be auto-deleted.
+            should_save = os.getenv(_SAVE_FLATC_ENV, "").strip() not in {"", "0"}
+            extra_message = ""
+            if should_save:
+                try:
+                    saved_dir = tempfile.mkdtemp(prefix="exir-saved-flatc-")
+                    for f in os.listdir(temp_dir):
+                        shutil.move(src=os.path.join(temp_dir, f), dst=saved_dir)
+                    extra_message += f" Moved input files to '{saved_dir}'."
+                except Exception as err2:
+                    extra_message += (
+                        f" (Failed to save input files for debugging: {err2})"
+                    )
+            else:
+                extra_message += (
+                    f" Set {_SAVE_FLATC_ENV}=1 to save input files on failure."
+                )
+
+            raise RuntimeError(
+                f"Failed to compile {json_path} to {output_path}." + extra_message
+            ) from err
         with open(output_path, "rb") as output_file:
             return _FlatbufferResult(
                 data=output_file.read(), max_alignment=schema_info.max_alignment

exir/_serialize/test/test_flatbuffer.py

@@ -6,13 +6,19 @@
 # LICENSE file in the root directory of this source tree.
 
 import os
+import re
+import shutil
 import tempfile
 import unittest
 from typing import Dict, Optional, Sequence
 from unittest.mock import patch
 
 from executorch.exir._serialize import _flatbuffer
-from executorch.exir._serialize._flatbuffer import _ResourceFiles, _SchemaInfo
+from executorch.exir._serialize._flatbuffer import (
+    _program_json_to_flatbuffer,
+    _ResourceFiles,
+    _SchemaInfo,
+)
 
 
 def read_file(dir: str, filename: str) -> bytes:
@@ -266,3 +272,60 @@ def test_bad_delegate_alignment_fails(self) -> None:
                             out_dir,
                             delegate_alignment=bad_alignment,
                         )
+
+
+class TestProgramJsonToFlatbuffer(unittest.TestCase):
+    @patch.dict(os.environ, {_flatbuffer._SAVE_FLATC_ENV: "1"})
+    def test_save_json_on_failure(self) -> None:
+        err_msg: Optional[str] = None
+        try:
+            _program_json_to_flatbuffer("} some bad json {")
+            self.fail("Should have raised an exception")
+        except RuntimeError as err:
+            err_msg = err.args[0]
+
+        self.assertIsNotNone(err_msg)
+        match = re.search(r"Moved input files to '(.*?)'", err_msg)
+        self.assertTrue(match, msg=f"Unexpected error message: {err_msg}")
+        path = match.group(1)
+
+        files = frozenset(os.listdir(path))
+        # Delete the files otherwise they'll accumulate every time the
+        # test is run.
+        shutil.rmtree(path)
+        # Check for a couple of the files that should be there.
+        self.assertIn("data.json", files)
+        self.assertIn("program.fbs", files)
+
+    @patch.dict(os.environ, {_flatbuffer._SAVE_FLATC_ENV: "1"})
+    def test_unable_to_save_json_on_failure(self) -> None:
+        err_msg: Optional[str] = None
+        try:
+            with patch.object(
+                _flatbuffer.shutil,
+                "move",
+                side_effect=Exception("shutil.move mock failure"),
+            ):
+                _program_json_to_flatbuffer("} some bad json {")
+            self.fail("Should have raised an exception")
+        except RuntimeError as err:
+            err_msg = err.args[0]
+
+        self.assertIsNotNone(err_msg)
+        self.assertIn("Failed to save input files", err_msg)
+
+    @patch.dict(os.environ, {_flatbuffer._SAVE_FLATC_ENV: ""})
+    def test_no_save_json_on_failure(self) -> None:
+        err_msg: Optional[str] = None
+        try:
+            _program_json_to_flatbuffer("} some bad json {")
+            self.fail("Should have raised an exception")
+        except RuntimeError as err:
+            err_msg = err.args[0]
+
+        self.assertIsNotNone(err_msg)
+        self.assertIn(
+            f"Set {_flatbuffer._SAVE_FLATC_ENV}=1 to save input files", err_msg
+        )
+        self.assertNotIn("Moved input files", err_msg)
+        self.assertNotIn("Failed to save input files", err_msg)