fix(sdk): filter captured module imports to those referenced by body (follow-up to #478)

[epappas]

Summary

Follow-up to #478 (v0.29.2). The prior fix captured every module-level import; that's too broad — it dragged import basilica (used only by the decorator) into the worker source, and the worker pod's pytorch/cuda image doesn't have basilica-sdk installed.

Runtime trace from basilica-backend#419 Stage 4 take-3 D2 verification:

File "/tmp/__basilica_source.py", line 3, in <module>
    import basilica
ModuleNotFoundError: No module named 'basilica'
E0517 19:58:18.832000 1 torch/distributed/elastic/multiprocessing/api.py:862] failed (exitcode: 1) local_rank: 0

Fix

Walk the function body's AST. Collect every free Name reference and the leftmost identifier of every Attribute chain (so os.environ.get(...) contributes os). Then filter the module's top-level imports down to those whose bound names appear in that set.

Binding rules:

• import foo binds foo
• import foo.bar binds foo (Python module-resolution semantics)
• import foo as bar binds bar
• from x import foo binds foo
• from x import foo as bar binds bar
• from x import * is kept verbatim (cannot statically enumerate exposed names)

Applied at all three call sites:

• DistributedFunction._extract_source (decorators.py)
• DeployedFunction._extract_source (decorators.py)
• SourcePackager.from_function (source.py)

Verification

Local repro against the exact ex20-shape decorator:

Pre-fix (v0.29.2) produced source:

import os
import sys
import time
import basilica           # <-- worker has no `basilica` package; fails at runtime
from basilica import ProviderFilter, WorldSize
def train() -> None:
    ...

Post-fix (v0.29.3) produced source:

import os
import time
def train() -> None:
    ...

Only the imports the body actually uses make it into the shipped source.

Tests

Adds test_extracted_source_filters_unused_module_imports to pin the contract: a body that only references os from module scope must NOT carry import basilica, import pytest, or from typing import Optional in its head (those are imported at module level in the test file but never referenced by the body).

Full SDK test suite: 139 passed, 0 failed.

Cross-repo

• Closes bug(sdk): @basilica.distributed decorator doesn't capture module-level imports — wrapped function fails with NameError #477 follow-up (the original v0.29.2 closure was incomplete in the runtime path)
• Refs: one-covenant/basilica-backend#419 Stage 4 take-3 Cell B + D2 runtime verification

Test plan

• New regression test passes
• Full SDK test suite green (139/139)
• ex20-shape extraction produces a head with import os only (verified local)
• Live cluster runtime verification (ex20 against api.basilica.ai). Will be exercised after CI green and merge.

[coderabbitai]

Warning

Rate limit exceeded

@epappas has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 31 minutes and 32 seconds before requesting another review.

You’ve run out of usage credits. Purchase more in the billing tab.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 3cc3b006-e93e-4725-a5ab-f9741854a798

📥 Commits

Reviewing files that changed from the base of the PR and between b9ece24 and 5cc81a7.

📒 Files selected for processing (6)

• crates/basilica-sdk-python/CHANGELOG.md
• crates/basilica-sdk-python/Cargo.toml
• crates/basilica-sdk-python/pyproject.toml
• crates/basilica-sdk-python/python/basilica/decorators.py
• crates/basilica-sdk-python/python/basilica/source.py
• crates/basilica-sdk-python/tests/test_decorator_source_extraction.py

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/477-decorator-imports-usage-filter

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}