Cache flatc binary and schema extraction to fix 3x fbpkg export slowdown (#19104)

Differential Revision: D102214303

Pull Request resolved: #19104

Author: navsud

backends/xnnpack/serialization/xnnpack_graph_serialize.py

@@ -10,6 +10,7 @@
 import logging
 import os
 import tempfile
+import threading
 from dataclasses import dataclass, fields, is_dataclass
 from typing import ClassVar, Literal, Optional
 
@@ -300,6 +301,29 @@ def pretty_print_xnngraph(xnnpack_graph_json: str, filename: Optional[str] = Non
 _delegate_instance_id = 0
 
 
+_cached_schema_bytes: Optional[bytes] = None
+# Lock protecting _cached_schema_bytes initialization. The race itself would
+# be benign (read_bytes() is idempotent), but a lock keeps this consistent
+# with the pattern used for caching the flatc binary in _flatbuffer.py.
+_schema_bytes_lock: threading.Lock = threading.Lock()
+
+
+def _get_schema_bytes() -> bytes:
+    """Returns the schema.fbs bytes, caching the result across calls."""
+    global _cached_schema_bytes
+    # Double-checked locking: fast path avoids the lock once cached.
+    if _cached_schema_bytes is not None:
+        return _cached_schema_bytes
+    with _schema_bytes_lock:
+        if _cached_schema_bytes is None:
+            _cached_schema_bytes = (
+                _resources.files(serialization_package)
+                .joinpath("schema.fbs")
+                .read_bytes()
+            )
+        return _cached_schema_bytes
+
+
 def convert_to_flatbuffer(xnnpack_graph: XNNGraph) -> bytes:
     global _delegate_instance_id
     sanity_check_xnngraph_dataclass(xnnpack_graph)
@@ -316,11 +340,7 @@ def convert_to_flatbuffer(xnnpack_graph: XNNGraph) -> bytes:
     with tempfile.TemporaryDirectory() as d:
         schema_path = os.path.join(d, "schema.fbs")
         with open(schema_path, "wb") as schema_file:
-            schema_file.write(
-                _resources.files(serialization_package)
-                .joinpath("schema.fbs")
-                .read_bytes()
-            )
+            schema_file.write(_get_schema_bytes())
         json_path = os.path.join(d, "schema.json")
         with open(json_path, "wb") as json_file:
             json_file.write(xnnpack_graph_json.encode("ascii"))

exir/_serialize/_flatbuffer.py

@@ -7,14 +7,16 @@
 
 # pyre-strict
 
+import atexit
+import contextlib
 import importlib.resources
 import os
 import re
 import shutil
 import stat
 import subprocess
-
 import tempfile
+import threading
 
 from dataclasses import dataclass
 from typing import Callable, Dict, List, Optional, Sequence
@@ -240,20 +242,41 @@ class _FlatbufferResult:
 # Name of an optional resource containing the `flatc` executable.
 _FLATC_RESOURCE_NAME: str = "flatbuffers-flatc"
 
-
-def _run_flatc(args: Sequence[str]) -> None:
-    """Runs the `flatc` command with the provided args.
-
-    If a resource matching _FLATC_RESOURCE_NAME exists, uses that executable.
-    Otherwise, expects the `flatc` tool to be available on the system path.
-    """
-    flatc_resource = importlib.resources.files(__package__).joinpath(
-        _FLATC_RESOURCE_NAME
-    )
-    if flatc_resource.is_file():
-        # Use the provided flatc binary.
-        with importlib.resources.as_file(flatc_resource) as flatc_path:
-            # Ensure the binary has execute permissions (needed for PAR files)
+# Cached flatc binary path. In PAR files, importlib.resources.as_file()
+# extracts the binary to a temp file on each call. With 200+ XNNPACK
+# partitions this adds ~30 min of overhead. Caching avoids re-extraction.
+# The ExitStack is registered with atexit so the extracted temp file is
+# cleaned up on normal process exit.
+#
+# Fork safety: the child inherits the parent's atexit registry and cached
+# path. Without _reset_flatc_cache_after_fork, the child's atexit would
+# run the inherited handler and unlink the parent's temp file. The
+# before/after_in_parent callbacks hold _flatc_lock across fork so the
+# child never inherits a half-initialized cache.
+_flatc_cached_path: Optional[str] = None
+_flatc_exit_stack: Optional[contextlib.ExitStack] = None
+_flatc_lock: threading.Lock = threading.Lock()
+
+
+def _get_flatc_path() -> str:
+    """Returns the path to the flatc executable, caching the result."""
+    global _flatc_cached_path, _flatc_exit_stack
+    # Double-checked locking: fast path avoids the lock once cached.
+    if _flatc_cached_path is not None:
+        return _flatc_cached_path
+
+    with _flatc_lock:
+        if _flatc_cached_path is not None:
+            return _flatc_cached_path
+
+        flatc_resource = importlib.resources.files(__package__).joinpath(
+            _FLATC_RESOURCE_NAME
+        )
+        if flatc_resource.is_file():
+            exit_stack = contextlib.ExitStack()
+            flatc_path = exit_stack.enter_context(
+                importlib.resources.as_file(flatc_resource)
+            )
             try:
                 current_mode = flatc_path.stat().st_mode
                 if not (current_mode & stat.S_IXUSR):
@@ -262,13 +285,49 @@ def _run_flatc(args: Sequence[str]) -> None:
                     )
             except OSError:
                 pass
-            subprocess.run([flatc_path] + list(args), check=True)
-    else:
-        # Expect the `flatc` tool to be on the system path or set as an env var.
-        flatc_path = os.getenv("FLATC_EXECUTABLE")
-        if not flatc_path:
-            flatc_path = "flatc"
-        subprocess.run([flatc_path] + list(args), check=True)
+            _flatc_exit_stack = exit_stack
+            # Clean up the extracted temp file on normal process exit.
+            atexit.register(exit_stack.close)
+            _flatc_cached_path = str(flatc_path)
+        else:
+            _flatc_cached_path = os.getenv("FLATC_EXECUTABLE", "flatc")
+
+        return _flatc_cached_path
+
+
+def _reset_flatc_cache_after_fork() -> None:
+    """Reset the flatc cache in the child after fork.
+
+    Unregister the inherited atexit handler (do NOT call .close() — the
+    parent still owns the file), clear the cached state so the child
+    re-extracts lazily, and replace the lock (the inherited one is held
+    by the `before` fork callback but the acquiring thread no longer
+    exists in the child).
+    """
+    global _flatc_cached_path, _flatc_exit_stack, _flatc_lock
+    if _flatc_exit_stack is not None:
+        atexit.unregister(_flatc_exit_stack.close)
+    _flatc_cached_path = None
+    _flatc_exit_stack = None
+    _flatc_lock = threading.Lock()
+
+
+# os.register_at_fork is Unix-only; guard for Windows importability.
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(
+        before=lambda: _flatc_lock.acquire(),
+        after_in_parent=lambda: _flatc_lock.release(),
+        after_in_child=_reset_flatc_cache_after_fork,
+    )
+
+
+def _run_flatc(args: Sequence[str]) -> None:
+    """Runs the `flatc` command with the provided args.
+
+    If a resource matching _FLATC_RESOURCE_NAME exists, uses that executable.
+    Otherwise, expects the `flatc` tool to be available on the system path.
+    """
+    subprocess.run([_get_flatc_path()] + list(args), check=True)
 
 
 def _flatc_compile(output_dir: str, schema_path: str, json_path: str) -> None: