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
```
This commit is contained in:
Mark Campbell 2025-05-29 17:53:45 +01:00 committed by GitHub
parent bfdd15d1fa
commit f0d8ceb242
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -107,6 +107,13 @@ def collect_template_dependencies(template_dir: Path) -> tuple[str | None, list[
return None, [] 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(): def main():
templates_dir = REPO_ROOT / "llama_stack" / "templates" templates_dir = REPO_ROOT / "llama_stack" / "templates"
change_tracker = ChangedPathTracker() change_tracker = ChangedPathTracker()
@ -118,6 +125,8 @@ def main():
template_dirs = list(find_template_dirs(templates_dir)) template_dirs = list(find_template_dirs(templates_dir))
task = progress.add_task("Processing distribution templates...", total=len(template_dirs)) task = progress.add_task("Processing distribution templates...", total=len(template_dirs))
pre_import_templates(template_dirs)
# Create a partial function with the progress bar # Create a partial function with the progress bar
process_func = partial(process_template, progress=progress, change_tracker=change_tracker) process_func = partial(process_template, progress=progress, change_tracker=change_tracker)