Skip unsupported JAX devices during module preload#1634
Conversation
Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe change centralizes JAX device module preloading, skips devices that cannot map to Warp, updates both FFI call paths, and adds coverage for mixed supported and unsupported devices. ChangesJAX FFI preload handling
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Greptile SummaryThis PR fixes a bug in the JAX FFI preload loop where catching an exception and using
Confidence Score: 5/5Safe to merge — the change is a narrow, well-tested bug fix with no behavioural regressions for the supported path. The only logic change is replacing except Exception: pass + fallthrough with except RuntimeError: continue, which is strictly safer: it eliminates the unbound/stale dev reference and aligns the caught type with the documented contract of device_from_jax. The helper is exercised by a new regression test that covers both device orderings. No other call sites are affected. No files require special attention. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[_preload_module_on_all_devices] --> B[iterate jax.local_devices]
B --> C{device_from_jax}
C -->|RuntimeError| D[continue - skip device]
D --> B
C -->|returns dev| E{dev.is_cuda?}
E -->|No| B
E -->|Yes| F[module.load dev]
F --> B
B -->|done| G[return]
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
A[_preload_module_on_all_devices] --> B[iterate jax.local_devices]
B --> C{device_from_jax}
C -->|RuntimeError| D[continue - skip device]
D --> B
C -->|returns dev| E{dev.is_cuda?}
E -->|No| B
E -->|Yes| F[module.load dev]
F --> B
B -->|done| G[return]
Reviews (2): Last reviewed commit: "Capitalize JAX in preload test message" | Re-trigger Greptile |
![coderabbitai[bot]](https://avatars.github.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
warp/_src/jax/ffi.py (1)
69-79: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a Google-style docstring for the helper.
This helper encodes an important exception contract—unsupported JAX platforms raise
RuntimeErrorand are skipped—but currently documents neither its arguments nor that behavior.As per coding guidelines, Python files must follow Google-style docstrings.
Proposed documentation
def _preload_module_on_all_devices(module, jax): + """Preload a module on local JAX devices that map to CUDA. + + Args: + module: Warp module to preload. + jax: JAX module-like object providing ``local_devices()``. + + A ``RuntimeError`` from JAX-to-Warp device conversion is treated as an + unsupported platform and skipped. + """ for d in jax.local_devices():🤖 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 `@warp/_src/jax/ffi.py` around lines 69 - 79, Add a Google-style docstring to `_preload_module_on_all_devices` documenting the `module` and `jax` arguments, the helper’s device-preloading behavior, and that `RuntimeError` from unsupported JAX platforms is caught and skipped.Source: Coding guidelines
🤖 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 `@warp/tests/interop/test_jax.py`:
- Line 865: Update the skip message in the `_jax_version` test decorator to use
the correctly capitalized name, changing “Jax version too old” to “JAX version
too old”.
---
Nitpick comments:
In `@warp/_src/jax/ffi.py`:
- Around line 69-79: Add a Google-style docstring to
`_preload_module_on_all_devices` documenting the `module` and `jax` arguments,
the helper’s device-preloading behavior, and that `RuntimeError` from
unsupported JAX platforms is caught and skipped.
🪄 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: Path: .coderabbit.yml
Review profile: CHILL
Plan: Enterprise
Run ID: 64d1395a-2a94-4107-ad07-7a4c12f2fb15
📒 Files selected for processing (2)
warp/_src/jax/ffi.pywarp/tests/interop/test_jax.py
Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>
1 participant
Summary
ALL_DEVICESmodule preloading forFfiKernelandFfiCallable.device_from_jax()raises its documentedRuntimeError.Root cause
The preload loops caught conversion failures but continued into
dev.is_cuda. That could read an unbound local on the first unsupported device or reuse the previous device on later failures.Validation
git diff --checkpass.Summary by CodeRabbit