diff --git a/python/sglang/kernels/jit/utils/compile/cache.py b/python/sglang/kernels/jit/utils/compile/cache.py index 78b664343..deba51e4c 100644 --- a/python/sglang/kernels/jit/utils/compile/cache.py +++ b/python/sglang/kernels/jit/utils/compile/cache.py @@ -298,10 +298,15 @@ def _normalize_text(text: str) -> str: return text -def build_key_dir(*, module_name: str, build_key: str) -> pathlib.Path: +def cache_root() -> pathlib.Path: configured = envs.SGLANG_JIT_CACHE_DIR.get() or "~/.cache/sglang/jit" - root = pathlib.Path(configured).expanduser() - return root / _target_tag() / module_name / f"{_BUILD_KEY_PREFIX}{build_key}" + return pathlib.Path(configured).expanduser() + + +def build_key_dir(*, module_name: str, build_key: str) -> pathlib.Path: + return ( + cache_root() / _target_tag() / module_name / f"{_BUILD_KEY_PREFIX}{build_key}" + ) # --------------------------------------------------------------------------- diff --git a/python/sglang/kernels/jit/utils/compile/loader.py b/python/sglang/kernels/jit/utils/compile/loader.py index dca05aba0..89fecce45 100644 --- a/python/sglang/kernels/jit/utils/compile/loader.py +++ b/python/sglang/kernels/jit/utils/compile/loader.py @@ -23,6 +23,7 @@ from sglang.kernels.jit.utils.compile.paths import ( ) from sglang.kernels.jit.utils.compile.spec import BuildSpec, resolve_sources from sglang.kernels.jit.utils.deps import REGISTERED_DEPENDENCIES +from sglang.srt.environ import envs if TYPE_CHECKING: from tvm_ffi import Module @@ -125,6 +126,16 @@ def load_jit( e, ) + # Before the lock: with the flag set no process ever publishes a build, + # so waiting on the lock cannot turn this miss into a hit. + if envs.SGLANG_CRASH_ON_JIT_COMPILE.get(): + raise RuntimeError( + f"JIT module '{spec.module_name}' has no usable cached build under " + f"{scope}, and SGLANG_CRASH_ON_JIT_COMPILE forbids compiling at " + "runtime. Seed the cache with prebuilt artifacts matching this " + "environment, or unset the flag." + ) + with _build_lock(scope): # Re-check: whoever held the lock before us has very likely just # published exactly what we were about to build. This is what turns N diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d858de885..d5bf1ccb1 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1095,6 +1095,10 @@ class Envs: # is what makes reverting an edit an instant hit instead of a rebuild; set # it to trade that away for disk (1 keeps only the most recent build). SGLANG_JIT_CACHE_KEEP = EnvInt(None) + # Raise instead of compiling when a module misses the cache, so a + # deployment that expects a pre-seeded cache fails loudly at startup + # rather than silently eating a cold compile. + SGLANG_CRASH_ON_JIT_COMPILE = EnvBool(False) # =================================================================== # Expert-parallel dispatch and MoE execution