[CI] Move existing unit tests into unit directory (#20631)
This commit is contained in:
Ke Bao
@@ -1,264 +0,0 @@
|
||||
# Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/test_batch_invariance.py
|
||||
import math
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.batch_invariant_ops import batch_invariant_ops
|
||||
from sglang.srt.batch_invariant_ops.batch_invariant_ops import set_batch_invariant_mode
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
# Note: MI300 (gfx942) has 64KB shared memory limit but kernel needs 66KB
|
||||
# MI35x (gfx950/CDNA4) may have different limits - testing on MI35x only
|
||||
register_cuda_ci(est_time=10, suite="nightly-1-gpu", nightly=True)
|
||||
register_amd_ci(est_time=10, suite="nightly-amd-1-gpu-mi35x", nightly=True)
|
||||
|
||||
device_type = getattr(torch.accelerator.current_accelerator(), "type", "cpu")
|
||||
torch.set_default_device(device_type)
|
||||
|
||||
# Just to get the logging out of the way
|
||||
with set_batch_invariant_mode(True):
|
||||
pass
|
||||
|
||||
|
||||
class TestBatchInvariantOps(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
batch_invariant_ops._ENABLE_MM_COMPARISON_TEST = True
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
batch_invariant_ops._ENABLE_MM_COMPARISON_TEST = False
|
||||
|
||||
def _test_batch_invariance(self, M, K, N, dtype):
|
||||
"""
|
||||
Test that matrix operations produce identical results for:
|
||||
- Method 1: Matrix-vector multiplication (batch size 1)
|
||||
- Method 2: Matrix-matrix multiplication, then slice (full batch)
|
||||
"""
|
||||
a = torch.linspace(-100, 100, M * K, dtype=dtype).reshape(M, K)
|
||||
|
||||
# Create non-contiguous tensor
|
||||
b = torch.linspace(-100, 100, K * N, dtype=dtype).reshape(N, K)
|
||||
b = b.transpose(0, 1)
|
||||
|
||||
# Method 1: Matrix-vector multiplication (batch size 1)
|
||||
out1 = torch.mm(a[:1], b)
|
||||
|
||||
# Method 2: Matrix-matrix multiplication, then slice (full batch)
|
||||
out2_pre = torch.mm(a, b)
|
||||
out2 = out2_pre[:1]
|
||||
|
||||
# Check if results are identical
|
||||
diff = (out1 - out2).abs().max()
|
||||
return diff.item()
|
||||
|
||||
def _run_multiple_iterations(self, iters, M, K, N, dtype):
|
||||
"""Run multiple iterations and collect diff statistics"""
|
||||
difflist = []
|
||||
for _ in range(iters):
|
||||
diff = self._test_batch_invariance(M, K, N, dtype)
|
||||
difflist.append(diff)
|
||||
return difflist
|
||||
|
||||
def _assert_batch_invariant_results(self, difflist, dtype, test_name):
|
||||
"""
|
||||
Assert that in batch-invariant mode:
|
||||
1. All diffs must not be NaN
|
||||
2. All diffs must be exactly 0
|
||||
3. Max, min, and diff of diffs must all be 0
|
||||
"""
|
||||
max_diff = max(difflist)
|
||||
min_diff = min(difflist)
|
||||
diff_range = max_diff - min_diff
|
||||
|
||||
# Check for NaN values
|
||||
self.assertFalse(
|
||||
math.isnan(max_diff), f"{test_name}: max_diff is NaN for {dtype}"
|
||||
)
|
||||
self.assertFalse(
|
||||
math.isnan(min_diff), f"{test_name}: min_diff is NaN for {dtype}"
|
||||
)
|
||||
self.assertFalse(
|
||||
math.isnan(diff_range), f"{test_name}: diff_range is NaN for {dtype}"
|
||||
)
|
||||
|
||||
# Check that all diffs are exactly 0
|
||||
self.assertEqual(
|
||||
max_diff,
|
||||
0.0,
|
||||
f"{test_name}: max_diff must be 0 in batch-invariant mode, got {max_diff} for {dtype}",
|
||||
)
|
||||
self.assertEqual(
|
||||
min_diff,
|
||||
0.0,
|
||||
f"{test_name}: min_diff must be 0 in batch-invariant mode, got {min_diff} for {dtype}",
|
||||
)
|
||||
self.assertEqual(
|
||||
diff_range,
|
||||
0.0,
|
||||
f"{test_name}: diff_range must be 0 in batch-invariant mode, got {diff_range} for {dtype}",
|
||||
)
|
||||
|
||||
def test_small_matrices(self):
|
||||
"""Test batch invariance with small matrix sizes"""
|
||||
test_cases = [
|
||||
("Small-1", 8, 64, 128),
|
||||
("Small-2", 16, 128, 256),
|
||||
("Small-3", 4, 32, 64),
|
||||
]
|
||||
|
||||
for name, M, K, N in test_cases:
|
||||
with self.subTest(name=name, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_multiple_iterations(
|
||||
iters=5, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_medium_matrices(self):
|
||||
"""Test batch invariance with medium matrix sizes"""
|
||||
test_cases = [
|
||||
("Medium-1", 32, 128, 1024),
|
||||
("Medium-2", 64, 512, 2048),
|
||||
("Medium-3", 24, 192, 768),
|
||||
]
|
||||
|
||||
for name, M, K, N in test_cases:
|
||||
with self.subTest(name=name, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_multiple_iterations(
|
||||
iters=5, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_large_matrices(self):
|
||||
"""Test batch invariance with large matrix sizes"""
|
||||
test_cases = [
|
||||
("Large-1", 128, 1024, 4096),
|
||||
("Large-2", 256, 2048, 8192),
|
||||
("Large-3", 96, 768, 3072),
|
||||
]
|
||||
|
||||
for name, M, K, N in test_cases:
|
||||
with self.subTest(name=name, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_multiple_iterations(
|
||||
iters=5, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_without_batch_invariant_mode(self):
|
||||
"""
|
||||
Test that without batch-invariant mode, results may differ.
|
||||
This test demonstrates the difference batch-invariant mode makes.
|
||||
"""
|
||||
M, K, N = 32, 128, 1024
|
||||
dtype = torch.float32
|
||||
|
||||
# Run without batch-invariant mode
|
||||
with set_batch_invariant_mode(False):
|
||||
difflist = self._run_multiple_iterations(
|
||||
iters=5, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
print(f"Without batch-invariant mode, we get diffs: {difflist}")
|
||||
|
||||
def _test_bmm_batch_invariance(self, B, M, K, N, dtype):
|
||||
"""
|
||||
Test that BMM operations produce identical results for:
|
||||
- Method 1: BMM with subset of batches
|
||||
- Method 2: BMM with all batches, then slice
|
||||
"""
|
||||
a = torch.linspace(-100, 100, B * M * K, dtype=dtype).reshape(B, M, K)
|
||||
b = torch.linspace(-100, 100, B * K * N, dtype=dtype).reshape(B, K, N)
|
||||
|
||||
# Method 1: BMM with subset (first 2 batches)
|
||||
subset_size = min(2, B)
|
||||
out1 = torch.bmm(a[:subset_size], b[:subset_size])
|
||||
|
||||
# Method 2: BMM with all batches, then slice
|
||||
out2_pre = torch.bmm(a, b)
|
||||
out2 = out2_pre[:subset_size]
|
||||
|
||||
# Check if results are identical
|
||||
diff = (out1 - out2).abs().max()
|
||||
return diff.item()
|
||||
|
||||
def _run_bmm_multiple_iterations(self, iters, B, M, K, N, dtype):
|
||||
"""Run multiple BMM iterations and collect diff statistics"""
|
||||
difflist = []
|
||||
for _ in range(iters):
|
||||
diff = self._test_bmm_batch_invariance(B, M, K, N, dtype)
|
||||
difflist.append(diff)
|
||||
return difflist
|
||||
|
||||
def test_bmm_small_matrices(self):
|
||||
"""Test BMM batch invariance with small matrix sizes"""
|
||||
test_cases = [
|
||||
("BMM-Small-1", 4, 8, 64, 128),
|
||||
("BMM-Small-2", 8, 16, 128, 256),
|
||||
("BMM-Small-3", 6, 4, 32, 64),
|
||||
]
|
||||
|
||||
for name, B, M, K, N in test_cases:
|
||||
with self.subTest(name=name, B=B, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_bmm_multiple_iterations(
|
||||
iters=5, B=B, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_bmm_medium_matrices(self):
|
||||
"""Test BMM batch invariance with medium matrix sizes"""
|
||||
test_cases = [
|
||||
("BMM-Medium-1", 8, 32, 128, 1024),
|
||||
("BMM-Medium-2", 16, 64, 512, 2048),
|
||||
("BMM-Medium-3", 12, 24, 192, 768),
|
||||
]
|
||||
|
||||
for name, B, M, K, N in test_cases:
|
||||
with self.subTest(name=name, B=B, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_bmm_multiple_iterations(
|
||||
iters=5, B=B, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
def test_bmm_large_matrices(self):
|
||||
"""Test BMM batch invariance with large matrix sizes"""
|
||||
test_cases = [
|
||||
("BMM-Large-1", 16, 128, 1024, 4096),
|
||||
("BMM-Large-2", 32, 256, 2048, 8192),
|
||||
("BMM-Large-3", 24, 96, 768, 3072),
|
||||
]
|
||||
|
||||
for name, B, M, K, N in test_cases:
|
||||
with self.subTest(name=name, B=B, M=M, K=K, N=N):
|
||||
for dtype in [torch.float32, torch.bfloat16]:
|
||||
with self.subTest(dtype=dtype):
|
||||
# Run with batch-invariant mode
|
||||
with set_batch_invariant_mode(True):
|
||||
difflist = self._run_bmm_multiple_iterations(
|
||||
iters=5, B=B, M=M, K=K, N=N, dtype=dtype
|
||||
)
|
||||
self._assert_batch_invariant_results(difflist, dtype, name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,581 +0,0 @@
|
||||
import copy
|
||||
import unittest
|
||||
|
||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||
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_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=8, suite="stage-b-test-large-1-gpu")
|
||||
register_amd_ci(est_time=8, suite="stage-b-test-small-1-gpu-amd")
|
||||
|
||||
|
||||
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,156 +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.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=6, suite="stage-b-test-small-1-gpu")
|
||||
register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")
|
||||
|
||||
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,336 +0,0 @@
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs, prepare_server_args
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
|
||||
CustomTestCase,
|
||||
)
|
||||
|
||||
register_cpu_ci(est_time=1, suite="stage-a-cpu-only")
|
||||
|
||||
# Mock get_device() so all tests run on CPU-only CI runners
|
||||
_mock_device = patch("sglang.srt.server_args.get_device", return_value="cuda")
|
||||
_mock_device.start()
|
||||
|
||||
|
||||
class TestPrepareServerArgs(CustomTestCase):
|
||||
def test_prepare_server_args(self):
|
||||
server_args = prepare_server_args(
|
||||
[
|
||||
"--model-path",
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
|
||||
"--json-model-override-args",
|
||||
'{"rope_scaling": {"factor": 2.0, "rope_type": "linear"}}',
|
||||
]
|
||||
)
|
||||
self.assertEqual(server_args.model_path, DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN)
|
||||
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.get_free_port")
|
||||
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
|
||||
def test_init_new_with_nccl_port_none(self, mock_temp_file, mock_get_free_port):
|
||||
"""Test that get_free_port() is called when nccl_port is None"""
|
||||
mock_temp_file.return_value.name = "temp_file"
|
||||
mock_get_free_port.return_value = 45678 # Mock ephemeral port
|
||||
|
||||
# Use MagicMock here to verify get_free_port is called
|
||||
server_args = MagicMock()
|
||||
server_args.nccl_port = None
|
||||
server_args.enable_dp_attention = False
|
||||
server_args.tokenizer_worker_num = 1
|
||||
|
||||
port_args = PortArgs.init_new(server_args)
|
||||
|
||||
# Verify get_free_port was called
|
||||
mock_get_free_port.assert_called_once()
|
||||
|
||||
# Verify the returned port is used
|
||||
self.assertEqual(port_args.nccl_port, 45678)
|
||||
|
||||
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
|
||||
def test_init_new_standard_case(self, mock_temp_file):
|
||||
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)
|
||||
|
||||
def test_init_new_with_single_node_dp_attention(self):
|
||||
|
||||
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)
|
||||
|
||||
def test_init_new_with_dp_rank(self):
|
||||
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)
|
||||
|
||||
def test_init_new_with_ipv4_address(self):
|
||||
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)
|
||||
|
||||
def test_init_new_with_malformed_ipv4_address(self):
|
||||
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(ValueError) as context:
|
||||
PortArgs.init_new(server_args)
|
||||
|
||||
self.assertIn("Missing port", str(context.exception))
|
||||
|
||||
def test_init_new_with_malformed_ipv4_address_invalid_port(self):
|
||||
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):
|
||||
PortArgs.init_new(server_args)
|
||||
|
||||
|
||||
class TestSSLArgs(unittest.TestCase):
|
||||
def test_default_ssl_fields_are_none(self):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
self.assertIsNone(server_args.ssl_keyfile)
|
||||
self.assertIsNone(server_args.ssl_certfile)
|
||||
self.assertIsNone(server_args.ssl_ca_certs)
|
||||
self.assertIsNone(server_args.ssl_keyfile_password)
|
||||
|
||||
def test_ssl_keyfile_without_certfile_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", ssl_keyfile="key.pem")
|
||||
self.assertIn("--ssl-certfile", str(context.exception))
|
||||
|
||||
def test_ssl_certfile_without_keyfile_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", ssl_certfile="cert.pem")
|
||||
self.assertIn("--ssl-keyfile", str(context.exception))
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_ssl_both_keyfile_and_certfile_accepted(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy", ssl_keyfile="key.pem", ssl_certfile="cert.pem"
|
||||
)
|
||||
self.assertEqual(server_args.ssl_keyfile, "key.pem")
|
||||
self.assertEqual(server_args.ssl_certfile, "cert.pem")
|
||||
|
||||
def test_url_returns_http_without_ssl(self):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
self.assertTrue(server_args.url().startswith("http://"))
|
||||
|
||||
def test_url_rewrites_all_interfaces_to_loopback(self):
|
||||
server_args = ServerArgs(model_path="dummy", host="0.0.0.0")
|
||||
self.assertEqual(server_args.url(), "http://127.0.0.1:30000")
|
||||
|
||||
def test_url_rewrites_empty_host_to_loopback(self):
|
||||
server_args = ServerArgs(model_path="dummy", host="")
|
||||
self.assertEqual(server_args.url(), "http://127.0.0.1:30000")
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_url_returns_https_with_ssl(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy", ssl_keyfile="key.pem", ssl_certfile="cert.pem"
|
||||
)
|
||||
self.assertTrue(server_args.url().startswith("https://"))
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_ssl_cli_args_parsed(self, _mock_isfile):
|
||||
server_args = prepare_server_args(
|
||||
[
|
||||
"--model-path",
|
||||
"dummy",
|
||||
"--ssl-keyfile",
|
||||
"key.pem",
|
||||
"--ssl-certfile",
|
||||
"cert.pem",
|
||||
"--ssl-ca-certs",
|
||||
"ca.pem",
|
||||
"--ssl-keyfile-password",
|
||||
"secret",
|
||||
]
|
||||
)
|
||||
self.assertEqual(server_args.ssl_keyfile, "key.pem")
|
||||
self.assertEqual(server_args.ssl_certfile, "cert.pem")
|
||||
self.assertEqual(server_args.ssl_ca_certs, "ca.pem")
|
||||
self.assertEqual(server_args.ssl_keyfile_password, "secret")
|
||||
|
||||
def test_ssl_verify_without_ssl(self):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
self.assertIs(server_args.ssl_verify(), True)
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_ssl_verify_with_ssl_no_ca(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy", ssl_keyfile="key.pem", ssl_certfile="cert.pem"
|
||||
)
|
||||
self.assertIs(server_args.ssl_verify(), False)
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_ssl_verify_with_ssl_and_ca(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile="key.pem",
|
||||
ssl_certfile="cert.pem",
|
||||
ssl_ca_certs="ca.pem",
|
||||
)
|
||||
self.assertEqual(server_args.ssl_verify(), "ca.pem")
|
||||
|
||||
def test_ssl_ca_certs_without_certfile_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", ssl_ca_certs="ca.pem")
|
||||
self.assertIn("--ssl-ca-certs", str(context.exception))
|
||||
|
||||
def test_ssl_keyfile_password_without_certfile_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", ssl_keyfile_password="secret")
|
||||
self.assertIn("--ssl-keyfile-password", str(context.exception))
|
||||
|
||||
def test_ssl_keyfile_not_found_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile="/nonexistent/key.pem",
|
||||
ssl_certfile="/nonexistent/cert.pem",
|
||||
)
|
||||
self.assertIn("not found", str(context.exception))
|
||||
|
||||
def test_ssl_certfile_not_found_raises(self):
|
||||
with tempfile.NamedTemporaryFile(suffix=".pem") as keyfile:
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile=keyfile.name,
|
||||
ssl_certfile="/nonexistent/cert.pem",
|
||||
)
|
||||
self.assertIn("SSL certificate file not found", str(context.exception))
|
||||
|
||||
def test_ssl_ca_certs_not_found_raises(self):
|
||||
with tempfile.NamedTemporaryFile(suffix=".pem") as keyfile:
|
||||
with tempfile.NamedTemporaryFile(suffix=".pem") as certfile:
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile=keyfile.name,
|
||||
ssl_certfile=certfile.name,
|
||||
ssl_ca_certs="/nonexistent/ca.pem",
|
||||
)
|
||||
self.assertIn(
|
||||
"SSL CA certificates file not found", str(context.exception)
|
||||
)
|
||||
|
||||
def test_enable_ssl_refresh_default_false(self):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
self.assertFalse(server_args.enable_ssl_refresh)
|
||||
|
||||
def test_enable_ssl_refresh_without_ssl_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", enable_ssl_refresh=True)
|
||||
self.assertIn("--enable-ssl-refresh", str(context.exception))
|
||||
self.assertIn("--ssl-certfile", str(context.exception))
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_enable_ssl_refresh_with_ssl_accepted(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile="key.pem",
|
||||
ssl_certfile="cert.pem",
|
||||
enable_ssl_refresh=True,
|
||||
)
|
||||
self.assertTrue(server_args.enable_ssl_refresh)
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_enable_ssl_refresh_cli_flag(self, _mock_isfile):
|
||||
server_args = prepare_server_args(
|
||||
[
|
||||
"--model-path",
|
||||
"dummy",
|
||||
"--ssl-keyfile",
|
||||
"key.pem",
|
||||
"--ssl-certfile",
|
||||
"cert.pem",
|
||||
"--enable-ssl-refresh",
|
||||
]
|
||||
)
|
||||
self.assertTrue(server_args.enable_ssl_refresh)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,165 +0,0 @@
|
||||
import asyncio
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from sglang.srt.entrypoints.ssl_utils import SSLCertRefresher
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=9, suite="stage-a-cpu-only")
|
||||
|
||||
|
||||
def _make_temp_pem(content: bytes) -> str:
|
||||
"""Create a temporary PEM file and return its path."""
|
||||
f = tempfile.NamedTemporaryFile(suffix=".pem", delete=False)
|
||||
f.write(content)
|
||||
f.flush()
|
||||
f.close()
|
||||
return f.name
|
||||
|
||||
|
||||
class TestSSLCertRefresher(CustomTestCase):
|
||||
"""Tests for the SSLCertRefresher class."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self._temp_files: list[str] = []
|
||||
|
||||
def tearDown(self):
|
||||
for path in self._temp_files:
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
super().tearDown()
|
||||
|
||||
def _track(self, path: str) -> str:
|
||||
"""Register a temp file for cleanup."""
|
||||
self._temp_files.append(path)
|
||||
return path
|
||||
|
||||
def _run_async(self, coro):
|
||||
"""Helper to run an async coroutine in tests."""
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
return loop.run_until_complete(coro)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
def test_reload_cert_key_on_file_change(self):
|
||||
"""SSLCertRefresher calls load_cert_chain when cert/key files change."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT_V1"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY_V1"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
with open(cert_path, "w") as f:
|
||||
f.write("CERT_V2")
|
||||
|
||||
await asyncio.sleep(1.5)
|
||||
refresher.stop()
|
||||
return mock_ctx
|
||||
|
||||
result_ctx = self._run_async(_test())
|
||||
result_ctx.load_cert_chain.assert_called_with(cert_path, key_path)
|
||||
|
||||
def test_reload_ca_on_file_change(self):
|
||||
"""SSLCertRefresher calls load_verify_locations when CA file changes."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
ca_path = self._track(_make_temp_pem(b"CA_V1"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path, ca_path)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
with open(ca_path, "w") as f:
|
||||
f.write("CA_V2")
|
||||
|
||||
await asyncio.sleep(1.5)
|
||||
refresher.stop()
|
||||
return mock_ctx
|
||||
|
||||
result_ctx = self._run_async(_test())
|
||||
result_ctx.load_verify_locations.assert_called_with(ca_path)
|
||||
|
||||
def test_stop_cancels_tasks(self):
|
||||
"""Calling stop() prevents further reloads."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
refresher.stop()
|
||||
|
||||
with open(cert_path, "w") as f:
|
||||
f.write("CERT_AFTER_STOP")
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
return mock_ctx
|
||||
|
||||
result_ctx = self._run_async(_test())
|
||||
result_ctx.load_cert_chain.assert_not_called()
|
||||
|
||||
def test_no_ca_watcher_when_ca_not_provided(self):
|
||||
"""No CA watcher task is created when ca_path is None."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path)
|
||||
self.assertEqual(len(refresher._tasks), 1)
|
||||
refresher.stop()
|
||||
|
||||
self._run_async(_test())
|
||||
|
||||
def test_ca_watcher_created_when_ca_provided(self):
|
||||
"""A CA watcher task is created when ca_path is provided."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
ca_path = self._track(_make_temp_pem(b"CA"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path, ca_path)
|
||||
self.assertEqual(len(refresher._tasks), 2)
|
||||
refresher.stop()
|
||||
|
||||
self._run_async(_test())
|
||||
|
||||
def test_reload_error_does_not_crash(self):
|
||||
"""A reload error is logged but doesn't crash the watcher."""
|
||||
mock_ctx = MagicMock()
|
||||
mock_ctx.load_cert_chain.side_effect = Exception("bad cert")
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
with open(cert_path, "w") as f:
|
||||
f.write("BAD_CERT")
|
||||
|
||||
await asyncio.sleep(1.5)
|
||||
|
||||
for task in refresher._tasks:
|
||||
self.assertFalse(task.done())
|
||||
|
||||
refresher.stop()
|
||||
|
||||
self._run_async(_test())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user