diff --git a/.github/workflows/python.yaml b/.github/workflows/python.yaml index 38539b5335..907c9c2fd6 100644 --- a/.github/workflows/python.yaml +++ b/.github/workflows/python.yaml @@ -383,6 +383,7 @@ jobs: xinference/model/llm/sglang/tests/test_pd.py \ xinference/model/llm/sglang/tests/test_xavier_gpu_lifecycle.py \ xinference/model/llm/sglang/tests/test_gc_lifecycle.py \ + xinference/model/llm/sglang/tests/test_runtime.py \ xinference/model/llm/vllm/tests/test_distributed_executor_v1.py \ xinference/model/llm/vllm/tests/test_executor_routing.py \ xinference/model/llm/vllm/tests/test_speculative_config.py \ diff --git a/xinference/model/llm/sglang/core.py b/xinference/model/llm/sglang/core.py index 99691f7dfd..551f106ed5 100644 --- a/xinference/model/llm/sglang/core.py +++ b/xinference/model/llm/sglang/core.py @@ -295,9 +295,11 @@ def load(self): # fork may cause sglang stuck, force set to spawn multiprocessing.set_start_method("spawn", force=True) - self._prepare_weight_cache() + from .runtime import create_runtime, prepare_jit_cache - from .runtime import create_runtime + # The weight daemon compiles JIT kernels before Runtime starts. + prepare_jit_cache() + self._prepare_weight_cache() logged_config = { key: "***" if key in ("api_key", "admin_api_key") else value diff --git a/xinference/model/llm/sglang/runtime.py b/xinference/model/llm/sglang/runtime.py index 49bc05a87f..1b24af04e9 100644 --- a/xinference/model/llm/sglang/runtime.py +++ b/xinference/model/llm/sglang/runtime.py @@ -1,14 +1,22 @@ # Copyright 2022-2026 Xinference Holdings Pte. Ltd # Licensed under the Apache License, Version 2.0. -"""Authenticate Runtime's startup metadata request as well as generation.""" +"""Prepare SGLang Runtime's cache path and authenticated startup requests.""" +import os import threading from functools import wraps _RUNTIME_START_LOCK = threading.Lock() +def prepare_jit_cache() -> None: + # SGLang compares resolved dependencies against its unresolved staging path. + cache_dir = os.environ.get("SGLANG_JIT_CACHE_DIR") or "~/.cache/sglang/jit" + os.environ["SGLANG_JIT_CACHE_DIR"] = os.path.realpath(os.path.expanduser(cache_dir)) + + def create_runtime(runtime_class, **config): + prepare_jit_cache() api_key = config.get("api_key") if not api_key: return runtime_class(**config) diff --git a/xinference/model/llm/sglang/tests/test_runtime.py b/xinference/model/llm/sglang/tests/test_runtime.py new file mode 100644 index 0000000000..85b12ef352 --- /dev/null +++ b/xinference/model/llm/sglang/tests/test_runtime.py @@ -0,0 +1,65 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. + +import os +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from ..runtime import create_runtime + + +@pytest.mark.parametrize("relative", [False, True]) +@pytest.mark.parametrize("authenticated", [False, True]) +def test_runtime_children_inherit_canonical_jit_cache( + tmp_path, monkeypatch, relative, authenticated +): + # SGLang records transient staging files as dependencies under symlinked roots. + physical = tmp_path / "physical" + physical.mkdir() + linked = tmp_path / "linked" + try: + linked.symlink_to(physical, target_is_directory=True) + except OSError: + pytest.skip("directory symlinks are unavailable") + monkeypatch.chdir(tmp_path) + configured = Path("linked/jit") if relative else linked / "jit" + monkeypatch.setenv("SGLANG_JIT_CACHE_DIR", str(configured)) + config = {"dtype": "float16"} + if authenticated: + config["api_key"] = "credential" + monkeypatch.setitem( + sys.modules, + "sglang.lang.backend", + SimpleNamespace(runtime_endpoint=SimpleNamespace(RuntimeEndpoint=object)), + ) + + def runtime(**kwargs): + assert kwargs == config + return subprocess.check_output( + [ + sys.executable, + "-c", + "import os; print(os.environ['SGLANG_JIT_CACHE_DIR'])", + ], + text=True, + timeout=30, + ).strip() + + assert create_runtime(runtime, **config) == str(physical / "jit") + assert not (physical / "jit").exists() + + +@pytest.mark.parametrize("configured", [None, "", "~/.cache/sglang/jit"]) +def test_runtime_default_jit_cache_is_expanded(monkeypatch, configured): + if configured is None: + # Track the absent key so teardown also restores direct runtime writes. + monkeypatch.setenv("SGLANG_JIT_CACHE_DIR", "unused") + monkeypatch.delenv("SGLANG_JIT_CACHE_DIR", raising=False) + else: + monkeypatch.setenv("SGLANG_JIT_CACHE_DIR", configured) + expected = str((Path.home() / ".cache" / "sglang" / "jit").resolve()) + assert create_runtime(lambda: os.environ.get("SGLANG_JIT_CACHE_DIR")) == expected diff --git a/xinference/model/llm/tests/test_weight_cache_engine_reload.py b/xinference/model/llm/tests/test_weight_cache_engine_reload.py index b57271f02f..070e481fd3 100644 --- a/xinference/model/llm/tests/test_weight_cache_engine_reload.py +++ b/xinference/model/llm/tests/test_weight_cache_engine_reload.py @@ -10,6 +10,7 @@ """ import asyncio +import os import sys from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -23,7 +24,9 @@ @pytest.mark.asyncio @pytest.mark.parametrize("backend", ["vllm", "sglang"]) -async def test_real_model_class_reloads_on_same_instance(monkeypatch, backend): +async def test_real_model_class_reloads_on_same_instance( + monkeypatch, tmp_path, backend +): family = match_llm("qwen2.5-instruct", "pytorch", "0_5", "none", "huggingface") engines = [] constructors = [] @@ -124,6 +127,16 @@ def create_engine_config(self): client_config=lambda: {}, engine_config=lambda config: config, stop=MagicMock() ) model._weight_cache = cache + if backend == "sglang": + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("SGLANG_JIT_CACHE_DIR", "jit-cache") + prepare_weight_cache = model._prepare_weight_cache + + def prepare(): + assert os.environ["SGLANG_JIT_CACHE_DIR"] == str(tmp_path / "jit-cache") + prepare_weight_cache() + + monkeypatch.setattr(model, "_prepare_weight_cache", prepare) try: await asyncio.to_thread(model.load) await asyncio.to_thread(model.wait_for_load)