diff --git a/python/pyproject.toml b/python/pyproject.toml index bb87cf64f..3340ecfdc 100755 --- a/python/pyproject.toml +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/ext/loader.py b/python/sglang/kernels/ops/diffusion/ext/loader.py index dd25a30d4..107f454fe 100644 --- a/python/sglang/kernels/ops/diffusion/ext/loader.py +++ b/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"] diff --git a/python/sglang/multimodal_gen/test/unit/test_cpp_extension_loader.py b/python/sglang/multimodal_gen/test/unit/test_cpp_extension_loader.py new file mode 100644 index 000000000..563757029 --- /dev/null +++ b/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