Fix stale lock file in CUDA loader — prevents infinite spin on crash recovery

torch.utils.cpp_extension.load creates a 'lock' file in the build
directory during compilation. If the compiling process is killed
(OOM, timeout, user interrupt), the lock file is never removed and
subsequent processes spin forever polling it (clock_nanosleep(100ms)
→ stat(lock) → repeat).

Fix: _cleanup_stale_lock() removes lock files older than 10 minutes
before any compilation attempt. This is the correct threshold — CUDA
kernel compilation should never take more than a few minutes, so a
10-minute-old lock is guaranteed stale.
This commit is contained in:
2026-06-02 21:34:58 +00:00
parent 0b6ca0df80
commit 845227c06c

View File

@@ -10,6 +10,7 @@ Usage:
result = mod.quantize_nvfp4_from_buffer(x, divisor)
"""
import os
import time
import hashlib
import torch
from torch.utils.cpp_extension import load
@@ -18,6 +19,34 @@ _KERNEL_DIR = os.path.dirname(os.path.abspath(__file__))
_CACHE_DIR = os.path.join(_KERNEL_DIR, "_build_cache")
_LOADED_MODULES = {}
# Maximum age of a stale lock file before we remove it (seconds).
# torch.utils.cpp_extension.load creates a lock file during compilation.
# If the process is killed during compilation, the lock remains and the
# next process spins forever polling it. This timeout prevents that.
_STALE_LOCK_TIMEOUT_S = 600 # 10 minutes
def _cleanup_stale_lock():
"""Remove stale lock files from the build cache directory.
torch.utils.cpp_extension.load creates a 'lock' file in the build
directory during compilation. If the compiling process is killed
(OOM, timeout, user interrupt), the lock file is never removed and
subsequent processes spin forever waiting for it.
This function checks if a lock file exists and is older than
_STALE_LOCK_TIMEOUT_S. If so, it removes it.
"""
lock_path = os.path.join(_CACHE_DIR, "lock")
if os.path.exists(lock_path):
try:
lock_age = time.time() - os.path.getmtime(lock_path)
if lock_age > _STALE_LOCK_TIMEOUT_S:
os.remove(lock_path)
print(f"[loader] Removed stale lock file (age={lock_age:.0f}s)", flush=True)
except OSError:
pass # Lock was removed between exists() and remove()
def get_cuda_module(name, sources, extra_cuda_cflags=None):
"""Load a CUDA kernel module, compiling once and caching forever.
@@ -33,6 +62,9 @@ def get_cuda_module(name, sources, extra_cuda_cflags=None):
if name in _LOADED_MODULES:
return _LOADED_MODULES[name]
# Clean up stale lock files from crashed previous compilations
_cleanup_stale_lock()
source_paths = [os.path.join(_KERNEL_DIR, s) for s in sources]
# Build a cache key from source file contents + compile flags