fix: let FLASHINFER_CUBIN_DIR env var override flashinfer-cubin package path#3062
fix: let FLASHINFER_CUBIN_DIR env var override flashinfer-cubin package path#3062ianliuy wants to merge 1 commit intoflashinfer-ai:mainfrom
Conversation
…ge path _get_cubin_dir() checked the flashinfer-cubin package before the FLASHINFER_CUBIN_DIR environment variable, making the env var silently ignored when the package was installed. This broke all non-root Kubernetes/OpenShift deployments where users need to redirect cubin storage to a writable directory. Swap the priority order so the env var is checked first, matching the standard Unix convention that explicit user overrides take precedence over automatic discovery. Default behavior (no env var set) is unchanged. Add regression tests for the priority logic in _get_cubin_dir(). Fixes flashinfer-ai#2976 Related: flashinfer-ai#2834
ianliuy
requested review from
aleozlx,
bkryu,
cyx-6,
jimmyzho,
kahyunnam,
nv-yunzheq,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
📝 WalkthroughWalkthroughThe pull request reorders the priority logic in Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the cubin directory resolution logic in flashinfer/jit/env.py to prioritize the FLASHINFER_CUBIN_DIR environment variable over the installed package. It also introduces a new test file, tests/test_env.py, to verify this priority and fallback behavior. Feedback was provided to improve test isolation by ensuring the module under test is correctly saved and restored in sys.modules within the test utility function.
|
|
||
| stubs["flashinfer.compilation_context"].CompilationContext = _Stub | ||
|
|
||
| saved = {k: sys.modules.get(k) for k in stubs} |
There was a problem hiding this comment.
The current implementation of _load_env_module does not save and restore the state of sys.modules["flashinfer.jit.env"]. This can lead to test isolation issues if other tests in the same process attempt to import this module, as they will receive the mocked version instead of the real one. It is safer to include the module being tested in the saved dictionary so it can be properly cleaned up in the finally block.
| saved = {k: sys.modules.get(k) for k in stubs} | |
| modules_to_save = list(stubs.keys()) + ["flashinfer.jit.env"] | |
| saved = {k: sys.modules.get(k) for k in modules_to_save} |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/test_env.py`:
- Around line 38-56: The helper _load_env_module mutates sys.modules by
inserting "flashinfer.jit.env" but never restores any prior entry, causing
cross-test contamination when _env is created at import time; update
_load_env_module to save the original value of "flashinfer.jit.env" before
assigning mod (e.g. original_env = sys.modules.get("flashinfer.jit.env") or
include that key in the saved dict) and in the finally block restore it exactly
as you do for other stubs (if original is None pop the key, else set it back),
ensuring "flashinfer.jit.env" is returned to its pre-call state and preventing
stale module reuse by _env.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 36af5a7d-350b-4068-98b8-988014915f26
📒 Files selected for processing (2)
flashinfer/jit/env.pytests/test_env.py
| saved = {k: sys.modules.get(k) for k in stubs} | ||
| sys.modules.update(stubs) | ||
| try: | ||
| spec = importlib.util.spec_from_file_location( | ||
| "flashinfer.jit.env", str(_REPO_ROOT / "flashinfer" / "jit" / "env.py") | ||
| ) | ||
| mod = importlib.util.module_from_spec(spec) | ||
| sys.modules["flashinfer.jit.env"] = mod | ||
| spec.loader.exec_module(mod) | ||
| return mod | ||
| finally: | ||
| for k, v in saved.items(): | ||
| if v is None: | ||
| sys.modules.pop(k, None) | ||
| else: | ||
| sys.modules[k] = v | ||
|
|
||
|
|
||
| _env = _load_env_module() |
There was a problem hiding this comment.
Prevent cross-test contamination from sys.modules mutation.
Line 45 installs flashinfer.jit.env into sys.modules, but _load_env_module() never restores any preexisting entry for that key. Because _env is created at import time (Line 56), later tests can accidentally reuse this stub-loaded module.
Proposed fix
def _load_env_module():
@@
- saved = {k: sys.modules.get(k) for k in stubs}
+ saved = {k: sys.modules.get(k) for k in stubs}
+ saved_env_module = sys.modules.get("flashinfer.jit.env")
sys.modules.update(stubs)
try:
spec = importlib.util.spec_from_file_location(
"flashinfer.jit.env", str(_REPO_ROOT / "flashinfer" / "jit" / "env.py")
)
mod = importlib.util.module_from_spec(spec)
sys.modules["flashinfer.jit.env"] = mod
spec.loader.exec_module(mod)
return mod
finally:
+ if saved_env_module is None:
+ sys.modules.pop("flashinfer.jit.env", None)
+ else:
+ sys.modules["flashinfer.jit.env"] = saved_env_module
for k, v in saved.items():
if v is None:
sys.modules.pop(k, None)
else:
sys.modules[k] = v🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_env.py` around lines 38 - 56, The helper _load_env_module mutates
sys.modules by inserting "flashinfer.jit.env" but never restores any prior
entry, causing cross-test contamination when _env is created at import time;
update _load_env_module to save the original value of "flashinfer.jit.env"
before assigning mod (e.g. original_env = sys.modules.get("flashinfer.jit.env")
or include that key in the saved dict) and in the finally block restore it
exactly as you do for other stubs (if original is None pop the key, else set it
back), ensuring "flashinfer.jit.env" is returned to its pre-call state and
preventing stale module reuse by _env.
1 participant
What's broken?
When
flashinfer-cubinis installed (all pre-built container images), settingFLASHINFER_CUBIN_DIRto redirect cubin storage to a writable directory is silently ignored. Non-root containers crash withPermissionErroron startup.Who is affected?
All non-root K8s/OpenShift deployments using pre-built images with
flashinfer-cubin(e.g., vLLM containers). No workaround exists.Why does it happen?
_get_cubin_dir()checks theflashinfer-cubinpackage before the env var. When the package is installed, the env var is never reached. Introduced in PR #1718 (intentional priority, but didn't consider non-root containers).Every other env var in
env.pyfollows "env var overrides automatic discovery." This was the sole exception.How did we fix it?
Swap priority: env var -> package -> default cache. Pure code-block reorder (~8 lines moved + docstring/comments updated). Zero new logic.
Before:
After:
Default behavior (no env var set) is unchanged.
How do we know it works?
Added
tests/test_env.pywith 4 regression tests covering the full priority matrix:All 4 pass. Pre-commit hooks (ruff, mypy, format) all pass.
Fixes #2976
Related: #2834
@aleozlx @yzh119
Summary by CodeRabbit
Bug Fixes
Tests