fix: handle race condition in ensure_symlink for multi-process TP workers - #2981
johnnynunez wants to merge 1 commit into
Conversation
…kers When using tensor parallelism (e.g. --tensor-parallel-size 2), multiple worker processes call `gen_trtllm_gen_fused_moe_sm100_module()` concurrently. Both workers check if the symlink exists (it doesn't), both call `symlink_to()`, and the second one crashes with `FileExistsError`. This is a classic TOCTOU (time-of-check-time-of-use) race. The fix wraps `symlink_to()` in a try/except that catches `FileExistsError` and verifies the existing symlink points to the correct target. If it does, the race is benign and we return successfully. If not, we re-raise. Fixes flashinfer-ai#2980 Made-with: Cursor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughModified Changes
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~3 minutes 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 modifies the ensure_symlink function in flashinfer/jit/cubin_loader.py to handle FileExistsError during symlink creation, aiming to support concurrent execution. However, the review feedback correctly identifies that the function is still not fully thread-safe because the preceding cleanup logic contains TOCTOU race conditions that could lead to FileNotFoundError when multiple processes attempt to remove the same file or directory.
| This function is safe to call concurrently from multiple processes (e.g. | ||
| tensor-parallel workers) that may race to create the same symlink. |
There was a problem hiding this comment.
The claim that this function is safe to call concurrently is not entirely accurate because the cleanup logic (lines 247-254) still contains TOCTOU (Time-of-Check to Time-of-Use) race conditions. If multiple processes attempt to remove a stale file or directory at the same time, one will likely fail with a FileNotFoundError when calling link.unlink() or shutil.rmtree(link). To make this truly safe for concurrent workers, these cleanup operations should be wrapped in a try-except FileNotFoundError block or use missing_ok=True for unlink (available in Python 3.8+).
There was a problem hiding this comment.
Code Review
This pull request updates the ensure_symlink function in flashinfer/jit/cubin_loader.py to support concurrent execution by handling FileExistsError during symlink creation. While the changes improve concurrency, the cleanup logic for stale links still contains a race condition where multiple processes might attempt to delete the same file, potentially leading to FileNotFoundError. Additionally, the verification logic for existing symlinks needs to be adjusted to correctly resolve relative paths against the symlink's parent directory rather than the current working directory.
| This function is safe to call concurrently from multiple processes (e.g. | ||
| tensor-parallel workers) that may race to create the same symlink. |
There was a problem hiding this comment.
The docstring states that this function is safe to call concurrently, but the cleanup logic for stale links (lines 251-254) still contains a race condition. If multiple processes attempt to unlink() or rmtree() the same stale path simultaneously, one will succeed while others may fail with a FileNotFoundError. To make this fully safe for concurrent workers, the cleanup phase should handle the case where the file has already been removed by another process (e.g., by catching FileNotFoundError or using missing_ok=True for unlink).
| link.symlink_to(target) | ||
| except FileExistsError: | ||
| # Another process created the symlink between our check and create. | ||
| if link.is_symlink() and link.resolve() == target.resolve(): |
There was a problem hiding this comment.
Comparing link.resolve() with target.resolve() can be unreliable if target is a relative path, as target.resolve() resolves against the current working directory instead of the symlink's parent directory. Resolving target relative to link.parent provides a more robust comparison that works correctly for both absolute and relative paths.
| if link.is_symlink() and link.resolve() == target.resolve(): | |
| if link.is_symlink() and link.resolve() == (link.parent / target).resolve(): |
|
closing in favor #2979 |
1 participant
Summary
ensure_symlink()that causesFileExistsErrorwhen multiple tensor-parallel workers concurrently create the same symlinkProblem
With
--tensor-parallel-size 2(or higher), multiple worker processes callgen_trtllm_gen_fused_moe_sm100_module()concurrently during startup. Both workers:Falselink.symlink_to(target)FileExistsError: [Errno 17] File existsThis crashes the vLLM server during initialization:
Fix
Wrap
symlink_to()in atry/except FileExistsError. If another process already created the correct symlink between our check and create, silently succeed. If the symlink points to the wrong target, re-raise.Test plan
--tensor-parallel-size 2 --moe-backend=flashinfer_trtllmon a model like Llama-4-Scout FP8Fixes #2980
Summary by CodeRabbit