fix(stubs): resolve in-function @remote imports for cross-endpoint dispatch#224
Merged
fix(stubs): resolve in-function @remote imports for cross-endpoint dispatch#224
Conversation
…spatch The dependency resolver only checked func.__globals__ (module-level namespace), missing @Remote functions imported inside function bodies. Pipeline functions like classify() that use `from gpu_worker import gpu_inference` inside the body would fail with ModuleNotFoundError on the remote worker. - Add resolve_in_function_imports() to discover @Remote deps via in-body imports - Add strip_remote_imports() to remove stale import statements from sent source - Update both LiveServerlessStub and LoadBalancerSlsStub to use new functions - Add 16 tests covering discovery, stripping, aliasing, and exec integration
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes a critical bug in the dependency resolver where @remote functions imported inside function bodies (rather than at module level) were not detected, causing ModuleNotFoundError on remote workers. The fix adds two new functions to discover and handle in-function imports:
Changes:
- Added
resolve_in_function_imports()to discover@remotedependencies via in-bodyfrom X import Ystatements - Added
strip_remote_imports()to remove stale import statements from source sent to workers - Updated both
LiveServerlessStubandLoadBalancerSlsStubto use the new resolution pipeline
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| src/runpod_flash/stubs/dependency_resolver.py | Adds resolve_in_function_imports() to detect in-function imports and strip_remote_imports() to remove them from source |
| src/runpod_flash/stubs/live_serverless.py | Integrates new resolution pipeline: augment globals → resolve deps → strip imports → build stubs |
| src/runpod_flash/stubs/load_balancer_sls.py | Integrates new resolution pipeline (same pattern as live_serverless.py) |
| tests/unit/test_dependency_resolver.py | Adds 16 comprehensive tests covering discovery, stripping, aliasing, and exec integration |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
… resolver Address PR review feedback: - Skip relative imports (module=None) in strip_remote_imports to avoid TypeError - Preserve first discovered import when same name imported from multiple modules - Add tests for both cases
KAJdev
approved these changes
Feb 26, 2026
|
📝 Documentation updates detected! New suggestion: Add code example for cross-endpoint dispatch Tip: Set up a GitHub Issues trigger in Projects to enable @Promptless mentions in issues 🐙 |
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
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
@remotefunctions imported inside function bodies (not just module-level)resolve_in_function_imports()to discover@remotedeps via in-bodyfrom X import Ystatementsstrip_remote_imports()to remove stale import statements from source sent to workersLiveServerlessStubandLoadBalancerSlsStubto use the new resolution pipelineRoot cause: Pipeline functions like
classify()usefrom gpu_worker import gpu_inferenceinside the function body. The dependency resolver only checkedfunc.__globals__(module-level namespace), missing these imports entirely. No dispatch stubs were generated, causingModuleNotFoundErroron the remote worker.Test plan
03_mixed_workersviaflash deploy --use-local-flash