[FEAT] Support fast engine recovery through weight cache (#27139)
Signed-off-by: Michael Qiu <qiudayu.qdy@antgroup.com> Co-authored-by: liusy58 <liusy58@linux.alibaba.com> Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
Mike Qiu
co-authored by
liusy58
Alex Nails
liusy58
Alex Nails
parent
bdd0698541
commit
f9c14e6bd4
@@ -0,0 +1,377 @@
|
||||
"""E2E test for WeightCacheDaemon with real model loading.
|
||||
|
||||
Launches TP daemons that load a real model, export IPC handles,
|
||||
and verifies the client can fetch and import them.
|
||||
|
||||
Usage:
|
||||
# With a small model (single GPU):
|
||||
python test/manual/test_weight_cache_e2e.py --model-path /path/to/model --tp-size 1
|
||||
|
||||
# With a large model (multi-GPU):
|
||||
python test/manual/test_weight_cache_e2e.py --model-path /path/to/model --tp-size 4
|
||||
|
||||
Requires GPUs and the model to be available on disk.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import socket
|
||||
import sys
|
||||
import time
|
||||
|
||||
|
||||
def find_free_port():
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def _global_rank(tp_size, pp_rank, tp_rank):
|
||||
return tp_size * pp_rank + tp_rank
|
||||
|
||||
|
||||
def _temp_path(tp_size, pp_rank, tp_rank, suffix):
|
||||
return f"/tmp/sglang_weight_cache_rank{_global_rank(tp_size, pp_rank, tp_rank)}{suffix}"
|
||||
|
||||
|
||||
def run_single_daemon(
|
||||
model_path,
|
||||
gpu_id,
|
||||
tp_size,
|
||||
tp_rank,
|
||||
pp_size,
|
||||
pp_rank,
|
||||
dist_init_method,
|
||||
socket_path,
|
||||
ready_path,
|
||||
done_path,
|
||||
load_format,
|
||||
dtype,
|
||||
quantization,
|
||||
trust_remote_code,
|
||||
):
|
||||
"""Run a single daemon process for one (pp_rank, tp_rank)."""
|
||||
import traceback
|
||||
|
||||
from sglang.srt.weight_cache.daemon import WeightCacheDaemon
|
||||
|
||||
daemon = WeightCacheDaemon(
|
||||
model_path=model_path,
|
||||
gpu_id=gpu_id,
|
||||
tp_size=tp_size,
|
||||
tp_rank=tp_rank,
|
||||
pp_size=pp_size,
|
||||
pp_rank=pp_rank,
|
||||
dp_size=1,
|
||||
load_format=load_format,
|
||||
dtype=dtype,
|
||||
quantization=quantization,
|
||||
trust_remote_code=trust_remote_code,
|
||||
dist_init_method=dist_init_method,
|
||||
)
|
||||
daemon.socket_path = socket_path
|
||||
daemon.ready_path = ready_path
|
||||
|
||||
try:
|
||||
daemon.load()
|
||||
except Exception as e:
|
||||
tb = traceback.format_exc()
|
||||
print(
|
||||
f"[Daemon gpu={gpu_id} pp_rank={pp_rank} tp_rank={tp_rank}] "
|
||||
f"LOAD ERROR: {e}\n{tb}",
|
||||
flush=True,
|
||||
)
|
||||
with open(ready_path, "w") as f:
|
||||
f.write(f"ERROR: {e}\n")
|
||||
return
|
||||
|
||||
with open(ready_path, "w") as f:
|
||||
f.write(f"pid={os.getpid()}\n")
|
||||
f.write(f"num_entries={len(daemon.state_entries)}\n")
|
||||
print(
|
||||
f"[Daemon gpu={gpu_id} pp_rank={pp_rank} tp_rank={tp_rank}] Ready with "
|
||||
f"{len(daemon.state_entries)} tensors",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
from sglang.srt.weight_cache.protocol import CacheConfig, recv_msg, send_msg
|
||||
|
||||
if os.path.exists(socket_path):
|
||||
os.unlink(socket_path)
|
||||
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
s.bind(socket_path)
|
||||
s.listen(5)
|
||||
s.settimeout(60)
|
||||
|
||||
try:
|
||||
while not os.path.exists(done_path):
|
||||
try:
|
||||
conn, _ = s.accept()
|
||||
try:
|
||||
req = recv_msg(conn)
|
||||
if req.get("type") == "query_config":
|
||||
send_msg(
|
||||
conn, {"status": "ok", "config": daemon.config.to_dict()}
|
||||
)
|
||||
elif req.get("type") == "fetch_state":
|
||||
engine_config = CacheConfig.from_dict(req["config"])
|
||||
if daemon.config.matches(engine_config):
|
||||
send_msg(
|
||||
conn,
|
||||
{
|
||||
"status": "ok",
|
||||
"config": daemon.config.to_dict(),
|
||||
"entries": daemon.state_entries,
|
||||
},
|
||||
)
|
||||
else:
|
||||
send_msg(
|
||||
conn,
|
||||
{
|
||||
"status": "mismatch",
|
||||
"daemon_config": daemon.config.to_dict(),
|
||||
},
|
||||
)
|
||||
elif req.get("type") == "ping":
|
||||
send_msg(conn, {"status": "ok"})
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
conn.close()
|
||||
except socket.timeout:
|
||||
continue
|
||||
finally:
|
||||
s.close()
|
||||
if os.path.exists(socket_path):
|
||||
os.unlink(socket_path)
|
||||
|
||||
print(
|
||||
f"[Daemon gpu={gpu_id} pp_rank={pp_rank} tp_rank={tp_rank}] Exiting",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="E2E test for WeightCacheDaemon")
|
||||
parser.add_argument("--model-path", required=True, help="Path to model weights")
|
||||
parser.add_argument("--tp-size", type=int, default=1, help="Tensor parallel size")
|
||||
parser.add_argument("--pp-size", type=int, default=1, help="Pipeline parallel size")
|
||||
parser.add_argument("--load-format", default="auto", help="Weight load format")
|
||||
parser.add_argument("--dtype", default="auto", help="Model dtype")
|
||||
parser.add_argument("--quantization", default=None, help="Quantization method")
|
||||
parser.add_argument("--trust-remote-code", action="store_true")
|
||||
parser.add_argument("--timeout", type=int, default=1800, help="Timeout in seconds")
|
||||
args = parser.parse_args()
|
||||
|
||||
tp_size = args.tp_size
|
||||
pp_size = args.pp_size
|
||||
total_ranks = tp_size * pp_size
|
||||
print(f"=== E2E Test: WeightCacheDaemon TP={tp_size} PP={pp_size} ===")
|
||||
print(f"Model: {args.model_path}")
|
||||
|
||||
dist_port = find_free_port()
|
||||
dist_init_method = f"tcp://127.0.0.1:{dist_port}"
|
||||
print(f"dist_init_method = {dist_init_method}")
|
||||
|
||||
# Clean up old files
|
||||
for pp_rank in range(pp_size):
|
||||
for tp_rank in range(tp_size):
|
||||
for suffix in [".sock", ".ready", ".done"]:
|
||||
p = _temp_path(tp_size, pp_rank, tp_rank, suffix)
|
||||
if os.path.exists(p):
|
||||
os.unlink(p)
|
||||
|
||||
mp.set_start_method("spawn", force=True)
|
||||
|
||||
# Launch all daemon processes
|
||||
procs = []
|
||||
for pp_rank in range(pp_size):
|
||||
for tp_rank in range(tp_size):
|
||||
gpu_id = pp_rank * tp_size + tp_rank
|
||||
sock_path = _temp_path(tp_size, pp_rank, tp_rank, ".sock")
|
||||
rdy_path = _temp_path(tp_size, pp_rank, tp_rank, ".ready")
|
||||
done_path = _temp_path(tp_size, pp_rank, tp_rank, ".done")
|
||||
|
||||
p = mp.Process(
|
||||
target=run_single_daemon,
|
||||
args=(
|
||||
args.model_path,
|
||||
gpu_id,
|
||||
tp_size,
|
||||
tp_rank,
|
||||
pp_size,
|
||||
pp_rank,
|
||||
dist_init_method,
|
||||
sock_path,
|
||||
rdy_path,
|
||||
done_path,
|
||||
args.load_format,
|
||||
args.dtype,
|
||||
args.quantization,
|
||||
args.trust_remote_code,
|
||||
),
|
||||
name=f"daemon_gpu{gpu_id}",
|
||||
)
|
||||
p.start()
|
||||
procs.append(p)
|
||||
print(
|
||||
f"Launched daemon gpu={gpu_id} pp_rank={pp_rank} "
|
||||
f"tp_rank={tp_rank} pid={p.pid}"
|
||||
)
|
||||
|
||||
# Wait for all daemons to be ready
|
||||
print("Waiting for all daemons to load model...")
|
||||
start = time.time()
|
||||
error_found = False
|
||||
for pp_rank in range(pp_size):
|
||||
for tp_rank in range(tp_size):
|
||||
ready_path = _temp_path(tp_size, pp_rank, tp_rank, ".ready")
|
||||
while not os.path.exists(ready_path):
|
||||
elapsed = time.time() - start
|
||||
if elapsed > args.timeout:
|
||||
print(
|
||||
f"ERROR: Daemon pp_rank={pp_rank} tp_rank={tp_rank} "
|
||||
f"timeout after {args.timeout}s"
|
||||
)
|
||||
error_found = True
|
||||
break
|
||||
for p in procs:
|
||||
if not p.is_alive() and not os.path.exists(ready_path):
|
||||
print(
|
||||
f"ERROR: Daemon pid={p.pid} exited with code {p.exitcode}"
|
||||
)
|
||||
error_found = True
|
||||
break
|
||||
if error_found:
|
||||
break
|
||||
time.sleep(2)
|
||||
|
||||
if error_found:
|
||||
break
|
||||
|
||||
with open(ready_path) as f:
|
||||
content = f.read()
|
||||
if content.startswith("ERROR:"):
|
||||
print(
|
||||
f"ERROR: Daemon pp_rank={pp_rank} tp_rank={tp_rank} "
|
||||
f"failed: {content.strip()}"
|
||||
)
|
||||
error_found = True
|
||||
break
|
||||
|
||||
print(
|
||||
f"Daemon pp_rank={pp_rank} tp_rank={tp_rank} ready "
|
||||
f"({time.time()-start:.0f}s)"
|
||||
)
|
||||
if error_found:
|
||||
break
|
||||
|
||||
if error_found:
|
||||
for p in procs:
|
||||
if p.is_alive():
|
||||
p.terminate()
|
||||
sys.exit(1)
|
||||
|
||||
print(
|
||||
f"\nAll {total_ranks} daemons ready! Total load time: {time.time()-start:.1f}s"
|
||||
)
|
||||
|
||||
# Query config from daemon (pp_rank=0, tp_rank=0)
|
||||
from sglang.srt.utils import MultiprocessingSerializer
|
||||
from sglang.srt.weight_cache.protocol import recv_msg, send_msg
|
||||
|
||||
socket_path_0 = _temp_path(tp_size, 0, 0, ".sock")
|
||||
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
s.settimeout(10)
|
||||
s.connect(socket_path_0)
|
||||
send_msg(s, {"type": "query_config"})
|
||||
result = recv_msg(s)
|
||||
s.close()
|
||||
daemon_config = result.get("config", {})
|
||||
print(
|
||||
f"\nDaemon (0,0) config: model={daemon_config.get('model_path')}, "
|
||||
f"arch={daemon_config.get('model_arch')}, "
|
||||
f"tp_size={daemon_config.get('tp_size')}, "
|
||||
f"dtype={daemon_config.get('dtype')}"
|
||||
)
|
||||
|
||||
# Fetch state from daemon (pp_rank=0, tp_rank=0)
|
||||
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
s.settimeout(300)
|
||||
s.connect(socket_path_0)
|
||||
send_msg(s, {"type": "fetch_state", "config": daemon_config})
|
||||
tic = time.perf_counter()
|
||||
result = recv_msg(s)
|
||||
fetch_time = time.perf_counter() - tic
|
||||
s.close()
|
||||
|
||||
print(
|
||||
f"\nFetch from daemon (0,0): status={result['status']}, time={fetch_time:.2f}s"
|
||||
)
|
||||
|
||||
if result["status"] == "ok":
|
||||
entries = result["entries"]
|
||||
print(f"Received {len(entries)} IPC handles from daemon (0,0)")
|
||||
|
||||
sample_names = list(entries.keys())[:5]
|
||||
for name in sample_names:
|
||||
entry = entries[name]
|
||||
imported = MultiprocessingSerializer.deserialize(entry["handle"])
|
||||
print(
|
||||
f" {name}: shape={tuple(imported.shape)}, "
|
||||
f"dtype={imported.dtype}, device={imported.device}"
|
||||
)
|
||||
del imported
|
||||
|
||||
print("\nIPC import OK!")
|
||||
else:
|
||||
print(f"ERROR: fetch_state returned {result}")
|
||||
for pp_rank in range(pp_size):
|
||||
for tp_rank in range(tp_size):
|
||||
done = _temp_path(tp_size, pp_rank, tp_rank, ".done")
|
||||
with open(done, "w") as f:
|
||||
f.write("done\n")
|
||||
for p in procs:
|
||||
if p.is_alive():
|
||||
p.terminate()
|
||||
sys.exit(1)
|
||||
|
||||
# Test config mismatch
|
||||
mismatch_config = dict(daemon_config)
|
||||
mismatch_config["tp_size"] = daemon_config.get("tp_size", 1) + 1
|
||||
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
s.settimeout(10)
|
||||
s.connect(socket_path_0)
|
||||
send_msg(s, {"type": "fetch_state", "config": mismatch_config})
|
||||
result = recv_msg(s)
|
||||
s.close()
|
||||
if result["status"] == "mismatch":
|
||||
print("Config mismatch detection: OK")
|
||||
else:
|
||||
print(f"WARNING: Expected mismatch, got {result['status']}")
|
||||
|
||||
# Signal all daemons done
|
||||
for pp_rank in range(pp_size):
|
||||
for tp_rank in range(tp_size):
|
||||
done = _temp_path(tp_size, pp_rank, tp_rank, ".done")
|
||||
with open(done, "w") as f:
|
||||
f.write("done\n")
|
||||
for p in procs:
|
||||
p.join(timeout=15)
|
||||
if p.is_alive():
|
||||
p.terminate()
|
||||
|
||||
# Clean up
|
||||
for pp_rank in range(pp_size):
|
||||
for tp_rank in range(tp_size):
|
||||
for suffix in [".sock", ".ready", ".done"]:
|
||||
path = _temp_path(tp_size, pp_rank, tp_rank, suffix)
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
print("\n=== E2E Test Passed! ===")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,338 @@
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
# A ~1B model keeps the daemon->client IPC handoff cheap to exercise on every
|
||||
# PR (fast download + load) while still covering the real block-load path; the
|
||||
# test asserts the IPC path ran, not any particular model's quality.
|
||||
DEFAULT_MODEL = "Qwen/Qwen3-0.6B"
|
||||
|
||||
# This file runs in two suites. The TP=2 class needs the 2-GPU runner (extra-a);
|
||||
# the TP=1 smoke class is always-on (base-b / 1-gpu-small) so the daemon->client
|
||||
# IPC handoff is exercised on every PR. Since the CI runner executes the whole
|
||||
# file per suite, TestWeightCacheDaemonTP2 self-skips when fewer than 2 GPUs are
|
||||
# visible (i.e. on the 1-gpu runner).
|
||||
register_cuda_ci(est_time=100, stage="extra-a", runner_config="2-gpu-large")
|
||||
register_cuda_ci(est_time=100, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
# Capture the client server's logs so test_loaded_via_ipc can assert the IPC
|
||||
# load path actually ran (and did not silently fall back to disk).
|
||||
STDOUT_FILENAME = "/tmp/test_weight_cache_daemon_stdout.log"
|
||||
STDERR_FILENAME = "/tmp/test_weight_cache_daemon_stderr.log"
|
||||
SMOKE_STDOUT_FILENAME = "/tmp/test_weight_cache_daemon_smoke_stdout.log"
|
||||
SMOKE_STDERR_FILENAME = "/tmp/test_weight_cache_daemon_smoke_stderr.log"
|
||||
|
||||
PROMPTS = [
|
||||
"The capital of France is",
|
||||
"Hello, my name is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
torch.cuda.device_count() < 2,
|
||||
"TP=2 weight cache daemon test requires >=2 GPUs (skipped on the 1-gpu runner)",
|
||||
)
|
||||
class TestWeightCacheDaemonTP2(CustomTestCase):
|
||||
"""E2E test: start weight cache daemons, then launch server in client mode with TP2."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.tp_size = 2
|
||||
|
||||
# Clean up stale ready/socket files from previous runs
|
||||
for rank in range(cls.tp_size):
|
||||
for suffix in (".ready", ".sock"):
|
||||
path = f"/tmp/sglang_weight_cache_rank{rank}{suffix}"
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
# Step 1: Launch weight cache daemons (blocks until all ranks are ready,
|
||||
# then monitors child processes)
|
||||
cls.daemon_process = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"sglang.srt.weight_cache.daemon",
|
||||
"--model-path",
|
||||
cls.model,
|
||||
"--tp-size",
|
||||
str(cls.tp_size),
|
||||
]
|
||||
)
|
||||
|
||||
# Step 2: Wait for all daemon ready files
|
||||
timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
|
||||
start = time.time()
|
||||
for rank in range(cls.tp_size):
|
||||
ready_path = f"/tmp/sglang_weight_cache_rank{rank}.ready"
|
||||
while not os.path.exists(ready_path):
|
||||
if time.time() - start > timeout:
|
||||
kill_process_tree(cls.daemon_process.pid)
|
||||
raise TimeoutError(
|
||||
f"Weight cache daemon rank {rank} not ready within {timeout}s"
|
||||
)
|
||||
if cls.daemon_process.poll() is not None:
|
||||
raise RuntimeError(
|
||||
f"Weight cache daemon exited prematurely "
|
||||
f"with code {cls.daemon_process.returncode}"
|
||||
)
|
||||
time.sleep(2)
|
||||
|
||||
# Step 3: Launch server in client mode — loads weights via IPC from daemons
|
||||
cls.stdout = open(STDOUT_FILENAME, "w")
|
||||
cls.stderr = open(STDERR_FILENAME, "w")
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--tp",
|
||||
str(cls.tp_size),
|
||||
"--weight-cache-mode",
|
||||
"client",
|
||||
],
|
||||
return_stdout_stderr=(cls.stdout, cls.stderr),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if hasattr(cls, "process") and cls.process:
|
||||
kill_process_tree(cls.process.pid)
|
||||
if hasattr(cls, "daemon_process") and cls.daemon_process:
|
||||
kill_process_tree(cls.daemon_process.pid)
|
||||
for stream in (getattr(cls, "stdout", None), getattr(cls, "stderr", None)):
|
||||
if stream is not None:
|
||||
try:
|
||||
stream.close()
|
||||
except OSError:
|
||||
pass
|
||||
for path in (STDOUT_FILENAME, STDERR_FILENAME):
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
for rank in range(getattr(cls, "tp_size", 2)):
|
||||
for suffix in (".ready", ".sock"):
|
||||
path = f"/tmp/sglang_weight_cache_rank{rank}{suffix}"
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def test_generate(self):
|
||||
for prompt in PROMPTS:
|
||||
resp = requests.post(
|
||||
f"{self.base_url}/v1/completions",
|
||||
json={
|
||||
"model": self.model,
|
||||
"prompt": prompt,
|
||||
"max_tokens": 32,
|
||||
"temperature": 0,
|
||||
},
|
||||
)
|
||||
self.assertEqual(resp.status_code, 200)
|
||||
data = resp.json()
|
||||
text = data["choices"][0]["text"]
|
||||
self.assertIsInstance(text, str)
|
||||
self.assertGreater(len(text), 0, f"Empty output for prompt: {prompt}")
|
||||
|
||||
def test_chat(self):
|
||||
resp = requests.post(
|
||||
f"{self.base_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": "What is 2+3?"}],
|
||||
"max_tokens": 32,
|
||||
"temperature": 0,
|
||||
},
|
||||
)
|
||||
self.assertEqual(resp.status_code, 200)
|
||||
data = resp.json()
|
||||
content = data["choices"][0]["message"]["content"]
|
||||
self.assertIsInstance(content, str)
|
||||
self.assertGreater(len(content), 0)
|
||||
|
||||
def test_loaded_via_ipc(self):
|
||||
"""Assert the server actually loaded weights over IPC.
|
||||
|
||||
Without this, the test would still pass if the IPC path silently
|
||||
regressed to disk loading (the daemon would just sit unused), because
|
||||
generation output looks identical either way. The daemon-side loader
|
||||
logs "[IpcModelLoader] Loaded model via IPC" on every rank, so its
|
||||
presence in the captured server logs is our proof the IPC path ran.
|
||||
"""
|
||||
for stream in (getattr(self, "stdout", None), getattr(self, "stderr", None)):
|
||||
if stream is not None:
|
||||
try:
|
||||
stream.flush()
|
||||
except OSError:
|
||||
pass
|
||||
logs = ""
|
||||
for path in (STDOUT_FILENAME, STDERR_FILENAME):
|
||||
if os.path.exists(path):
|
||||
with open(path, errors="replace") as f:
|
||||
logs += f.read()
|
||||
self.assertIn(
|
||||
"Loaded model via IPC",
|
||||
logs,
|
||||
"Expected the client server to load weights via IPC, but the IPC "
|
||||
"load log line was not found — the loader likely fell back to disk.",
|
||||
)
|
||||
|
||||
|
||||
class TestWeightCacheDaemonTP1Smoke(CustomTestCase):
|
||||
"""Always-on TP=1 smoke: start a single weight cache daemon, launch a server
|
||||
in client mode, and confirm it loads weights via IPC and generates.
|
||||
|
||||
This is the fast single-GPU sanity check (small model) that runs on every PR
|
||||
in the base-b / 1-gpu-small suite; the heavier TP=2 case above only runs on
|
||||
the 2-GPU runner.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.tp_size = 1
|
||||
|
||||
# Clean up stale ready/socket files from previous runs.
|
||||
for rank in range(cls.tp_size):
|
||||
for suffix in (".ready", ".sock"):
|
||||
path = f"/tmp/sglang_weight_cache_rank{rank}{suffix}"
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
# Step 1: Launch the weight cache daemon (blocks until the rank is
|
||||
# ready, then monitors the child process).
|
||||
cls.daemon_process = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"sglang.srt.weight_cache.daemon",
|
||||
"--model-path",
|
||||
cls.model,
|
||||
"--tp-size",
|
||||
str(cls.tp_size),
|
||||
]
|
||||
)
|
||||
|
||||
# Step 2: Wait for the daemon ready file.
|
||||
timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
|
||||
start = time.time()
|
||||
for rank in range(cls.tp_size):
|
||||
ready_path = f"/tmp/sglang_weight_cache_rank{rank}.ready"
|
||||
while not os.path.exists(ready_path):
|
||||
if time.time() - start > timeout:
|
||||
kill_process_tree(cls.daemon_process.pid)
|
||||
raise TimeoutError(
|
||||
f"Weight cache daemon rank {rank} not ready within {timeout}s"
|
||||
)
|
||||
if cls.daemon_process.poll() is not None:
|
||||
raise RuntimeError(
|
||||
f"Weight cache daemon exited prematurely "
|
||||
f"with code {cls.daemon_process.returncode}"
|
||||
)
|
||||
time.sleep(2)
|
||||
|
||||
# Step 3: Launch server in client mode — loads weights via IPC.
|
||||
cls.stdout = open(SMOKE_STDOUT_FILENAME, "w")
|
||||
cls.stderr = open(SMOKE_STDERR_FILENAME, "w")
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--tp",
|
||||
str(cls.tp_size),
|
||||
"--weight-cache-mode",
|
||||
"client",
|
||||
],
|
||||
return_stdout_stderr=(cls.stdout, cls.stderr),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if hasattr(cls, "process") and cls.process:
|
||||
kill_process_tree(cls.process.pid)
|
||||
if hasattr(cls, "daemon_process") and cls.daemon_process:
|
||||
kill_process_tree(cls.daemon_process.pid)
|
||||
for stream in (getattr(cls, "stdout", None), getattr(cls, "stderr", None)):
|
||||
if stream is not None:
|
||||
try:
|
||||
stream.close()
|
||||
except OSError:
|
||||
pass
|
||||
for path in (SMOKE_STDOUT_FILENAME, SMOKE_STDERR_FILENAME):
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
for rank in range(getattr(cls, "tp_size", 1)):
|
||||
for suffix in (".ready", ".sock"):
|
||||
path = f"/tmp/sglang_weight_cache_rank{rank}{suffix}"
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def test_generate(self):
|
||||
resp = requests.post(
|
||||
f"{self.base_url}/v1/completions",
|
||||
json={
|
||||
"model": self.model,
|
||||
"prompt": "The capital of France is",
|
||||
"max_tokens": 32,
|
||||
"temperature": 0,
|
||||
},
|
||||
)
|
||||
self.assertEqual(resp.status_code, 200)
|
||||
data = resp.json()
|
||||
text = data["choices"][0]["text"]
|
||||
self.assertIsInstance(text, str)
|
||||
self.assertGreater(len(text), 0, "Empty generation output")
|
||||
|
||||
def test_loaded_via_ipc(self):
|
||||
"""Assert the server actually loaded weights over IPC (see the TP=2
|
||||
variant for why this guard matters)."""
|
||||
for stream in (getattr(self, "stdout", None), getattr(self, "stderr", None)):
|
||||
if stream is not None:
|
||||
try:
|
||||
stream.flush()
|
||||
except OSError:
|
||||
pass
|
||||
logs = ""
|
||||
for path in (SMOKE_STDOUT_FILENAME, SMOKE_STDERR_FILENAME):
|
||||
if os.path.exists(path):
|
||||
with open(path, errors="replace") as f:
|
||||
logs += f.read()
|
||||
self.assertIn(
|
||||
"Loaded model via IPC",
|
||||
logs,
|
||||
"Expected the client server to load weights via IPC, but the IPC "
|
||||
"load log line was not found — the loader likely fell back to disk.",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,366 @@
|
||||
"""
|
||||
CPU-only unit tests for the weight cache protocol layer.
|
||||
|
||||
These cover the pure-Python logic that the GPU end-to-end test
|
||||
(test_weight_cache_daemon.py) cannot exercise cheaply:
|
||||
|
||||
- length-prefixed socket framing (send_msg/recv_msg) over socketpair()
|
||||
- CacheConfig fingerprint matching / (de)serialization
|
||||
- quant-config hashing and method-name extraction
|
||||
- the daemon rank formula and socket/ready path derivation
|
||||
- the IPC quantization allowlist (the gate that keeps silently-wrong
|
||||
quant methods off the zero-copy path)
|
||||
- stale-vs-live daemon file cleanup
|
||||
|
||||
They intentionally require no CUDA, no model download, and no daemon
|
||||
process, so they run in the fast CPU suite and would catch a regression
|
||||
in any of these branches before it reaches the expensive GPU path.
|
||||
"""
|
||||
|
||||
import os
|
||||
import socket
|
||||
import struct
|
||||
import unittest
|
||||
|
||||
from sglang.srt.weight_cache.protocol import (
|
||||
IPC_QUANT_ALLOWLIST,
|
||||
CacheConfig,
|
||||
UnsupportedQuantForIPCError,
|
||||
check_ipc_quant_support,
|
||||
cleanup_stale_daemon_files,
|
||||
compute_global_rank,
|
||||
compute_local_gpu_id,
|
||||
get_quant_method_name,
|
||||
get_ready_path,
|
||||
get_socket_path,
|
||||
hash_quant_config,
|
||||
is_ipc_quant_supported,
|
||||
recv_msg,
|
||||
send_msg,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _make_cache_config(**overrides) -> CacheConfig:
|
||||
base = dict(
|
||||
model_path="/models/demo",
|
||||
model_arch="LlamaForCausalLM",
|
||||
tp_size=2,
|
||||
tp_rank=0,
|
||||
pp_size=1,
|
||||
pp_rank=0,
|
||||
dp_size=1,
|
||||
ep_size=1,
|
||||
quant_method="",
|
||||
quant_config_hash="",
|
||||
dtype="torch.float16",
|
||||
revision="",
|
||||
device_capability="8.0",
|
||||
torch_version="2.5.1",
|
||||
)
|
||||
base.update(overrides)
|
||||
return CacheConfig(**base)
|
||||
|
||||
|
||||
class TestProtocolFraming(CustomTestCase):
|
||||
"""Length-prefixed pickle framing over a real socket pair."""
|
||||
|
||||
def test_round_trip(self):
|
||||
a, b = socket.socketpair()
|
||||
try:
|
||||
payload = {"handles": [1, 2, 3], "meta": ("x", 4.5), "flag": True}
|
||||
send_msg(a, payload)
|
||||
self.assertEqual(recv_msg(b), payload)
|
||||
finally:
|
||||
a.close()
|
||||
b.close()
|
||||
|
||||
def test_multiple_messages_are_framed_independently(self):
|
||||
a, b = socket.socketpair()
|
||||
try:
|
||||
send_msg(a, {"n": 1})
|
||||
send_msg(a, {"n": 2})
|
||||
self.assertEqual(recv_msg(b), {"n": 1})
|
||||
self.assertEqual(recv_msg(b), {"n": 2})
|
||||
finally:
|
||||
a.close()
|
||||
b.close()
|
||||
|
||||
def test_connection_closed_mid_header_raises(self):
|
||||
a, b = socket.socketpair()
|
||||
try:
|
||||
# Peer sends only a partial header, then hangs up.
|
||||
a.sendall(struct.pack("!I", 128)[:2])
|
||||
a.close()
|
||||
with self.assertRaises(ConnectionError):
|
||||
recv_msg(b)
|
||||
finally:
|
||||
b.close()
|
||||
|
||||
def test_connection_closed_mid_body_raises(self):
|
||||
a, b = socket.socketpair()
|
||||
try:
|
||||
# Full header promising 128 bytes, but no body follows.
|
||||
a.sendall(struct.pack("!I", 128))
|
||||
a.close()
|
||||
with self.assertRaises(ConnectionError):
|
||||
recv_msg(b)
|
||||
finally:
|
||||
b.close()
|
||||
|
||||
|
||||
class TestCacheConfig(CustomTestCase):
|
||||
def test_identical_configs_match(self):
|
||||
self.assertTrue(_make_cache_config().matches(_make_cache_config()))
|
||||
|
||||
def test_any_field_difference_breaks_match(self):
|
||||
base = _make_cache_config()
|
||||
for field, value in (
|
||||
("tp_rank", 1),
|
||||
("dtype", "torch.bfloat16"),
|
||||
("quant_method", "fp8"),
|
||||
("model_path", "/models/other"),
|
||||
("revision", "v2"),
|
||||
("device_capability", "9.0"),
|
||||
("torch_version", "2.4.0"),
|
||||
):
|
||||
self.assertFalse(
|
||||
base.matches(_make_cache_config(**{field: value})),
|
||||
msg=f"{field} difference should break match",
|
||||
)
|
||||
|
||||
def test_dict_round_trip(self):
|
||||
cfg = _make_cache_config(quant_method="fp8", quant_config_hash="abc123")
|
||||
restored = CacheConfig.from_dict(cfg.to_dict())
|
||||
self.assertTrue(cfg.matches(restored))
|
||||
self.assertEqual(cfg.to_dict(), restored.to_dict())
|
||||
|
||||
|
||||
class TestQuantConfigHashing(CustomTestCase):
|
||||
def test_none_hashes_to_empty(self):
|
||||
self.assertEqual(hash_quant_config(None), "")
|
||||
|
||||
def test_dict_hash_is_deterministic_and_order_insensitive(self):
|
||||
h1 = hash_quant_config({"bits": 8, "group_size": 128})
|
||||
h2 = hash_quant_config({"group_size": 128, "bits": 8})
|
||||
self.assertEqual(h1, h2)
|
||||
self.assertNotEqual(h1, hash_quant_config({"bits": 4, "group_size": 128}))
|
||||
|
||||
def test_hash_is_not_truncated(self):
|
||||
# A correctness gate must use the full SHA-256 digest, not a 16-char prefix.
|
||||
self.assertEqual(len(hash_quant_config({"bits": 8})), 64)
|
||||
|
||||
def test_hash_does_not_embed_object_address(self):
|
||||
# Two distinct instances with identical public attrs must hash equal,
|
||||
# otherwise configs would mismatch across processes (the bug the
|
||||
# docstring warns about).
|
||||
class _Q:
|
||||
def __init__(self):
|
||||
self.bits = 8
|
||||
self.method = "fp8"
|
||||
|
||||
self.assertEqual(hash_quant_config(_Q()), hash_quant_config(_Q()))
|
||||
|
||||
def test_get_quant_method_name_variants(self):
|
||||
self.assertEqual(get_quant_method_name(None), "")
|
||||
self.assertEqual(get_quant_method_name("fp8"), "fp8")
|
||||
|
||||
class _WithGetName:
|
||||
def get_name(self):
|
||||
return "gptq_marlin"
|
||||
|
||||
class _WithName:
|
||||
name = "awq"
|
||||
|
||||
self.assertEqual(get_quant_method_name(_WithGetName()), "gptq_marlin")
|
||||
self.assertEqual(get_quant_method_name(_WithName()), "awq")
|
||||
|
||||
|
||||
class TestGlobalRankAndPaths(CustomTestCase):
|
||||
def test_compute_global_rank_formula(self):
|
||||
self.assertEqual(compute_global_rank(tp_size=4, pp_rank=0, tp_rank=3), 3)
|
||||
self.assertEqual(compute_global_rank(tp_size=4, pp_rank=1, tp_rank=0), 4)
|
||||
self.assertEqual(compute_global_rank(tp_size=4, pp_rank=2, tp_rank=1), 9)
|
||||
|
||||
def test_socket_and_ready_paths_are_unique_per_rank(self):
|
||||
self.assertNotEqual(get_socket_path(0), get_socket_path(1))
|
||||
self.assertTrue(get_socket_path(3).endswith("rank3.sock"))
|
||||
self.assertTrue(get_ready_path(3).endswith("rank3.ready"))
|
||||
|
||||
def test_compute_local_gpu_id_honors_base_and_step(self):
|
||||
# Single-node TP=4: identity mapping rank -> gpu.
|
||||
self.assertEqual(
|
||||
compute_local_gpu_id(0, 2, pp_size_per_node=1, tp_size_per_node=4),
|
||||
2,
|
||||
)
|
||||
# base_gpu_id offsets every rank; gpu_id_step strides between them.
|
||||
self.assertEqual(
|
||||
compute_local_gpu_id(
|
||||
0, 2, pp_size_per_node=1, tp_size_per_node=4, base_gpu_id=4
|
||||
),
|
||||
6,
|
||||
)
|
||||
self.assertEqual(
|
||||
compute_local_gpu_id(
|
||||
0, 2, pp_size_per_node=1, tp_size_per_node=4, gpu_id_step=2
|
||||
),
|
||||
4,
|
||||
)
|
||||
|
||||
|
||||
class TestIpcQuantAllowlist(CustomTestCase):
|
||||
def test_unquantized_is_supported(self):
|
||||
self.assertTrue(is_ipc_quant_supported("", None))
|
||||
|
||||
def test_block_fp8_supported_but_per_tensor_fp8_rejected(self):
|
||||
self.assertTrue(
|
||||
is_ipc_quant_supported("fp8", {"weight_block_size": [128, 128]})
|
||||
)
|
||||
# Per-tensor FP8 (no weight_block_size) transposes the weight during
|
||||
# post-processing -> not reproducible by the meta-init client.
|
||||
self.assertFalse(is_ipc_quant_supported("fp8", {}))
|
||||
self.assertFalse(is_ipc_quant_supported("fp8", None))
|
||||
|
||||
def test_unknown_method_rejected(self):
|
||||
self.assertFalse(is_ipc_quant_supported("gptq_marlin", None))
|
||||
self.assertFalse(is_ipc_quant_supported("awq", None))
|
||||
|
||||
def test_check_raises_on_unsupported(self):
|
||||
with self.assertRaises(UnsupportedQuantForIPCError):
|
||||
check_ipc_quant_support("awq", None, where="client")
|
||||
# Per-tensor FP8 must also raise even though "fp8" is a known key.
|
||||
with self.assertRaises(UnsupportedQuantForIPCError):
|
||||
check_ipc_quant_support("fp8", {}, where="daemon")
|
||||
|
||||
def test_check_passes_on_supported(self):
|
||||
# Should not raise.
|
||||
check_ipc_quant_support("", None, where="daemon")
|
||||
check_ipc_quant_support(
|
||||
"fp8", {"weight_block_size": [128, 128]}, where="daemon"
|
||||
)
|
||||
|
||||
def test_allowlist_registry_shape(self):
|
||||
# Guard against accidentally widening the allowlist without review.
|
||||
self.assertEqual(set(IPC_QUANT_ALLOWLIST), {"", "fp8"})
|
||||
|
||||
|
||||
class TestCleanupStaleDaemonFiles(CustomTestCase):
|
||||
# Use a rank far outside any realistic tp*pp layout so we never collide
|
||||
# with a daemon that might actually be running on the test host.
|
||||
RANK = 987654
|
||||
|
||||
def _paths(self):
|
||||
return get_ready_path(self.RANK), get_socket_path(self.RANK)
|
||||
|
||||
def tearDown(self):
|
||||
for path in self._paths():
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
def test_no_files_is_noop(self):
|
||||
# Neither file present: must return quietly, not raise.
|
||||
cleanup_stale_daemon_files(self.RANK)
|
||||
|
||||
def test_stale_files_without_live_pid_are_removed(self):
|
||||
ready_path, socket_path = self._paths()
|
||||
# A .ready file whose PID is unreadable is treated as a crashed-daemon
|
||||
# leftover and cleaned up.
|
||||
with open(ready_path, "w") as f:
|
||||
f.write("stale contents, no pid line\n")
|
||||
open(socket_path, "w").close()
|
||||
|
||||
cleanup_stale_daemon_files(self.RANK)
|
||||
|
||||
self.assertFalse(os.path.exists(ready_path))
|
||||
self.assertFalse(os.path.exists(socket_path))
|
||||
|
||||
def test_live_daemon_pid_blocks_cleanup(self):
|
||||
ready_path, socket_path = self._paths()
|
||||
# Our own PID is alive -> cleanup must refuse and leave files intact.
|
||||
with open(ready_path, "w") as f:
|
||||
f.write(f"pid={os.getpid()}\n")
|
||||
open(socket_path, "w").close()
|
||||
|
||||
with self.assertRaises(RuntimeError):
|
||||
cleanup_stale_daemon_files(self.RANK)
|
||||
|
||||
self.assertTrue(os.path.exists(ready_path))
|
||||
self.assertTrue(os.path.exists(socket_path))
|
||||
|
||||
def test_force_takes_over_from_live_pid(self):
|
||||
ready_path, socket_path = self._paths()
|
||||
# Spawn a real child we are allowed to kill, point the ready file at it,
|
||||
# then force-takeover: the child must be killed and the files removed.
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
|
||||
try:
|
||||
with open(ready_path, "w") as f:
|
||||
f.write(f"pid={child.pid}\n")
|
||||
open(socket_path, "w").close()
|
||||
|
||||
cleanup_stale_daemon_files(self.RANK, force=True)
|
||||
|
||||
self.assertFalse(os.path.exists(ready_path))
|
||||
self.assertFalse(os.path.exists(socket_path))
|
||||
# The daemon holding the rank must have been killed.
|
||||
self.assertEqual(child.wait(timeout=5), -9)
|
||||
finally:
|
||||
if child.poll() is None:
|
||||
child.kill()
|
||||
child.wait(timeout=5)
|
||||
|
||||
|
||||
class TestDaemonModeRefusesDiskLoad(CustomTestCase):
|
||||
"""In daemon mode the engine and daemon share a GPU, so a missing daemon
|
||||
must be a hard error — NOT a silent disk-load that would OOM the shared
|
||||
device. This exercises that contract without a GPU or a live daemon by
|
||||
pointing the loader at a socket path that does not exist.
|
||||
"""
|
||||
|
||||
RANK = 987655
|
||||
|
||||
def _model_config(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
# Minimal stand-in: the loader only reads hf_config.quantization_config,
|
||||
# quantization, and (unreached here) hf_config.architectures.
|
||||
hf_config = SimpleNamespace(
|
||||
architectures=["LlamaForCausalLM"], quantization_config=None
|
||||
)
|
||||
return SimpleNamespace(
|
||||
model_path="/models/demo",
|
||||
hf_config=hf_config,
|
||||
quantization=None,
|
||||
revision=None,
|
||||
dtype="torch.float16",
|
||||
)
|
||||
|
||||
def test_daemon_mode_missing_daemon_raises_instead_of_disk_load(self):
|
||||
from sglang.srt.configs.load_config import LoadConfig, LoadFormat
|
||||
from sglang.srt.weight_cache.ipc_loader import IpcModelLoader
|
||||
|
||||
missing_socket = get_socket_path(self.RANK)
|
||||
if os.path.exists(missing_socket):
|
||||
os.unlink(missing_socket)
|
||||
|
||||
loader = IpcModelLoader(
|
||||
load_config=LoadConfig(load_format=LoadFormat.IPC_CACHE),
|
||||
socket_path=missing_socket,
|
||||
weight_cache_mode="daemon",
|
||||
fallback_load_format="auto",
|
||||
)
|
||||
|
||||
with self.assertRaises(RuntimeError) as ctx:
|
||||
loader.load_model(model_config=self._model_config(), device_config=None)
|
||||
# The error must be about the missing daemon, proving we did not quietly
|
||||
# fall through to a disk load.
|
||||
self.assertIn("daemon", str(ctx.exception).lower())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user