Skip to content

Fix: Prevent race condition in cubin loader when file is being consumed #1852

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix: Prevent race condition in cubin loader when file is being consumed #1852

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions flashinfer/jit/cubin_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,9 @@ def download_file(
# Handle local file copy
if os.path.exists(source):
try:
shutil.copy(source, local_path)
temp_path = f"{local_path}.tmp"
shutil.copy(source, temp_path)
os.replace(temp_path, local_path) # Atomic rename
Comment on lines +83 to +85
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

While using a temporary file and atomic rename is a good solution for the race condition, this implementation can leave behind .tmp files if an error occurs during the copy or rename operation. This could lead to disk space exhaustion over time. To make this more robust, I suggest wrapping the file operations in a try...finally block to ensure the temporary file is always cleaned up.

Suggested change
temp_path = f"{local_path}.tmp"
shutil.copy(source, temp_path)
os.replace(temp_path, local_path) # Atomic rename
temp_path = f"{local_path}.tmp"
try:
shutil.copy(source, temp_path)
os.replace(temp_path, local_path) # Atomic rename
finally:
if os.path.exists(temp_path):
os.remove(temp_path)

logger.info(f"File copied successfully: {local_path}")
return True
except Exception as e:
Expand All @@ -93,9 +95,13 @@ def download_file(
response = session.get(source, timeout=timeout)
response.raise_for_status()

with open(local_path, "wb") as file:
temp_path = f"{local_path}.tmp"
with open(temp_path, "wb") as file:
file.write(response.content)

# Atomic rename to prevent readers from seeing partial writes
os.replace(temp_path, local_path)
Comment on lines +98 to +103
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Similar to the local file copy, the URL download logic can leave temporary .tmp files if an error occurs during file writing or the atomic rename. This can happen, for example, if there are file permission issues or if the disk is full. Adding a try...finally block will ensure these temporary files are cleaned up, making the download process more robust.

Suggested change
temp_path = f"{local_path}.tmp"
with open(temp_path, "wb") as file:
file.write(response.content)
# Atomic rename to prevent readers from seeing partial writes
os.replace(temp_path, local_path)
temp_path = f"{local_path}.tmp"
try:
with open(temp_path, "wb") as file:
file.write(response.content)
# Atomic rename to prevent readers from seeing partial writes
os.replace(temp_path, local_path)
finally:
if os.path.exists(temp_path):
os.remove(temp_path)


logger.info(
f"File downloaded successfully: {source} -> {local_path}"
)
Expand Down