Skip to content

fix(cpp): add cross-process file lock for JIT kernel builds - #3019

Merged
Qubitium merged 1 commit into
ModelCloud:mainfrom
Ndgandhi23:fix/concurrent-jit-build-cache-lock
Aug 19, 2026
Merged

fix(cpp): add cross-process file lock for JIT kernel builds#3019
Qubitium merged 1 commit into
ModelCloud:mainfrom
Ndgandhi23:fix/concurrent-jit-build-cache-lock

Conversation

@Ndgandhi23

@Ndgandhi23 Ndgandhi23 commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Summary

Two quantize() processes on one host share the same JIT build cache, but the lock guarding it is process-local. torch serializes builds with a lock file that has no timeout or owner check. If a process is killed mid-build, that file is left behind and every later load() waits on it forever at 0% cpu. Reproduced locally without a GPU. Surfaced while investigating #3015.

The fix is an fcntl.flock around the build step. The kernel releases the lock automatically when the holding process dies, so a killed process can never leave the cache locked. Once the lock is held, any leftover torch lock file is removed. The wait is also bounded, so a stuck holder degrades the load into the existing non-JIT fallback instead of hanging a worker.

What Changed

  • gptqmodel/utils/cpp.py:
    • New _cross_process_build_lock context manager using stdlib fcntl, with a no-op fallback where fcntl is unavailable. The lock file lives outside the build dir so cache clears never delete it.
    • TorchOpsJitExtension.load() runs its cache-probe and compile section under this lock, and removes any stale torch lock file after acquiring it.
    • The lock wait is bounded by GPTQMODEL_TORCH_OPS_LOCK_TIMEOUT (default 600s, or 5x the kernel's compile baseline if larger). On timeout the load fails cleanly into the non-JIT path.
    • clear_cache() and force-rebuild delete the cache under the same lock, so a wipe cannot land under another process's in-flight build.
    • Note: most of the diff line count is the existing load() tail reindented into the new with block. It is byte-identical apart from the stale-lock removal call.
  • tests/test_torch_ops_jit_extension.py: three new tests — a stale torch lock file is removed before compiling, a SIGKILLed lock holder releases the lock so load proceeds, and a live holder hits the bounded wait instead of hanging. The subprocess tests skip where fcntl is unavailable.
  • README.md: new "JIT kernel cache and multi-process quantization" section documenting the cache dir and both env vars.

Tests

  • I added a new simple/fast unit test for this change, or documented why that is not applicable.
  • I ran the new targeted test locally before opening this PR.
  • I ran any other directly relevant local tests.

Paste the exact test commands and results here:

pytest tests/test_torch_ops_jit_extension.py tests/test_extension_load_api.py tests/test_cpp_jit_progress.py tests/test_cpp_nvcc_flags.py
# 57 passed in 4.18s

cd format && ruff check --config ruff.toml \
  ../gptqmodel/utils/cpp.py \
  ../tests/test_torch_ops_jit_extension.py
# All checks passed!

git diff --check
# no output

Review Requirements

  • I personally reviewed every file in this diff.
  • I checked that the code matches existing project structure, APIs, and conventions.
  • I avoided unnecessary monkeypatching and used the project's normal extension points where possible.

Notes

flock on NFS-mounted cache dirs has weaker advisory semantics; worst case is today's behavior.

Assisted with Claude Code; every change personally designed and reviewed.

@Ndgandhi23
Ndgandhi23 marked this pull request as draft August 19, 2026 03:22
@Ndgandhi23
Ndgandhi23 force-pushed the fix/concurrent-jit-build-cache-lock branch from d1a22b8 to 2c19be5 Compare August 19, 2026 03:35
Two quantize() processes on one host share the same JIT build cache,
but the existing lock is process-local. If a process is killed mid
build, torch's lock file is left behind and every later load() waits
on it forever at 0% cpu (ModelCloud#3015).

Serialize builds with an fcntl.flock that the kernel auto-releases
when the holding process dies. Once the lock is held, any leftover
torch lock file is removed. The wait is bounded by
GPTQMODEL_TORCH_OPS_LOCK_TIMEOUT and falls back to non-JIT paths on
timeout. Cache clears run under the same lock so they cannot wipe
another process's in-flight build.

Ref ModelCloud#3015
@Ndgandhi23
Ndgandhi23 force-pushed the fix/concurrent-jit-build-cache-lock branch from 2c19be5 to 4a1ee5a Compare August 19, 2026 03:53
@Ndgandhi23
Ndgandhi23 marked this pull request as ready for review August 19, 2026 03:54
@Qubitium
Qubitium merged commit 9dcefe4 into ModelCloud:main Aug 19, 2026
@Qubitium

Copy link
Copy Markdown
Collaborator

@Ndgandhi23 LGTM. Thank you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants