131 lines
3.7 KiB
Python
131 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
import torch
|
|
|
|
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _get_build_directory(name: str) -> Path:
|
|
try:
|
|
from torch.utils.cpp_extension import _get_build_directory
|
|
|
|
return Path(_get_build_directory(name, False))
|
|
except (ImportError, AttributeError):
|
|
from torch.utils.cpp_extension import get_default_build_root
|
|
|
|
root = os.environ.get("TORCH_EXTENSIONS_DIR") or get_default_build_root()
|
|
if "TORCH_EXTENSIONS_DIR" not in os.environ:
|
|
cu_str = (
|
|
"cpu"
|
|
if torch.version.cuda is None
|
|
else f"cu{torch.version.cuda.replace('.', '')}"
|
|
)
|
|
py_str = (
|
|
f"py{sys.version_info.major}{sys.version_info.minor}"
|
|
f"{getattr(sys, 'abiflags', '')}"
|
|
)
|
|
root = os.path.join(root, f"{py_str}_{cu_str}")
|
|
return Path(root) / name
|
|
|
|
|
|
def _is_recoverable_load_error(
|
|
exc: BaseException, name: str, build_directory: Path
|
|
) -> bool:
|
|
message = str(exc).lower()
|
|
current = exc.__cause__ or exc.__context__
|
|
while current is not None:
|
|
message += f"\n{current}".lower()
|
|
current = current.__cause__ or current.__context__
|
|
|
|
if any(
|
|
marker in message
|
|
for marker in (
|
|
"error building extension",
|
|
"error compiling objects for extension",
|
|
"ninja",
|
|
"nvcc",
|
|
"gcc",
|
|
"g++",
|
|
"fatal error:",
|
|
"compilation terminated",
|
|
)
|
|
):
|
|
return False
|
|
|
|
if not any(
|
|
marker in message
|
|
for marker in (str(build_directory / f"{name}.so").lower(), f"{name}.so")
|
|
):
|
|
return False
|
|
|
|
return any(
|
|
marker in message
|
|
for marker in (
|
|
"undefined symbol",
|
|
"cannot open shared object file",
|
|
"no such file or directory",
|
|
"file too short",
|
|
"invalid elf header",
|
|
"wrong elf class",
|
|
"elf load command",
|
|
"dlopen",
|
|
"version `glibcxx",
|
|
)
|
|
)
|
|
|
|
|
|
def load_extension_with_recovery(
|
|
name: str,
|
|
sources: Sequence[str],
|
|
extra_cflags: Sequence[str] | None = None,
|
|
extra_cuda_cflags: Sequence[str] | None = None,
|
|
verbose: bool = False,
|
|
) -> 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
|
|
|
|
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)
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
__all__ = ["load_extension_with_recovery"]
|