[diffusion] fix: fix hunyuan3d stale extension lock hangs (#35989)

python/pyproject.toml

@@ -76,6 +76,7 @@ dependencies = [
   "tiktoken",
   "tilelang==0.1.12",
   "timm==1.0.16",
+  "tokenizers==0.22.2", # 0.23.0rc0 is incompatible with transformers' CLIPTokenizer.
   "tokenspeed_mla==0.1.8",
   "tomli ; python_version < '3.11'",
   "torch==2.13.0",

python/sglang/kernels/ops/diffusion/ext/loader.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import fcntl
 import logging
 import os
 import shutil
 import sys
+from contextlib import contextmanager
 from pathlib import Path
-from typing import Any, Sequence
+from typing import Any, Iterator, Sequence
 
 import torch
 
@@ -81,6 +83,28 @@ def _is_recoverable_load_error(
     )
 
 
+@contextmanager
+def _extension_build_lock(build_directory: Path) -> Iterator[None]:
+    """Serialize builds and discard PyTorch lock files left by dead processes."""
+    build_directory.parent.mkdir(parents=True, exist_ok=True)
+    lock_path = build_directory.parent / f".{build_directory.name}.sglang.lock"
+    with lock_path.open("a+") as lock_file:
+        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+        try:
+            torch_lock_path = build_directory / "lock"
+            if torch_lock_path.exists():
+                logger.warning(
+                    "Removing stale PyTorch extension lock for %s at %s",
+                    build_directory.name,
+                    torch_lock_path,
+                )
+                torch_lock_path.unlink(missing_ok=True)
+            build_directory.mkdir(parents=True, exist_ok=True)
+            yield
+        finally:
+            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
+
+
 def load_extension_with_recovery(
     name: str,
     sources: Sequence[str],
@@ -90,40 +114,36 @@ def load_extension_with_recovery(
 ) -> Any:
     from torch.utils.cpp_extension import load
 
-    try:
-        return load(
-            name=name,
-            sources=list(sources),
-            extra_cflags=None if extra_cflags is None else list(extra_cflags),
-            extra_cuda_cflags=(
-                None if extra_cuda_cflags is None else list(extra_cuda_cflags)
-            ),
-            verbose=verbose,
-        )
-    except Exception as exc:
-        build_directory = _get_build_directory(name)
-        if not _is_recoverable_load_error(exc, name, build_directory):
-            raise
+    build_directory = _get_build_directory(name)
+    load_kwargs = {
+        "name": name,
+        "sources": list(sources),
+        "extra_cflags": None if extra_cflags is None else list(extra_cflags),
+        "extra_cuda_cflags": (
+            None if extra_cuda_cflags is None else list(extra_cuda_cflags)
+        ),
+        "build_directory": str(build_directory),
+        "verbose": verbose,
+    }
 
-        logger.warning(
-            "Detected a stale or broken JIT extension for %s at %s; clearing "
-            "its cache and retrying once.",
-            name,
-            build_directory,
-        )
-        sys.modules.pop(name, None)
-        if build_directory.exists():
-            shutil.rmtree(build_directory)
+    with _extension_build_lock(build_directory):
+        try:
+            return load(**load_kwargs)
+        except Exception as exc:
+            if not _is_recoverable_load_error(exc, name, build_directory):
+                raise
 
-        return load(
-            name=name,
-            sources=list(sources),
-            extra_cflags=None if extra_cflags is None else list(extra_cflags),
-            extra_cuda_cflags=(
-                None if extra_cuda_cflags is None else list(extra_cuda_cflags)
-            ),
-            verbose=verbose,
-        )
+            logger.warning(
+                "Detected a stale or broken JIT extension for %s at %s; clearing "
+                "its cache and retrying once.",
+                name,
+                build_directory,
+            )
+            sys.modules.pop(name, None)
+            if build_directory.exists():
+                shutil.rmtree(build_directory)
+            build_directory.mkdir(parents=True)
+            return load(**load_kwargs)
 
 
 __all__ = ["load_extension_with_recovery"]

python/sglang/multimodal_gen/test/unit/test_cpp_extension_loader.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from pathlib import Path
+from unittest.mock import patch
+
+from sglang.kernels.ops.diffusion.ext.loader import load_extension_with_recovery
+
+
+def test_stale_torch_lock_is_removed_before_loading(tmp_path: Path):
+    build_directory = tmp_path / "test_extension"
+    build_directory.mkdir()
+    torch_lock_path = build_directory / "lock"
+    torch_lock_path.touch()
+
+    expected = object()
+    with (
+        patch(
+            "sglang.kernels.ops.diffusion.ext.loader._get_build_directory",
+            return_value=build_directory,
+        ),
+        patch("torch.utils.cpp_extension.load", return_value=expected) as load,
+    ):
+        result = load_extension_with_recovery("test_extension", ["source.cpp"])
+
+    assert result is expected
+    assert not torch_lock_path.exists()
+    load.assert_called_once_with(
+        name="test_extension",
+        sources=["source.cpp"],
+        extra_cflags=None,
+        extra_cuda_cflags=None,
+        build_directory=str(build_directory),
+        verbose=False,
+    )
+
+
+def test_broken_extension_is_rebuilt_under_the_same_lock(tmp_path: Path):
+    build_directory = tmp_path / "test_extension"
+    build_directory.mkdir()
+    expected = object()
+    load_error = OSError(f"{build_directory}/test_extension.so: file too short")
+
+    with (
+        patch(
+            "sglang.kernels.ops.diffusion.ext.loader._get_build_directory",
+            return_value=build_directory,
+        ),
+        patch(
+            "torch.utils.cpp_extension.load",
+            side_effect=[load_error, expected],
+        ) as load,
+    ):
+        result = load_extension_with_recovery("test_extension", ["source.cpp"])
+
+    assert result is expected
+    assert build_directory.is_dir()
+    assert load.call_count == 2