243 lines
9.0 KiB
Python
243 lines
9.0 KiB
Python
import unittest
|
|
from typing import Dict, List
|
|
|
|
import requests
|
|
from prometheus_client.parser import text_string_to_metric_families
|
|
from prometheus_client.samples import Sample
|
|
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.metrics.collector import (
|
|
ROUTING_KEY_REQ_COUNT_BUCKET_BOUNDS,
|
|
compute_routing_key_stats,
|
|
)
|
|
from sglang.srt.utils import kill_process_tree
|
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
DEFAULT_URL_FOR_TEST,
|
|
CustomTestCase,
|
|
is_in_ci,
|
|
popen_launch_server,
|
|
)
|
|
|
|
register_cuda_ci(est_time=32, suite="stage-b-test-small-1-gpu")
|
|
register_amd_ci(est_time=32, suite="stage-b-test-small-1-gpu-amd")
|
|
|
|
_MODEL_NAME = "Qwen/Qwen3-0.6B"
|
|
|
|
|
|
class TestEnableMetrics(CustomTestCase):
|
|
def test_metrics_1gpu(self):
|
|
"""Test that metrics endpoint returns data when enabled"""
|
|
self._execute_core(
|
|
other_args=[],
|
|
verify_metrics_extra=None,
|
|
)
|
|
|
|
def test_metrics_2gpu(self):
|
|
# TODO enable when we have 2-gpu runner in nightly CI
|
|
if is_in_ci():
|
|
print("Skip test_metrics_2gpu since in 1-gpu CI")
|
|
return
|
|
|
|
def _verify_metrics_extra(metrics):
|
|
metrics_to_check = [
|
|
(
|
|
"sglang:dp_cooperation_realtime_tokens_total",
|
|
{"mode": "prefill_compute"},
|
|
),
|
|
(
|
|
"sglang:dp_cooperation_realtime_tokens_total",
|
|
{"mode": "decode"},
|
|
),
|
|
(
|
|
"sglang:dp_cooperation_gpu_execution_seconds_total",
|
|
{"category": "forward_extend"},
|
|
),
|
|
(
|
|
"sglang:dp_cooperation_gpu_execution_seconds_total",
|
|
{"category": "forward_decode"},
|
|
),
|
|
]
|
|
_check_metrics_positive(self, metrics, metrics_to_check)
|
|
|
|
num_prefill_ranks_values = {
|
|
s.labels["num_prefill_ranks"]
|
|
for s in metrics["sglang:dp_cooperation_realtime_tokens_total"]
|
|
}
|
|
self.assertIn("0", num_prefill_ranks_values)
|
|
self.assertIn("1", num_prefill_ranks_values)
|
|
|
|
self._execute_core(
|
|
other_args=["--tp", "2", "--dp", "2", "--enable-dp-attention"],
|
|
verify_metrics_extra=_verify_metrics_extra,
|
|
)
|
|
|
|
def _execute_core(self, other_args, verify_metrics_extra):
|
|
with (
|
|
envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.override(True),
|
|
envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.override(True),
|
|
envs.SGLANG_TEST_RETRACT.override(True),
|
|
):
|
|
process = popen_launch_server(
|
|
_MODEL_NAME,
|
|
DEFAULT_URL_FOR_TEST,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=["--enable-metrics", "--cuda-graph-max-bs", 2, *other_args],
|
|
)
|
|
|
|
try:
|
|
# Make some requests to generate some metrics
|
|
response = requests.get(f"{DEFAULT_URL_FOR_TEST}/health_generate")
|
|
self.assertEqual(response.status_code, 200)
|
|
|
|
response = requests.post(
|
|
f"{DEFAULT_URL_FOR_TEST}/generate",
|
|
json={
|
|
"text": ["The capital of France is"] * 20,
|
|
"sampling_params": {
|
|
"temperature": 0,
|
|
"max_new_tokens": 50,
|
|
},
|
|
"stream": True,
|
|
"ignore_eos": True,
|
|
},
|
|
stream=True,
|
|
)
|
|
for _ in response.iter_lines(decode_unicode=False):
|
|
pass
|
|
|
|
response = requests.post(
|
|
f"{DEFAULT_URL_FOR_TEST}/generate",
|
|
json={
|
|
"text": "Hello",
|
|
"sampling_params": {"temperature": 0, "max_new_tokens": 5},
|
|
},
|
|
headers={"x-smg-routing-key": "test-key"},
|
|
)
|
|
self.assertEqual(response.status_code, 200)
|
|
|
|
# Get metrics
|
|
metrics_response = requests.get(f"{DEFAULT_URL_FOR_TEST}/metrics")
|
|
self.assertEqual(metrics_response.status_code, 200)
|
|
metrics_text = metrics_response.text
|
|
|
|
print(f"metrics_text=\n{metrics_text}")
|
|
|
|
metrics = _parse_prometheus_metrics(metrics_text)
|
|
self._verify_metrics_common(metrics_text, metrics)
|
|
if verify_metrics_extra is not None:
|
|
verify_metrics_extra(metrics)
|
|
finally:
|
|
kill_process_tree(process.pid)
|
|
|
|
def _verify_metrics_common(self, metrics_text, metrics):
|
|
essential_metrics = [
|
|
"sglang:num_running_reqs",
|
|
"sglang:num_used_tokens",
|
|
"sglang:token_usage",
|
|
"sglang:gen_throughput",
|
|
"sglang:num_queue_reqs",
|
|
"sglang:num_grammar_queue_reqs",
|
|
"sglang:cache_hit_rate",
|
|
"sglang:spec_accept_length",
|
|
"sglang:prompt_tokens_total",
|
|
"sglang:generation_tokens_total",
|
|
"sglang:cached_tokens_total",
|
|
"sglang:num_requests_total",
|
|
"sglang:time_to_first_token_seconds",
|
|
"sglang:inter_token_latency_seconds",
|
|
"sglang:e2e_request_latency_seconds",
|
|
"sglang:http_requests_active",
|
|
"sglang:routing_keys_active",
|
|
"sglang:num_unique_running_routing_keys",
|
|
"sglang:routing_key_running_req_count",
|
|
"sglang:routing_key_all_req_count",
|
|
]
|
|
for metric in essential_metrics:
|
|
self.assertIn(metric, metrics_text, f"Missing metric: {metric}")
|
|
|
|
# Verify routing key GaugeHistogram buckets
|
|
expected_buckets = len(ROUTING_KEY_REQ_COUNT_BUCKET_BOUNDS) + 1
|
|
for metric_name in [
|
|
"sglang:routing_key_running_req_count",
|
|
"sglang:routing_key_all_req_count",
|
|
]:
|
|
gt_le_pairs = set()
|
|
for sample in metrics.get(metric_name, []):
|
|
gt_le_pairs.add((sample.labels.get("gt"), sample.labels.get("le")))
|
|
self.assertEqual(
|
|
len(gt_le_pairs),
|
|
expected_buckets,
|
|
f"{metric_name}: Expected {expected_buckets} buckets, got {len(gt_le_pairs)}",
|
|
)
|
|
|
|
self.assertIn(f'model_name="{_MODEL_NAME}"', metrics_text)
|
|
self.assertIn("_sum{", metrics_text)
|
|
self.assertIn("_count{", metrics_text)
|
|
self.assertIn("_bucket{", metrics_text)
|
|
|
|
metrics_to_check = [
|
|
("sglang:realtime_tokens_total", {"mode": "prefill_compute"}),
|
|
("sglang:realtime_tokens_total", {"mode": "decode"}),
|
|
("sglang:gpu_execution_seconds_total", {"category": "forward_extend"}),
|
|
("sglang:gpu_execution_seconds_total", {"category": "forward_decode"}),
|
|
("sglang:process_cpu_seconds_total", {"component": "tokenizer"}),
|
|
]
|
|
_check_metrics_positive(self, metrics, metrics_to_check)
|
|
|
|
|
|
def _parse_prometheus_metrics(metrics_text: str) -> Dict[str, List[Sample]]:
|
|
result = {}
|
|
for family in text_string_to_metric_families(metrics_text):
|
|
for sample in family.samples:
|
|
if sample.name not in result:
|
|
result[sample.name] = []
|
|
result[sample.name].append(sample)
|
|
return result
|
|
|
|
|
|
def _get_sample_value_by_labels(samples: List[Sample], labels: Dict[str, str]) -> float:
|
|
for sample in samples:
|
|
if all(sample.labels.get(k) == v for k, v in labels.items()):
|
|
return sample.value
|
|
raise KeyError(f"No sample found with labels {labels}")
|
|
|
|
|
|
def _check_metrics_positive(test_case, metrics, metrics_to_check):
|
|
for metric_name, labels in metrics_to_check:
|
|
value = _get_sample_value_by_labels(metrics[metric_name], labels)
|
|
test_case.assertGreater(value, 0, f"{metric_name} {labels}")
|
|
|
|
|
|
class TestComputeRoutingKeyStats(unittest.TestCase):
|
|
def test_empty(self):
|
|
num_unique, req_counts = compute_routing_key_stats([])
|
|
self.assertEqual(num_unique, 0)
|
|
self.assertEqual(req_counts, [])
|
|
|
|
def test_all_none(self):
|
|
num_unique, req_counts = compute_routing_key_stats([None, None, None])
|
|
self.assertEqual(num_unique, 0)
|
|
self.assertEqual(req_counts, [])
|
|
|
|
def test_with_none(self):
|
|
num_unique, req_counts = compute_routing_key_stats([None, "key1", None])
|
|
self.assertEqual(num_unique, 1)
|
|
self.assertEqual(req_counts, [1])
|
|
|
|
def test_single_key_multiple_reqs(self):
|
|
num_unique, req_counts = compute_routing_key_stats(["key1"] * 5)
|
|
self.assertEqual(num_unique, 1)
|
|
self.assertEqual(req_counts, [5])
|
|
|
|
def test_distribution(self):
|
|
routing_keys = ["key1"] * 5 + ["key2"] * 1 + ["key3"] * 15 + ["key4"] * 250
|
|
num_unique, req_counts = compute_routing_key_stats(routing_keys)
|
|
self.assertEqual(num_unique, 4)
|
|
self.assertEqual(sorted(req_counts), [1, 5, 15, 250])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|