fix(sdk): @basilica.distributed decorator captures module-level imports (fixes #477)#478
Merged
Merged
Conversation
…evel imports Closes #477. The decorators previously shipped only the function body via inspect.getsource(fn); module-level imports were not in scope on the worker pod, so any reference to a module-level name (e.g. the `import os` in examples/20_distributed_diloco.py) raised `NameError: name 'os' is not defined` at runtime. Fix walks the defining module's AST and prepends every top-level `Import` / `ImportFrom` node to the extracted function source. Only import statements are captured; other top-level state is not leaked into the shipped source. Applied to: - DistributedFunction._extract_source (decorators.py) - DeployedFunction._extract_source (decorators.py) - SourcePackager.from_function (source.py) Adds tests/test_decorator_source_extraction.py with 6 cases pinning the import-capture contract: simple imports, from-imports, no-imports no-op, and a NameError exec-check that mirrors the worker-pod failure mode from one-covenant/basilica-backend#419 Stage 4 take-3 Cell B. Bumps version to 0.29.2.
Contributor
WalkthroughThis PR fixes a bug where ChangesModule-level import capture for distributed and packaging paths
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related issues
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
crates/basilica-sdk-python/python/basilica/decorators.py (1)
22-65: ⚡ Quick winConsider extracting this helper to a shared module.
_extract_module_level_importsis duplicated nearly verbatim insource.py(lines 36-70). Consider moving it to a shared internal utility module (e.g.,_utils.pyor_source_utils.py) to avoid divergence.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/basilica-sdk-python/python/basilica/decorators.py` around lines 22 - 65, The function _extract_module_level_imports is duplicated between decorators.py and source.py; extract it into a shared internal module (e.g., create _source_utils.py or _utils.py) and export a single function (keep the name _extract_module_level_imports for minimal churn), then import and use that shared function from both decorators.py and source.py (replace the local implementation in each file with "from ._source_utils import _extract_module_level_imports"); ensure tests/imports still resolve and preserve the current behavior and docstring.crates/basilica-sdk-python/tests/test_decorator_source_extraction.py (2)
81-93: 💤 Low valueConsider a more robust approach for extracting the import header.
The pattern
source.split("def ")[0]is repeated across multiple test methods and could be fragile if the string "def " appears in comments, docstrings, or import statements in the extracted source. While this works for the current implementation, a more robust approach might use a helper method or regex to locate the first function definition.♻️ Example of a more robust approach
def _extract_import_header(source: str) -> str: """Extract everything before the first function definition.""" import re match = re.search(r'^def\s+\w+', source, re.MULTILINE) return source[:match.start()] if match else sourceThen use:
head = _extract_import_header(source)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/basilica-sdk-python/tests/test_decorator_source_extraction.py` around lines 81 - 93, The tests repeatedly use fragile head = source.split("def ")[0] to capture the module import header; replace this with a small helper (e.g., _extract_import_header) used by test functions that takes the extracted source from _train_with_aliased_import._extract_source() and returns everything before the first actual function definition by searching for the first multiline match of '^def\\s+\\w+' (or returning whole source if no match). Update each test in test_decorator_source_extraction.py to call this helper instead of using split to make import-header extraction robust to "def " appearing in comments/docstrings.
147-168: 💤 Low valueConsider simplifying the line trimming logic.
The loop at lines 157-163 that removes the trailing function call is complex and could be more readable. The current approach iteratively pops lines with multiple conditions, which makes the intent less clear.
♻️ Simpler alternative
# Find the last line of the function definition (excluding trailing call) lines = source.splitlines() # Remove the decorator-appended call (last non-empty line) func_name = _serve_with_module_imports.__name__ while lines and (lines[-1].strip() == "" or lines[-1].strip() == f"{func_name}()"): lines.pop() trimmed = "\n".join(lines) + "\n"Or even simpler using rsplit:
# Split on the trailing call and take everything before it func_name = _serve_with_module_imports.__name__ trimmed = source.rsplit(f"\n{func_name}()", 1)[0] + "\n"Note: The
exec()call on line 167 is flagged by static analysis (S102) but is a false positive—this is legitimate test code validating source execution.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/basilica-sdk-python/tests/test_decorator_source_extraction.py` around lines 147 - 168, The trimming loop that removes the decorator-appended call is more complex than needed; simplify it by removing trailing blank lines and any trailing line equal to f"{_serve_with_module_imports.__name__}()" in one clear step before joining into trimmed. Locate the test function test_extracted_source_executes_without_name_error and replace the while/pop logic that manipulates lines with a concise loop or rsplit approach that strips empty lines and the exact trailing call, keeping the rest of the source intact so compile(trimmed, "<test-source>", "exec") and exec(compiled, ns) behave the same.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/basilica-sdk-python/CHANGELOG.md`:
- Around line 10-23: Update the changelog footer links to include a link target
for [0.29.2] and correct the comparison target for [Unreleased]; replace the
current reference to v0.14.0 with the appropriate v0.29.2 comparison (and add a
new `[0.29.2]` link target) so the `[Unreleased]` vs release links and the
`[0.29.2]` entry resolve correctly instead of pointing to v0.14.0.
---
Nitpick comments:
In `@crates/basilica-sdk-python/python/basilica/decorators.py`:
- Around line 22-65: The function _extract_module_level_imports is duplicated
between decorators.py and source.py; extract it into a shared internal module
(e.g., create _source_utils.py or _utils.py) and export a single function (keep
the name _extract_module_level_imports for minimal churn), then import and use
that shared function from both decorators.py and source.py (replace the local
implementation in each file with "from ._source_utils import
_extract_module_level_imports"); ensure tests/imports still resolve and preserve
the current behavior and docstring.
In `@crates/basilica-sdk-python/tests/test_decorator_source_extraction.py`:
- Around line 81-93: The tests repeatedly use fragile head = source.split("def
")[0] to capture the module import header; replace this with a small helper
(e.g., _extract_import_header) used by test functions that takes the extracted
source from _train_with_aliased_import._extract_source() and returns everything
before the first actual function definition by searching for the first multiline
match of '^def\\s+\\w+' (or returning whole source if no match). Update each
test in test_decorator_source_extraction.py to call this helper instead of using
split to make import-header extraction robust to "def " appearing in
comments/docstrings.
- Around line 147-168: The trimming loop that removes the decorator-appended
call is more complex than needed; simplify it by removing trailing blank lines
and any trailing line equal to f"{_serve_with_module_imports.__name__}()" in one
clear step before joining into trimmed. Locate the test function
test_extracted_source_executes_without_name_error and replace the while/pop
logic that manipulates lines with a concise loop or rsplit approach that strips
empty lines and the exact trailing call, keeping the rest of the source intact
so compile(trimmed, "<test-source>", "exec") and exec(compiled, ns) behave the
same.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 5a1f5540-212f-458b-8ee5-a4fac96c0bd3
📒 Files selected for processing (6)
crates/basilica-sdk-python/CHANGELOG.mdcrates/basilica-sdk-python/Cargo.tomlcrates/basilica-sdk-python/pyproject.tomlcrates/basilica-sdk-python/python/basilica/decorators.pycrates/basilica-sdk-python/python/basilica/source.pycrates/basilica-sdk-python/tests/test_decorator_source_extraction.py
Comment on lines
+10
to
+23
| ## [0.29.2] - 2026-05-16 | ||
|
|
||
| ### Fixed | ||
| - `@basilica.distributed` and `@basilica.deployment` now capture the | ||
| defining module's top-level `import` and `from ... import ...` | ||
| statements and prepend them to the source shipped to the worker pod. | ||
| Before this fix, only the function body was shipped; module-level | ||
| names referenced inside the body (e.g. the `import os` in | ||
| `examples/20_distributed_diloco.py`) raised `NameError` at worker | ||
| runtime. Closes #477. Cross-repo reference: | ||
| `one-covenant/basilica-backend#419` Stage 4 take-3 Cell B. The same | ||
| capture is applied in `SourcePackager.from_function()` for the | ||
| lower-level packaging path. | ||
|
|
Contributor
There was a problem hiding this comment.
Update changelog reference links for the new release.
After adding 0.29.2, the footer links were not updated: Line 320 still compares [Unreleased] from v0.14.0, and there is no [0.29.2] link target. Please add/update those references so changelog navigation is correct.
📝 Suggested patch
-[Unreleased]: https://github.com/one-covenant/basilica/compare/basilica-sdk-python-v0.14.0...HEAD
+[Unreleased]: https://github.com/one-covenant/basilica/compare/basilica-sdk-python-v0.29.2...HEAD
+[0.29.2]: https://github.com/one-covenant/basilica/compare/basilica-sdk-python-v0.29.1...basilica-sdk-python-v0.29.2
[0.14.0]: https://github.com/one-covenant/basilica/compare/basilica-sdk-python-v0.13.0...basilica-sdk-python-v0.14.0🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/basilica-sdk-python/CHANGELOG.md` around lines 10 - 23, Update the
changelog footer links to include a link target for [0.29.2] and correct the
comparison target for [Unreleased]; replace the current reference to v0.14.0
with the appropriate v0.29.2 comparison (and add a new `[0.29.2]` link target)
so the `[Unreleased]` vs release links and the `[0.29.2]` entry resolve
correctly instead of pointing to v0.14.0.
Merged
4 tasks
epappas
added a commit
that referenced
this pull request
May 17, 2026
…sage-filter fix(sdk): filter captured module imports to those referenced by body (follow-up to #478)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes #477. Fixes the canonical
NameError: name 'os' is not definedfailure mode inexamples/20_distributed_diloco.pydocumented atone-covenant/basilica-backend#419Stage 4 take-3 Cell B.Bug
The decorator-based deploy path ships only the wrapped function body via
inspect.getsource(fn):DistributedFunction._extract_sourceincrates/basilica-sdk-python/python/basilica/decorators.py:313-337(pre-fix)DeployedFunction._extract_sourceincrates/basilica-sdk-python/python/basilica/decorators.py:92-115(pre-fix)SourcePackager.from_functionincrates/basilica-sdk-python/python/basilica/source.py:336-389(pre-fix)Module-level
importstatements in the user's source file are not in scope inside the body when the body is exec'd on the worker pod. Any reference to a module-level name raisesNameError. The failure was visible end-to-end on every worker rank of ex20:Fix
New helper
_extract_module_level_imports(func)walks the AST ofsys.modules[func.__module__]and emits every top-levelast.Import/ast.ImportFromnode viaast.unparse. Only import statements are captured; other top-level state (function defs, dataclasses, constants) is NOT leaked into the shipped source.Applied at all three extraction sites. Each produces source whose first lines are now:
Tests
New
tests/test_decorator_source_extraction.pypins the import-capture contract with 6 cases:@basilica.distributedcapturesimport os/import timefrom module scope@basilica.distributedcapturesfrom typing import Optional(from-imports)@basilica.distributedwith a body that uses no module-level names still produces a valid call@basilica.distributedbody can beexec'd in a clean namespace with__name__ == "__main__"without raisingNameError(the canonical pre-fix failure mode)@basilica.deploymentcapturesimport osfrom module scope (same code path)@basilica.deploymentbody can beexec'd in a clean namespacePre-fix:
NameError: name 'os' is not defined— exactly the worker-pod failure mode fromone-covenant/basilica-backend#419Cell BPost-fix:
Test plan
NameError)import basilicastill works post-fix/tmp/__basilica_source.pythat includesimport os, and the NameError is gone). Will be exercised againstapi.basilica.aiafter CI green and merge.Cross-repo trace
one-covenant/basilica-backend#419Stage 4 take-3 Cell B (data/evidence/419-stage-4-take-3/B-example20-diloco/)Summary by CodeRabbit