[diffusion] CI: fix hunyuan3d JIT cache (#20773)

Co-authored-by: daiweitao <dwti614707404@163.com>
2026-03-20 17:51:55 +08:00
parent a0a4dae67f
commit fa89d152c0
6 changed files with 198 additions and 68 deletions
--- a/python/sglang/multimodal_gen/csrc/render/init.py
+++ b/python/sglang/multimodal_gen/csrc/render/init.py
@@ -0,0 +1,130 @@
+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"]
--- a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/init.py
+++ b/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/init.py
@@ -12,6 +12,7 @@ import os
 from typing import List, Tuple

 import torch
+from sglang.multimodal_gen.csrc.render import load_extension_with_recovery

 _abs_path = os.path.dirname(os.path.abspath(__file__))
 _custom_rasterizer_kernel = None
@@ -24,9 +25,7 @@ def _load_custom_rasterizer():
    if _custom_rasterizer_kernel is not None:
        return _custom_rasterizer_kernel

-    from torch.utils.cpp_extension import load
-
-    _custom_rasterizer_kernel = load(
+    _custom_rasterizer_kernel = load_extension_with_recovery(
        name="custom_rasterizer_kernel",
        sources=[
            f"{_abs_path}/rasterizer.cpp",
--- a/python/sglang/multimodal_gen/csrc/render/mesh_processor/init.py
+++ b/python/sglang/multimodal_gen/csrc/render/mesh_processor/init.py
@@ -12,6 +12,7 @@ import os
 from typing import Tuple

 import numpy as np
+from sglang.multimodal_gen.csrc.render import load_extension_with_recovery

 _abs_path = os.path.dirname(os.path.abspath(__file__))
 _mesh_processor_kernel = None
@@ -24,9 +25,7 @@ def _load_mesh_processor():
    if _mesh_processor_kernel is not None:
        return _mesh_processor_kernel

-    from torch.utils.cpp_extension import load
-
-    _mesh_processor_kernel = load(
+    _mesh_processor_kernel = load_extension_with_recovery(
        name="mesh_processor_kernel",
        sources=[
            f"{_abs_path}/mesh_processor.cpp",
--- a/python/sglang/multimodal_gen/test/server/perf_baselines.json
+++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json
@@ -1988,69 +1988,69 @@
        },
        "hunyuan3d_shape_gen": {
            "stages_ms": {
-                "Hunyuan3DShapeBeforeDenoisingStage": 31.42,
-                "Hunyuan3DShapeDenoisingStage": 3259.83,
-                "Hunyuan3DShapeExportStage": 8735.55,
-                "Hunyuan3DShapeSaveStage": 981.64,
-                "Hunyuan3DPaintPreprocessStage": 226071.67,
-                "Hunyuan3DPaintTexGenStage": 11083.05,
-                "Hunyuan3DPaintPostprocessStage": 7469.29
+                "Hunyuan3DShapeBeforeDenoisingStage": 235.65,
+                "Hunyuan3DShapeDenoisingStage": 3452.51,
+                "Hunyuan3DShapeExportStage": 8819.6,
+                "Hunyuan3DShapeSaveStage": 752.4,
+                "Hunyuan3DPaintPreprocessStage": 218136.45,
+                "Hunyuan3DPaintTexGenStage": 10259.21,
+                "Hunyuan3DPaintPostprocessStage": 6387.55
            },
            "denoise_step_ms": {
-                "0": 32.26,
-                "1": 63.34,
-                "2": 65.44,
-                "3": 65.44,
-                "4": 65.6,
-                "5": 65.81,
-                "6": 65.82,
-                "7": 65.48,
-                "8": 65.9,
-                "9": 65.77,
-                "10": 65.54,
-                "11": 65.68,
-                "12": 65.85,
-                "13": 65.77,
-                "14": 65.7,
-                "15": 65.78,
-                "16": 66.0,
-                "17": 66.15,
-                "18": 65.91,
-                "19": 66.5,
-                "20": 65.76,
-                "21": 66.08,
-                "22": 66.06,
-                "23": 66.23,
-                "24": 65.79,
-                "25": 65.58,
-                "26": 65.88,
-                "27": 65.67,
-                "28": 65.87,
-                "29": 66.09,
-                "30": 65.81,
-                "31": 65.91,
-                "32": 66.18,
-                "33": 65.93,
-                "34": 66.26,
-                "35": 66.26,
-                "36": 66.27,
-                "37": 65.57,
-                "38": 66.02,
-                "39": 66.19,
-                "40": 65.23,
-                "41": 66.11,
-                "42": 66.18,
-                "43": 65.86,
-                "44": 65.86,
-                "45": 65.92,
-                "46": 65.65,
-                "47": 65.78,
-                "48": 66.01,
-                "49": 66.08
+                "0": 150.72,
+                "1": 26.65,
+                "2": 65.91,
+                "3": 68.09,
+                "4": 76.12,
+                "5": 68.16,
+                "6": 61.19,
+                "7": 68.26,
+                "8": 67.92,
+                "9": 68.26,
+                "10": 68.06,
+                "11": 68.22,
+                "12": 68.11,
+                "13": 68.19,
+                "14": 69.58,
+                "15": 66.91,
+                "16": 68.03,
+                "17": 68.36,
+                "18": 68.49,
+                "19": 67.69,
+                "20": 68.19,
+                "21": 69.1,
+                "22": 67.78,
+                "23": 68.36,
+                "24": 68.19,
+                "25": 68.26,
+                "26": 68.06,
+                "27": 68.25,
+                "28": 68.39,
+                "29": 68.26,
+                "30": 68.05,
+                "31": 68.27,
+                "32": 68.2,
+                "33": 68.19,
+                "34": 68.02,
+                "35": 68.3,
+                "36": 68.2,
+                "37": 68.48,
+                "38": 68.23,
+                "39": 68.36,
+                "40": 67.9,
+                "41": 75.76,
+                "42": 62.04,
+                "43": 66.78,
+                "44": 67.85,
+                "45": 68.11,
+                "46": 67.92,
+                "47": 68.15,
+                "48": 67.89,
+                "49": 68.3
            },
-            "expected_e2e_ms": 257696.97,
-            "expected_avg_denoise_ms": 65.16,
-            "expected_median_denoise_ms": 65.86
+            "expected_e2e_ms": 248171.5,
+            "expected_avg_denoise_ms": 68.98,
+            "expected_median_denoise_ms": 68.19
        },
        "wan2_1_t2v_1.3b_frame_interp_2x": {
            "stages_ms": {
--- a/python/sglang/multimodal_gen/test/server/test_server_common.py
+++ b/python/sglang/multimodal_gen/test/server/test_server_common.py
@@ -96,8 +96,8 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
    if server_args.lora_path:
        extra_args += f" --lora-path {server_args.lora_path}"

-    # default warmup
-    extra_args += f" --warmup"
+    if server_args.enable_warmup:
+        extra_args += " --warmup"

    for arg in server_args.extras:
        extra_args += f" {arg}"
--- a/python/sglang/multimodal_gen/test/server/testcase_configs.py
+++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py
@@ -198,6 +198,7 @@ class DiffusionServerArgs:
    dit_offload_prefetch_size: int | float | None = None
    enable_cache_dit: bool = False
    text_encoder_cpu_offload: bool = False
+    enable_warmup: bool = True

    extras: list[str] = field(default_factory=lambda: [])

@@ -726,6 +727,7 @@ if not current_platform.is_hip():
            DiffusionServerArgs(
                model_path="tencent/Hunyuan3D-2",
                modality="3d",
+                enable_warmup=False,
            ),
            HUNYUAN3D_SHAPE_sampling_params,
        ),