metrics: support customer buckets for prompt/generation_tokens_histogram (#9634)

test/srt/run_suite.py

@@ -80,6 +80,7 @@ suites = {
         TestFile("test_io_struct.py", 8),
         TestFile("test_jinja_template_utils.py", 1),
         TestFile("test_metrics.py", 32),
+        TestFile("test_metrics_utils.py", 1),
         TestFile("test_mla.py", 167),
         TestFile("test_mla_deepseek_v3.py", 700),
         TestFile("test_mla_int8_deepseek_v3.py", 429),
@@ -214,6 +215,7 @@ suite_amd = {
         TestFile("test_io_struct.py", 8),
         TestFile("test_jinja_template_utils.py", 1),
         TestFile("test_metrics.py", 32),
+        TestFile("test_metrics_utils.py", 1),
         TestFile("test_mla.py", 242),
         TestFile("test_mla_deepseek_v3.py", 221),
         TestFile("test_no_chunked_prefill.py", 108),

test/srt/test_metrics_utils.py

@@ -0,0 +1,137 @@
+import unittest
+
+from sglang.srt.metrics.utils import generate_buckets, two_sides_exponential_buckets
+
+
+class TestMetricsUtils(unittest.TestCase):
+    """Test cases for metrics utility functions."""
+
+    def test_two_sides_exponential_buckets_basic(self):
+        """Test basic functionality of two_sides_exponential_buckets."""
+        # Test with simple parameters
+        count = 5
+        buckets = two_sides_exponential_buckets(middle=10.0, base=2.0, count=count)
+
+        # Should contain the middle value
+        self.assertIn(10.0, buckets)
+
+        # Should be sorted
+        self.assertEqual(buckets, sorted(buckets))
+
+        # Should have unique values (no duplicates)
+        self.assertEqual(len(buckets), len(set(buckets)))
+
+        # Should have reasonable number of buckets (not exactly count due to ceiling and deduplication)
+        self.assertGreaterEqual(len(buckets), 3)
+        self.assertLessEqual(len(buckets), count + 2)
+
+    def test_two_sides_exponential_buckets_specific_values(self):
+        """Test specific values for two_sides_exponential_buckets."""
+        buckets = two_sides_exponential_buckets(middle=100.0, base=2.0, count=4)
+        expected_values = [96.0, 98.0, 100.0, 102.0, 104.0]
+        self.assertEqual(buckets, expected_values)
+
+    def test_two_sides_exponential_buckets_negative_values(self):
+        """Test two_sides_exponential_buckets with values that could go negative."""
+        buckets = two_sides_exponential_buckets(middle=5.0, base=3.0, count=4)
+
+        # Should not contain negative values (max(0, middle - distance))
+        for bucket in buckets:
+            self.assertGreaterEqual(bucket, 0.0)
+
+        # Should contain the middle value
+        self.assertIn(5.0, buckets)
+
+    def test_two_sides_exponential_buckets_edge_cases(self):
+        """Test edge cases for two_sides_exponential_buckets."""
+        # Count = 1
+        buckets = two_sides_exponential_buckets(middle=10.0, base=2.0, count=1)
+        self.assertIn(10.0, buckets)
+
+        # Very small middle value
+        buckets = two_sides_exponential_buckets(middle=0.1, base=2.0, count=2)
+        self.assertIn(0.1, buckets)
+        for bucket in buckets:
+            self.assertGreaterEqual(bucket, 0.0)
+
+    def test_generate_buckets_default(self):
+        """Test generate_buckets with default rule."""
+        default_buckets = [1.0, 5.0, 10.0, 50.0, 100.0]
+
+        # Test with "default" rule
+        result = generate_buckets(["default"], default_buckets)
+        self.assertEqual(result, default_buckets)
+
+        # Test with None (should default to "default")
+        result = generate_buckets(None, default_buckets)
+        self.assertEqual(result, default_buckets)
+
+        # Test with empty (should default to "default")
+        result = generate_buckets(None, default_buckets)
+        self.assertEqual(result, default_buckets)
+
+    def test_generate_buckets_tse(self):
+        """Test generate_buckets with tse (two sides exponential) rule."""
+        default_buckets = [1.0, 5.0, 10.0]
+
+        # Test with "tse" rule
+        result = generate_buckets(["tse", "10", "2.0", "4"], default_buckets)
+
+        # Should return the same as calling two_sides_exponential_buckets directly
+        expected = two_sides_exponential_buckets(10.0, 2.0, 4)
+        self.assertEqual(result, expected)
+
+    def test_generate_buckets_customer(self):
+        """Test generate_buckets with customer rule."""
+        default_buckets = [1.0, 5.0, 10.0]
+
+        # Test with "customer" rule
+        result = generate_buckets(
+            ["customer", "1.5", "3.2", "7.8", "15.6"], default_buckets
+        )
+        expected = [1.5, 3.2, 7.8, 15.6]
+        self.assertEqual(result, expected)
+
+    def test_generate_buckets_customer_with_integers(self):
+        """Test generate_buckets with customer rule using integer strings."""
+        default_buckets = [1.0, 5.0, 10.0]
+
+        # Test with integer strings
+        result = generate_buckets(["customer", "1", "5", "10", "50"], default_buckets)
+        expected = [1.0, 5.0, 10.0, 50.0]
+        self.assertEqual(result, expected)
+
+    def test_generate_buckets_preserves_order_and_type(self):
+        """Test that generate_buckets preserves order and returns floats."""
+        default_buckets = [1, 5, 10, 50, 100]  # integers
+
+        # Test default rule
+        result = generate_buckets(["default"], default_buckets)
+        self.assertEqual(result, default_buckets)
+        self.assertIsInstance(result, list)
+
+        # Test customer rule with proper float conversion
+        result = generate_buckets(
+            ["customer", "100", "50", "10", "5", "1"], default_buckets
+        )
+        expected = [1.0, 5.0, 10.0, 50.0, 100.0]
+        self.assertEqual(result, expected)
+
+        # All values should be floats
+        for value in result:
+            self.assertIsInstance(value, float)
+
+    def test_integration_tse_through_generate_buckets(self):
+        """Test integration of TSE buckets through generate_buckets function."""
+        default_buckets = [1.0, 10.0, 100.0]
+
+        # Generate buckets using both methods
+        direct_result = two_sides_exponential_buckets(50.0, 1.5, 6)
+        indirect_result = generate_buckets(["tse", "50.0", "1.5", "6"], default_buckets)
+
+        # Results should be identical
+        self.assertEqual(direct_result, indirect_result)
+
+
+if __name__ == "__main__":
+    unittest.main()