Fix gpt_oss_common import path and migrate core tests (#16426)
This commit is contained in:
Alison Shao
+1
-24
@@ -9,24 +9,12 @@ from sglang.test.ci.ci_utils import TestFile, run_unittest_files
|
||||
# NOTE: please sort the test cases alphabetically by the test file name
|
||||
suites = {
|
||||
"per-commit-1-gpu": [
|
||||
TestFile("test_deterministic.py", 228),
|
||||
TestFile("test_evs.py", 20),
|
||||
TestFile("test_external_models.py", 30),
|
||||
TestFile("test_gpt_oss_1gpu.py", 402),
|
||||
TestFile("test_hidden_states.py", 55),
|
||||
TestFile("test_input_embeddings.py", 38),
|
||||
TestFile("test_io_struct.py", 8),
|
||||
TestFile("test_jinja_template_utils.py", 7),
|
||||
TestFile("test_model_hooks.py", 6),
|
||||
TestFile("test_modelopt_loader.py", 11),
|
||||
TestFile("test_multi_tokenizer.py", 230),
|
||||
TestFile("test_page_size.py", 60),
|
||||
TestFile("test_request_queue_validation.py", 47),
|
||||
TestFile("test_score_api.py", 260),
|
||||
TestFile("test_server_args.py", 9),
|
||||
TestFile("test_skip_tokenizer_init.py", 77),
|
||||
TestFile("test_srt_endpoint.py", 127),
|
||||
TestFile("test_srt_engine.py", 252),
|
||||
TestFile("test_utils_update_weights.py", 29),
|
||||
TestFile("test_video_utils.py", 5),
|
||||
TestFile("test_modelopt_export.py", 9),
|
||||
@@ -98,7 +86,6 @@ suites = {
|
||||
TestFile("test_bench_one_batch.py"),
|
||||
TestFile("test_bench_serving.py"),
|
||||
TestFile("test_eval_accuracy_large.py"),
|
||||
TestFile("test_gpt_oss_common.py"),
|
||||
TestFile("test_moe_eval_accuracy_large.py"),
|
||||
TestFile("test_profile_v2.py"),
|
||||
TestFile("models/test_ministral3_models.py"),
|
||||
@@ -120,18 +107,10 @@ suite_amd = {
|
||||
# TestFile("lora/test_lora_qwen3.py", 97), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
|
||||
TestFile("test_bench_typebaseddispatcher.py", 10),
|
||||
TestFile("test_external_models.py", 45),
|
||||
TestFile("test_input_embeddings.py", 38),
|
||||
TestFile("test_io_struct.py", 8),
|
||||
TestFile("test_jinja_template_utils.py", 1),
|
||||
TestFile("test_model_hooks.py", 10),
|
||||
TestFile("test_multi_tokenizer.py", 345),
|
||||
TestFile("test_page_size.py", 60),
|
||||
TestFile("test_request_queue_validation.py", 70),
|
||||
TestFile("test_rope_rocm.py", 3),
|
||||
TestFile("test_server_args.py", 1),
|
||||
TestFile("test_skip_tokenizer_init.py", 117),
|
||||
TestFile("test_srt_endpoint.py", 130),
|
||||
TestFile("test_srt_engine.py", 261),
|
||||
# TestFile("test_torch_compile_moe.py", 210), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
|
||||
TestFile("test_type_based_dispatcher.py", 10),
|
||||
TestFile("test_video_utils.py", 8),
|
||||
@@ -143,9 +122,7 @@ suite_amd = {
|
||||
# TestFile("test_vision_chunked_prefill.py", 175), # Disabled temporarily and track in #7701
|
||||
# TestFile("test_wave_attention_backend.py", 150), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/11127
|
||||
],
|
||||
"per-commit-amd-mi35x": [
|
||||
TestFile("test_gpt_oss_1gpu.py", 750),
|
||||
],
|
||||
"per-commit-amd-mi35x": [],
|
||||
"per-commit-2-gpu-amd": [
|
||||
TestFile("test_data_parallelism.py", 73),
|
||||
TestFile("test_load_weights_from_remote_instance.py", 72),
|
||||
|
||||
@@ -1,61 +0,0 @@
|
||||
"""
|
||||
Usage:
|
||||
cd test/srt
|
||||
python3 -m unittest test_deterministic.TestDeterministic.TESTCASE
|
||||
|
||||
Note that there is also `python/sglang/test/test_deterministic.py` as an interactive test. We are converting that
|
||||
test into unit tests so that's easily reproducible in CI.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.test_deterministic_utils import (
|
||||
COMMON_SERVER_ARGS,
|
||||
TestDeterministicBase,
|
||||
)
|
||||
|
||||
|
||||
class TestFlashinferDeterministic(TestDeterministicBase):
|
||||
# Test with flashinfer attention backend
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = COMMON_SERVER_ARGS
|
||||
args.extend(
|
||||
[
|
||||
"--attention-backend",
|
||||
"flashinfer",
|
||||
]
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
class TestFa3Deterministic(TestDeterministicBase):
|
||||
# Test with fa3 attention backend
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = COMMON_SERVER_ARGS
|
||||
args.extend(
|
||||
[
|
||||
"--attention-backend",
|
||||
"fa3",
|
||||
]
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
class TestTritonDeterministic(TestDeterministicBase):
|
||||
# Test with triton attention backend
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = COMMON_SERVER_ARGS
|
||||
args.extend(
|
||||
[
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
]
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,31 +0,0 @@
|
||||
import unittest
|
||||
|
||||
from test_gpt_oss_common import BaseTestGptOss
|
||||
|
||||
|
||||
class TestGptOss1Gpu(BaseTestGptOss):
|
||||
def test_mxfp4_20b(self):
|
||||
self.run_test(
|
||||
model_variant="20b",
|
||||
quantization="mxfp4",
|
||||
expected_score_of_reasoning_effort={
|
||||
"low": 0.34,
|
||||
"medium": 0.34,
|
||||
"high": 0.27, # TODO investigate
|
||||
},
|
||||
)
|
||||
|
||||
def test_bf16_20b(self):
|
||||
self.run_test(
|
||||
model_variant="20b",
|
||||
quantization="bf16",
|
||||
expected_score_of_reasoning_effort={
|
||||
"low": 0.34,
|
||||
"medium": 0.34,
|
||||
"high": 0.27, # TODO investigate
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,6 +1,6 @@
|
||||
import unittest
|
||||
|
||||
from test_gpt_oss_common import BaseTestGptOss
|
||||
from sglang.test.gpt_oss_common import BaseTestGptOss
|
||||
|
||||
|
||||
class TestGptOss4Gpu(BaseTestGptOss):
|
||||
|
||||
@@ -1,152 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from types import SimpleNamespace
|
||||
from typing import Dict, List, Literal, Optional
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import is_hip, kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
_base_url = DEFAULT_URL_FOR_TEST
|
||||
_is_hip = is_hip()
|
||||
|
||||
|
||||
class BaseTestGptOss(CustomTestCase):
|
||||
def run_test(
|
||||
self,
|
||||
model_variant: Literal["20b", "120b"],
|
||||
quantization: Literal["mxfp4", "bf16"],
|
||||
expected_score_of_reasoning_effort: Dict[str, float],
|
||||
other_args: Optional[List[str]] = None,
|
||||
):
|
||||
if other_args is None:
|
||||
other_args = []
|
||||
|
||||
model = {
|
||||
("20b", "bf16"): "lmsys/gpt-oss-20b-bf16",
|
||||
("120b", "bf16"): "lmsys/gpt-oss-120b-bf16",
|
||||
("20b", "mxfp4"): "openai/gpt-oss-20b",
|
||||
("120b", "mxfp4"): "openai/gpt-oss-120b",
|
||||
}[(model_variant, quantization)]
|
||||
|
||||
if model_variant == "20b":
|
||||
other_args += ["--cuda-graph-max-bs", "600"]
|
||||
if _is_hip:
|
||||
os.environ["SGLANG_USE_AITER"] = "0"
|
||||
self._run_test_raw(
|
||||
model=model,
|
||||
expected_score_of_reasoning_effort=expected_score_of_reasoning_effort,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
def _run_test_raw(
|
||||
self,
|
||||
model: str,
|
||||
expected_score_of_reasoning_effort: Dict[str, float],
|
||||
other_args: List[str],
|
||||
):
|
||||
process = popen_launch_server(
|
||||
model,
|
||||
_base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
try:
|
||||
self._check_streaming_responses_api_request(model)
|
||||
|
||||
# run multiple tests in parallel since we are mostly bound by the longest generate sequence
|
||||
# instead of the number of questions
|
||||
with ThreadPoolExecutor(max_workers=4) as executor:
|
||||
list(
|
||||
executor.map(
|
||||
lambda d: self._run_one_eval(**d),
|
||||
[
|
||||
dict(
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
expected_score=expected_score,
|
||||
)
|
||||
for reasoning_effort, expected_score in expected_score_of_reasoning_effort.items()
|
||||
],
|
||||
)
|
||||
)
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
def _check_streaming_responses_api_request(self, model):
|
||||
# Use requests to verify /v1/responses streaming
|
||||
url = f"{_base_url}/v1/responses"
|
||||
payload = {
|
||||
"model": model,
|
||||
"input": "What is 1 + 1?",
|
||||
"stream": True,
|
||||
"temperature": 0,
|
||||
}
|
||||
|
||||
response = requests.post(url, json=payload, stream=True)
|
||||
if response.status_code != 200:
|
||||
print(f"Response API failed: {response.text}")
|
||||
response.raise_for_status()
|
||||
|
||||
content = ""
|
||||
for line in response.iter_lines():
|
||||
if line:
|
||||
decoded_line = line.decode("utf-8")
|
||||
if decoded_line.startswith("data: "):
|
||||
data_str = decoded_line[6:]
|
||||
if data_str.strip() == "[DONE]":
|
||||
break
|
||||
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
if data.get("type") == "response.output_text.delta":
|
||||
delta = data.get("delta", "")
|
||||
content += delta
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
print(f"Streaming check response: {content}")
|
||||
self.assertTrue(len(content) > 0)
|
||||
self.assertIn("2", content)
|
||||
|
||||
def _run_one_eval(self, model, reasoning_effort, expected_score):
|
||||
args = SimpleNamespace(
|
||||
base_url=_base_url,
|
||||
model=model,
|
||||
eval_name="gpqa",
|
||||
num_examples=198,
|
||||
# use enough threads to allow parallelism
|
||||
num_threads=198,
|
||||
# TODO 4k is still not enough, we need e.g. 64k token, but that is super slow
|
||||
# otherwise a lot of questions are not answered
|
||||
max_tokens=4096,
|
||||
# simple-evals by default use 0.5 and is better than 0.0 temperature
|
||||
# but here for reproducibility, we use 0.1
|
||||
temperature=0.1,
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
|
||||
setup = f"model={model} reasoning_effort={reasoning_effort} expected_score={expected_score}"
|
||||
|
||||
print(f"Evaluation start: {setup}")
|
||||
metrics = run_eval(args)
|
||||
print(f"Evaluation end: {setup} {metrics=}")
|
||||
self.assertGreaterEqual(metrics["score"], expected_score)
|
||||
|
||||
if is_in_ci():
|
||||
write_github_step_summary(
|
||||
f"### test_gpt_oss_common\n"
|
||||
f"Setup: {setup}\n"
|
||||
f"Score: {metrics['score']:.2f}\n"
|
||||
)
|
||||
@@ -1,139 +0,0 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
|
||||
|
||||
|
||||
class TestHiddenState(CustomTestCase):
|
||||
def test_return_hidden_states(self):
|
||||
prompts = ["Today is", "Today is a sunny day and I like"]
|
||||
model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
input_ids = tokenizer(prompts).input_ids
|
||||
|
||||
sampling_params = {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 8,
|
||||
}
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=model_path,
|
||||
random_seed=42,
|
||||
skip_tokenizer_init=True,
|
||||
enable_return_hidden_states=True,
|
||||
)
|
||||
outputs = engine.generate(
|
||||
input_ids=input_ids,
|
||||
sampling_params=sampling_params,
|
||||
return_hidden_states=True,
|
||||
)
|
||||
engine.shutdown()
|
||||
|
||||
for output in outputs:
|
||||
self.assertEqual(len(output["meta_info"]["hidden_states"]), 8)
|
||||
for i in range(len(output["meta_info"]["hidden_states"])):
|
||||
assert isinstance(output["meta_info"]["hidden_states"][i], list)
|
||||
output["meta_info"]["hidden_states"][i] = torch.tensor(
|
||||
output["meta_info"]["hidden_states"][i], dtype=torch.bfloat16
|
||||
)
|
||||
# Checks that splicing of the batch was done correctly
|
||||
self.assertGreater(
|
||||
outputs[1]["meta_info"]["hidden_states"][0].shape[0],
|
||||
outputs[0]["meta_info"]["hidden_states"][0].shape[0],
|
||||
)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path, torch_dtype=torch.bfloat16, device_map="cuda"
|
||||
)
|
||||
|
||||
for input_id, output in zip(input_ids, outputs):
|
||||
with torch.inference_mode():
|
||||
hf_out = model(
|
||||
torch.tensor(
|
||||
[input_id + output["output_ids"][:-1]], device=model.device
|
||||
),
|
||||
output_hidden_states=True,
|
||||
)
|
||||
print("=== HF Hiddens ===")
|
||||
print(hf_out["hidden_states"][-1][0])
|
||||
sg_hidden_states = torch.cat(
|
||||
[
|
||||
i.unsqueeze(0) if len(i.shape) == 1 else i
|
||||
for i in output["meta_info"]["hidden_states"]
|
||||
]
|
||||
).to("cuda")
|
||||
print("=== SRT Hiddens ===")
|
||||
print(sg_hidden_states)
|
||||
|
||||
print(
|
||||
f"Max diff: {torch.max(torch.abs(hf_out['hidden_states'][-1][0] - sg_hidden_states))}"
|
||||
)
|
||||
|
||||
atol = 0.8
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
hf_out["hidden_states"][-1][0],
|
||||
sg_hidden_states,
|
||||
atol=atol,
|
||||
rtol=0,
|
||||
)
|
||||
)
|
||||
|
||||
def test_repeatedly_changes_hidden_states(self):
|
||||
prompts = ["Today is", "Today is a sunny day and I like"]
|
||||
model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
input_ids = tokenizer(prompts).input_ids
|
||||
|
||||
sampling_params = {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 8,
|
||||
}
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=model_path,
|
||||
random_seed=42,
|
||||
skip_tokenizer_init=True,
|
||||
enable_return_hidden_states=True,
|
||||
)
|
||||
outputs_completion_first_round = engine.generate(
|
||||
input_ids=input_ids,
|
||||
sampling_params=sampling_params,
|
||||
return_hidden_states=True,
|
||||
)
|
||||
outputs_hidden_state = engine.generate(
|
||||
input_ids=input_ids,
|
||||
sampling_params=sampling_params,
|
||||
return_hidden_states=False,
|
||||
)
|
||||
|
||||
outputs_completion_last_round = engine.generate(
|
||||
input_ids=input_ids,
|
||||
sampling_params=sampling_params,
|
||||
return_hidden_states=True,
|
||||
)
|
||||
engine.shutdown()
|
||||
|
||||
for (
|
||||
output_completion_first_round,
|
||||
output_hidden_state,
|
||||
output_completion_last_round,
|
||||
) in zip(
|
||||
outputs_completion_first_round,
|
||||
outputs_hidden_state,
|
||||
outputs_completion_last_round,
|
||||
):
|
||||
self.assertEqual(
|
||||
len(output_completion_first_round["meta_info"]["hidden_states"]), 8
|
||||
)
|
||||
self.assertNotIn("hidden_states", output_hidden_state["meta_info"])
|
||||
self.assertEqual(
|
||||
len(output_completion_last_round["meta_info"]["hidden_states"]), 8
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,151 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestInputEmbeds(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.tokenizer = AutoTokenizer.from_pretrained(cls.model)
|
||||
cls.ref_model = AutoModelForCausalLM.from_pretrained(cls.model)
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--disable-radix", "--cuda-graph-max-bs", 4],
|
||||
)
|
||||
cls.texts = [
|
||||
"The capital of France is",
|
||||
"What is the best time of year to visit Japan for cherry blossoms?",
|
||||
]
|
||||
|
||||
def generate_input_embeddings(self, text):
|
||||
"""Generate input embeddings for a given text."""
|
||||
input_ids = self.tokenizer(text, return_tensors="pt")["input_ids"]
|
||||
embeddings = self.ref_model.get_input_embeddings()(input_ids)
|
||||
return embeddings.squeeze().tolist() # Convert tensor to a list for API use
|
||||
|
||||
def send_request(self, payload):
|
||||
"""Send a POST request to the /generate endpoint and return the response."""
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json=payload,
|
||||
timeout=30, # Set a reasonable timeout for the API request
|
||||
)
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
return {
|
||||
"error": f"Request failed with status {response.status_code}: {response.text}"
|
||||
}
|
||||
|
||||
def send_file_request(self, file_path):
|
||||
"""Send a POST request to the /generate_from_file endpoint with a file."""
|
||||
with open(file_path, "rb") as f:
|
||||
response = requests.post(
|
||||
self.base_url + "/generate_from_file",
|
||||
files={"file": f},
|
||||
timeout=30, # Set a reasonable timeout for the API request
|
||||
)
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
return {
|
||||
"error": f"Request failed with status {response.status_code}: {response.text}"
|
||||
}
|
||||
|
||||
def test_text_based_response(self):
|
||||
"""Test and print API responses using text-based input."""
|
||||
for text in self.texts:
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"text": text,
|
||||
"sampling_params": {"temperature": 0, "max_new_tokens": 50},
|
||||
}
|
||||
response = self.send_request(payload)
|
||||
print(
|
||||
f"Text Input: {text}\nResponse: {json.dumps(response, indent=2)}\n{'-' * 80}"
|
||||
)
|
||||
|
||||
def test_embedding_based_response(self):
|
||||
"""Test and print API responses using input embeddings."""
|
||||
for text in self.texts:
|
||||
embeddings = self.generate_input_embeddings(text)
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"input_embeds": embeddings,
|
||||
"sampling_params": {"temperature": 0, "max_new_tokens": 50},
|
||||
}
|
||||
response = self.send_request(payload)
|
||||
print(
|
||||
f"Embeddings Input (for text '{text}'):\nResponse: {json.dumps(response, indent=2)}\n{'-' * 80}"
|
||||
)
|
||||
|
||||
def test_compare_text_vs_embedding(self):
|
||||
"""Test and compare responses for text-based and embedding-based inputs."""
|
||||
for text in self.texts:
|
||||
# Text-based payload
|
||||
text_payload = {
|
||||
"model": self.model,
|
||||
"text": text,
|
||||
"sampling_params": {"temperature": 0, "max_new_tokens": 50},
|
||||
}
|
||||
# Embedding-based payload
|
||||
embeddings = self.generate_input_embeddings(text)
|
||||
embed_payload = {
|
||||
"model": self.model,
|
||||
"input_embeds": embeddings,
|
||||
"sampling_params": {"temperature": 0, "max_new_tokens": 50},
|
||||
}
|
||||
# Get responses
|
||||
text_response = self.send_request(text_payload)
|
||||
embed_response = self.send_request(embed_payload)
|
||||
# Print responses
|
||||
print(
|
||||
f"Text Input: {text}\nText-Based Response: {json.dumps(text_response, indent=2)}\n"
|
||||
)
|
||||
print(
|
||||
f"Embeddings Input (for text '{text}'):\nEmbedding-Based Response: {json.dumps(embed_response, indent=2)}\n{'-' * 80}"
|
||||
)
|
||||
# This is flaky, so we skip this temporarily
|
||||
# self.assertEqual(text_response["text"], embed_response["text"])
|
||||
|
||||
def test_generate_from_file(self):
|
||||
"""Test the /generate_from_file endpoint using tokenized embeddings."""
|
||||
for text in self.texts:
|
||||
embeddings = self.generate_input_embeddings(text)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w", suffix=".json", delete=False
|
||||
) as tmp_file:
|
||||
json.dump(embeddings, tmp_file)
|
||||
tmp_file_path = tmp_file.name
|
||||
|
||||
try:
|
||||
response = self.send_file_request(tmp_file_path)
|
||||
print(
|
||||
f"Text Input: {text}\nResponse from /generate_from_file: {json.dumps(response, indent=2)}\n{'-' * 80}"
|
||||
)
|
||||
finally:
|
||||
# Ensure the temporary file is deleted
|
||||
os.remove(tmp_file_path)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,577 +0,0 @@
|
||||
import copy
|
||||
import unittest
|
||||
|
||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
)
|
||||
|
||||
|
||||
class TestGenerateReqInputNormalization(CustomTestCase):
|
||||
"""Test the normalization of GenerateReqInput for batch processing and different input formats."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
def setUp(self):
|
||||
# Common setup for all tests
|
||||
self.base_req = GenerateReqInput(
|
||||
text=["Hello", "World"],
|
||||
sampling_params=[{}, {}],
|
||||
rid=["id1", "id2"],
|
||||
)
|
||||
|
||||
def test_single_image_to_list_of_lists(self):
|
||||
"""Test that a single image is converted to a list of single-image lists."""
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.image_data = "single_image.jpg" # A single image (non-list)
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should be converted to [[image], [image]]
|
||||
self.assertEqual(len(req.image_data), 2)
|
||||
self.assertEqual(len(req.image_data[0]), 1)
|
||||
self.assertEqual(len(req.image_data[1]), 1)
|
||||
self.assertEqual(req.image_data[0][0], "single_image.jpg")
|
||||
self.assertEqual(req.image_data[1][0], "single_image.jpg")
|
||||
|
||||
# Check modalities
|
||||
self.assertEqual(req.modalities, ["image", "image"])
|
||||
|
||||
def test_list_of_images_to_list_of_lists(self):
|
||||
"""Test that a list of images is converted to a list of single-image lists."""
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.image_data = ["image1.jpg", "image2.jpg"] # List of images
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should be converted to [[image1], [image2]]
|
||||
self.assertEqual(len(req.image_data), 2)
|
||||
self.assertEqual(len(req.image_data[0]), 1)
|
||||
self.assertEqual(len(req.image_data[1]), 1)
|
||||
self.assertEqual(req.image_data[0][0], "image1.jpg")
|
||||
self.assertEqual(req.image_data[1][0], "image2.jpg")
|
||||
|
||||
# Check modalities
|
||||
self.assertEqual(req.modalities, ["image", "image"])
|
||||
|
||||
def test_list_of_lists_with_different_modalities(self):
|
||||
"""Test handling of list of lists of images with different modalities."""
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.image_data = [
|
||||
["image1.jpg"], # Single image (image modality)
|
||||
["image2.jpg", "image3.jpg"], # Multiple images (multi-images modality)
|
||||
]
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Structure should remain the same
|
||||
self.assertEqual(len(req.image_data), 2)
|
||||
self.assertEqual(len(req.image_data[0]), 1)
|
||||
self.assertEqual(len(req.image_data[1]), 2)
|
||||
|
||||
# Check modalities
|
||||
self.assertEqual(req.modalities, ["image", "multi-images"])
|
||||
|
||||
def test_list_of_lists_with_none_values(self):
|
||||
"""Test handling of list of lists with None values."""
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.image_data = [
|
||||
[None], # None value
|
||||
["image.jpg"], # Single image
|
||||
]
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Structure should remain the same
|
||||
self.assertEqual(len(req.image_data), 2)
|
||||
self.assertEqual(len(req.image_data[0]), 1)
|
||||
self.assertEqual(len(req.image_data[1]), 1)
|
||||
|
||||
# Check modalities
|
||||
self.assertEqual(req.modalities, [None, "image"])
|
||||
|
||||
def test_expanding_parallel_sample_correlation(self):
|
||||
"""Test that when expanding with parallel samples, prompts, images and modalities are properly correlated."""
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.text = ["Prompt 1", "Prompt 2"]
|
||||
req.image_data = [
|
||||
["image1.jpg"],
|
||||
["image2.jpg", "image3.jpg"],
|
||||
]
|
||||
req.sampling_params = {"n": 3} # All prompts get 3 samples
|
||||
|
||||
# Define expected values before normalization
|
||||
expected_text = req.text * 3
|
||||
expected_images = req.image_data * 3
|
||||
expected_modalities = ["image", "multi-images"] * 3
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should be expanded to 6 items (2 original * 3 parallel)
|
||||
self.assertEqual(len(req.image_data), 6)
|
||||
|
||||
# Check that images are properly expanded
|
||||
self.assertEqual(req.image_data, expected_images)
|
||||
|
||||
# Check modalities
|
||||
self.assertEqual(req.modalities, expected_modalities)
|
||||
|
||||
# Ensure that text items are properly duplicated too
|
||||
self.assertEqual(req.text, expected_text)
|
||||
|
||||
def test_specific_parallel_n_per_sample(self):
|
||||
"""Test parallel expansion when different samples have different n values."""
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.text = ["Prompt 1", "Prompt 2"]
|
||||
req.image_data = [
|
||||
["image1.jpg"],
|
||||
["image2.jpg", "image3.jpg"],
|
||||
]
|
||||
req.sampling_params = [
|
||||
{"n": 2},
|
||||
{"n": 2},
|
||||
] # First prompt gets 2 samples, second prompt gets 2 samples
|
||||
|
||||
expected_images = req.image_data * 2
|
||||
expected_modalities = ["image", "multi-images"] * 2
|
||||
expected_text = req.text * 2
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should be expanded to 4 items (2 original * 2 parallel)
|
||||
self.assertEqual(len(req.image_data), 4)
|
||||
|
||||
# Check that the first 2 are copies for the first prompt
|
||||
self.assertEqual(req.image_data, expected_images)
|
||||
|
||||
# Check modalities
|
||||
self.assertEqual(req.modalities, expected_modalities)
|
||||
|
||||
# Check text expansion
|
||||
self.assertEqual(req.text, expected_text)
|
||||
|
||||
def test_mixed_none_and_images_with_parallel_samples(self):
|
||||
"""Test that when some batch items have images and others None, parallel expansion works correctly."""
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.text = ["Prompt 1", "Prompt 2", "Prompt 3"]
|
||||
req.rid = ["id1", "id2", "id3"]
|
||||
req.image_data = [
|
||||
["image1.jpg"],
|
||||
None,
|
||||
["image3_1.jpg", "image3_2.jpg"],
|
||||
]
|
||||
req.sampling_params = {"n": 2} # All prompts get 2 samples
|
||||
|
||||
expected_images = req.image_data * 2
|
||||
expected_modalities = ["image", None, "multi-images"] * 2
|
||||
expected_text = req.text * 2
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should be expanded to 6 items (3 original * 2 parallel)
|
||||
self.assertEqual(len(req.image_data), 6)
|
||||
|
||||
# Check image data
|
||||
self.assertEqual(req.image_data, expected_images)
|
||||
|
||||
# Check modalities
|
||||
self.assertEqual(req.modalities, expected_modalities)
|
||||
|
||||
# Check text expansion
|
||||
self.assertEqual(req.text, expected_text)
|
||||
|
||||
def test_correlation_with_sampling_params(self):
|
||||
"""Test that sampling parameters are correctly correlated with prompts during expansion."""
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.text = ["Prompt 1", "Prompt 2"]
|
||||
req.image_data = [
|
||||
["image1.jpg"],
|
||||
["image2.jpg"],
|
||||
]
|
||||
req.sampling_params = [
|
||||
{"temperature": 0.7, "n": 2},
|
||||
{"temperature": 0.9, "n": 2},
|
||||
]
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Check sampling params expansion
|
||||
self.assertEqual(len(req.sampling_params), 4)
|
||||
self.assertEqual(req.sampling_params[0]["temperature"], 0.7)
|
||||
self.assertEqual(req.sampling_params[1]["temperature"], 0.9)
|
||||
self.assertEqual(req.sampling_params[2]["temperature"], 0.7)
|
||||
self.assertEqual(req.sampling_params[3]["temperature"], 0.9)
|
||||
|
||||
# Should be expanded to 4 items (2 original * 2 parallel)
|
||||
self.assertEqual(len(req.image_data), 4)
|
||||
|
||||
# Check correlation with images
|
||||
self.assertEqual(req.image_data[0], ["image1.jpg"])
|
||||
self.assertEqual(req.image_data[1], ["image2.jpg"])
|
||||
self.assertEqual(req.image_data[2], ["image1.jpg"])
|
||||
self.assertEqual(req.image_data[3], ["image2.jpg"])
|
||||
|
||||
def test_single_example_with_image(self):
|
||||
"""Test handling of single example with image."""
|
||||
req = GenerateReqInput(
|
||||
text="Hello",
|
||||
image_data="single_image.jpg",
|
||||
)
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# For single examples, image_data doesn't get processed into lists
|
||||
self.assertEqual(req.image_data, "single_image.jpg")
|
||||
self.assertIsNone(req.modalities) # Modalities isn't set for single examples
|
||||
|
||||
def test_single_to_batch_with_parallel_sampling(self):
|
||||
"""Test single example converted to batch with parallel sampling."""
|
||||
req = GenerateReqInput(
|
||||
text="Hello",
|
||||
image_data="single_image.jpg",
|
||||
sampling_params={"n": 3}, # parallel_sample_num = 3
|
||||
)
|
||||
|
||||
# Define expected values before normalization
|
||||
expected_text = ["Hello"] * 3
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should be converted to batch with text=["Hello"]
|
||||
self.assertEqual(req.text, expected_text)
|
||||
|
||||
# Image should be automatically wrapped to list of lists with length 1*3=3
|
||||
self.assertEqual(len(req.image_data), 3)
|
||||
self.assertEqual(req.image_data[0][0], "single_image.jpg")
|
||||
self.assertEqual(req.image_data[1][0], "single_image.jpg")
|
||||
self.assertEqual(req.image_data[2][0], "single_image.jpg")
|
||||
|
||||
# Modalities should be set for all 3 examples
|
||||
self.assertEqual(req.modalities, ["image", "image", "image"])
|
||||
|
||||
def test_audio_data_handling(self):
|
||||
"""Test handling of audio_data."""
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.audio_data = "audio.mp3" # Single audio
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should be converted to ["audio.mp3", "audio.mp3"]
|
||||
self.assertEqual(len(req.audio_data), 2)
|
||||
self.assertEqual(req.audio_data[0], "audio.mp3")
|
||||
self.assertEqual(req.audio_data[1], "audio.mp3")
|
||||
|
||||
# Test with list
|
||||
req = copy.deepcopy(self.base_req)
|
||||
req.audio_data = ["audio1.mp3", "audio2.mp3"]
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should remain the same
|
||||
self.assertEqual(len(req.audio_data), 2)
|
||||
self.assertEqual(req.audio_data[0], "audio1.mp3")
|
||||
self.assertEqual(req.audio_data[1], "audio2.mp3")
|
||||
|
||||
def test_input_ids_normalization(self):
|
||||
"""Test normalization of input_ids instead of text."""
|
||||
# Test single input_ids
|
||||
req = GenerateReqInput(input_ids=[1, 2, 3])
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertTrue(req.is_single)
|
||||
self.assertEqual(req.batch_size, 1)
|
||||
|
||||
# Test batch input_ids
|
||||
req = GenerateReqInput(input_ids=[[1, 2, 3], [4, 5, 6]])
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertFalse(req.is_single)
|
||||
self.assertEqual(req.batch_size, 2)
|
||||
|
||||
# Test with parallel sampling
|
||||
req = GenerateReqInput(
|
||||
input_ids=[[1, 2, 3], [4, 5, 6]], sampling_params={"n": 2}
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(len(req.input_ids), 4) # 2 original * 2 parallel
|
||||
|
||||
def test_input_embeds_normalization(self):
|
||||
"""Test normalization of input_embeds."""
|
||||
# Test single input_embeds
|
||||
req = GenerateReqInput(input_embeds=[[0.1, 0.2], [0.3, 0.4]])
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertTrue(req.is_single)
|
||||
self.assertEqual(req.batch_size, 1)
|
||||
|
||||
# Test batch input_embeds
|
||||
req = GenerateReqInput(input_embeds=[[[0.1, 0.2]], [[0.3, 0.4]]])
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertFalse(req.is_single)
|
||||
self.assertEqual(req.batch_size, 2)
|
||||
|
||||
def test_input_embeds_with_parallel_sampling(self):
|
||||
"""Test input_embeds normalization with parallel sampling (n > 1)."""
|
||||
# Test single input_embeds with parallel sampling
|
||||
req = GenerateReqInput(
|
||||
input_embeds=[[0.1, 0.2]], # single embedding vector
|
||||
sampling_params={"n": 2},
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should be converted from single to batch and then expanded
|
||||
self.assertFalse(req.is_single)
|
||||
self.assertEqual(len(req.input_embeds), 2)
|
||||
# Both should be the same input_embeds
|
||||
self.assertEqual(req.input_embeds[0], [[0.1, 0.2]])
|
||||
self.assertEqual(req.input_embeds[1], [[0.1, 0.2]])
|
||||
|
||||
# Test batch input_embeds with parallel sampling
|
||||
req = GenerateReqInput(
|
||||
input_embeds=[[[0.1, 0.2]], [[0.3, 0.4]]], sampling_params={"n": 3}
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should be expanded
|
||||
self.assertFalse(req.is_single)
|
||||
self.assertEqual(len(req.input_embeds), 6)
|
||||
|
||||
# Check that the expansion is correct
|
||||
expected_embeds = [[[0.1, 0.2]], [[0.3, 0.4]]] * 3
|
||||
self.assertEqual(req.input_embeds, expected_embeds)
|
||||
|
||||
# Test with different n values per sample (should raise error)
|
||||
req = GenerateReqInput(
|
||||
input_embeds=[[[0.1, 0.2]], [[0.3, 0.4]]],
|
||||
sampling_params=[{"n": 2}, {"n": 3}],
|
||||
)
|
||||
with self.assertRaises(ValueError):
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
def test_input_embeds_single_to_batch_conversion(self):
|
||||
"""Test that single input_embeds are properly converted to batch when using parallel sampling."""
|
||||
# Test the specific case that was fixed: single input_embeds with n > 1
|
||||
req = GenerateReqInput(
|
||||
input_embeds=[[0.1, 0.2, 0.3]], sampling_params={"n": 2} # Single embedding
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Should convert single to batch and then expand
|
||||
self.assertFalse(req.is_single)
|
||||
self.assertEqual(len(req.input_embeds), 2)
|
||||
|
||||
# Both should be the same single embedding
|
||||
self.assertEqual(req.input_embeds[0], [[0.1, 0.2, 0.3]])
|
||||
self.assertEqual(req.input_embeds[1], [[0.1, 0.2, 0.3]])
|
||||
|
||||
# Test with higher n value
|
||||
req = GenerateReqInput(input_embeds=[[0.1, 0.2, 0.3]], sampling_params={"n": 5})
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
self.assertFalse(req.is_single)
|
||||
self.assertEqual(len(req.input_embeds), 5)
|
||||
|
||||
# All should be the same
|
||||
for i in range(5):
|
||||
self.assertEqual(req.input_embeds[i], [[0.1, 0.2, 0.3]])
|
||||
|
||||
def test_lora_path_normalization(self):
|
||||
"""Test normalization of lora_path."""
|
||||
# Test single lora_path with batch input
|
||||
req = GenerateReqInput(text=["Hello", "World"], lora_path="path/to/lora")
|
||||
|
||||
# Define expected lora_paths before normalization
|
||||
expected_lora_paths = ["path/to/lora", "path/to/lora"]
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req.lora_path, expected_lora_paths)
|
||||
|
||||
# Test list of lora_paths
|
||||
req = GenerateReqInput(text=["Hello", "World"], lora_path=["path1", "path2"])
|
||||
|
||||
# Define expected lora_paths before normalization
|
||||
expected_lora_paths = ["path1", "path2"]
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req.lora_path, expected_lora_paths)
|
||||
|
||||
# Test with parallel sampling
|
||||
req = GenerateReqInput(
|
||||
text=["Hello", "World"],
|
||||
lora_path=["path1", "path2"],
|
||||
sampling_params={"n": 2},
|
||||
)
|
||||
|
||||
# Define expected lora_paths before normalization
|
||||
expected_lora_paths = ["path1", "path2"] * 2
|
||||
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req.lora_path, expected_lora_paths)
|
||||
|
||||
def test_logprob_parameters_normalization(self):
|
||||
"""Test normalization of logprob-related parameters."""
|
||||
# Test single example
|
||||
req = GenerateReqInput(
|
||||
text="Hello",
|
||||
return_logprob=True,
|
||||
logprob_start_len=10,
|
||||
top_logprobs_num=5,
|
||||
token_ids_logprob=[7, 8, 9],
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req.return_logprob, True)
|
||||
self.assertEqual(req.logprob_start_len, 10)
|
||||
self.assertEqual(req.top_logprobs_num, 5)
|
||||
self.assertEqual(req.token_ids_logprob, [7, 8, 9])
|
||||
|
||||
# Test batch with scalar values
|
||||
req = GenerateReqInput(
|
||||
text=["Hello", "World"],
|
||||
return_logprob=True,
|
||||
logprob_start_len=10,
|
||||
top_logprobs_num=5,
|
||||
token_ids_logprob=[7, 8, 9],
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req.return_logprob, [True, True])
|
||||
self.assertEqual(req.logprob_start_len, [10, 10])
|
||||
self.assertEqual(req.top_logprobs_num, [5, 5])
|
||||
self.assertEqual(req.token_ids_logprob, [[7, 8, 9], [7, 8, 9]])
|
||||
|
||||
# Test batch with list values
|
||||
req = GenerateReqInput(
|
||||
text=["Hello", "World"],
|
||||
return_logprob=[True, False],
|
||||
logprob_start_len=[10, 5],
|
||||
top_logprobs_num=[5, 3],
|
||||
token_ids_logprob=[[7, 8, 9], [4, 5, 6]],
|
||||
return_hidden_states=[False, False, True],
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req.return_logprob, [True, False])
|
||||
self.assertEqual(req.logprob_start_len, [10, 5])
|
||||
self.assertEqual(req.top_logprobs_num, [5, 3])
|
||||
self.assertEqual(req.token_ids_logprob, [[7, 8, 9], [4, 5, 6]])
|
||||
self.assertEqual(req.return_hidden_states, [False, False, True])
|
||||
|
||||
def test_custom_logit_processor_normalization(self):
|
||||
"""Test normalization of custom_logit_processor."""
|
||||
# Test single processor
|
||||
req = GenerateReqInput(
|
||||
text=["Hello", "World"], custom_logit_processor="serialized_processor"
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(
|
||||
req.custom_logit_processor, ["serialized_processor", "serialized_processor"]
|
||||
)
|
||||
|
||||
# Test list of processors
|
||||
req = GenerateReqInput(
|
||||
text=["Hello", "World"], custom_logit_processor=["processor1", "processor2"]
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req.custom_logit_processor, ["processor1", "processor2"])
|
||||
|
||||
def test_session_params_handling(self):
|
||||
"""Test handling of session_params."""
|
||||
# Test with dict
|
||||
req = GenerateReqInput(
|
||||
text=["Hello", "World"], session_params={"id": "session1", "offset": 10}
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req.session_params, {"id": "session1", "offset": 10})
|
||||
|
||||
# Test with list of dicts
|
||||
req = GenerateReqInput(
|
||||
text=["Hello", "World"],
|
||||
session_params=[{"id": "session1"}, {"id": "session2"}],
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req.session_params, [{"id": "session1"}, {"id": "session2"}])
|
||||
|
||||
def test_getitem_method(self):
|
||||
"""Test the __getitem__ method."""
|
||||
req = GenerateReqInput(
|
||||
text=["Hello", "World"],
|
||||
image_data=[["img1.jpg"], ["img2.jpg"]],
|
||||
audio_data=["audio1.mp3", "audio2.mp3"],
|
||||
sampling_params=[{"temp": 0.7}, {"temp": 0.8}],
|
||||
rid=["id1", "id2"],
|
||||
return_logprob=[True, False],
|
||||
logprob_start_len=[10, 5],
|
||||
top_logprobs_num=[5, 3],
|
||||
token_ids_logprob=[[7, 8, 9], [4, 5, 6]],
|
||||
stream=True,
|
||||
log_metrics=True,
|
||||
modalities=["image", "image"],
|
||||
lora_path=["path1", "path2"],
|
||||
custom_logit_processor=["processor1", "processor2"],
|
||||
return_hidden_states=True,
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Get the first item
|
||||
item0 = req[0]
|
||||
self.assertEqual(item0.text, "Hello")
|
||||
self.assertEqual(item0.image_data, ["img1.jpg"])
|
||||
self.assertEqual(item0.audio_data, "audio1.mp3")
|
||||
self.assertEqual(item0.sampling_params, {"temp": 0.7})
|
||||
self.assertEqual(item0.rid, "id1")
|
||||
self.assertEqual(item0.return_logprob, True)
|
||||
self.assertEqual(item0.logprob_start_len, 10)
|
||||
self.assertEqual(item0.top_logprobs_num, 5)
|
||||
self.assertEqual(item0.token_ids_logprob, [7, 8, 9])
|
||||
self.assertEqual(item0.stream, True)
|
||||
self.assertEqual(item0.log_metrics, True)
|
||||
self.assertEqual(item0.modalities, "image")
|
||||
self.assertEqual(item0.lora_path, "path1")
|
||||
self.assertEqual(item0.custom_logit_processor, "processor1")
|
||||
self.assertEqual(item0.return_hidden_states, True)
|
||||
|
||||
def test_regenerate_rid(self):
|
||||
"""Test the regenerate_rid method."""
|
||||
req = GenerateReqInput(text="Hello")
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
original_rid = req.rid
|
||||
new_rid = req.regenerate_rid()
|
||||
|
||||
self.assertNotEqual(original_rid, new_rid)
|
||||
self.assertEqual(req.rid, new_rid)
|
||||
|
||||
def test_error_cases(self):
|
||||
"""Test various error cases."""
|
||||
# Test when neither text, input_ids, nor input_embeds is provided
|
||||
with self.assertRaises(ValueError):
|
||||
req = GenerateReqInput()
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
# Test when all of text, input_ids, and input_embeds are provided
|
||||
with self.assertRaises(ValueError):
|
||||
req = GenerateReqInput(
|
||||
text="Hello", input_ids=[1, 2, 3], input_embeds=[[0.1, 0.2]]
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
|
||||
def test_multiple_input_formats(self):
|
||||
"""Test different combinations of input formats."""
|
||||
# Test with text only
|
||||
req = GenerateReqInput(text="Hello")
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertTrue(req.is_single)
|
||||
|
||||
# Test with input_ids only
|
||||
req = GenerateReqInput(input_ids=[1, 2, 3])
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertTrue(req.is_single)
|
||||
|
||||
# Test with input_embeds only
|
||||
req = GenerateReqInput(input_embeds=[[0.1, 0.2]])
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertTrue(req.is_single)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,152 +0,0 @@
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.srt.model_executor.hook_manager import register_forward_hooks
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
HOOK_CALLS = []
|
||||
|
||||
|
||||
def dummy_hook_factory(config):
|
||||
"""Factory that returns a forward hook capturing a tag from config."""
|
||||
tag = config.get("tag", "default")
|
||||
|
||||
def hook(module, inputs, output):
|
||||
HOOK_CALLS.append(
|
||||
{
|
||||
"module_type": type(module).__name__,
|
||||
"tag": tag,
|
||||
"shape": tuple(output.shape),
|
||||
}
|
||||
)
|
||||
return output
|
||||
|
||||
return hook
|
||||
|
||||
|
||||
class TinyModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.inner = nn.Sequential(
|
||||
nn.Linear(4, 2),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self.outer = nn.Sequential(
|
||||
nn.Linear(4, 4),
|
||||
nn.ReLU(),
|
||||
self.inner,
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.outer(x)
|
||||
|
||||
|
||||
class TestAttachHooks(CustomTestCase):
|
||||
"""Tests for register_forward_hooks / resolve_callable integration."""
|
||||
|
||||
def setUp(self):
|
||||
HOOK_CALLS.clear()
|
||||
|
||||
def test_hook_is_attached(self):
|
||||
"""Hook from a factory string is registered and fired."""
|
||||
hook_specs = [
|
||||
{
|
||||
"target_modules": ["outer.0", "outer.1"],
|
||||
"hook_factory": "test_model_hooks:dummy_hook_factory",
|
||||
"config": {"tag": "forward-ok"},
|
||||
},
|
||||
{
|
||||
"target_modules": ["inner.*"],
|
||||
"hook_factory": "test_model_hooks:dummy_hook_factory",
|
||||
"config": {"tag": "forward-ok"},
|
||||
},
|
||||
]
|
||||
|
||||
model = TinyModel()
|
||||
register_forward_hooks(model, hook_specs)
|
||||
|
||||
x = torch.randn(3, 4)
|
||||
_ = model(x)
|
||||
|
||||
self.assertEqual(
|
||||
len(HOOK_CALLS),
|
||||
4,
|
||||
"Forward hook was not called correct number of times",
|
||||
)
|
||||
tags = {call["tag"] for call in HOOK_CALLS}
|
||||
self.assertIn("forward-ok", tags)
|
||||
|
||||
def test_no_matching_modules_does_not_crash(self):
|
||||
"""Hook spec with no matching modules should not crash."""
|
||||
model = TinyModel()
|
||||
hook_specs = [
|
||||
{
|
||||
"name": "no_match",
|
||||
"target_modules": ["does_not_exist.*"],
|
||||
"hook_factory": "test_model_hooks:dummy_hook_factory",
|
||||
"config": {"tag": "unused"},
|
||||
}
|
||||
]
|
||||
|
||||
register_forward_hooks(model, hook_specs)
|
||||
|
||||
x = torch.randn(3, 4)
|
||||
_ = model(x)
|
||||
|
||||
# No hooks should have fired
|
||||
self.assertEqual(len(HOOK_CALLS), 0)
|
||||
|
||||
def test_cli_hooks_reach_model(self):
|
||||
"""
|
||||
Ensure that when hooks are provided via CLI, they are parsed into
|
||||
ServerArgs, passed to register_forward_hooks, and actually
|
||||
run during a forward pass.
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
ServerArgs.add_cli_args(parser)
|
||||
|
||||
hooks_spec = [
|
||||
{
|
||||
"name": "outer_and_inner_from_cli",
|
||||
"target_modules": ["outer.0", "outer.1", "inner.*"],
|
||||
"hook_factory": "test_model_hooks:dummy_hook_factory",
|
||||
"config": {"tag": "cli-hook"},
|
||||
}
|
||||
]
|
||||
|
||||
cli_args = [
|
||||
"--model-path",
|
||||
"Qwen/Qwen2-7B-Instruct", # Dummy value; not used in this test
|
||||
"--forward-hooks",
|
||||
json.dumps(hooks_spec),
|
||||
]
|
||||
|
||||
args = parser.parse_args(cli_args)
|
||||
server_args = ServerArgs.from_cli_args(args)
|
||||
|
||||
self.assertEqual(server_args.forward_hooks, hooks_spec)
|
||||
|
||||
model = TinyModel()
|
||||
register_forward_hooks(model, server_args.forward_hooks)
|
||||
|
||||
x = torch.randn(3, 4)
|
||||
_ = model(x)
|
||||
|
||||
# We expect hooks on outer.0, outer.1, inner.0, inner.1 => 4 calls
|
||||
self.assertEqual(
|
||||
len(HOOK_CALLS),
|
||||
4,
|
||||
"CLI-configured hooks did not fire expected number of times",
|
||||
)
|
||||
|
||||
tags = {call["tag"] for call in HOOK_CALLS}
|
||||
self.assertEqual(tags, {"cli-hook"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
# unittest.main()
|
||||
@@ -1,47 +0,0 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestPageSize(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
os.environ["SGLANG_DEBUG_MEMORY_POOL"] = "1"
|
||||
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--page-size", 4, "--chunked-prefill-size", 128],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], 0.65)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,88 +0,0 @@
|
||||
import asyncio
|
||||
import os
|
||||
import re
|
||||
import unittest
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
STDERR_FILENAME,
|
||||
STDOUT_FILENAME,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
send_concurrent_generate_requests,
|
||||
send_generate_requests,
|
||||
)
|
||||
|
||||
|
||||
class TestMaxQueuedRequests(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
cls.stdout = open(STDOUT_FILENAME, "w")
|
||||
cls.stderr = open(STDERR_FILENAME, "w")
|
||||
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=(
|
||||
"--max-running-requests", # Enforce max request concurrency is 1
|
||||
"1",
|
||||
"--max-queued-requests", # Enforce max queued request number is 1
|
||||
"1",
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
),
|
||||
return_stdout_stderr=(cls.stdout, cls.stderr),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
cls.stdout.close()
|
||||
cls.stderr.close()
|
||||
os.remove(STDOUT_FILENAME)
|
||||
os.remove(STDERR_FILENAME)
|
||||
|
||||
def test_max_queued_requests_validation_with_serial_requests(self):
|
||||
"""Verify request is not throttled when the max concurrency is 1."""
|
||||
status_codes = send_generate_requests(
|
||||
self.base_url,
|
||||
num_requests=10,
|
||||
)
|
||||
|
||||
for status_code in status_codes:
|
||||
assert status_code == 200 # request shouldn't be throttled
|
||||
|
||||
def test_max_queued_requests_validation_with_concurrent_requests(self):
|
||||
"""Verify request throttling with concurrent requests."""
|
||||
status_codes = asyncio.run(
|
||||
send_concurrent_generate_requests(self.base_url, num_requests=10)
|
||||
)
|
||||
self.assertLessEqual(status_codes.count(200), 2)
|
||||
|
||||
# expected_status_codes = [200, 200, 503, 503, 503, 503, 503, 503, 503, 503]
|
||||
# self.assertEqual(status_codes, expected_status_codes)
|
||||
|
||||
def test_max_running_requests_and_max_queued_request_validation(self):
|
||||
"""Verify running request and queued request numbers based on server logs."""
|
||||
rr_pattern = re.compile(r"#running-req:\s*(\d+)")
|
||||
qr_pattern = re.compile(r"#queue-req:\s*(\d+)")
|
||||
|
||||
with open(STDERR_FILENAME) as lines:
|
||||
for line in lines:
|
||||
rr_match, qr_match = rr_pattern.search(line), qr_pattern.search(line)
|
||||
if rr_match:
|
||||
assert int(rr_match.group(1)) <= 1
|
||||
if qr_match:
|
||||
assert int(qr_match.group(1)) <= 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,590 +0,0 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from sglang.srt.entrypoints.engine import Engine
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
|
||||
|
||||
TEST_MODEL_NAME = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
|
||||
class TestScoreAPI(CustomTestCase):
|
||||
"""Test the scoring API functionality."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up each test case."""
|
||||
self.engine = Engine(model_path=TEST_MODEL_NAME)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after each test case."""
|
||||
if self.engine is not None:
|
||||
self.engine.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def compute_hf_scores(
|
||||
self, query, items, label_token_ids, apply_softmax=False, item_first=False
|
||||
):
|
||||
"""Compute scores using direct HuggingFace model inference.
|
||||
Returns probabilities for each token ID, optionally normalized with softmax.
|
||||
|
||||
Args:
|
||||
query: The query text
|
||||
items: List of item texts
|
||||
label_token_ids: List of token IDs to compute probabilities for
|
||||
apply_softmax: Whether to normalize probabilities using softmax
|
||||
item_first: If True, prepend items to query. Otherwise append items to query.
|
||||
"""
|
||||
# Initialize HF model and tokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
TEST_MODEL_NAME, trust_remote_code=True
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
TEST_MODEL_NAME, trust_remote_code=True
|
||||
)
|
||||
|
||||
try:
|
||||
scores = []
|
||||
for item in items:
|
||||
# Construct full text based on item_first parameter
|
||||
full_text = f"{item}{query}" if item_first else f"{query}{item}"
|
||||
inputs = tokenizer(full_text, return_tensors="pt").to(model.device)
|
||||
|
||||
# Get logits for the last token
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
last_token_logits = outputs.logits[0, -1]
|
||||
|
||||
# Get logits for just our target tokens
|
||||
target_logits = last_token_logits[label_token_ids]
|
||||
|
||||
# Apply softmax over just the target tokens
|
||||
target_probs = torch.softmax(target_logits, dim=-1)
|
||||
|
||||
# Convert to list of probabilities in order of label_token_ids
|
||||
probs = [target_probs[i].item() for i in range(len(label_token_ids))]
|
||||
|
||||
scores.append(probs)
|
||||
|
||||
return scores
|
||||
finally:
|
||||
# Clean up HF resources
|
||||
model.cpu()
|
||||
del model
|
||||
del tokenizer
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def _get_token_ids(self, tokens):
|
||||
"""Helper method to get token IDs for a list of tokens."""
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
TEST_MODEL_NAME, trust_remote_code=True
|
||||
)
|
||||
try:
|
||||
label_token_ids = []
|
||||
for token in tokens:
|
||||
encoding = tokenizer.encode_plus(token, add_special_tokens=False)
|
||||
token_ids = encoding["input_ids"]
|
||||
label_token_ids.append(token_ids[0])
|
||||
return label_token_ids
|
||||
finally:
|
||||
del tokenizer
|
||||
|
||||
def _compare_scores(self, hf_scores, sglang_scores, label_token_ids, case_name=""):
|
||||
"""Helper method to compare scores between HF and SGLang using relative tolerance."""
|
||||
self.assertEqual(
|
||||
len(hf_scores),
|
||||
len(sglang_scores),
|
||||
f"Score lengths don't match for {case_name}",
|
||||
)
|
||||
|
||||
# Use a relative tolerance of 1%
|
||||
TOLERANCE = 0.01
|
||||
|
||||
for hf_score_list, sglang_score_list in zip(hf_scores, sglang_scores):
|
||||
self.assertEqual(
|
||||
len(hf_score_list),
|
||||
len(sglang_score_list),
|
||||
f"Score list lengths don't match for {case_name}",
|
||||
)
|
||||
|
||||
for hf_score, sglang_score in zip(hf_score_list, sglang_score_list):
|
||||
diff = abs(hf_score - sglang_score)
|
||||
self.assertLessEqual(
|
||||
diff,
|
||||
TOLERANCE,
|
||||
msg=f"Scores differ by {diff:.2%} ({case_name}): "
|
||||
f"HF={hf_score:.6f}, SGLang={sglang_score:.6f}",
|
||||
)
|
||||
|
||||
self.assertGreaterEqual(
|
||||
sglang_score, 0, f"SGLang score {sglang_score:.6f} not in [0,1]"
|
||||
)
|
||||
self.assertLessEqual(
|
||||
sglang_score, 1, f"SGLang score {sglang_score:.6f} not in [0,1]"
|
||||
)
|
||||
|
||||
self.assertAlmostEqual(
|
||||
sum(sglang_score_list),
|
||||
1.0,
|
||||
places=6,
|
||||
msg=f"SGLang scores don't sum to 1 ({case_name}): {sum(sglang_score_list):.6f}",
|
||||
)
|
||||
|
||||
def test_score_consistency(self):
|
||||
"""Test that SGLang scoring matches direct HuggingFace model scoring."""
|
||||
# Define test cases
|
||||
test_cases = [
|
||||
{
|
||||
"name": "default case",
|
||||
"query": "I pledge allegiance",
|
||||
"items": ["", " to"],
|
||||
"item_first": False,
|
||||
},
|
||||
{
|
||||
"name": "item_first case",
|
||||
"query": " is a city",
|
||||
"items": ["Tokyo", "Japan"],
|
||||
"item_first": True,
|
||||
},
|
||||
]
|
||||
|
||||
# Common tokens to test for all cases
|
||||
tokens = [" to", " the"]
|
||||
label_token_ids = self._get_token_ids(tokens)
|
||||
|
||||
# Run each test case
|
||||
for case in test_cases:
|
||||
# Get scores from SGLang
|
||||
sglang_scores = self.engine.score(
|
||||
query=case["query"],
|
||||
items=case["items"],
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
item_first=case["item_first"],
|
||||
)
|
||||
|
||||
# Get scores from HuggingFace using the same parameters
|
||||
hf_scores = self.compute_hf_scores(
|
||||
query=case["query"],
|
||||
items=case["items"],
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
item_first=case["item_first"],
|
||||
)
|
||||
|
||||
# Compare scores
|
||||
self._compare_scores(
|
||||
hf_scores, sglang_scores, label_token_ids, case["name"]
|
||||
)
|
||||
|
||||
def test_score_batch_handling(self):
|
||||
"""Test that batch scoring works correctly."""
|
||||
# Test with different batch sizes
|
||||
batch_sizes = [1, 2, 4, 8]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
for batch_size in batch_sizes:
|
||||
texts = [f"test {i}" for i in range(batch_size)]
|
||||
scores = self.engine.score(
|
||||
query="The test was",
|
||||
items=texts,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
len(scores),
|
||||
batch_size,
|
||||
f"Expected {batch_size} scores, got {len(scores)}",
|
||||
)
|
||||
|
||||
# Verify each score list has the correct length
|
||||
for score_list in scores:
|
||||
self.assertEqual(
|
||||
len(score_list),
|
||||
len(label_token_ids),
|
||||
f"Score list length {len(score_list)} doesn't match label_token_ids length {len(label_token_ids)}",
|
||||
)
|
||||
self.assertTrue(
|
||||
all(isinstance(v, float) for v in score_list),
|
||||
"All scores should be floats",
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
1.0, sum(score_list), 6, "Scores should sum to 1"
|
||||
)
|
||||
|
||||
def test_score_request_construction(self):
|
||||
"""Test that scoring requests are constructed to avoid decode phase."""
|
||||
from unittest.mock import patch
|
||||
|
||||
# Capture the internal request to verify optimization
|
||||
captured_requests = []
|
||||
original_gen = self.engine.tokenizer_manager.generate_request
|
||||
|
||||
async def mock_generate_request(req, request=None):
|
||||
captured_requests.append(req)
|
||||
async for result in original_gen(req, request):
|
||||
yield result
|
||||
|
||||
# Patch the generate_request method
|
||||
with patch.object(
|
||||
self.engine.tokenizer_manager,
|
||||
"generate_request",
|
||||
side_effect=mock_generate_request,
|
||||
):
|
||||
# Run a scoring request
|
||||
query = "What is the capital of"
|
||||
items = ["France", "Germany"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
# Verify we got results
|
||||
self.assertEqual(len(scores), len(items))
|
||||
|
||||
# Verify the captured request has decode-avoiding properties
|
||||
self.assertEqual(len(captured_requests), 1)
|
||||
request = captured_requests[0]
|
||||
|
||||
# Key assertions for decode phase avoidance:
|
||||
# 1. max_new_tokens should be 0 (prevents token generation)
|
||||
# Handle both single and batch request cases
|
||||
if isinstance(request.sampling_params, dict):
|
||||
max_new_tokens = request.sampling_params.get("max_new_tokens", 0)
|
||||
elif isinstance(request.sampling_params, list):
|
||||
# For batch requests, check the first item
|
||||
max_new_tokens = request.sampling_params[0].get("max_new_tokens", 0)
|
||||
else:
|
||||
max_new_tokens = getattr(request.sampling_params, "max_new_tokens", 0)
|
||||
|
||||
self.assertEqual(
|
||||
max_new_tokens, 0, "max_new_tokens should be 0 to avoid decode phase"
|
||||
)
|
||||
|
||||
# 2. Should have token_ids_logprob for scoring
|
||||
# Handle both single and batch request cases
|
||||
if (
|
||||
isinstance(request.token_ids_logprob, list)
|
||||
and len(request.token_ids_logprob) > 0
|
||||
and isinstance(request.token_ids_logprob[0], list)
|
||||
):
|
||||
# Batch case: token_ids_logprob is a list of lists
|
||||
# Each item in the batch should have the same label_token_ids
|
||||
for item_token_ids in request.token_ids_logprob:
|
||||
self.assertEqual(
|
||||
item_token_ids,
|
||||
label_token_ids,
|
||||
"Each batch item should have label_token_ids for scoring",
|
||||
)
|
||||
else:
|
||||
# Single request case
|
||||
self.assertEqual(
|
||||
request.token_ids_logprob,
|
||||
label_token_ids,
|
||||
"Should have label_token_ids for scoring",
|
||||
)
|
||||
|
||||
# 3. Should request logprobs but not stream
|
||||
self.assertTrue(
|
||||
request.return_logprob, "Should request logprobs for scoring"
|
||||
)
|
||||
self.assertFalse(request.stream, "Scoring requests should not stream")
|
||||
|
||||
def test_multi_item_scoring_basic(self):
|
||||
"""Test basic multi-item scoring functionality."""
|
||||
# Test with a simple query and items
|
||||
query = "What is the capital of California? Answer Yes or No for each of the following options:"
|
||||
items = ["Sacramento", "San Jose", "San Francisco"]
|
||||
label_token_ids = [9454, 2753] # "Yes" and "No" tokens
|
||||
|
||||
# Get scores using SGLang
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
# Verify we get the expected number of scores
|
||||
self.assertEqual(len(scores), len(items), "Should get one score list per item")
|
||||
|
||||
# Verify each score list has the correct length
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(
|
||||
len(score_list),
|
||||
len(label_token_ids),
|
||||
f"Item {i} should have {len(label_token_ids)} scores",
|
||||
)
|
||||
# Verify scores are probabilities (sum to 1)
|
||||
self.assertAlmostEqual(
|
||||
sum(score_list),
|
||||
1.0,
|
||||
places=6,
|
||||
msg=f"Scores for item {i} should sum to 1",
|
||||
)
|
||||
# Verify all scores are non-negative
|
||||
for j, score in enumerate(score_list):
|
||||
self.assertGreaterEqual(
|
||||
score, 0, f"Score {j} for item {i} should be non-negative"
|
||||
)
|
||||
|
||||
def test_multi_item_scoring_consistency(self):
|
||||
"""Test that multi-item scoring gives consistent results."""
|
||||
query = "Choose the best option:"
|
||||
items = ["Option A", "Option B", "Option C"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
# Run the same test multiple times
|
||||
scores1 = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
scores2 = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
# Results should be identical (deterministic)
|
||||
self.assertEqual(len(scores1), len(scores2), "Should get same number of items")
|
||||
for i, (s1, s2) in enumerate(zip(scores1, scores2)):
|
||||
self.assertEqual(
|
||||
len(s1), len(s2), f"Item {i} should have same number of scores"
|
||||
)
|
||||
for j, (score1, score2) in enumerate(zip(s1, s2)):
|
||||
self.assertAlmostEqual(
|
||||
score1,
|
||||
score2,
|
||||
places=6,
|
||||
msg=f"Score {j} for item {i} should be identical",
|
||||
)
|
||||
|
||||
def test_multi_item_scoring_different_sizes(self):
|
||||
"""Test multi-item scoring with different numbers of items."""
|
||||
query = "Rate each option:"
|
||||
label_token_ids = [1, 2, 3, 4, 5]
|
||||
|
||||
# Test with different numbers of items
|
||||
test_cases = [
|
||||
["Single item"],
|
||||
["Item 1", "Item 2"],
|
||||
["A", "B", "C", "D"],
|
||||
["X", "Y", "Z", "W", "V", "U"],
|
||||
]
|
||||
|
||||
for items in test_cases:
|
||||
with self.subTest(items=items):
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
len(scores), len(items), f"Should get {len(items)} score lists"
|
||||
)
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(
|
||||
len(score_list),
|
||||
len(label_token_ids),
|
||||
f"Item {i} should have {len(label_token_ids)} scores",
|
||||
)
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_empty_items(self):
|
||||
"""Test multi-item scoring with empty items list."""
|
||||
query = "Test query"
|
||||
items = []
|
||||
label_token_ids = [1, 2]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
self.assertEqual(len(scores), 0, "Should return empty list for empty items")
|
||||
|
||||
def test_multi_item_scoring_single_item(self):
|
||||
"""Test multi-item scoring with single item (should work like regular scoring)."""
|
||||
query = "Complete this sentence: The capital of France is"
|
||||
items = ["Paris"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
self.assertEqual(len(scores), 1, "Should get one score list")
|
||||
self.assertEqual(
|
||||
len(scores[0]), len(label_token_ids), "Should have correct number of scores"
|
||||
)
|
||||
self.assertAlmostEqual(sum(scores[0]), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_different_queries(self):
|
||||
"""Test multi-item scoring with different types of queries."""
|
||||
items = ["Yes", "No"]
|
||||
label_token_ids = [1, 2]
|
||||
|
||||
test_queries = [
|
||||
"Is this true?",
|
||||
"Choose the correct answer:",
|
||||
"What is the best option?",
|
||||
"Select all that apply:",
|
||||
"", # Empty query
|
||||
]
|
||||
|
||||
for query in test_queries:
|
||||
with self.subTest(query=query):
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
len(scores),
|
||||
len(items),
|
||||
f"Should get {len(items)} score lists for query: '{query}'",
|
||||
)
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), len(label_token_ids))
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_different_label_tokens(self):
|
||||
"""Test multi-item scoring with different label token sets."""
|
||||
query = "Choose the best option:"
|
||||
items = ["Option A", "Option B"]
|
||||
|
||||
test_label_tokens = [
|
||||
[1, 2], # Two tokens
|
||||
[1, 2, 3, 4], # Four tokens
|
||||
[1], # Single token
|
||||
[1, 2, 3, 4, 5, 6, 7, 8], # Many tokens
|
||||
]
|
||||
|
||||
for label_token_ids in test_label_tokens:
|
||||
with self.subTest(label_tokens=label_token_ids):
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(
|
||||
len(score_list),
|
||||
len(label_token_ids),
|
||||
f"Item {i} should have {len(label_token_ids)} scores",
|
||||
)
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_without_softmax(self):
|
||||
"""Test multi-item scoring without softmax normalization."""
|
||||
query = "Rate each option:"
|
||||
items = ["Good", "Bad", "Neutral"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=False, # No softmax
|
||||
)
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), len(label_token_ids))
|
||||
# Without softmax, scores don't need to sum to 1
|
||||
# But they should still be valid logits/probabilities
|
||||
for j, score in enumerate(score_list):
|
||||
self.assertIsInstance(
|
||||
score, (int, float), f"Score {j} for item {i} should be numeric"
|
||||
)
|
||||
|
||||
def test_multi_item_scoring_large_batch(self):
|
||||
"""Test multi-item scoring with a large number of items."""
|
||||
query = "Classify each item:"
|
||||
items = [f"Item {i}" for i in range(20)] # 20 items
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
self.assertEqual(len(scores), len(items), "Should handle large batches")
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), len(label_token_ids))
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_unicode(self):
|
||||
"""Test multi-item scoring with unicode characters."""
|
||||
query = "选择最佳选项："
|
||||
items = ["选项A", "选项B", "选项C"]
|
||||
label_token_ids = [1, 2, 3]
|
||||
|
||||
scores = self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
self.assertEqual(len(scores), len(items))
|
||||
|
||||
for i, score_list in enumerate(scores):
|
||||
self.assertEqual(len(score_list), len(label_token_ids))
|
||||
self.assertAlmostEqual(sum(score_list), 1.0, places=6)
|
||||
|
||||
def test_multi_item_scoring_error_handling(self):
|
||||
"""Test multi-item scoring error handling."""
|
||||
query = "Test query"
|
||||
items = ["Item 1", "Item 2"]
|
||||
label_token_ids = [1, 2]
|
||||
|
||||
# Test with invalid label_token_ids
|
||||
with self.assertRaises((ValueError, TypeError)):
|
||||
self.engine.score(
|
||||
query=query,
|
||||
items=items,
|
||||
label_token_ids="invalid", # Should be list of ints
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
# Test with None items
|
||||
with self.assertRaises((ValueError, TypeError)):
|
||||
self.engine.score(
|
||||
query=query,
|
||||
items=None,
|
||||
label_token_ids=label_token_ids,
|
||||
apply_softmax=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,288 +0,0 @@
|
||||
import json
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs, prepare_server_args
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestPrepareServerArgs(CustomTestCase):
|
||||
def test_prepare_server_args(self):
|
||||
server_args = prepare_server_args(
|
||||
[
|
||||
"--model-path",
|
||||
"meta-llama/Meta-Llama-3.1-8B-Instruct",
|
||||
"--json-model-override-args",
|
||||
'{"rope_scaling": {"factor": 2.0, "rope_type": "linear"}}',
|
||||
]
|
||||
)
|
||||
self.assertEqual(
|
||||
server_args.model_path, "meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
)
|
||||
self.assertEqual(
|
||||
json.loads(server_args.json_model_override_args),
|
||||
{"rope_scaling": {"factor": 2.0, "rope_type": "linear"}},
|
||||
)
|
||||
|
||||
|
||||
class TestLoadBalanceMethod(unittest.TestCase):
|
||||
def test_non_pd_defaults_to_round_robin(self):
|
||||
server_args = ServerArgs(model_path="dummy", disaggregation_mode="null")
|
||||
self.assertEqual(server_args.load_balance_method, "round_robin")
|
||||
|
||||
def test_pd_prefill_defaults_to_follow_bootstrap_room(self):
|
||||
server_args = ServerArgs(model_path="dummy", disaggregation_mode="prefill")
|
||||
self.assertEqual(server_args.load_balance_method, "follow_bootstrap_room")
|
||||
|
||||
def test_pd_decode_defaults_to_round_robin(self):
|
||||
server_args = ServerArgs(model_path="dummy", disaggregation_mode="decode")
|
||||
self.assertEqual(server_args.load_balance_method, "round_robin")
|
||||
|
||||
|
||||
class TestPortArgs(unittest.TestCase):
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
|
||||
def test_init_new_standard_case(self, mock_temp_file, mock_is_port_available):
|
||||
mock_is_port_available.return_value = True
|
||||
mock_temp_file.return_value.name = "temp_file"
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
server_args.enable_dp_attention = False
|
||||
|
||||
port_args = PortArgs.init_new(server_args)
|
||||
|
||||
self.assertTrue(port_args.tokenizer_ipc_name.startswith("ipc://"))
|
||||
self.assertTrue(port_args.scheduler_input_ipc_name.startswith("ipc://"))
|
||||
self.assertTrue(port_args.detokenizer_ipc_name.startswith("ipc://"))
|
||||
self.assertIsInstance(port_args.nccl_port, int)
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
def test_init_new_with_single_node_dp_attention(self, mock_is_port_available):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 1
|
||||
server_args.dist_init_addr = None
|
||||
|
||||
port_args = PortArgs.init_new(server_args)
|
||||
|
||||
self.assertTrue(port_args.tokenizer_ipc_name.startswith("tcp://127.0.0.1:"))
|
||||
self.assertTrue(
|
||||
port_args.scheduler_input_ipc_name.startswith("tcp://127.0.0.1:")
|
||||
)
|
||||
self.assertTrue(port_args.detokenizer_ipc_name.startswith("tcp://127.0.0.1:"))
|
||||
self.assertIsInstance(port_args.nccl_port, int)
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
def test_init_new_with_dp_rank(self, mock_is_port_available):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 1
|
||||
server_args.dist_init_addr = "192.168.1.1:25000"
|
||||
|
||||
worker_ports = [25006, 25007, 25008, 25009]
|
||||
port_args = PortArgs.init_new(server_args, dp_rank=2, worker_ports=worker_ports)
|
||||
|
||||
self.assertTrue(port_args.scheduler_input_ipc_name.endswith(":25008"))
|
||||
|
||||
self.assertTrue(port_args.tokenizer_ipc_name.startswith("tcp://192.168.1.1:"))
|
||||
self.assertTrue(port_args.detokenizer_ipc_name.startswith("tcp://192.168.1.1:"))
|
||||
self.assertIsInstance(port_args.nccl_port, int)
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
def test_init_new_with_ipv4_address(self, mock_is_port_available):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
|
||||
server_args.nccl_port = None
|
||||
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 2
|
||||
server_args.dist_init_addr = "192.168.1.1:25000"
|
||||
|
||||
port_args = PortArgs.init_new(server_args)
|
||||
|
||||
self.assertTrue(port_args.tokenizer_ipc_name.startswith("tcp://192.168.1.1:"))
|
||||
self.assertTrue(
|
||||
port_args.scheduler_input_ipc_name.startswith("tcp://192.168.1.1:")
|
||||
)
|
||||
self.assertTrue(port_args.detokenizer_ipc_name.startswith("tcp://192.168.1.1:"))
|
||||
self.assertIsInstance(port_args.nccl_port, int)
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
def test_init_new_with_malformed_ipv4_address(self, mock_is_port_available):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 2
|
||||
server_args.dist_init_addr = "192.168.1.1"
|
||||
|
||||
with self.assertRaises(AssertionError) as context:
|
||||
PortArgs.init_new(server_args)
|
||||
|
||||
self.assertIn(
|
||||
"please provide --dist-init-addr as host:port", str(context.exception)
|
||||
)
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
def test_init_new_with_malformed_ipv4_address_invalid_port(
|
||||
self, mock_is_port_available
|
||||
):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 2
|
||||
server_args.dist_init_addr = "192.168.1.1:abc"
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
PortArgs.init_new(server_args)
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
@patch("sglang.srt.server_args.is_valid_ipv6_address", return_value=True)
|
||||
def test_init_new_with_ipv6_address(
|
||||
self, mock_is_valid_ipv6, mock_is_port_available
|
||||
):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 2
|
||||
server_args.dist_init_addr = "[2001:db8::1]:25000"
|
||||
|
||||
port_args = PortArgs.init_new(server_args)
|
||||
|
||||
self.assertTrue(port_args.tokenizer_ipc_name.startswith("tcp://[2001:db8::1]:"))
|
||||
self.assertTrue(
|
||||
port_args.scheduler_input_ipc_name.startswith("tcp://[2001:db8::1]:")
|
||||
)
|
||||
self.assertTrue(
|
||||
port_args.detokenizer_ipc_name.startswith("tcp://[2001:db8::1]:")
|
||||
)
|
||||
self.assertIsInstance(port_args.nccl_port, int)
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
@patch("sglang.srt.server_args.is_valid_ipv6_address", return_value=False)
|
||||
def test_init_new_with_invalid_ipv6_address(
|
||||
self, mock_is_valid_ipv6, mock_is_port_available
|
||||
):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 2
|
||||
server_args.dist_init_addr = "[invalid-ipv6]:25000"
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
PortArgs.init_new(server_args)
|
||||
|
||||
self.assertIn("invalid IPv6 address", str(context.exception))
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
def test_init_new_with_malformed_ipv6_address_missing_bracket(
|
||||
self, mock_is_port_available
|
||||
):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 2
|
||||
server_args.dist_init_addr = "[2001:db8::1:25000"
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
PortArgs.init_new(server_args)
|
||||
|
||||
self.assertIn("invalid IPv6 address format", str(context.exception))
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
@patch("sglang.srt.server_args.is_valid_ipv6_address", return_value=True)
|
||||
def test_init_new_with_malformed_ipv6_address_missing_port(
|
||||
self, mock_is_valid_ipv6, mock_is_port_available
|
||||
):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 2
|
||||
server_args.dist_init_addr = "[2001:db8::1]"
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
PortArgs.init_new(server_args)
|
||||
|
||||
self.assertIn(
|
||||
"a port must be specified in IPv6 address", str(context.exception)
|
||||
)
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
@patch("sglang.srt.server_args.is_valid_ipv6_address", return_value=True)
|
||||
def test_init_new_with_malformed_ipv6_address_invalid_port(
|
||||
self, mock_is_valid_ipv6, mock_is_port_available
|
||||
):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 2
|
||||
server_args.dist_init_addr = "[2001:db8::1]:abcde"
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
PortArgs.init_new(server_args)
|
||||
|
||||
self.assertIn("invalid port in IPv6 address", str(context.exception))
|
||||
|
||||
@patch("sglang.srt.server_args.is_port_available")
|
||||
@patch("sglang.srt.server_args.is_valid_ipv6_address", return_value=True)
|
||||
def test_init_new_with_malformed_ipv6_address_wrong_separator(
|
||||
self, mock_is_valid_ipv6, mock_is_port_available
|
||||
):
|
||||
mock_is_port_available.return_value = True
|
||||
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
server_args.port = 30000
|
||||
server_args.nccl_port = None
|
||||
|
||||
server_args.enable_dp_attention = True
|
||||
server_args.nnodes = 2
|
||||
server_args.dist_init_addr = "[2001:db8::1]#25000"
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
PortArgs.init_new(server_args)
|
||||
|
||||
self.assertIn("expected ':' after ']'", str(context.exception))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,743 +0,0 @@
|
||||
"""
|
||||
python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_simple_decode
|
||||
python3 -m unittest test_srt_endpoint.TestSRTEndpoint.test_logprob_with_chunked_prefill
|
||||
python3 -m unittest test_srt_endpoint.TestTokenizeDetokenize
|
||||
"""
|
||||
|
||||
import json
|
||||
import random
|
||||
import time
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from functools import partial
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
|
||||
from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
run_logprob_check,
|
||||
)
|
||||
|
||||
|
||||
class TestSRTEndpoint(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=(
|
||||
"--enable-custom-logit-processor",
|
||||
"--mem-fraction-static",
|
||||
"0.7",
|
||||
"--cuda-graph-max-bs",
|
||||
"8",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def run_decode(
|
||||
self,
|
||||
return_logprob=False,
|
||||
top_logprobs_num=0,
|
||||
return_text=False,
|
||||
n=1,
|
||||
stream=False,
|
||||
batch=False,
|
||||
):
|
||||
if batch:
|
||||
text = ["The capital of France is"]
|
||||
else:
|
||||
text = "The capital of France is"
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": text,
|
||||
"sampling_params": {
|
||||
"temperature": 0 if n == 1 else 0.5,
|
||||
"max_new_tokens": 16,
|
||||
"n": n,
|
||||
},
|
||||
"stream": stream,
|
||||
"return_logprob": return_logprob,
|
||||
"top_logprobs_num": top_logprobs_num,
|
||||
"return_text_in_logprobs": return_text,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
)
|
||||
if not stream:
|
||||
response_json = response.json()
|
||||
else:
|
||||
response_json = []
|
||||
for line in response.iter_lines():
|
||||
if line.startswith(b"data: ") and line[6:] != b"[DONE]":
|
||||
response_json.append(json.loads(line[6:]))
|
||||
|
||||
print(json.dumps(response_json, indent=2))
|
||||
print("=" * 100)
|
||||
|
||||
def test_simple_decode(self):
|
||||
self.run_decode()
|
||||
|
||||
def test_simple_decode_batch(self):
|
||||
self.run_decode(batch=True)
|
||||
|
||||
def test_parallel_sample(self):
|
||||
self.run_decode(n=3)
|
||||
|
||||
def test_parallel_sample_stream(self):
|
||||
self.run_decode(n=3, stream=True)
|
||||
|
||||
def test_logprob(self):
|
||||
self.run_decode(
|
||||
return_logprob=True,
|
||||
top_logprobs_num=5,
|
||||
return_text=True,
|
||||
)
|
||||
|
||||
def test_logprob_start_len(self):
|
||||
logprob_start_len = 4
|
||||
new_tokens = 4
|
||||
prompts = [
|
||||
"I have a very good idea on",
|
||||
"Today is a sunndy day and",
|
||||
]
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompts,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": new_tokens,
|
||||
},
|
||||
"return_logprob": True,
|
||||
"top_logprobs_num": 5,
|
||||
"return_text_in_logprobs": True,
|
||||
"logprob_start_len": logprob_start_len,
|
||||
},
|
||||
)
|
||||
response_json = response.json()
|
||||
print(json.dumps(response_json, indent=2))
|
||||
|
||||
for i, res in enumerate(response_json):
|
||||
self.assertEqual(
|
||||
res["meta_info"]["prompt_tokens"],
|
||||
logprob_start_len + len(res["meta_info"]["input_token_logprobs"]),
|
||||
)
|
||||
assert prompts[i].endswith(
|
||||
"".join([x[-1] for x in res["meta_info"]["input_token_logprobs"]])
|
||||
)
|
||||
|
||||
self.assertEqual(res["meta_info"]["completion_tokens"], new_tokens)
|
||||
self.assertEqual(len(res["meta_info"]["output_token_logprobs"]), new_tokens)
|
||||
self.assertEqual(
|
||||
res["text"],
|
||||
"".join([x[-1] for x in res["meta_info"]["output_token_logprobs"]]),
|
||||
)
|
||||
|
||||
def test_logprob_with_chunked_prefill(self):
|
||||
"""Test a long prompt that requests output logprobs will not hit OOM."""
|
||||
new_tokens = 4
|
||||
prompts = "I have a very good idea on this. " * 8000
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompts,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": new_tokens,
|
||||
},
|
||||
"return_logprob": True,
|
||||
"logprob_start_len": -1,
|
||||
"top_logprobs_num": 5,
|
||||
},
|
||||
)
|
||||
response_json = response.json()
|
||||
# print(json.dumps(response_json, indent=2))
|
||||
|
||||
res = response_json
|
||||
self.assertEqual(res["meta_info"]["completion_tokens"], new_tokens)
|
||||
|
||||
# Test the number of tokens are correct
|
||||
self.assertEqual(len(res["meta_info"]["output_token_logprobs"]), new_tokens)
|
||||
self.assertEqual(len(res["meta_info"]["output_top_logprobs"]), new_tokens)
|
||||
|
||||
# Test the top-1 tokens are the same as output tokens (because temp = 0.0)
|
||||
for i in range(new_tokens):
|
||||
self.assertListEqual(
|
||||
res["meta_info"]["output_token_logprobs"][i],
|
||||
res["meta_info"]["output_top_logprobs"][i][0],
|
||||
)
|
||||
self.assertEqual(len(res["meta_info"]["output_top_logprobs"][i]), 5)
|
||||
|
||||
def test_logprob_match(self):
|
||||
"""Test the output logprobs are close to the input logprobs if we run a prefill again."""
|
||||
|
||||
def run_generate(
|
||||
prompt, return_logprob=False, max_new_tokens=512, logprob_start_len=-1
|
||||
):
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt_kwargs = {"text": prompt}
|
||||
else:
|
||||
prompt_kwargs = {"input_ids": prompt}
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
**prompt_kwargs,
|
||||
"sampling_params": {
|
||||
"temperature": 1.0,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
"return_logprob": return_logprob,
|
||||
"return_text_in_logprobs": True,
|
||||
"logprob_start_len": logprob_start_len,
|
||||
},
|
||||
)
|
||||
return response.json()
|
||||
|
||||
prompt = "I have a very good idea on how to"
|
||||
|
||||
gen = run_generate(prompt, return_logprob=True, logprob_start_len=0)
|
||||
output_logprobs = np.array(
|
||||
[x[0] for x in gen["meta_info"]["output_token_logprobs"]]
|
||||
)
|
||||
num_prompts_tokens = gen["meta_info"]["prompt_tokens"]
|
||||
|
||||
input_tokens = [x[1] for x in gen["meta_info"]["input_token_logprobs"]]
|
||||
output_tokens = [x[1] for x in gen["meta_info"]["output_token_logprobs"]]
|
||||
|
||||
new_prompt = input_tokens + output_tokens
|
||||
score = run_generate(
|
||||
new_prompt, return_logprob=True, logprob_start_len=0, max_new_tokens=0
|
||||
)
|
||||
output_logprobs_score = np.array(
|
||||
[
|
||||
x[0]
|
||||
for x in score["meta_info"]["input_token_logprobs"][num_prompts_tokens:]
|
||||
]
|
||||
)
|
||||
|
||||
print(f"{output_logprobs[-10:]=}")
|
||||
print(f"{output_logprobs_score[-10:]=}")
|
||||
|
||||
diff = np.abs(output_logprobs - output_logprobs_score)
|
||||
max_diff = np.max(diff)
|
||||
self.assertLess(max_diff, 0.35)
|
||||
|
||||
def test_logprob_mixed(self):
|
||||
args = []
|
||||
temperature = 0
|
||||
# input_len, output_len, temperature, logprob_start_len, return_logprob, top_logprobs_num
|
||||
for input_len in [1000, 5000, 10000, 50000]:
|
||||
for output_len in [4, 8]:
|
||||
for logprob_start_len in [0, 500, 2500, 5000, 25000]:
|
||||
for return_logprob in [True, False]:
|
||||
for top_logprobs_num in [0, 5]:
|
||||
|
||||
if logprob_start_len >= input_len:
|
||||
continue
|
||||
|
||||
args.append(
|
||||
(
|
||||
input_len,
|
||||
output_len,
|
||||
temperature,
|
||||
logprob_start_len,
|
||||
return_logprob,
|
||||
top_logprobs_num,
|
||||
)
|
||||
)
|
||||
|
||||
random.shuffle(args)
|
||||
|
||||
func = partial(run_logprob_check, self)
|
||||
with ThreadPoolExecutor(8) as executor:
|
||||
list(executor.map(func, args))
|
||||
|
||||
def test_logprob_grammar(self):
|
||||
prompts = "Question: Is Paris the Capital of France? Answer:"
|
||||
allowed_tokens = [" Yes", " No"]
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompts,
|
||||
"sampling_params": {
|
||||
"temperature": 1.0,
|
||||
"max_new_tokens": 1,
|
||||
"regex": "( Yes| No)",
|
||||
},
|
||||
"return_logprob": True,
|
||||
"top_logprobs_num": 5, # The grammar constraint allows all prefix tokens so we need to use a larger top_k.
|
||||
"return_text_in_logprobs": True,
|
||||
},
|
||||
)
|
||||
response_json = response.json()
|
||||
output_top_logprobs = response_json["meta_info"]["output_top_logprobs"][0]
|
||||
print(f"{output_top_logprobs=}")
|
||||
|
||||
# Parse results
|
||||
# This is because the grammar constraint allows all prefix tokens
|
||||
logprobs = [None] * 2
|
||||
for i in range(len(output_top_logprobs)):
|
||||
try:
|
||||
idx = allowed_tokens.index(output_top_logprobs[i][2])
|
||||
except ValueError:
|
||||
# Not found
|
||||
continue
|
||||
logprobs[idx] = output_top_logprobs[i][0]
|
||||
|
||||
self.assertTrue(all(x is not None for x in logprobs))
|
||||
|
||||
def run_custom_logit_processor(self, target_token_id: Optional[int] = None):
|
||||
"""Test custom logit processor with custom params.
|
||||
|
||||
If target_token_id is None, the custom logit processor won't be passed in.
|
||||
"""
|
||||
|
||||
custom_params = {"token_id": target_token_id}
|
||||
|
||||
class DeterministicLogitProcessor(CustomLogitProcessor):
|
||||
"""A dummy logit processor that changes the logits to always
|
||||
sample the given token id.
|
||||
"""
|
||||
|
||||
def __call__(self, logits, custom_param_list):
|
||||
assert logits.shape[0] == len(custom_param_list)
|
||||
key = "token_id"
|
||||
|
||||
for i, param_dict in enumerate(custom_param_list):
|
||||
# Mask all other tokens
|
||||
logits[i, :] = -float("inf")
|
||||
# Assign highest probability to the specified token
|
||||
logits[i, param_dict[key]] = 0.0
|
||||
return logits
|
||||
|
||||
prompts = "Question: Is Paris the Capital of France? Answer:"
|
||||
|
||||
# Base case json data to be posted to the server.
|
||||
base_json = {
|
||||
"text": prompts,
|
||||
"sampling_params": {"temperature": 0.0},
|
||||
"return_logprob": True,
|
||||
}
|
||||
|
||||
# Custom json data with custom logit processor and params.
|
||||
custom_json = base_json.copy()
|
||||
# Only set the custom logit processor if target_token_id is not None.
|
||||
if target_token_id is not None:
|
||||
custom_json["custom_logit_processor"] = DeterministicLogitProcessor.to_str()
|
||||
custom_json["sampling_params"]["custom_params"] = custom_params
|
||||
|
||||
custom_response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json=custom_json,
|
||||
).json()
|
||||
|
||||
output_token_logprobs = custom_response["meta_info"]["output_token_logprobs"]
|
||||
sampled_tokens = [x[1] for x in output_token_logprobs]
|
||||
|
||||
# The logit processor should always sample the given token as the logits is deterministic.
|
||||
if target_token_id is not None:
|
||||
self.assertTrue(
|
||||
all(x == custom_params["token_id"] for x in sampled_tokens),
|
||||
# Print the detailed test case info if the test fails.
|
||||
f"{target_token_id=}\n{sampled_tokens=}\n{custom_response=}",
|
||||
)
|
||||
|
||||
def run_stateful_custom_logit_processor(
|
||||
self, first_token_id: int | None, delay: int = 2
|
||||
):
|
||||
"""Test custom logit processor with custom params and state.
|
||||
|
||||
Should sample the first `delay` tokens normally, then output first_token_id and consecutive tokens after that.
|
||||
If first_token_id is None, the custom logit processor won't be passed in.
|
||||
"""
|
||||
custom_params = {"token_id": first_token_id, "delay": 2}
|
||||
|
||||
class DeterministicStatefulLogitProcessor(CustomLogitProcessor):
|
||||
"""A dummy logit processor that changes the logits to always
|
||||
sample the given token id.
|
||||
"""
|
||||
|
||||
def __call__(self, logits, custom_param_list):
|
||||
assert logits.shape[0] == len(custom_param_list)
|
||||
|
||||
for i, param_dict in enumerate(custom_param_list):
|
||||
if param_dict["delay"] > 0:
|
||||
param_dict["delay"] -= 1
|
||||
continue
|
||||
if param_dict["delay"] == 0:
|
||||
param_dict["delay"] -= 1
|
||||
force_token = param_dict["token_id"]
|
||||
else:
|
||||
output_ids = param_dict["__req__"].output_ids
|
||||
force_token = output_ids[-1] + 1
|
||||
# Mask all other tokens
|
||||
logits[i, :] = -float("inf")
|
||||
# Assign highest probability to the specified token
|
||||
logits[i, force_token] = 0.0
|
||||
return logits
|
||||
|
||||
prompts = "Question: Is Paris the Capital of France? Answer:"
|
||||
|
||||
# Base case json data to be posted to the server.
|
||||
base_json = {
|
||||
"text": prompts,
|
||||
"sampling_params": {"temperature": 0.0},
|
||||
"return_logprob": True,
|
||||
}
|
||||
|
||||
# Custom json data with custom logit processor and params.
|
||||
custom_json = base_json.copy()
|
||||
# Only set the custom logit processor if target_token_id is not None.
|
||||
if first_token_id is not None:
|
||||
custom_json["custom_logit_processor"] = (
|
||||
DeterministicStatefulLogitProcessor().to_str()
|
||||
)
|
||||
custom_json["sampling_params"]["custom_params"] = custom_params
|
||||
|
||||
custom_response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json=custom_json,
|
||||
).json()
|
||||
|
||||
output_token_logprobs = custom_response["meta_info"]["output_token_logprobs"]
|
||||
sampled_tokens = [x[1] for x in output_token_logprobs]
|
||||
# The logit processor should always sample the given token as the logits is deterministic.
|
||||
if first_token_id is not None:
|
||||
self.assertTrue(
|
||||
all(
|
||||
x == custom_params["token_id"] + k
|
||||
for k, x in enumerate(sampled_tokens[custom_params["delay"] :])
|
||||
),
|
||||
# Print the detailed test case info if the test fails.
|
||||
f"{first_token_id=}\n{sampled_tokens=}\n{custom_response=}",
|
||||
)
|
||||
|
||||
def test_custom_logit_processor(self):
|
||||
"""Test custom logit processor with a single request."""
|
||||
self.run_custom_logit_processor(target_token_id=5)
|
||||
|
||||
def test_custom_logit_processor_batch_mixed(self):
|
||||
"""Test a batch of requests mixed of requests with and without custom logit processor."""
|
||||
target_token_ids = list(range(32)) + [None] * 16
|
||||
random.shuffle(target_token_ids)
|
||||
with ThreadPoolExecutor(len(target_token_ids)) as executor:
|
||||
list(executor.map(self.run_custom_logit_processor, target_token_ids))
|
||||
|
||||
@unittest.skip("Skip this test because this feature has a bug. See comments below.")
|
||||
def test_stateful_custom_logit_processor(self):
|
||||
"""Test custom logit processor with a single request."""
|
||||
|
||||
"""
|
||||
NOTE: This feature has a race condition bug.
|
||||
This line https://github.com/sgl-project/sglang/blob/ef8ec07b2ce4c70c2a33ec5acda4ce529bc3cda4/test/srt/test_srt_endpoint.py#L395-L396 can be accessed by two concurrent threads at the same time. The access order is not guaranteed.
|
||||
In sglang, we use two python threads to overlap the GPU computation and CPU scheduling.
|
||||
Thread 1 (the CPU scheduling thread) will update the `param_dict["__req__"].output_ids`.
|
||||
Thread 2 (the GPU computation thread) will call `DeterministicStatefulLogitProcessor` because sampling is considered as GPU computation.
|
||||
We can fix this by moving the call of DeterministicStatefulLogitProcessor to the CPU scheduling thread.
|
||||
"""
|
||||
|
||||
self.run_stateful_custom_logit_processor(first_token_id=5)
|
||||
|
||||
@unittest.skip("Skip this test because this feature has a bug. See comments above.")
|
||||
def test_stateful_custom_logit_processor_batch_mixed(self):
|
||||
"""Test a batch of requests mixed of requests with and without custom logit processor."""
|
||||
target_token_ids = list(range(32)) + [None] * 16
|
||||
random.shuffle(target_token_ids)
|
||||
with ThreadPoolExecutor(len(target_token_ids)) as executor:
|
||||
list(
|
||||
executor.map(self.run_stateful_custom_logit_processor, target_token_ids)
|
||||
)
|
||||
|
||||
def test_cache_tokens(self):
|
||||
for _ in range(2):
|
||||
time.sleep(1)
|
||||
response = requests.post(self.base_url + "/flush_cache")
|
||||
assert response.status_code == 200
|
||||
|
||||
def send_and_check_cached_tokens(input_ids):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"input_ids": list(input_ids),
|
||||
"sampling_params": {
|
||||
"max_new_tokens": 1,
|
||||
},
|
||||
},
|
||||
)
|
||||
response_json = response.json()
|
||||
return response_json["meta_info"]["cached_tokens"]
|
||||
|
||||
self.assertEqual(send_and_check_cached_tokens(range(0, 100)), 0)
|
||||
self.assertEqual(send_and_check_cached_tokens(range(0, 10000)), 100)
|
||||
self.assertEqual(send_and_check_cached_tokens(range(0, 10000)), 9999)
|
||||
self.assertEqual(send_and_check_cached_tokens(range(0, 1000)), 999)
|
||||
self.assertEqual(send_and_check_cached_tokens(range(0, 11000)), 10000)
|
||||
|
||||
def test_get_server_info(self):
|
||||
response = requests.get(self.base_url + "/get_server_info")
|
||||
response_json = response.json()
|
||||
|
||||
max_total_num_tokens = response_json["max_total_num_tokens"]
|
||||
self.assertIsInstance(max_total_num_tokens, int)
|
||||
|
||||
version = response_json["version"]
|
||||
self.assertIsInstance(version, str)
|
||||
|
||||
def test_logit_bias(self):
|
||||
"""Test that a very high logit bias forces sampling of a specific token."""
|
||||
# Choose a token ID to bias (using 5 as an example)
|
||||
target_token_id = 60704 # Paris for meta-llama/Llama-3.2-1B-Instruct, DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
logit_bias = {str(target_token_id): 100.0} # Very high positive bias
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 1.0, # Use high temperature to encourage exploration
|
||||
"max_new_tokens": 4,
|
||||
"logit_bias": logit_bias,
|
||||
},
|
||||
"return_logprob": True,
|
||||
},
|
||||
)
|
||||
response_json = response.json()
|
||||
|
||||
# Extract the sampled token IDs from the output
|
||||
output_token_logprobs = response_json["meta_info"]["output_token_logprobs"]
|
||||
sampled_tokens = [x[1] for x in output_token_logprobs]
|
||||
|
||||
# Verify that all sampled tokens are the target token
|
||||
self.assertTrue(
|
||||
all(x == target_token_id for x in sampled_tokens),
|
||||
f"Expected all tokens to be {target_token_id}, but got {sampled_tokens}",
|
||||
)
|
||||
|
||||
def test_forbidden_token(self):
|
||||
"""Test that a forbidden token (very negative logit bias) doesn't appear in the output."""
|
||||
# Choose a token ID to forbid (using 10 as an example)
|
||||
forbidden_token_id = 23994 # rice for meta-llama/Llama-3.2-1B-Instruct, DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
logit_bias = {
|
||||
str(forbidden_token_id): -100.0
|
||||
} # Very negative bias to forbid the token
|
||||
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": "Only output 'rice' exactly like this, in lowercase ONLY: rice",
|
||||
"sampling_params": {
|
||||
"temperature": 1.0, # Use high temperature to encourage diverse output
|
||||
"max_new_tokens": 50, # Generate enough tokens to likely include numbers
|
||||
"logit_bias": logit_bias,
|
||||
},
|
||||
"return_logprob": True,
|
||||
},
|
||||
)
|
||||
response_json = response.json()
|
||||
|
||||
# Extract the sampled token IDs from the output
|
||||
output_token_logprobs = response_json["meta_info"]["output_token_logprobs"]
|
||||
sampled_tokens = [x[1] for x in output_token_logprobs]
|
||||
|
||||
# Verify that the forbidden token doesn't appear in the output
|
||||
self.assertNotIn(
|
||||
forbidden_token_id,
|
||||
sampled_tokens,
|
||||
f"Expected forbidden token {forbidden_token_id} not to be present, but it was found",
|
||||
)
|
||||
|
||||
def test_logit_bias_isolation(self):
|
||||
"""Test that logit_bias applied to one request doesn't affect other requests in batch."""
|
||||
# Choose a token ID to bias in first request only
|
||||
biased_token_id = 60704 # Paris for meta-llama/Llama-3.2-1B-Instruct, DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
# Prepare batch requests - one with logit_bias and one without
|
||||
requests_data = [
|
||||
{
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 1.0,
|
||||
"max_new_tokens": 4,
|
||||
"logit_bias": {str(biased_token_id): 100.0}, # Strong bias
|
||||
},
|
||||
"return_logprob": True,
|
||||
},
|
||||
{
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 1.0,
|
||||
"max_new_tokens": 4,
|
||||
},
|
||||
"return_logprob": True,
|
||||
},
|
||||
]
|
||||
|
||||
# Send both requests
|
||||
responses = []
|
||||
for req in requests_data:
|
||||
response = requests.post(self.base_url + "/generate", json=req)
|
||||
responses.append(response.json())
|
||||
|
||||
# Extract token IDs from each response
|
||||
biased_tokens = [
|
||||
x[1] for x in responses[0]["meta_info"]["output_token_logprobs"]
|
||||
]
|
||||
unbiased_tokens = [
|
||||
x[1] for x in responses[1]["meta_info"]["output_token_logprobs"]
|
||||
]
|
||||
|
||||
# Verify first response contains only biased tokens
|
||||
self.assertTrue(
|
||||
all(x == biased_token_id for x in biased_tokens),
|
||||
f"Expected all tokens to be {biased_token_id} in first response, but got {biased_tokens}",
|
||||
)
|
||||
|
||||
# Verify second response contains at least some different tokens
|
||||
# (We can't guarantee exactly what tokens will be generated, but they shouldn't all be the biased token)
|
||||
self.assertTrue(
|
||||
any(x != biased_token_id for x in unbiased_tokens),
|
||||
f"Expected some tokens to be different from {biased_token_id} in second response, but got {unbiased_tokens}",
|
||||
)
|
||||
|
||||
def test_get_server_info_concurrent(self):
|
||||
"""Make sure the concurrent get_server_info doesn't crash the server."""
|
||||
tp = ThreadPoolExecutor(max_workers=30)
|
||||
|
||||
def s():
|
||||
server_info = requests.get(self.base_url + "/get_server_info")
|
||||
server_info.json()
|
||||
|
||||
futures = []
|
||||
for _ in range(4):
|
||||
futures.append(tp.submit(s))
|
||||
|
||||
for f in futures:
|
||||
f.result()
|
||||
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# /tokenize & /detokenize Test Class: TestTokenizeDetokenize
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTokenizeDetokenize(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.tokenize_url = f"{cls.base_url}/tokenize"
|
||||
cls.detokenize_url = f"{cls.base_url}/detokenize"
|
||||
cls.session = requests.Session()
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
cls.session.close()
|
||||
|
||||
def _post_json(self, url, payload):
|
||||
r = self.session.post(url, json=payload)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
def test_tokenize_various_inputs(self):
|
||||
single = "Hello SGLang world! 123 😊, ಪರ್ವತದ ಮೇಲೆ ಹಿಮ."
|
||||
multi = ["First sentence.", "Second, with 中文."]
|
||||
scenarios = [
|
||||
{"prompt": single, "add_special_tokens": True},
|
||||
{"prompt": single, "add_special_tokens": False},
|
||||
{"prompt": multi, "add_special_tokens": True},
|
||||
{"prompt": multi, "add_special_tokens": False},
|
||||
{"prompt": "", "add_special_tokens": False},
|
||||
]
|
||||
for case in scenarios:
|
||||
payload = {"model": self.model, "prompt": case["prompt"]}
|
||||
if "add_special_tokens" in case:
|
||||
payload["add_special_tokens"] = case["add_special_tokens"]
|
||||
resp = self._post_json(self.tokenize_url, payload)
|
||||
tokens = resp["tokens"]
|
||||
count = resp["count"]
|
||||
self.assertIsInstance(tokens, list)
|
||||
if not tokens:
|
||||
self.assertEqual(count, 0)
|
||||
else:
|
||||
if isinstance(tokens[0], list):
|
||||
total = sum(len(t) for t in tokens)
|
||||
expected = sum(count) if isinstance(count, list) else count
|
||||
else:
|
||||
total = len(tokens)
|
||||
expected = count
|
||||
self.assertEqual(total, expected)
|
||||
|
||||
def test_tokenize_invalid_type(self):
|
||||
r = self.session.post(
|
||||
self.tokenize_url, json={"model": self.model, "prompt": 12345}
|
||||
)
|
||||
self.assertEqual(r.status_code, 400)
|
||||
|
||||
def test_detokenize_roundtrip(self):
|
||||
text = "Verify detokenization round trip. यह डिटोकेनाइजेशन है"
|
||||
t0 = self._post_json(
|
||||
self.tokenize_url,
|
||||
{"model": self.model, "prompt": text, "add_special_tokens": False},
|
||||
)["tokens"]
|
||||
t1 = self._post_json(
|
||||
self.tokenize_url,
|
||||
{"model": self.model, "prompt": text, "add_special_tokens": True},
|
||||
)["tokens"]
|
||||
cases = [
|
||||
{"tokens": t0, "skip_special_tokens": True, "expected": text},
|
||||
{"tokens": t1, "skip_special_tokens": True, "expected": text},
|
||||
{"tokens": t1, "skip_special_tokens": False, "expected": None},
|
||||
{"tokens": [], "skip_special_tokens": True, "expected": ""},
|
||||
]
|
||||
for case in cases:
|
||||
payload = {"model": self.model, "tokens": case["tokens"]}
|
||||
if "skip_special_tokens" in case:
|
||||
payload["skip_special_tokens"] = case["skip_special_tokens"]
|
||||
resp = self._post_json(self.detokenize_url, payload)
|
||||
text_out = resp["text"]
|
||||
if case["expected"] is not None:
|
||||
self.assertEqual(text_out, case["expected"])
|
||||
else:
|
||||
self.assertIsInstance(text_out, str)
|
||||
|
||||
def test_detokenize_invalid_tokens(self):
|
||||
r = self.session.post(
|
||||
self.detokenize_url, json={"model": self.model, "tokens": ["a", "b"]}
|
||||
)
|
||||
self.assertEqual(r.status_code, 400)
|
||||
r2 = self.session.post(
|
||||
self.detokenize_url, json={"model": self.model, "tokens": [1, -1, 2]}
|
||||
)
|
||||
self.assertEqual(r2.status_code, 500)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,219 +0,0 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_srt_engine.TestSRTEngine.test_4_sync_async_stream_combination
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.bench_offline_throughput import BenchArgs, throughput_test
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.few_shot_gsm8k_engine import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
CustomTestCase,
|
||||
)
|
||||
|
||||
|
||||
class TestSRTEngine(CustomTestCase):
|
||||
|
||||
def test_1_engine_runtime_consistency(self):
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
sampling_params = {"temperature": 0, "max_new_tokens": 8}
|
||||
|
||||
engine = sgl.Engine(model_path=model_path, random_seed=42)
|
||||
out1 = engine.generate(prompt, sampling_params)["text"]
|
||||
engine.shutdown()
|
||||
|
||||
runtime = sgl.Runtime(model_path=model_path, random_seed=42)
|
||||
out2 = json.loads(runtime.generate(prompt, sampling_params))["text"]
|
||||
runtime.shutdown()
|
||||
|
||||
print("==== Answer 1 ====")
|
||||
print(out1)
|
||||
|
||||
print("==== Answer 2 ====")
|
||||
print(out2)
|
||||
self.assertEqual(out1, out2)
|
||||
|
||||
def test_2_engine_runtime_encode_consistency(self):
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST
|
||||
|
||||
engine = sgl.Engine(model_path=model_path, is_embedding=True, random_seed=42)
|
||||
out1 = torch.tensor(engine.encode(prompt)["embedding"])
|
||||
engine.shutdown()
|
||||
|
||||
runtime = sgl.Runtime(model_path=model_path, is_embedding=True, random_seed=42)
|
||||
out2 = torch.tensor(json.loads(runtime.encode(prompt))["embedding"])
|
||||
runtime.shutdown()
|
||||
|
||||
self.assertTrue(torch.allclose(out1, out2, atol=1e-5, rtol=1e-3))
|
||||
|
||||
def test_3_engine_token_ids_consistency(self):
|
||||
# just to ensure there is no issue running multiple generate calls
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
sampling_params = {"temperature": 0, "max_new_tokens": 8}
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=model_path, random_seed=42, disable_radix_cache=True
|
||||
)
|
||||
out1 = engine.generate(prompt, sampling_params)["text"]
|
||||
|
||||
tokenizer = get_tokenizer(model_path)
|
||||
token_ids = tokenizer.encode(prompt)
|
||||
out2 = engine.generate(input_ids=token_ids, sampling_params=sampling_params)[
|
||||
"text"
|
||||
]
|
||||
|
||||
engine.shutdown()
|
||||
|
||||
print("==== Answer 1 ====")
|
||||
print(out1)
|
||||
|
||||
print("==== Answer 2 ====")
|
||||
print(out2)
|
||||
self.assertEqual(out1, out2)
|
||||
|
||||
def test_4_sync_async_stream_combination(self):
|
||||
prompt = "AI safety is"
|
||||
sampling_params = {"temperature": 0.8, "top_p": 0.95}
|
||||
|
||||
# Create an LLM.
|
||||
llm = sgl.Engine(
|
||||
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
)
|
||||
|
||||
if True:
|
||||
# 1. sync + non streaming
|
||||
print("\n\n==== 1. sync + non streaming ====")
|
||||
output = llm.generate(prompt, sampling_params)
|
||||
print(output["text"])
|
||||
|
||||
# 2. sync + streaming
|
||||
print("\n\n==== 2. sync + streaming ====")
|
||||
output_generator = llm.generate(prompt, sampling_params, stream=True)
|
||||
offset = 0
|
||||
for output in output_generator:
|
||||
print(output["text"][offset:], end="", flush=True)
|
||||
offset = len(output["text"])
|
||||
print()
|
||||
|
||||
if True:
|
||||
loop = asyncio.get_event_loop()
|
||||
# 3. async + non_streaming
|
||||
print("\n\n==== 3. async + non streaming ====")
|
||||
output = loop.run_until_complete(
|
||||
llm.async_generate(prompt, sampling_params)
|
||||
)
|
||||
print(output["text"])
|
||||
|
||||
# 4. async + streaming
|
||||
async def async_streaming(engine):
|
||||
generator = await engine.async_generate(
|
||||
prompt, sampling_params, stream=True
|
||||
)
|
||||
|
||||
offset = 0
|
||||
async for output in generator:
|
||||
print(output["text"][offset:], end="", flush=True)
|
||||
offset = len(output["text"])
|
||||
print()
|
||||
|
||||
print("\n\n==== 4. async + streaming ====")
|
||||
loop.run_until_complete(async_streaming(llm))
|
||||
|
||||
llm.shutdown()
|
||||
|
||||
def test_5_gsm8k(self):
|
||||
|
||||
args = SimpleNamespace(
|
||||
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
local_data_path=None,
|
||||
num_shots=5,
|
||||
num_questions=1400,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreater(metrics["accuracy"], 0.33)
|
||||
|
||||
def test_6_engine_cpu_offload(self):
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
sampling_params = {"temperature": 0, "max_new_tokens": 8}
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=model_path,
|
||||
random_seed=42,
|
||||
max_total_tokens=128,
|
||||
)
|
||||
out1 = engine.generate(prompt, sampling_params)["text"]
|
||||
engine.shutdown()
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=model_path,
|
||||
random_seed=42,
|
||||
max_total_tokens=128,
|
||||
cpu_offload_gb=3,
|
||||
)
|
||||
out2 = engine.generate(prompt, sampling_params)["text"]
|
||||
engine.shutdown()
|
||||
|
||||
print("==== Answer 1 ====")
|
||||
print(out1)
|
||||
|
||||
print("==== Answer 2 ====")
|
||||
print(out2)
|
||||
self.assertEqual(out1, out2)
|
||||
|
||||
def test_7_engine_offline_throughput(self):
|
||||
server_args = ServerArgs(
|
||||
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
)
|
||||
bench_args = BenchArgs(num_prompts=10)
|
||||
result = throughput_test(server_args=server_args, bench_args=bench_args)
|
||||
self.assertGreater(result["total_throughput"], 3000)
|
||||
|
||||
def test_8_engine_async_encode_consistency(self):
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=model_path,
|
||||
is_embedding=True,
|
||||
random_seed=42,
|
||||
disable_radix_cache=True,
|
||||
)
|
||||
|
||||
# Get sync and async embeddings
|
||||
out1 = torch.tensor(engine.encode(prompt)["embedding"])
|
||||
loop = asyncio.get_event_loop()
|
||||
out2 = torch.tensor(
|
||||
loop.run_until_complete(engine.async_encode(prompt))["embedding"]
|
||||
)
|
||||
|
||||
engine.shutdown()
|
||||
|
||||
print("\n==== Shapes ====")
|
||||
print(f"sync shape: {out1.shape}")
|
||||
print(f"async shape: {out2.shape}")
|
||||
|
||||
self.assertTrue(
|
||||
torch.allclose(out1, out2, atol=1e-5, rtol=1e-3),
|
||||
"Sync and async embeddings are not equal within tolerance",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user