[CICD] [prefill-only] Consolidate prefill-only model E2E tests (#22405)
This commit is contained in:
Sundara Raman Ramachandran
@@ -0,0 +1,155 @@
|
||||
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.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
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,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=38, suite="stage-b-test-1-gpu-small")
|
||||
register_amd_ci(est_time=38, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Regression tests for input_embeds shape-mismatch bugs.
|
||||
|
||||
Covers two bugs with the same crash signature
|
||||
(RuntimeError: shape mismatch in set_kv_buffer) but opposite polarity:
|
||||
|
||||
- Chunked prefill truncation (#20376): PrefillAdder truncates fill_ids and
|
||||
extend_input_len on chunk overflow but not input_embeds, so the full array
|
||||
flows through while out_cache_loc is sized for the truncated length.
|
||||
Polarity: cache_k > loc.
|
||||
|
||||
- Retraction with output_ids (#14110): after retraction, fill_ids includes
|
||||
accumulated output_ids but input_embeds only covers origin_input_ids.
|
||||
Polarity: cache_k < loc.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
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,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=45, suite="stage-b-test-1-gpu-small")
|
||||
|
||||
CHUNKED_PREFILL_SIZE = 256
|
||||
|
||||
# Shared reference model — loaded once per process, not per test class.
|
||||
_MODEL = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
_tokenizer = None
|
||||
_ref_model = None
|
||||
|
||||
|
||||
def _load_ref():
|
||||
global _tokenizer, _ref_model
|
||||
if _tokenizer is None:
|
||||
_tokenizer = AutoTokenizer.from_pretrained(_MODEL)
|
||||
_ref_model = AutoModelForCausalLM.from_pretrained(_MODEL)
|
||||
|
||||
|
||||
def _embeds_for(text: str) -> list[list[float]]:
|
||||
_load_ref()
|
||||
ids = _tokenizer(text, return_tensors="pt")["input_ids"]
|
||||
embeds = _ref_model.get_input_embeddings()(ids)
|
||||
return embeds.squeeze(0).to(torch.float32).tolist()
|
||||
|
||||
|
||||
def _generate(base_url, input_embeds, max_new_tokens, ignore_eos=False, timeout=120):
|
||||
resp = requests.post(
|
||||
f"{base_url}/generate",
|
||||
json={
|
||||
"input_embeds": input_embeds,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"ignore_eos": ignore_eos,
|
||||
},
|
||||
},
|
||||
timeout=timeout,
|
||||
)
|
||||
return resp
|
||||
|
||||
|
||||
class TestInputEmbedsChunkedAndRetract(CustomTestCase):
|
||||
"""Single server launch covering both bugs.
|
||||
|
||||
Both tests require --disable-radix-cache (for input_embeds). The chunked
|
||||
prefill test needs a small --chunked-prefill-size. The retraction test
|
||||
uses SGLANG_TEST_RETRACT to deterministically force retraction every few
|
||||
scheduler iterations regardless of KV pressure.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
# SGLANG_TEST_RETRACT forces retraction periodically; this is
|
||||
# deterministic and doesn't require guessing KV budgets.
|
||||
with envs.SGLANG_TEST_RETRACT.override(True):
|
||||
cls.process = popen_launch_server(
|
||||
_MODEL,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--disable-radix-cache",
|
||||
"--chunked-prefill-size",
|
||||
str(CHUNKED_PREFILL_SIZE),
|
||||
"--cuda-graph-max-bs",
|
||||
"4",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def _assert_server_alive(self):
|
||||
self.assertIsNone(self.process.poll(), "server process crashed")
|
||||
|
||||
def test_chunked_prefill_truncation_and_continuation(self):
|
||||
"""Regression test for #20376.
|
||||
|
||||
A single request longer than chunked_prefill_size deterministically
|
||||
exercises both (a) first-chunk truncation and (b) chunk continuation,
|
||||
without any concurrent-timing dependency. Pre-fix this crashes in
|
||||
set_kv_buffer on both chunks.
|
||||
"""
|
||||
# ~80 tokens each repetition; 6 repetitions exceeds CHUNKED_PREFILL_SIZE
|
||||
# comfortably. Token count is model-dependent so assert it.
|
||||
text = "The quick brown fox jumps over the lazy dog. " * 40
|
||||
embeds = _embeds_for(text)
|
||||
self.assertGreater(
|
||||
len(embeds),
|
||||
CHUNKED_PREFILL_SIZE,
|
||||
f"prompt must exceed chunked_prefill_size={CHUNKED_PREFILL_SIZE} "
|
||||
f"to trigger chunking; got {len(embeds)} tokens",
|
||||
)
|
||||
|
||||
resp = _generate(self.base_url, embeds, max_new_tokens=8)
|
||||
self.assertEqual(resp.status_code, 200, resp.text[:300])
|
||||
body = resp.json()
|
||||
self.assertIn("text", body)
|
||||
self.assertIsInstance(body["text"], str)
|
||||
self._assert_server_alive()
|
||||
|
||||
def test_chunked_prefill_batch_truncation(self):
|
||||
"""Regression test for #20376 — multi-request batch case.
|
||||
|
||||
A batch POST with total tokens > chunked_prefill_size goes through a
|
||||
single ZMQ send, so all requests land in the same scheduler iteration
|
||||
and the PrefillAdder is forced to truncate at least one. This matches
|
||||
the original thundering-herd trigger without HTTP timing races.
|
||||
"""
|
||||
text = "The quick brown fox jumps over the lazy dog. " * 8
|
||||
embeds = _embeds_for(text)
|
||||
seq_len = len(embeds)
|
||||
|
||||
# Enough batched requests to overflow the chunk budget.
|
||||
n = max(4, CHUNKED_PREFILL_SIZE // seq_len + 2)
|
||||
self.assertGreater(n * seq_len, CHUNKED_PREFILL_SIZE)
|
||||
|
||||
resp = _generate(self.base_url, [embeds] * n, max_new_tokens=8)
|
||||
self.assertEqual(resp.status_code, 200, resp.text[:300])
|
||||
results = resp.json()
|
||||
self.assertEqual(len(results), n)
|
||||
for r in results:
|
||||
self.assertIn("text", r)
|
||||
self._assert_server_alive()
|
||||
|
||||
def test_retraction_with_output_ids(self):
|
||||
"""Regression test for #14110.
|
||||
|
||||
SGLANG_TEST_RETRACT forces retraction every few scheduler iterations.
|
||||
Combined with ignore_eos and a reasonable max_new_tokens, at least one
|
||||
request is retracted mid-decode with non-empty output_ids, then
|
||||
re-prefilled. Pre-#14110 this crashes (cache_k < loc) because fill_ids
|
||||
includes output_ids but input_embeds does not.
|
||||
"""
|
||||
text = "The quick brown fox jumps over the lazy dog. " * 4
|
||||
embeds = _embeds_for(text)
|
||||
|
||||
# Batch of requests with enough decode steps that SGLANG_TEST_RETRACT
|
||||
# (interval=3 by default) fires mid-decode.
|
||||
n = 4
|
||||
resp = _generate(
|
||||
self.base_url,
|
||||
[embeds] * n,
|
||||
max_new_tokens=32,
|
||||
ignore_eos=True,
|
||||
)
|
||||
self.assertEqual(resp.status_code, 200, resp.text[:300])
|
||||
results = resp.json()
|
||||
self.assertEqual(len(results), n)
|
||||
for r in results:
|
||||
self.assertIn("text", r)
|
||||
self._assert_server_alive()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user