minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

weight cache: key daemon paths by GPU UUID (#36101)

Browse Source 
Co-authored-by: siyu <liusy58@linux.alibaba.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
Tarang Khanna
2026-08-31 01:44:45 -07:00
committed by GitHub
co-authored by siyu Alex Nails
parent 3865efc9f7
commit 6580d5cd9a
11 changed files with 176 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/model_loading/test_weight_cache_daemon.py
+25 -18
View File
@@ -7,7 +7,9 @@ import unittest
import requests
import torch
from sglang.srt.platforms import current_platform
from sglang.srt.utils import kill_process_tree
from sglang.srt.weight_cache.protocol import get_ready_path, get_socket_path
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_TARGET_MODEL_EAGLE_DP_ATTN,
@@ -44,6 +46,11 @@ PROMPTS = [
]
def _gpu_uuids(tp_size: int) -> list:
# Single-node, default base_gpu_id/gpu_id_step: rank i runs on physical GPU i.
return [current_platform.get_device_uuid(i) for i in range(tp_size)]
@unittest.skipIf(
torch.cuda.device_count() < 2,
"TP=2 weight cache daemon test requires >=2 GPUs (skipped on the 1-gpu runner)",
@@ -60,11 +67,11 @@ class TestWeightCacheDaemonTP2(CustomTestCase):
cls.model = cls.model_override or DEFAULT_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.tp_size = 2
cls.gpu_uuids = _gpu_uuids(cls.tp_size)
# 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}"
for device_uuid in cls.gpu_uuids:
for path in (get_ready_path(device_uuid), get_socket_path(device_uuid)):
if os.path.exists(path):
os.unlink(path)
@@ -86,13 +93,14 @@ class TestWeightCacheDaemonTP2(CustomTestCase):
# 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"
for device_uuid in cls.gpu_uuids:
ready_path = get_ready_path(device_uuid)
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"
f"Weight cache daemon for GPU {device_uuid} not ready "
f"within {timeout}s"
)
if cls.daemon_process.poll() is not None:
raise RuntimeError(
@@ -136,9 +144,8 @@ class TestWeightCacheDaemonTP2(CustomTestCase):
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}"
for device_uuid in getattr(cls, "gpu_uuids", ()):
for path in (get_ready_path(device_uuid), get_socket_path(device_uuid)):
if os.path.exists(path):
try:
os.unlink(path)
@@ -220,11 +227,11 @@ class TestWeightCacheDaemonTP1Smoke(CustomTestCase):
cls.model = DEFAULT_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.tp_size = 1
cls.gpu_uuids = _gpu_uuids(cls.tp_size)
# 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}"
for device_uuid in cls.gpu_uuids:
for path in (get_ready_path(device_uuid), get_socket_path(device_uuid)):
if os.path.exists(path):
os.unlink(path)
@@ -245,13 +252,14 @@ class TestWeightCacheDaemonTP1Smoke(CustomTestCase):
# 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"
for device_uuid in cls.gpu_uuids:
ready_path = get_ready_path(device_uuid)
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"
f"Weight cache daemon for GPU {device_uuid} not ready "
f"within {timeout}s"
)
if cls.daemon_process.poll() is not None:
raise RuntimeError(
@@ -294,9 +302,8 @@ class TestWeightCacheDaemonTP1Smoke(CustomTestCase):
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}"
for device_uuid in getattr(cls, "gpu_uuids", ()):
for path in (get_ready_path(device_uuid), get_socket_path(device_uuid)):
if os.path.exists(path):
try:
os.unlink(path)