Tiny support http headers in bench serving (#16606)

test/registered/bench_fn/test_bench_serving_functionality.py

@@ -1,10 +1,12 @@
 import json
 import tempfile
+import threading
 import time
 import unittest
+from http.server import BaseHTTPRequestHandler, HTTPServer
 from pathlib import Path
 
-from sglang.bench_serving import run_benchmark
+from sglang.bench_serving import parse_custom_headers, run_benchmark
 from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import (
@@ -96,5 +98,87 @@ class TestBenchServingFunctionality(CustomTestCase):
         )
 
 
+class TestBenchServingCustomHeaders(CustomTestCase):
+    def test_parse_custom_headers(self):
+        headers = parse_custom_headers(["MyHeader=MY_VALUE", "Another=value=hello"])
+        self.assertEqual(headers, {"MyHeader": "MY_VALUE", "Another": "value=hello"})
+
+        headers = parse_custom_headers(["InvalidNoEquals"])
+        self.assertEqual(headers, {})
+
+        headers = parse_custom_headers(["=NoKey"])
+        self.assertEqual(headers, {})
+
+    # TODO: Using well-implemented mock server, e.g. the on in sgl-router
+    def test_custom_headers_sent_to_server(self):
+        import queue
+
+        received_requests = queue.Queue()
+
+        class HeaderEchoHandler(BaseHTTPRequestHandler):
+            def _handle(self):
+                received_requests.put(
+                    {
+                        "method": self.command,
+                        "path": self.path,
+                        "headers": dict(self.headers),
+                    }
+                )
+                self.send_response(200)
+                self.send_header("Content-Type", "application/json")
+                self.end_headers()
+                if self.path == "/v1/models":
+                    self.wfile.write(json.dumps({"data": [{"id": "gpt2"}]}).encode())
+                elif self.path == "/generate":
+                    self.wfile.write(
+                        json.dumps(
+                            {"text": "ok", "meta_info": {"completion_tokens": 1}}
+                        ).encode()
+                    )
+                else:
+                    self.wfile.write(json.dumps({}).encode())
+
+            do_GET = do_POST = _handle
+
+        server = HTTPServer(("127.0.0.1", 0), HeaderEchoHandler)
+        port = server.server_address[1]
+        server_thread = threading.Thread(target=server.serve_forever)
+        server_thread.daemon = True
+        server_thread.start()
+
+        try:
+            args = get_benchmark_args(
+                base_url=f"http://127.0.0.1:{port}",
+                backend="sglang",
+                dataset_name="random",
+                tokenizer="gpt2",
+                num_prompts=1,
+                random_input_len=8,
+                random_output_len=8,
+                header=["X-Custom-Test=TestValue123", "X-Another=AnotherVal"],
+            )
+            args.warmup_requests = 0
+            args.disable_tqdm = True
+            run_benchmark(args)
+        except Exception:
+            pass
+        finally:
+            server.shutdown()
+
+        all_reqs = []
+        while not received_requests.empty():
+            all_reqs.append(received_requests.get_nowait())
+
+        generate_reqs = [r for r in all_reqs if r["path"] == "/generate"]
+        self.assertGreater(
+            len(generate_reqs),
+            0,
+            f"No /generate request. All: {[r['path'] for r in all_reqs]}",
+        )
+        headers = generate_reqs[0]["headers"]
+        self.assertEqual(headers.get("X-Custom-Test"), "TestValue123")
+        self.assertEqual(headers.get("X-Another"), "AnotherVal")
+
+
 if __name__ == "__main__":
     unittest.main()