SGLang GKE test

PiperOrigin-RevId: 832193475

Author: AnilAltinay

test/gpu/sglang/sglang.go

@@ -547,6 +547,43 @@ func (r *Response) NumTokens() int {
 	return len(r.data)
 }
 
+// TimeToFirstToken returns the time it took between the request starting
+// and the first token being received by the client.
+func (r *Response) TimeToFirstToken() time.Duration {
+	if !r.Done() {
+		return -1
+	}
+	return r.metrics.FirstByteRead.Sub(r.metrics.RequestSent)
+}
+
+// TimeToLastToken returns the time it took between the request starting
+// and the last token being received by the client.
+func (r *Response) TimeToLastToken() time.Duration {
+	if !r.Done() {
+		return -1
+	}
+	return r.metrics.LastByteRead.Sub(r.metrics.RequestSent)
+}
+
+// OutputTokensPerSecond computes the average number of output tokens
+// generated per second.
+func (r *Response) OutputTokensPerSecond() float64 {
+	if !r.Done() {
+		return -1
+	}
+	return float64(len(r.data)-1) / r.data[len(r.data)-1].MetaInfo.E2ELatency
+}
+
+// E2ELatency returns the elapsed time between when start_time was recorded and
+// the current moment in seconds.
+// See https://github.com/sgl-project/sglang/blob/4a2768a86b2905b9b7f19d415261b9d4af639e19/sgl-router/src/routers/grpc/regular/streaming.rs#L904
+func (r *Response) E2ELatency() float64 {
+	if len(r.data) == 0 {
+		return 0
+	}
+	return r.data[len(r.data)-1].MetaInfo.E2ELatency
+}
+
 // String returns the response text, if it is done.
 func (r *Response) String() string {
 	if len(r.data) == 0 {

test/kubernetes/benchmarks/BUILD

@@ -397,6 +397,46 @@ go_test(
     ],
 )
 
+go_library(
+    name = "sglang",
+    testonly = True,
+    srcs = ["sglang.go"],
+    embedsrcs = [
+        "//test/kubernetes/benchmarks/resources:files",  # keep
+    ],
+    nogo = False,
+    deps = [
+        "//pkg/sync",
+        "//test/gpu/sglang",
+        "//test/kubernetes",
+        "//test/kubernetes/benchmarks/profiling",
+        "//test/kubernetes/benchmetric",
+        "//test/kubernetes/k8sctx",
+        "//test/kubernetes/testcluster",
+        "@io_k8s_api//core/v1:go_default_library",
+        "@io_k8s_apimachinery//pkg/api/resource:go_default_library",
+        "@io_k8s_apimachinery//pkg/apis/meta/v1:go_default_library",
+        "@io_k8s_apimachinery//pkg/util/intstr:go_default_library",
+    ],
+)
+
+go_test(
+    name = "sglang_test",
+    srcs = ["sglang_test.go"],
+    library = ":sglang",
+    nogo = False,
+    tags = [
+        "local",
+        "noguitar",
+        "notap",
+    ],
+    deps = [
+        "//test/kubernetes/k8sctx",
+        "//test/kubernetes/k8sctx/kubectlctx",
+        "//test/kubernetes/testcluster",
+    ],
+)
+
 go_library(
     name = "stablediffusion",
     testonly = True,