84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
"""Unit tests for func_timer.py — no server, no model loading."""
|
|
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=6, suite="base-a-test-cpu")
|
|
register_cpu_ci(est_time=7, suite="base-b-test-cpu")
|
|
|
|
import asyncio
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import sglang.srt.observability.func_timer as func_timer
|
|
from sglang.srt.observability.func_timer import enable_func_timer, time_func_latency
|
|
|
|
|
|
class TestFuncTimer(unittest.TestCase):
|
|
def setUp(self):
|
|
self.orig_enable = func_timer.enable_metrics
|
|
self.orig_latency = func_timer.FUNC_LATENCY
|
|
|
|
def tearDown(self):
|
|
func_timer.enable_metrics = self.orig_enable
|
|
func_timer.FUNC_LATENCY = self.orig_latency
|
|
|
|
@patch("prometheus_client.Histogram")
|
|
def test_enable_func_timer(self, MockHistogram):
|
|
"""Sets enable_metrics and creates FUNC_LATENCY histogram."""
|
|
enable_func_timer()
|
|
self.assertTrue(func_timer.enable_metrics)
|
|
self.assertIs(func_timer.FUNC_LATENCY, MockHistogram.return_value)
|
|
MockHistogram.assert_called_once()
|
|
|
|
def test_sync_disabled(self):
|
|
"""Sync function passes through when metrics disabled."""
|
|
func_timer.enable_metrics = False
|
|
|
|
@time_func_latency
|
|
def add(a, b):
|
|
return a + b
|
|
|
|
self.assertEqual(add(2, 3), 5)
|
|
|
|
def test_sync_enabled(self):
|
|
"""Sync function timed with custom name when metrics enabled."""
|
|
mock_histogram = MagicMock()
|
|
func_timer.enable_metrics = True
|
|
func_timer.FUNC_LATENCY = mock_histogram
|
|
|
|
@time_func_latency(name="custom_op")
|
|
def add(a, b):
|
|
return a + b
|
|
|
|
self.assertEqual(add(2, 3), 5)
|
|
mock_histogram.labels.assert_called_with(name="custom_op")
|
|
mock_histogram.labels().observe.assert_called_once()
|
|
|
|
def test_async_disabled(self):
|
|
"""Async function passes through when metrics disabled."""
|
|
func_timer.enable_metrics = False
|
|
|
|
@time_func_latency
|
|
async def add(a, b):
|
|
return a + b
|
|
|
|
self.assertEqual(asyncio.run(add(2, 3)), 5)
|
|
|
|
def test_async_enabled(self):
|
|
"""Async function timed with default name when metrics enabled."""
|
|
mock_histogram = MagicMock()
|
|
func_timer.enable_metrics = True
|
|
func_timer.FUNC_LATENCY = mock_histogram
|
|
|
|
@time_func_latency
|
|
async def add(a, b):
|
|
return a + b
|
|
|
|
self.assertEqual(asyncio.run(add(2, 3)), 5)
|
|
mock_histogram.labels.assert_called_with(name="add")
|
|
mock_histogram.labels().observe.assert_called_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|