diff --git a/python/sglang/srt/entrypoints/grpc_server.py b/python/sglang/srt/entrypoints/grpc_server.py index 3a879cb6a..11424676f 100644 --- a/python/sglang/srt/entrypoints/grpc_server.py +++ b/python/sglang/srt/entrypoints/grpc_server.py @@ -210,7 +210,10 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer) sampling_params.normalize(tokenizer=None) # Create health check request - is_generation = self.scheduler_info.get("is_generation", True) + is_generation = self.scheduler_info.get("is_generation") + if is_generation is None: + is_generation = not self.server_args.is_embedding + if is_generation: health_req = TokenizedGenerateReqInput( rid=rid, @@ -229,6 +232,7 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer) health_req.bootstrap_host = FAKE_BOOTSTRAP_HOST health_req.bootstrap_room = 0 else: + sampling_params.max_new_tokens = 0 health_req = TokenizedEmbeddingReqInput( rid=rid, input_text="", diff --git a/sgl-model-gateway/py_test/e2e_grpc/basic/test_embedding_server.py b/sgl-model-gateway/py_test/e2e_grpc/basic/test_embedding_server.py new file mode 100644 index 000000000..71b6da5f4 --- /dev/null +++ b/sgl-model-gateway/py_test/e2e_grpc/basic/test_embedding_server.py @@ -0,0 +1,109 @@ +""" +gRPC Router E2E Test - Embedding Server +Test the embedding functionality of the gRPC router. +""" + +import sys +import unittest +from pathlib import Path + +import openai + +_TEST_DIR = Path(__file__).parent +sys.path.insert(0, str(_TEST_DIR.parent)) +from fixtures import popen_launch_workers_and_router +from util import ( + DEFAULT_EMBEDDING_MODEL_PATH, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + kill_process_tree, +) + + +class TestEmbeddingServer(CustomTestCase): + """ + Test Embedding API through gRPC router. + """ + + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_EMBEDDING_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + cls.api_key = "sk-123456" + + # Launch workers with --is-embedding flag + cls.cluster = popen_launch_workers_and_router( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + num_workers=1, + tp_size=1, + policy="round_robin", + api_key=cls.api_key, + worker_args=["--is-embedding"], + ) + + cls.base_url += "/v1" + + @classmethod + def tearDownClass(cls): + # Cleanup router and workers + kill_process_tree(cls.cluster["router"].pid) + for worker in cls.cluster.get("workers", []): + kill_process_tree(worker.pid) + + def test_embedding(self): + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + + input_text = "Hello world" + response = client.embeddings.create( + model=self.model, + input=input_text, + ) + + assert response.object == "list" + assert len(response.data) == 1 + embedding = response.data[0] + assert embedding.object == "embedding" + assert embedding.index == 0 + assert len(embedding.embedding) > 0 + assert isinstance(embedding.embedding[0], float) + + # Verify usage statistics + assert response.usage.prompt_tokens > 0 + assert response.usage.total_tokens == response.usage.prompt_tokens + + def test_embedding_batch(self): + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + + input_texts = ["Hello world", "SGLang is fast"] + response = client.embeddings.create( + model=self.model, + input=input_texts, + ) + + assert len(response.data) == 1 + assert response.data[0].index == 0 + assert len(response.data[0].embedding) > 0 + + def test_embedding_dimensions(self): + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + + response1 = client.embeddings.create( + model=self.model, + input="A short text", + ) + dim1 = len(response1.data[0].embedding) + + response2 = client.embeddings.create( + model=self.model, + input="A much longer text to ensure dimensions match", + ) + dim2 = len(response2.data[0].embedding) + + assert dim1 == dim2 + + +if __name__ == "__main__": + unittest.main() diff --git a/sgl-model-gateway/py_test/e2e_grpc/util.py b/sgl-model-gateway/py_test/e2e_grpc/util.py index 3f85cd67d..bebe75ea4 100644 --- a/sgl-model-gateway/py_test/e2e_grpc/util.py +++ b/sgl-model-gateway/py_test/e2e_grpc/util.py @@ -92,6 +92,9 @@ DEFAULT_MISTRAL_FUNCTION_CALLING_MODEL_PATH = _get_model_path( # GPT-OSS models DEFAULT_GPT_OSS_MODEL_PATH = _get_model_path("openai/gpt-oss-20b") +# Embedding models +DEFAULT_EMBEDDING_MODEL_PATH = _get_model_path("intfloat/e5-mistral-7b-instruct") + # ============================================================================ # Process Management diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 95595e723..86cf5d3cd 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -1,3 +1,4 @@ +/// test use std::{ sync::{ atomic::{AtomicBool, Ordering},