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

Merge /get_load into /v1/loads (#23010)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-04-17 13:36:51 -07:00
committed by GitHub
parent 44e67c6835
commit 573e12a7fc
13 changed files with 235 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/ascend/interface/test_npu_api.py
+12 -9
View File
@@ -30,7 +30,7 @@ class TestNpuApi(CustomTestCase):
"""Testcase: Verify that the basic functions of the API interfaces work properly and the returned parameters are consistent with the configurations.
[Test Category] Interface
[Test Target] /health; /health_generate; /ping; /model_info; /server_info; /get_load; /v1/models; /v1/models/{model:path}; /generate
[Test Target] /health; /health_generate; /ping; /model_info; /server_info; /v1/loads; /v1/models; /v1/models/{model:path}; /generate
"""
@classmethod
@@ -90,15 +90,18 @@ class TestNpuApi(CustomTestCase):
self.assertEqual(response.json()["model_path"], self.model)
self.assertEqual(response.json()["tokenizer_path"], self.model)
def test_api_get_load(self):
response = requests.get(f"{self.base_url}/get_load")
def test_api_v1_loads(self):
response = requests.get(f"{self.base_url}/v1/loads")
self.assertEqual(response.status_code, 200)
self.assertIsNone(response.json()[0]["rid"])
self.assertIsNone(response.json()[0]["http_worker_ipc"])
self.assertIsNone(response.json()[0]["dp_rank"])
self.assertGreaterEqual(response.json()[0]["num_reqs"], 0)
self.assertGreaterEqual(response.json()[0]["num_waiting_reqs"], 0)
self.assertGreaterEqual(response.json()[0]["num_tokens"], 0)
body = response.json()
self.assertIn("loads", body)
self.assertIn("aggregate", body)
self.assertGreaterEqual(len(body["loads"]), 1)
load = body["loads"][0]
self.assertGreaterEqual(load["num_running_reqs"], 0)
self.assertGreaterEqual(load["num_waiting_reqs"], 0)
self.assertGreaterEqual(load["num_used_tokens"], 0)
self.assertGreaterEqual(load["num_total_tokens"], 0)
def test_api_v1_models(self):
response = requests.get(f"{self.base_url}/v1/models")
test/registered/unit/entrypoints/test_v1_loads_aggregate.py
+76
View File
@@ -0,0 +1,76 @@
"""Unit tests for /v1/loads _compute_aggregate.
Narrow scope: lock in the semantic of new aggregate keys added by this PR
(total_used_tokens vs total_tokens). Trivial helpers (dict filtering,
zero-init branch) are not covered — they would just restate Python.
"""
import unittest
from sglang.srt.entrypoints.v1_loads import _compute_aggregate
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
def _load(
*,
dp_rank=0,
running=0,
waiting=0,
used=0,
total=0,
token_usage=0.0,
throughput=0.0,
utilization=0.0,
):
return {
"dp_rank": dp_rank,
"num_running_reqs": running,
"num_waiting_reqs": waiting,
"num_used_tokens": used,
"num_total_tokens": total,
"token_usage": token_usage,
"gen_throughput": throughput,
"utilization": utilization,
}
class TestComputeAggregate(CustomTestCase):
def test_multi_dp_rank_sums(self):
agg = _compute_aggregate(
[
_load(dp_rank=0, running=3, waiting=1, used=50, total=70),
_load(dp_rank=1, running=5, waiting=2, used=80, total=100),
_load(dp_rank=2, running=0, waiting=4, used=0, total=40),
]
)
self.assertEqual(agg["total_running_reqs"], 8)
self.assertEqual(agg["total_waiting_reqs"], 7)
self.assertEqual(agg["total_reqs"], 15)
self.assertEqual(agg["total_used_tokens"], 130)
self.assertEqual(agg["total_tokens"], 210)
def test_averages_over_dp_count(self):
agg = _compute_aggregate(
[
_load(token_usage=0.6, throughput=100.0, utilization=0.5),
_load(token_usage=0.8, throughput=200.0, utilization=0.7),
]
)
self.assertAlmostEqual(agg["avg_token_usage"], 0.7)
self.assertAlmostEqual(agg["avg_throughput"], 150.0)
self.assertAlmostEqual(agg["avg_utilization"], 0.6)
def test_total_tokens_differs_from_total_used_tokens(self):
# Regression: total_tokens sums num_total_tokens, NOT num_used_tokens.
# Gateway reads aggregate.total_tokens for DP load estimation, so a
# silent swap would under-report load.
agg = _compute_aggregate([_load(used=10, total=30), _load(used=20, total=45)])
self.assertEqual(agg["total_used_tokens"], 30)
self.assertEqual(agg["total_tokens"], 75)
if __name__ == "__main__":
unittest.main()
test/registered/unit/managers/test_dp_budget.py
+91
View File
@@ -0,0 +1,91 @@
"""Unit tests for DPBudget — field mapping regression guard.
This PR changed DPBudget.update_budget to read num_running_reqs +
num_waiting_reqs and num_total_tokens from the new GetLoadsReqOutput.
These tests lock in that mapping. Pre-existing dispatch logic is not
retested here — it's covered by DP balance integration tests.
"""
import dataclasses
import unittest
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.managers.data_parallel_controller import DPBudget
from sglang.srt.managers.io_struct import GetLoadsReqOutput, WatchLoadUpdateReq
register_cpu_ci(est_time=2, suite="stage-a-test-cpu")
_BASE_LOAD = GetLoadsReqOutput(
dp_rank=0,
timestamp=0.0,
num_running_reqs=0,
num_waiting_reqs=0,
num_used_tokens=0,
num_total_tokens=0,
max_total_num_tokens=4096,
token_usage=0.0,
gen_throughput=0.0,
cache_hit_rate=0.0,
utilization=0.0,
max_running_requests=128,
)
def _load(**overrides) -> GetLoadsReqOutput:
return dataclasses.replace(_BASE_LOAD, **overrides)
class TestDPBudgetUpdateBudget(CustomTestCase):
def test_maps_running_plus_waiting_to_total_requests(self):
budget = DPBudget(dp_size=2)
budget.update_budget(
WatchLoadUpdateReq(
loads=[
_load(dp_rank=0, num_running_reqs=3, num_waiting_reqs=2),
_load(dp_rank=1, num_running_reqs=5, num_waiting_reqs=1),
]
)
)
self.assertEqual(budget.total_requests, [5, 6])
def test_maps_num_total_tokens_not_num_used_tokens(self):
# Reads num_total_tokens (used + pending prefill), NOT num_used_tokens.
# A silent swap here would break DP balance for long-prompt workloads.
budget = DPBudget(dp_size=2)
budget.update_budget(
WatchLoadUpdateReq(
loads=[
_load(dp_rank=0, num_used_tokens=100, num_total_tokens=150),
_load(dp_rank=1, num_used_tokens=80, num_total_tokens=80),
]
)
)
self.assertEqual(budget.total_tokens, [150, 80])
def test_partial_update_only_affects_reported_rank(self):
budget = DPBudget(dp_size=3)
budget.total_requests = [10, 20, 30]
budget.total_tokens = [100, 200, 300]
budget.update_budget(
WatchLoadUpdateReq(
loads=[
_load(
dp_rank=1,
num_running_reqs=1,
num_waiting_reqs=1,
num_total_tokens=50,
)
]
)
)
self.assertEqual(budget.total_requests, [10, 2, 30])
self.assertEqual(budget.total_tokens, [100, 50, 300])
if __name__ == "__main__":
unittest.main()
test/registered/utils/test_type_based_dispatcher.py
+3 -3
View File
@@ -33,7 +33,7 @@ class TestTypeBasedDispatcher(unittest.TestCase):
FlushCacheReqInput,
FreezeGCReq,
GetInternalStateReq,
GetLoadReqInput,
GetLoadsReqInput,
GetWeightsByNameReqInput,
InitWeightsSendGroupForRemoteInstanceReqInput,
InitWeightsUpdateGroupReqInput,
@@ -113,7 +113,7 @@ class TestTypeBasedDispatcher(unittest.TestCase):
(ExpertDistributionReq, lambda req: "expert_distribution_handled"),
(LoadLoRAAdapterReqInput, lambda req: "load_lora_adapter_handled"),
(UnloadLoRAAdapterReqInput, lambda req: "unload_lora_adapter_handled"),
(GetLoadReqInput, lambda req: "get_load_handled"),
(GetLoadsReqInput, lambda req: "get_loads_handled"),
]
# Create requests that conforms to the real distribution
@@ -204,7 +204,7 @@ class TestTypeBasedDispatcher(unittest.TestCase):
test_requests.append(GetWeightsByNameReqInput(name=""))
test_requests.append(ReleaseMemoryOccupationReqInput())
test_requests.append(RpcReqInput(method=""))
test_requests.append(GetLoadReqInput())
test_requests.append(GetLoadsReqInput())
dispatcher = TypeBasedDispatcher(mapping)