Import-time cpp_extension JIT can hang forever on a stale build lock (FileBaton spin-wait); please add a timeout or lock-staleness check

[JinukMoon]

Summary

import openequivariance (v0.4.1) JIT-builds two torch extensions at import time in openequivariance/extlib/__init__.py:

torch_module = torch.utils.cpp_extension.load("libtorch_tp_jit", ...)
generic_module = torch.utils.cpp_extension.load("generic_module", ...)

torch.utils.cpp_extension.load has two properties that combine badly on shared/limited hosts:

1. ninja parallelism defaults to nproc (only bounded by the MAX_JOBS env var), and
2. concurrent/subsequent builds serialize on a lock file in the extension build dir via torch's FileBaton, which spin-waits indefinitely with no timeout.

If a first build is killed part-way (OOM-killer, cgroup pids limit, user Ctrl-C, batch scheduler), the stale lock file remains, and every later import openequivariance hangs forever, silently — no traceback, no log, nothing on the GPU.

Observed in the field

On a login node with a per-user cgroup pids.max=100 (Fedora 36, RTX A4500, torch 2.x cu128, openequivariance 0.4.1, installed as the enable_OpenEquivariance modifier path of nequip-compile):

• the first import's ninja swarm (-j <nproc>) exceeded the pids limit and died;
• the stale lock remained in ~/.cache/torch_extensions/<tag>/libtorch_tp_jit/;
• every subsequent nequip-compile ... --modifiers enable_OpenEquivariance then hung deterministically at model load, reproducibly, with zero diagnostics.

The same workload with prebuilt-kernel backends (cuEquivariance) on the same host was unaffected, which is what isolated the JIT path.

Suggested fixes (any of these would prevent the silent hang)

1. Wrap the two cpp_extension.load calls with a bounded wait: if the baton does not clear within N minutes, raise with a message naming the lock path instead of spinning forever.
2. Detect a stale lock proactively (lock exists but no live ninja/compiler process) and remove it before building.
3. Default MAX_JOBS to a small value (e.g. 4) for these two import-time builds unless the user overrides — an import-time side effect should not fan out to nproc compile jobs.
4. At minimum, document MAX_JOBS and the stale-lock failure mode in the install docs.

Our downstream workaround (for reference): set MAX_JOBS=4 for any process importing oeq, and pre-clean libtorch_tp_jit/lock + generic_module/lock under $TORCH_EXTENSIONS_DIR when no ninja process is alive.

Thanks for the great kernels — happy to test a patch.

[asglover]

Hello, thank you for the report, I'll take a look and see what can be done to make this less likely to happen in the future

[vbharadwaj-bk]

Hi @JinukMoon - can you try upgrading to the latest version (0.6.x or above) - we do ahead-of-time compilation to avoid this exact issue.

[asglover]

I was just typing that out, yeah, it's a duplicate of this issue:
#156
But maybe we can look at putting in the time out or lowering the max jobs to prevent the issue. Lets discuss further

[JinukMoon]

Thanks @vbharadwaj-bk @asglover — tested 0.6.8 on our stack (torch 2.8.0+cu128, cp311, RTX 4060 Ti / sm89). Results:

1. AOT confirmed — the hang class is gone in 0.6.x. import openequivariance completes in ~1.6 s with no torch.utils.cpp_extension JIT (0.4.1 took 30–60 s of ninja JIT on a cold cache, with the lock-file failure mode described above). So the upgrade indeed removes this exact issue. 👍

2. One install friction worth knowing: the 0.6.8 sdist's CMake build requires a full CUDA toolkit for CUDA::nvrtc — a conda cuda-nvcc-only toolchain (compiler, no nvrtc dev libs) fails at target_link_libraries ... CUDA::nvrtc but the target was not found. Setting CUDAToolkit_ROOT=/usr/local/cuda-12.8 (full toolkit) fixed it. Might be worth a line in the install docs.

3. Heads-up for downstreams pinned on 0.4.x artifacts: models AOT-compiled under 0.4.1 (e.g. nequip-compile ... --modifiers enable_OpenEquivariance → AOTI .pt2) fail to load under 0.6.8 with Couldn't resolve type '__torch__.torch.classes.libtorch_tp_jit.TorchJITConv' — i.e. upgrading requires recompiling those artifacts, it's not drop-in. For projects that can't immediately recompile (ours pins 0.4.1 until we re-verify the recompiled artifacts), the timeout / lower-MAX_JOBS mitigation @asglover mentioned would still be valuable on the 0.4.x line. And agreed this is a duplicate of #156 — feel free to close in favor of it.

[vbharadwaj-bk]

This is valuable feedback - I think if you are using these old classes, @asglover's fix is valuable. Now, no guarantee that it will still work - I hope we haven't changed the serialization format- but I think adding serialization stability is actually a good regression test, since now this library has been use for some time. Let's keep this issue open and start fixing.

In the meantime, I would suggest an upgrade and recompile of your artifacts, if possible, with the knowledge that things will be more stable in the near-future and you won't have to repeat this. Of course, only if it does not consume too much of your resources.

[vbharadwaj-bk]

I've taken a closer look at this - and we've removed some of the infrastructure we used in 0.4.x, so deserializing on 0.6.x will indeed be impossible. But adding a ser/de test for forward compatibility seems like a good idea.