diff --git a/docs/blog/fixing-import-loops.mdx b/docs/blog/fixing-import-loops.mdx new file mode 100644 index 000000000..7633fc6ae --- /dev/null +++ b/docs/blog/fixing-import-loops.mdx @@ -0,0 +1,159 @@ +--- +title: "Import Loops in PyTorch" +icon: "arrows-rotate" +iconType: "solid" +description: "Identifying and visualizing import loops in the PyTorch codebase" +--- + +In this post, we will visualize all import loops in the [PyTorch](https://github.com/pytorch/pytorch) codebase, propose a fix for one potentially unstable case, and use Codegen to refactor that fix. + + +You can find the complete jupyter notebook in our [examples repository](https://github.com/codegen-sh/codegen-examples/tree/main/examples/removing_import_loops_in_pytorch). + + +Import loops (or circular dependencies) occur when two or more Python modules depend on each other, creating a cycle. For example: + +```python +# module_a.py +from module_b import function_b + +# module_b.py +from module_a import function_a +``` + +While Python can handle some import cycles through its import machinery, they can lead to runtime errors, import deadlocks, or initialization order problems. + +Debugging import cycle errors can be a challenge, especially when they occur in large codebases. However, Codegen allows us to identify these loops through our visualization tools and fix them very deterministically and at scale. + + + + + + + +## Visualize Import Loops in PyTorch + +Using Codegen, we discovered several import cycles in PyTorch's codebase. The code to gather and visualize these loops is as follows: + +```python +G = nx.MultiDiGraph() + +# Add all edges to the graph +for imp in codebase.imports: + if imp.from_file and imp.to_file: + edge_color = "red" if imp.is_dynamic else "black" + edge_label = "dynamic" if imp.is_dynamic else "static" + + # Store the import statement and its metadata + G.add_edge( + imp.to_file.filepath, + imp.from_file.filepath, + color=edge_color, + label=edge_label, + is_dynamic=imp.is_dynamic, + import_statement=imp, # Store the whole import object + key=id(imp.import_statement), + ) +# Find strongly connected components +cycles = [scc for scc in nx.strongly_connected_components(G) if len(scc) > 1] + +print(f" Found {len(cycles)} import cycles:") +for i, cycle in enumerate(cycles, 1): + print(f"\nCycle #{i}:") + print(f"Size: {len(cycle)} files") + + # Create subgraph for this cycle to count edges + cycle_subgraph = G.subgraph(cycle) + + # Count total edges + total_edges = cycle_subgraph.number_of_edges() + print(f"Total number of imports in cycle: {total_edges}") + + # Count dynamic and static imports separately + dynamic_imports = sum(1 for u, v, data in cycle_subgraph.edges(data=True) if data.get("color") == "red") + static_imports = sum(1 for u, v, data in cycle_subgraph.edges(data=True) if data.get("color") == "black") + + print(f"Number of dynamic imports: {dynamic_imports}") + print(f"Number of static imports: {static_imports}") +``` + +Here is one example visualized ⤵️ + + + + + +Not all import cycles are problematic! Some cycles using dynamic imports can work perfectly fine: + + + Valid import loop example + + + +PyTorch prevents most circular import issues through dynamic imports which can be seen through the `import_symbol.is_dynamic` property. If any edge in a strongly connected component is dynamic, runtime conflicts are typically resolved. + +However, we discovered an import loop worth investigating between [`flex_decoding.py`](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/kernel/flex_decoding.py) and [`flex_attention.py`](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/kernel/flex_attention.py): + +Invalid import loop example + +`flex_decoding.py` imports `flex_attention.py` *twice* — once dynamically and once at top-level. This mixed static/dynamic import pattern from the same module creates potential runtime instability. + +*Thus, we propose the following refactoring using Codegen*: + +## Move Shared Code to a Separate `utils.py` File + +```python +# Create new utils file +utils_file = codebase.create_file("torch/_inductor/kernel/flex_utils.py") + +# Get the two files involved in the import cycle +decoding_file = codebase.get_file("torch/_inductor/kernel/flex_decoding.py") +attention_file = codebase.get_file("torch/_inductor/kernel/flex_attention.py") +attention_file_path = "torch/_inductor/kernel/flex_attention.py" +decoding_file_path = "torch/_inductor/kernel/flex_decoding.py" + +# Track symbols to move +symbols_to_move = set() + +# Find imports from flex_attention in flex_decoding +for imp in decoding_file.imports: + if imp.from_file and imp.from_file.filepath == attention_file_path: + # Get the actual symbol from flex_attention + if imp.imported_symbol: + symbols_to_move.add(imp.imported_symbol) + +# Move identified symbols to utils file +for symbol in symbols_to_move: + symbol.move_to_file(utils_file) + +print(f" Moved {len(symbols_to_move)} symbols to flex_utils.py") +for symbol in symbols_to_move: + print(symbol.name) +``` + +Running this codemod will move all the shared symbols to a separate `utils.py` as well as resolve the imports from both files to point to the newly created file solving this potential unpredictable error that could lead issues later on. + + +## Conclusion + +Import loops are a common challenge in large Python codebases. Using Codegen, no matter the repo size, you will gain some new insights into your codebase's import structure and be able to perform deterministic manipulations saving developer hours and future runtime errors. + +Want to try it yourself? Check out our [complete example](https://github.com/codegen-sh/codegen-examples/tree/main/examples/removing_import_loops_in_pytorch) of fixing import loops using Codegen. diff --git a/docs/mint.json b/docs/mint.json index 2a8786f26..ab61d473e 100644 --- a/docs/mint.json +++ b/docs/mint.json @@ -159,7 +159,7 @@ }, { "group": "Blog", - "pages": ["blog/posts", "blog/act-via-code"] + "pages": ["blog/posts", "blog/act-via-code", "blog/fixing-import-loops"] }, { "group": "API Reference", diff --git a/docs/tutorials/fixing-import-loops-in-pytorch.mdx b/docs/tutorials/fixing-import-loops-in-pytorch.mdx index e196f72f1..3eb3ed2d4 100644 --- a/docs/tutorials/fixing-import-loops-in-pytorch.mdx +++ b/docs/tutorials/fixing-import-loops-in-pytorch.mdx @@ -101,12 +101,14 @@ def some_func(): A dynamic import is an import defined inside of a function, method or any executable body of code which delays the import execution until that function, method or body of code is called. -You can use `imp.is_dynamic` to check if the import is dynamic allowing you to investigate imports that are handled more intentionally. +You can use [Import.is_dynamic](/api-reference/core/Import#is-dynamic) to check if the import is dynamic allowing you to investigate imports that are handled more intentionally. # Step 2: Visualize Import Loops - Create a new subgraph to visualize one cycle - color and label the edges based on their type (dynamic/static) -- visualize the cycle graph using `codebase.visualize(graph)` +- visualize the cycle graph using [codebase.visualize(graph)](/api-reference/core/Codebase#visualize) + +Learn more about codebase visualization [here](/building-with-codegen/codebase-visualization) ```python cycle = cycles[0] @@ -217,7 +219,9 @@ problematic_cycles = find_problematic_import_loops(G, cycles) ``` # Step 4: Fix the loop by moving the shared symbols to a separate `utils.py` file -One common fix to this problem to break this cycle is to move all the shared symbols to a separate `utils.py` file. We can do this using the method `symbol.move_to_file`: +One common fix to this problem to break this cycle is to move all the shared symbols to a separate `utils.py` file. We can do this using the method [symbol.move_to_file](/api-reference/core/Symbol#move-to-file): + +Learn more about moving symbols [here](/building-with-codegen/moving-symbols) ```python # Create new utils file @@ -246,15 +250,27 @@ for symbol in symbols_to_move: print(f"🔄 Moved {len(symbols_to_move)} symbols to flex_utils.py") for symbol in symbols_to_move: print(symbol.name) -``` -```python -# run this command to have the changes take effect in the codebase +# Commit changes codebase.commit() ``` -Next Steps -Verify all tests pass after the migration and fix other problematic import loops using the suggested strategies: - 1. Move the shared symbols to a separate file - 2. If a module needs imports only for type hints, consider using `if TYPE_CHECKING` from the `typing` module - 3. Use lazy imports using `importlib` to load imports dynamically \ No newline at end of file +# Conclusions & Next Steps + +Import loops can be tricky to identify and fix, but Codegen provides powerful tools to help manage them: + +- Use `codebase.imports` to analyze import relationships across your project +- Visualize import cycles to better understand dependencies +- Distinguish between static and dynamic imports using `Import.is_dynamic` +- Move shared symbols to break cycles using `symbol.move_to_file` + +Here are some next steps you can take: + +1. **Analyze Your Codebase**: Run similar analysis on your own codebase to identify potential import cycles +2. **Create Import Guidelines**: Establish best practices for your team around when to use static vs dynamic imports +3. **Automate Fixes**: Create scripts to automatically detect and fix problematic import patterns +4. **Monitor Changes**: Set up CI checks to prevent new problematic import cycles from being introduced + + +For more examples of codebase analysis and refactoring, check out our other [tutorials](/tutorials/at-a-glance). + \ No newline at end of file