From f0d8ceb2422247b1c68bfda9d92f9561012310df Mon Sep 17 00:00:00 2001
From: Mark Campbell <mcampbel@redhat.com>
Date: Thu, 29 May 2025 17:53:45 +0100
Subject: [PATCH] chore: fix flaky distro_codegen script (#2305)

# What does this PR do?
<!-- Provide a short summary of what this PR does and why. Link to
relevant issues if applicable. -->
Adds an import for all of the template modules before the executor to
prevent deadlock
<!-- If resolving an issue, uncomment and update the line below -->
Closes #2278

## Test Plan
<!-- Describe the tests you ran to verify your changes with result
summaries. *Provide clear instructions so the plan can be easily
re-executed.* -->
```
# Run the pre-commit multiple times and verify the deadlock doesn't occur
for i in {1..10}; do pre-commit run --all-files; done
```
---
 scripts/distro_codegen.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/scripts/distro_codegen.py b/scripts/distro_codegen.py
index 8820caf55..d33c5de67 100755
--- a/scripts/distro_codegen.py
+++ b/scripts/distro_codegen.py
@@ -107,6 +107,13 @@ def collect_template_dependencies(template_dir: Path) -> tuple[str | None, list[
     return None, []
 
 
+def pre_import_templates(template_dirs: list[Path]) -> None:
+    # Pre-import all template modules to avoid deadlocks.
+    for template_dir in template_dirs:
+        module_name = f"llama_stack.templates.{template_dir.name}"
+        importlib.import_module(module_name)
+
+
 def main():
     templates_dir = REPO_ROOT / "llama_stack" / "templates"
     change_tracker = ChangedPathTracker()
@@ -118,6 +125,8 @@ def main():
         template_dirs = list(find_template_dirs(templates_dir))
         task = progress.add_task("Processing distribution templates...", total=len(template_dirs))
 
+        pre_import_templates(template_dirs)
+
         # Create a partial function with the progress bar
         process_func = partial(process_template, progress=progress, change_tracker=change_tracker)