Add filelock to ensure_symlink - #2979
Conversation
📝 WalkthroughWalkthroughAdds per-link file locking to serialize symlink creation in Changes
Sequence Diagram(s)sequenceDiagram
participant ProcA as Process A
participant ProcB as Process B
participant Lock as FileLock (`<link>.lock`)
participant FS as Filesystem
ProcA->>Lock: acquire (timeout 60s)
alt granted
Lock-->>ProcA: granted
ProcA->>FS: check symlink exists & target
FS-->>ProcA: result
alt stale or wrong
ProcA->>FS: remove stale file/dir
FS-->>ProcA: removed
end
ProcA->>FS: create parent dirs
FS-->>ProcA: dirs created
ProcA->>FS: create symlink -> target
FS-->>ProcA: symlink created
ProcA->>Lock: release
else timeout
Lock-->>ProcA: timeout error
end
ProcB->>Lock: acquire (queued/blocked until release)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related issues
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 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 introduces a file lock mechanism in the ensure_symlink function to prevent race conditions during symlink creation. A review comment points out that the FileLock constructor may fail if the parent directory does not yet exist, suggesting that the directory creation should be moved before the lock initialization.
|
|
||
| lock_path = str(link) + ".lock" | ||
| lock = filelock.FileLock(lock_path, timeout=60) | ||
| with lock: | ||
| if link.is_symlink() or link.exists(): | ||
| if link.is_symlink() and link.resolve() == target.resolve(): | ||
| return # already correct | ||
| # Stale symlink or directory from a previous version; remove it. | ||
| if link.is_symlink() or link.is_file(): | ||
| link.unlink() | ||
| else: | ||
| shutil.rmtree(link) | ||
| link.parent.mkdir(parents=True, exist_ok=True) | ||
| link.symlink_to(target) |
There was a problem hiding this comment.
The filelock.FileLock constructor will fail with a FileNotFoundError if the parent directory of the lock_path does not exist. Since ensure_symlink is responsible for creating the directory structure, the mkdir call should be moved before the lock creation to ensure the lock file can be created successfully.
Additionally, once the parent directory is created and the lock is acquired, the redundant mkdir call inside the lock can be removed, as the parent directory is guaranteed to exist (and the open lock file prevents the parent directory from being removed by standard directory removal operations).
| lock_path = str(link) + ".lock" | |
| lock = filelock.FileLock(lock_path, timeout=60) | |
| with lock: | |
| if link.is_symlink() or link.exists(): | |
| if link.is_symlink() and link.resolve() == target.resolve(): | |
| return # already correct | |
| # Stale symlink or directory from a previous version; remove it. | |
| if link.is_symlink() or link.is_file(): | |
| link.unlink() | |
| else: | |
| shutil.rmtree(link) | |
| link.parent.mkdir(parents=True, exist_ok=True) | |
| link.symlink_to(target) | |
| link.parent.mkdir(parents=True, exist_ok=True) | |
| lock_path = str(link) + ".lock" | |
| lock = filelock.FileLock(lock_path, timeout=60) | |
| with lock: | |
| if link.is_symlink() or link.exists(): | |
| if link.is_symlink() and link.resolve() == target.resolve(): | |
| return # already correct | |
| # Stale symlink or directory from a previous version; remove it. | |
| if link.is_symlink() or link.is_file(): | |
| link.unlink() | |
| else: | |
| shutil.rmtree(link) | |
| link.symlink_to(target) |
There was a problem hiding this comment.
fixed. can you check again?
![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 `@flashinfer/jit/cubin_loader.py`:
- Around line 246-257: The ensure_symlink function currently creates a
filelock.FileLock(lock_path, timeout=60) and enters with lock but does not
handle filelock.Timeout; wrap the with-lock block in a try/except catching
filelock.Timeout (and optionally filelock.AlreadyLocked) around the context
manager usage for the lock created at lock_path, and inside the except check
whether the symlink (link) already exists and resolves to the desired target
(target.resolve() == link.resolve()); if so, return normally, otherwise raise a
new, clearer RuntimeError (or re-raise a wrapped exception) indicating lock
acquisition timed out while attempting to create the symlink for the JIT cubin,
so callers like moe_utils and fused_moe get a descriptive error instead of an
uncaught filelock.Timeout.
🪄 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: 3d6871ae-f473-461c-817e-c041e58ba34f
📒 Files selected for processing (1)
flashinfer/jit/cubin_loader.py
| lock = filelock.FileLock(lock_path, timeout=60) | ||
| with lock: | ||
| if link.is_symlink() or link.exists(): | ||
| if link.is_symlink() and link.resolve() == target.resolve(): | ||
| return # already correct | ||
| # Stale symlink or directory from a previous version; remove it. | ||
| if link.is_symlink() or link.is_file(): | ||
| link.unlink() | ||
| else: | ||
| shutil.rmtree(link) | ||
| link.parent.mkdir(parents=True, exist_ok=True) | ||
| link.symlink_to(target) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify direct ensure_symlink call sites and lack of local timeout handling context.
set -euo pipefail
rg -n -C3 '\bensure_symlink\s*\(' --type=py
rg -n -C3 'except\s+filelock\.Timeout|except\s+Timeout' --type=pyRepository: flashinfer-ai/flashinfer
Length of output: 2036
🏁 Script executed:
# Get the full ensure_symlink function and surrounding context
sed -n '233,260p' flashinfer/jit/cubin_loader.pyRepository: flashinfer-ai/flashinfer
Length of output: 1127
🏁 Script executed:
# Get context around line 134 to understand which function it's in
sed -n '120,150p' flashinfer/jit/cubin_loader.pyRepository: flashinfer-ai/flashinfer
Length of output: 1269
🏁 Script executed:
# Check the full context of ensure_symlink calls in moe_utils.py
sed -n '55,70p' flashinfer/jit/moe_utils.pyRepository: flashinfer-ai/flashinfer
Length of output: 609
🏁 Script executed:
# Check the full context of ensure_symlink calls in fused_moe.py
sed -n '270,285p' flashinfer/jit/fused_moe.pyRepository: flashinfer-ai/flashinfer
Length of output: 565
Handle filelock.Timeout inside ensure_symlink to prevent uncaught exceptions during JIT module generation under contention.
The function creates a filelock.FileLock with 60-second timeout (line 247) but lacks exception handling. If lock acquisition times out, the exception propagates uncaught to callers like moe_utils.py:63 and fused_moe.py:277, which do not catch it. Under concurrent access, module generation may fail abruptly.
Wrap the lock block with try-except, check if the symlink was created correctly by another process, and raise a clearer error if the timeout is genuine.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@flashinfer/jit/cubin_loader.py` around lines 246 - 257, The ensure_symlink
function currently creates a filelock.FileLock(lock_path, timeout=60) and enters
with lock but does not handle filelock.Timeout; wrap the with-lock block in a
try/except catching filelock.Timeout (and optionally filelock.AlreadyLocked)
around the context manager usage for the lock created at lock_path, and inside
the except check whether the symlink (link) already exists and resolves to the
desired target (target.resolve() == link.resolve()); if so, return normally,
otherwise raise a new, clearer RuntimeError (or re-raise a wrapped exception)
indicating lock acquisition timed out while attempting to create the symlink for
the JIT cubin, so callers like moe_utils and fused_moe get a descriptive error
instead of an uncaught filelock.Timeout.
|
/bot run |
|
@johnnynunez seems that you are running into this too. |
yes, vllm team... trying to upgrade to have the fix #2903 |
|
@johnnynunez Can you check if this fix resolves the issue? |
|
Added |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/utils/test_gen_module_symlink_race_condition.py (2)
82-84: Usepytest.skip()for proper test skipping.Using
print()andreturncauses pytest to report this test as PASSED on unsupported hardware rather than SKIPPED, masking test coverage gaps.♻️ Proposed fix
Add
pytestto imports at line 17:import os +import pytest import tempfileThen update the skip logic:
if not (is_sm100a_supported(device) or is_sm12x_supported(device)): - print("Skipping: gen_trtllm_gen_fused_moe_sm100_module requires SM100 or SM12x") - return + pytest.skip("gen_trtllm_gen_fused_moe_sm100_module requires SM100 or SM12x")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/utils/test_gen_module_symlink_race_condition.py` around lines 82 - 84, Replace the print+return skip pattern with pytest.skip so pytest records the test as SKIPPED: add pytest to the module imports and change the conditional in the test (the block using is_sm100a_supported/is_sm12x_supported) to call pytest.skip("gen_trtllm_gen_fused_moe_sm100_module requires SM100 or SM12x") when the condition is met instead of printing and returning; update the test function that contains the is_sm100a_supported and is_sm12x_supported check accordingly.
118-121: Consider replacing×withxfor ASCII compatibility.Static analysis flagged the multiplication sign
×as ambiguous. While it renders correctly in most environments, using ASCIIxavoids potential display issues.🔤 Proposed fix
print( f"\nAll gen_fused_moe symlink race tests passed: " - f"{num_iterations} iterations × {num_processes} processes" + f"{num_iterations} iterations x {num_processes} processes" )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/utils/test_gen_module_symlink_race_condition.py` around lines 118 - 121, Replace the non-ASCII multiplication character in the test print with an ASCII 'x' to ensure compatibility: update the f-string in the print call that composes the message (the print(...) using f"\nAll gen_fused_moe symlink race tests passed: {num_iterations} iterations × {num_processes} processes") to use 'x' instead of '×' while keeping the same variables num_iterations and num_processes and message formatting.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/utils/test_gen_module_symlink_race_condition.py`:
- Around line 82-84: Replace the print+return skip pattern with pytest.skip so
pytest records the test as SKIPPED: add pytest to the module imports and change
the conditional in the test (the block using
is_sm100a_supported/is_sm12x_supported) to call
pytest.skip("gen_trtllm_gen_fused_moe_sm100_module requires SM100 or SM12x")
when the condition is met instead of printing and returning; update the test
function that contains the is_sm100a_supported and is_sm12x_supported check
accordingly.
- Around line 118-121: Replace the non-ASCII multiplication character in the
test print with an ASCII 'x' to ensure compatibility: update the f-string in the
print call that composes the message (the print(...) using f"\nAll gen_fused_moe
symlink race tests passed: {num_iterations} iterations × {num_processes}
processes") to use 'x' instead of '×' while keeping the same variables
num_iterations and num_processes and message formatting.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: bcdb802d-2796-4b7a-bb80-c33fabed237b
📒 Files selected for processing (1)
tests/utils/test_gen_module_symlink_race_condition.py
|
[SUCCESS] Pipeline #47728558: 10/20 passed |
Super thanks! |
|
/bot run |
|
/bot stop |
|
The GitLab CI pipeline #47780170 has been cancelled. |
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/utils/test_gen_module_symlink_race_condition.py (2)
20-20: Please verify the default pool start method after the CUDA probe.Line 85-86 performs the GPU capability check before Line 100 constructs a plain
multiprocessing.Pool. On Linux that meansfork, which is the multiprocessing mode PyTorch/CUDA tends to be most sensitive to. If this worker path ever starts touching CUDA,get_context("spawn").Pool(...)will be safer and more portable.Possible hardening change
-from multiprocessing import Pool +from multiprocessing import get_context ... - with Pool(processes=num_processes) as pool: + with get_context("spawn").Pool(processes=num_processes) as pool:Also applies to: 85-100
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/utils/test_gen_module_symlink_race_condition.py` at line 20, The test constructs a plain multiprocessing.Pool after performing the CUDA probe which may leave the default start method as "fork" (risky for CUDA); update the Pool creation to explicitly use the "spawn" context by replacing multiprocessing.Pool(...) with multiprocessing.get_context("spawn").Pool(...) (or call multiprocessing.get_context("spawn").Pool in the same scope where the Pool is constructed) to ensure a safe start method after the GPU capability check and verify the default via multiprocessing.get_start_method() if needed.
90-110: Prime the shared artifact cache once before the stress loop.
gen_trtllm_gen_fused_moe_sm100_module()does checksum/header fetches before it reachesensure_symlink(), so the first iteration is still mixing cache population with the symlink race. A one-time serial warm-up againsttemp_dirwould keep this regression tighter and less flaky.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/utils/test_gen_module_symlink_race_condition.py` around lines 90 - 110, The test mixes cache population with the symlink race on the first parallel iteration; perform a one-time serial warm-up against the shared temp_dir before spawning the Pool so checksum/header fetches in gen_trtllm_gen_fused_moe_sm100_module complete up-front. Call the worker path used by the pool (e.g., invoke gen_fused_moe_worker_process with temp_dir once, or directly call gen_trtllm_gen_fused_moe_sm100_module against temp_dir) before the for-loop/Pool creation so subsequent iterations only exercise ensure_symlink race conditions.
🤖 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/utils/test_gen_module_symlink_race_condition.py`:
- Around line 81-88: The test uses a bare return when GPU support checks fail
which hides skipped tests; in the test function
test_gen_fused_moe_symlink_race_condition replace the early return after calling
is_sm100a_supported(device) or is_sm12x_supported(device) with pytest.skip(...)
so the test is reported as skipped (import pytest at top if needed) and include
a clear message like "Skipping: gen_trtllm_gen_fused_moe_sm100_module requires
SM100 or SM12x" to match existing patterns.
---
Nitpick comments:
In `@tests/utils/test_gen_module_symlink_race_condition.py`:
- Line 20: The test constructs a plain multiprocessing.Pool after performing the
CUDA probe which may leave the default start method as "fork" (risky for CUDA);
update the Pool creation to explicitly use the "spawn" context by replacing
multiprocessing.Pool(...) with multiprocessing.get_context("spawn").Pool(...)
(or call multiprocessing.get_context("spawn").Pool in the same scope where the
Pool is constructed) to ensure a safe start method after the GPU capability
check and verify the default via multiprocessing.get_start_method() if needed.
- Around line 90-110: The test mixes cache population with the symlink race on
the first parallel iteration; perform a one-time serial warm-up against the
shared temp_dir before spawning the Pool so checksum/header fetches in
gen_trtllm_gen_fused_moe_sm100_module complete up-front. Call the worker path
used by the pool (e.g., invoke gen_fused_moe_worker_process with temp_dir once,
or directly call gen_trtllm_gen_fused_moe_sm100_module against temp_dir) before
the for-loop/Pool creation so subsequent iterations only exercise ensure_symlink
race conditions.
🪄 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: 6891581d-4b41-4e03-abe8-d5b569e6ca87
📒 Files selected for processing (1)
tests/utils/test_gen_module_symlink_race_condition.py
| import shutil | ||
| import torch | ||
| from flashinfer.utils import is_sm100a_supported, is_sm12x_supported | ||
|
|
||
| device = torch.device("cuda") | ||
| if not (is_sm100a_supported(device) or is_sm12x_supported(device)): | ||
| print("Skipping: gen_trtllm_gen_fused_moe_sm100_module requires SM100 or SM12x") | ||
| return |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n tests/utils/test_gen_module_symlink_race_condition.py | head -100Repository: flashinfer-ai/flashinfer
Length of output: 4138
🏁 Script executed:
# Check how other tests in the repository handle skipping on unsupported architectures
rg "pytest\.skip\(|is_sm.*_supported" tests/ -A 2 -B 2 | head -100Repository: flashinfer-ai/flashinfer
Length of output: 7658
🏁 Script executed:
# Check if there are other test functions in the file that might show patterns
rg "def test_|def gen_" tests/utils/test_gen_module_symlink_race_condition.pyRepository: flashinfer-ai/flashinfer
Length of output: 196
Use pytest.skip() instead of bare return.
A bare return reports this as passed on unsupported GPUs, causing CI to lose visibility that the test never ran. The codebase pattern (e.g., tests/attention/test_trtllm_gen_attention.py) consistently uses pytest.skip() for this purpose. Since test_gen_fused_moe_symlink_race_condition() is a pytest test function, apply the same pattern:
Suggested change
import shutil
+ import pytest
import torch
from flashinfer.utils import is_sm100a_supported, is_sm12x_supported
device = torch.device("cuda")
if not (is_sm100a_supported(device) or is_sm12x_supported(device)):
- print("Skipping: gen_trtllm_gen_fused_moe_sm100_module requires SM100 or SM12x")
- return
+ pytest.skip(
+ "gen_trtllm_gen_fused_moe_sm100_module requires SM100 or SM12x"
+ )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/utils/test_gen_module_symlink_race_condition.py` around lines 81 - 88,
The test uses a bare return when GPU support checks fail which hides skipped
tests; in the test function test_gen_fused_moe_symlink_race_condition replace
the early return after calling is_sm100a_supported(device) or
is_sm12x_supported(device) with pytest.skip(...) so the test is reported as
skipped (import pytest at top if needed) and include a clear message like
"Skipping: gen_trtllm_gen_fused_moe_sm100_module requires SM100 or SM12x" to
match existing patterns.
|
filed an issue to track future work |
|
once https://gitlab-master.nvidia.com/dl/flashinfer/flashinfer-ci/-/pipelines/47787607 clears the PR, we can go forth with a post3 fix |
4 participants
📌 Description
Add filelock to
ensure_symlink.Previously hitting issues on v0.6.7.post1
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Bug Fixes
Tests