Skip to content

Blog for Import loops #207

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Feb 5, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 159 additions & 0 deletions docs/blog/fixing-import-loops.mdx
Original file line number Diff line number Diff line change
@@ -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.

<Info>
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).
</Info>

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.

<Frame caption="Import loop in pytorch/torchgen/model.py">
<iframe
width="100%"
height="500px"
scrolling="no"
src={`https://www.codegen.sh/embedded/graph/?id=8b575318-ff94-41f1-94df-6e21d9de45d1&zoom=1&targetNodeName=model`}
className="rounded-xl"
style={{
backgroundColor: "#15141b",
}}
></iframe>
</Frame>



## 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 ⤵️

<Frame caption="Import loops in pytorch/torchgen/model.py">
<iframe
width="100%"
height="500px"
scrolling="no"
src={`https://www.codegen.sh/embedded/graph/?id=8b575318-ff94-41f1-94df-6e21d9de45d1&zoom=1&targetNodeName=model`}
className="rounded-xl"
style={{
backgroundColor: "#15141b",
}}
></iframe>
</Frame>

Not all import cycles are problematic! Some cycles using dynamic imports can work perfectly fine:

<Frame>
<img src="/images/valid-import-loop.png" alt="Valid import loop example" />
</Frame>


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):

<img src="/images/problematic-import-loop.png" alt="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.
2 changes: 1 addition & 1 deletion docs/mint.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
38 changes: 27 additions & 11 deletions docs/tutorials/fixing-import-loops-in-pytorch.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -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)

<Tip>Learn more about codebase visualization [here](/building-with-codegen/codebase-visualization)</Tip>

```python
cycle = cycles[0]
Expand Down Expand Up @@ -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):

<Tip>Learn more about moving symbols [here](/building-with-codegen/moving-symbols)</Tip>

```python
# Create new utils file
Expand Down Expand Up @@ -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
# 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

<Info>
For more examples of codebase analysis and refactoring, check out our other [tutorials](/tutorials/at-a-glance).
</Info>
Loading