[CI] Move existing unit tests into unit directory (#20631)
This commit is contained in:
Ke Bao
@@ -0,0 +1,581 @@
|
||||
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()
|
||||
@@ -0,0 +1,447 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.srt.managers.schedule_batch import Req
|
||||
from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder
|
||||
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=1, suite="stage-b-test-small-1-gpu")
|
||||
register_amd_ci(est_time=2, suite="stage-b-test-small-1-gpu-amd")
|
||||
|
||||
|
||||
class TestPrefillAdder(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.mock_tree_cache = self.create_tree_cache()
|
||||
self.mock_token_allocator = self.create_token_allocator()
|
||||
patcher = patch(
|
||||
"sglang.srt.managers.schedule_policy.is_nsa_prefill_cp_in_seq_split",
|
||||
return_value=False,
|
||||
)
|
||||
self.mock_is_nsa = patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
def create_tree_cache(
|
||||
self,
|
||||
*,
|
||||
full_evictable_size: int = 0,
|
||||
swa_evictable_size: int = 0,
|
||||
evictable_size: int = 0,
|
||||
) -> MagicMock:
|
||||
tree_cache = MagicMock()
|
||||
tree_cache.full_evictable_size.return_value = full_evictable_size
|
||||
tree_cache.swa_evictable_size.return_value = swa_evictable_size
|
||||
tree_cache.evictable_size.return_value = evictable_size
|
||||
tree_cache.disable = False
|
||||
tree_cache.inc_lock_ref.return_value = None
|
||||
tree_cache.dec_lock_ref.return_value = None
|
||||
return tree_cache
|
||||
|
||||
def create_token_allocator(
|
||||
self,
|
||||
*,
|
||||
full_available_size: int = 0,
|
||||
swa_available_size: int = 0,
|
||||
available_size: int = 0,
|
||||
) -> MagicMock:
|
||||
allocator = MagicMock()
|
||||
allocator.full_available_size.return_value = full_available_size
|
||||
allocator.swa_available_size.return_value = swa_available_size
|
||||
allocator.available_size.return_value = available_size
|
||||
return allocator
|
||||
|
||||
def create_running_batch(self, reqs=None) -> MagicMock:
|
||||
batch = MagicMock()
|
||||
batch.reqs = list(reqs or [])
|
||||
batch.release_req.return_value = None
|
||||
batch.filter_batch.return_value = None
|
||||
return batch
|
||||
|
||||
def create_server_args(
|
||||
self, *, schedule_low_priority_values_first: bool
|
||||
) -> MagicMock:
|
||||
server_args = MagicMock()
|
||||
server_args.schedule_low_priority_values_first = (
|
||||
schedule_low_priority_values_first
|
||||
)
|
||||
return server_args
|
||||
|
||||
def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time=0):
|
||||
req = MagicMock(spec=Req)
|
||||
req.rid = str(rid)
|
||||
req.priority = priority
|
||||
req.extend_input_len = 0
|
||||
req.extend_logprob_start_len = 0
|
||||
req.output_ids = [0] * output_len
|
||||
req.sampling_params = SimpleNamespace(max_new_tokens=max_new_tokens)
|
||||
req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time)
|
||||
req.finished.return_value = False
|
||||
return req
|
||||
|
||||
def create_adder(self, running_batch, **kwargs):
|
||||
defaults = dict(
|
||||
page_size=1,
|
||||
tree_cache=self.mock_tree_cache,
|
||||
token_to_kv_pool_allocator=self.mock_token_allocator,
|
||||
running_batch=running_batch,
|
||||
new_token_ratio=1.0,
|
||||
rem_input_tokens=10000,
|
||||
rem_chunk_tokens=None,
|
||||
mixed_with_decode_tokens=0,
|
||||
priority_scheduling_preemption_threshold=0,
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return PrefillAdder(**defaults)
|
||||
|
||||
def test_preempt_success_high_priority_values_first(self):
|
||||
params = [
|
||||
("run1", 0, 50),
|
||||
("run2", 1, 75),
|
||||
("run3", 2, 100),
|
||||
]
|
||||
running_reqs = [
|
||||
self.create_mock_req(rid, priority, max_new_tokens)
|
||||
for rid, priority, max_new_tokens in params
|
||||
]
|
||||
mock_server_args = self.create_server_args(
|
||||
schedule_low_priority_values_first=False
|
||||
)
|
||||
running_batch = self.create_running_batch(running_reqs)
|
||||
adder = self.create_adder(running_batch)
|
||||
|
||||
self.assertEqual(adder.rem_total_token_offset, 225)
|
||||
|
||||
self.mock_token_allocator.full_available_size.return_value = (
|
||||
225 # full occupation of GRam
|
||||
)
|
||||
self.mock_token_allocator.available_size.return_value = 225
|
||||
|
||||
new_req = self.create_mock_req("new1", priority=1, max_new_tokens=49)
|
||||
|
||||
success = adder.preempt_to_schedule(new_req, mock_server_args)
|
||||
|
||||
self.assertTrue(success)
|
||||
self.assertIn(running_reqs[0], adder.preempt_list)
|
||||
self.assertEqual(adder.rem_total_token_offset, 175) # 50 + 75 + 100 - 50 = 175
|
||||
running_batch.release_req.assert_called_once()
|
||||
|
||||
def test_preempt_success_low_priority_values_first(self):
|
||||
params = [
|
||||
("run1", 0, 50),
|
||||
("run2", 1, 75),
|
||||
("run3", 2, 100),
|
||||
]
|
||||
running_reqs = [
|
||||
self.create_mock_req(rid, priority, max_new_tokens)
|
||||
for rid, priority, max_new_tokens in params
|
||||
]
|
||||
mock_server_args = self.create_server_args(
|
||||
schedule_low_priority_values_first=True
|
||||
)
|
||||
running_batch = self.create_running_batch(running_reqs)
|
||||
adder = self.create_adder(running_batch)
|
||||
|
||||
self.assertEqual(adder.rem_total_token_offset, 225)
|
||||
|
||||
self.mock_token_allocator.full_available_size.return_value = (
|
||||
225 # full occupation of GRam
|
||||
)
|
||||
self.mock_token_allocator.available_size.return_value = 225
|
||||
|
||||
new_req = self.create_mock_req("new1", priority=1, max_new_tokens=49)
|
||||
|
||||
success = adder.preempt_to_schedule(new_req, mock_server_args)
|
||||
|
||||
self.assertTrue(success)
|
||||
self.assertIn(running_reqs[2], adder.preempt_list)
|
||||
self.assertEqual(adder.rem_total_token_offset, 125) # 50 + 75 + 100 - 100 = 125
|
||||
running_batch.release_req.assert_called_once()
|
||||
|
||||
def test_preempt_fail_low_priority_values_first(self):
|
||||
params = [
|
||||
("run1", 0, 50),
|
||||
("run2", 1, 75),
|
||||
("run3", 2, 100),
|
||||
]
|
||||
running_reqs = [
|
||||
self.create_mock_req(rid, priority, max_new_tokens)
|
||||
for rid, priority, max_new_tokens in params
|
||||
]
|
||||
mock_server_args = self.create_server_args(
|
||||
schedule_low_priority_values_first=True
|
||||
)
|
||||
running_batch = self.create_running_batch(running_reqs)
|
||||
adder = self.create_adder(running_batch)
|
||||
|
||||
self.assertEqual(adder.rem_total_token_offset, 225)
|
||||
|
||||
self.mock_token_allocator.full_available_size.return_value = (
|
||||
225 # full occupation of GRam
|
||||
)
|
||||
self.mock_token_allocator.available_size.return_value = 225
|
||||
|
||||
new_req_fail_by_priority_check = self.create_mock_req(
|
||||
"new1", priority=2, max_new_tokens=49
|
||||
)
|
||||
|
||||
success_by_priority_check = adder.preempt_to_schedule(
|
||||
new_req_fail_by_priority_check, mock_server_args
|
||||
)
|
||||
self.assertFalse(success_by_priority_check)
|
||||
|
||||
new_req_fail_by_priority_check = self.create_mock_req(
|
||||
"new2", priority=1, max_new_tokens=110
|
||||
)
|
||||
success_by_capacity_check = adder.preempt_to_schedule(
|
||||
new_req_fail_by_priority_check, mock_server_args
|
||||
)
|
||||
self.assertFalse(success_by_capacity_check)
|
||||
|
||||
def test_preempt_fail_high_priority_values_first(self):
|
||||
params = [
|
||||
("run1", 0, 50),
|
||||
("run2", 1, 75),
|
||||
("run3", 2, 100),
|
||||
]
|
||||
running_reqs = [
|
||||
self.create_mock_req(rid, priority, max_new_tokens)
|
||||
for rid, priority, max_new_tokens in params
|
||||
]
|
||||
mock_server_args = self.create_server_args(
|
||||
schedule_low_priority_values_first=False
|
||||
)
|
||||
running_batch = self.create_running_batch(running_reqs)
|
||||
adder = self.create_adder(running_batch)
|
||||
|
||||
self.assertEqual(adder.rem_total_token_offset, 225)
|
||||
|
||||
self.mock_token_allocator.full_available_size.return_value = (
|
||||
225 # full occupation of GRam
|
||||
)
|
||||
self.mock_token_allocator.available_size.return_value = 225
|
||||
|
||||
new_req_fail_by_priority_check = self.create_mock_req(
|
||||
"new1", priority=0, max_new_tokens=49
|
||||
)
|
||||
|
||||
success_by_priority_check = adder.preempt_to_schedule(
|
||||
new_req_fail_by_priority_check, mock_server_args
|
||||
)
|
||||
self.assertFalse(success_by_priority_check)
|
||||
|
||||
new_req_fail_by_priority_check = self.create_mock_req(
|
||||
"new2", priority=-1, max_new_tokens=110
|
||||
)
|
||||
success_by_capacity_check = adder.preempt_to_schedule(
|
||||
new_req_fail_by_priority_check, mock_server_args
|
||||
)
|
||||
self.assertFalse(success_by_capacity_check)
|
||||
|
||||
def test_preempt_skip_already_preempted_request(self):
|
||||
params = [
|
||||
("req_prio_0", 0, 50),
|
||||
("req_prio_1", 1, 75),
|
||||
("req_prio_2", 2, 100),
|
||||
]
|
||||
running_reqs = [
|
||||
self.create_mock_req(rid, priority, max_new_tokens)
|
||||
for rid, priority, max_new_tokens in params
|
||||
]
|
||||
mock_server_args = self.create_server_args(
|
||||
schedule_low_priority_values_first=False
|
||||
)
|
||||
running_batch = self.create_running_batch(running_reqs)
|
||||
adder = self.create_adder(running_batch)
|
||||
|
||||
self.assertEqual(adder.rem_total_token_offset, 225)
|
||||
|
||||
self.mock_token_allocator.full_available_size.return_value = 225
|
||||
self.mock_token_allocator.available_size.return_value = 225
|
||||
|
||||
# New request preempts req_prio_0
|
||||
first_req = self.create_mock_req(
|
||||
"new_req_prio_1", priority=1, max_new_tokens=49
|
||||
)
|
||||
first_success = adder.preempt_to_schedule(first_req, mock_server_args)
|
||||
self.assertTrue(first_success)
|
||||
self.assertIn(running_reqs[0], adder.preempt_list)
|
||||
self.assertEqual(adder.rem_total_token_offset, 175)
|
||||
running_batch.release_req.assert_called_once()
|
||||
|
||||
# Second call needs more tokens than currently free, so it would need to
|
||||
# preempt req_prio_0 again if already-preempted requests were not filtered out.
|
||||
second_req = self.create_mock_req(
|
||||
"second_new_req_prio_1", priority=1, max_new_tokens=76
|
||||
)
|
||||
second_success = adder.preempt_to_schedule(second_req, mock_server_args)
|
||||
|
||||
self.assertFalse(second_success)
|
||||
self.assertEqual(adder.rem_total_token_offset, 175)
|
||||
self.assertEqual(adder.preempt_list.count(running_reqs[0]), 1)
|
||||
running_batch.release_req.assert_called_once()
|
||||
|
||||
def test_preempt_success_low_priority_values_first_exact_once(self):
|
||||
params = [
|
||||
("run1", 0, 50),
|
||||
("run2", 1, 75),
|
||||
("run3", 2, 100),
|
||||
("run4", 2, 125),
|
||||
("run4", 2, 125),
|
||||
]
|
||||
running_reqs = [
|
||||
self.create_mock_req(rid, priority, max_new_tokens)
|
||||
for rid, priority, max_new_tokens in params
|
||||
]
|
||||
mock_server_args = self.create_server_args(
|
||||
schedule_low_priority_values_first=True
|
||||
)
|
||||
running_batch = self.create_running_batch(running_reqs)
|
||||
adder = self.create_adder(running_batch)
|
||||
|
||||
self.assertEqual(adder.rem_total_token_offset, 475)
|
||||
|
||||
self.mock_token_allocator.full_available_size.return_value = (
|
||||
475 # full occupation of GRam
|
||||
)
|
||||
self.mock_token_allocator.available_size.return_value = 475
|
||||
|
||||
new_req = self.create_mock_req("new1", priority=1, max_new_tokens=75)
|
||||
|
||||
success = adder.preempt_to_schedule(new_req, mock_server_args)
|
||||
self.assertTrue(success)
|
||||
self.assertIn(running_reqs[2], adder.preempt_list)
|
||||
self.assertEqual(
|
||||
adder.rem_total_token_offset, 375
|
||||
) # 50 + 75 + 100 + 125 + 125 - 100 = 375
|
||||
running_batch.release_req.assert_called_once()
|
||||
|
||||
def test_preempt_success_low_priority_values_first_exact_twice(self):
|
||||
params = [
|
||||
("run1", 0, 50),
|
||||
("run2", 1, 75),
|
||||
("run3", 2, 100),
|
||||
("run4", 2, 125),
|
||||
("run4", 2, 125),
|
||||
]
|
||||
running_reqs = [
|
||||
self.create_mock_req(rid, priority, max_new_tokens)
|
||||
for rid, priority, max_new_tokens in params
|
||||
]
|
||||
mock_server_args = self.create_server_args(
|
||||
schedule_low_priority_values_first=True
|
||||
)
|
||||
running_batch = self.create_running_batch(running_reqs)
|
||||
adder = self.create_adder(running_batch)
|
||||
|
||||
self.assertEqual(adder.rem_total_token_offset, 475)
|
||||
|
||||
self.mock_token_allocator.full_available_size.return_value = (
|
||||
475 # full occupation of GRam
|
||||
)
|
||||
self.mock_token_allocator.available_size.return_value = 475
|
||||
|
||||
new_req = self.create_mock_req("new1", priority=1, max_new_tokens=200)
|
||||
|
||||
success = adder.preempt_to_schedule(new_req, mock_server_args)
|
||||
self.assertTrue(success)
|
||||
self.assertIn(running_reqs[2], adder.preempt_list)
|
||||
self.assertIn(running_reqs[3], adder.preempt_list)
|
||||
self.assertEqual(
|
||||
adder.rem_total_token_offset, 250
|
||||
) # 50 + 75 + 100 + 125 + 125 - 100 - 125 = 250
|
||||
self.assertEqual(running_batch.release_req.call_count, 2)
|
||||
|
||||
def test_mixed_chunk_prefill_budgets(self):
|
||||
self.mock_token_allocator.available_size.return_value = 1000
|
||||
|
||||
decode_reqs = [
|
||||
self.create_mock_req(f"decode_{i}", priority=0, max_new_tokens=50)
|
||||
for i in range(8)
|
||||
]
|
||||
running_batch = self.create_running_batch(decode_reqs)
|
||||
|
||||
adder = self.create_adder(
|
||||
running_batch,
|
||||
rem_input_tokens=200,
|
||||
rem_chunk_tokens=64,
|
||||
mixed_with_decode_tokens=len(decode_reqs),
|
||||
)
|
||||
|
||||
self.assertEqual(adder.rem_input_tokens, 192) # 200 - 8
|
||||
self.assertEqual(adder.rem_chunk_tokens, 56) # 64 - 8
|
||||
self.assertEqual(adder.rem_total_token_offset, 408) # 8 + 8 * 50
|
||||
self.assertEqual(adder.cur_rem_token_offset, 8)
|
||||
self.assertEqual(adder.budget_state(), AddReqResult.CONTINUE)
|
||||
|
||||
# Add a prefill that exactly consumes the chunk budget
|
||||
req1 = self.create_mock_req("req1", priority=0, max_new_tokens=64)
|
||||
req1.extend_input_len = 56
|
||||
req1.host_hit_length = 0
|
||||
req1.prefix_indices = []
|
||||
req1.fill_ids = list(range(56))
|
||||
req1.last_node = MagicMock()
|
||||
req1.sampling_params.ignore_eos = False
|
||||
|
||||
result1 = adder.add_one_req(
|
||||
req1, has_chunked_req=False, truncation_align_size=None
|
||||
)
|
||||
|
||||
self.assertEqual(len(adder.can_run_list), 1)
|
||||
self.assertEqual(adder.rem_chunk_tokens, 0) # 56 - 56
|
||||
self.assertEqual(adder.rem_input_tokens, 136) # 192 - 56
|
||||
self.assertEqual(result1, AddReqResult.OTHER)
|
||||
|
||||
# 3 decode requests finished
|
||||
remaining_decode_reqs = decode_reqs[3:]
|
||||
running_batch2 = self.create_running_batch(remaining_decode_reqs)
|
||||
|
||||
adder2 = self.create_adder(
|
||||
running_batch2,
|
||||
rem_input_tokens=200,
|
||||
rem_chunk_tokens=64,
|
||||
mixed_with_decode_tokens=len(remaining_decode_reqs),
|
||||
)
|
||||
|
||||
self.assertEqual(adder2.rem_input_tokens, 195) # 200 - 5
|
||||
self.assertEqual(adder2.rem_chunk_tokens, 59) # 64 - 5
|
||||
self.assertEqual(adder2.rem_total_token_offset, 255) # 5 + 5 * 50
|
||||
self.assertEqual(adder2.budget_state(), AddReqResult.CONTINUE)
|
||||
|
||||
# Same prefill no longer exhausts the chunk budget
|
||||
req2 = self.create_mock_req("req2", priority=0, max_new_tokens=64)
|
||||
req2.extend_input_len = 56
|
||||
req2.host_hit_length = 0
|
||||
req2.prefix_indices = []
|
||||
req2.fill_ids = list(range(56))
|
||||
req2.last_node = MagicMock()
|
||||
req2.sampling_params.ignore_eos = False
|
||||
|
||||
result2 = adder2.add_one_req(
|
||||
req2, has_chunked_req=False, truncation_align_size=None
|
||||
)
|
||||
|
||||
self.assertEqual(len(adder2.can_run_list), 1)
|
||||
self.assertEqual(adder2.rem_chunk_tokens, 3) # 59 - 56 = 3 remaining
|
||||
self.assertEqual(result2, AddReqResult.CONTINUE)
|
||||
|
||||
# Fit last small prefill request
|
||||
req3 = self.create_mock_req("req3", priority=0, max_new_tokens=16)
|
||||
req3.extend_input_len = 3
|
||||
req3.host_hit_length = 0
|
||||
req3.prefix_indices = []
|
||||
req3.fill_ids = list(range(3))
|
||||
req3.last_node = MagicMock()
|
||||
req3.sampling_params.ignore_eos = False
|
||||
|
||||
result3 = adder2.add_one_req(
|
||||
req3, has_chunked_req=False, truncation_align_size=None
|
||||
)
|
||||
|
||||
self.assertEqual(len(adder2.can_run_list), 2)
|
||||
self.assertEqual(adder2.rem_chunk_tokens, 0) # 3 - 3 = 0
|
||||
self.assertEqual(result3, AddReqResult.OTHER)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,166 @@
|
||||
import json
|
||||
import unittest
|
||||
|
||||
from sglang.srt.managers.io_struct import ProfileReqInput
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=9, suite="stage-b-test-small-1-gpu")
|
||||
register_amd_ci(est_time=9, suite="stage-b-test-small-1-gpu-amd")
|
||||
|
||||
|
||||
class TestProfileMergerHTTPAPI(unittest.TestCase):
|
||||
def test_profile_req_input_merge_profiles_json_serialization(self):
|
||||
# Test with merge_profiles=True
|
||||
req_input = ProfileReqInput(
|
||||
output_dir="/tmp/test",
|
||||
num_steps=5,
|
||||
activities=["CPU", "GPU"],
|
||||
profile_by_stage=True,
|
||||
merge_profiles=True,
|
||||
)
|
||||
|
||||
# Convert to dict (as would happen in HTTP request)
|
||||
req_dict = {
|
||||
"output_dir": req_input.output_dir,
|
||||
"num_steps": req_input.num_steps,
|
||||
"activities": req_input.activities,
|
||||
"profile_by_stage": req_input.profile_by_stage,
|
||||
"merge_profiles": req_input.merge_profiles,
|
||||
}
|
||||
|
||||
# Test JSON serialization
|
||||
json_str = json.dumps(req_dict)
|
||||
parsed_data = json.loads(json_str)
|
||||
|
||||
self.assertTrue(parsed_data["merge_profiles"])
|
||||
self.assertEqual(parsed_data["output_dir"], "/tmp/test")
|
||||
self.assertEqual(parsed_data["num_steps"], 5)
|
||||
self.assertEqual(parsed_data["activities"], ["CPU", "GPU"])
|
||||
self.assertTrue(parsed_data["profile_by_stage"])
|
||||
|
||||
def test_profile_req_input_merge_profiles_json_deserialization(self):
|
||||
# Test JSON data as would come from HTTP request
|
||||
json_data = {
|
||||
"output_dir": "/tmp/test",
|
||||
"num_steps": 10,
|
||||
"activities": ["CPU", "GPU", "MEM"],
|
||||
"profile_by_stage": False,
|
||||
"merge_profiles": True,
|
||||
}
|
||||
|
||||
# Create ProfileReqInput from dict (as HTTP server would do)
|
||||
req_input = ProfileReqInput(**json_data)
|
||||
|
||||
self.assertTrue(req_input.merge_profiles)
|
||||
self.assertEqual(req_input.output_dir, "/tmp/test")
|
||||
self.assertEqual(req_input.num_steps, 10)
|
||||
self.assertEqual(req_input.activities, ["CPU", "GPU", "MEM"])
|
||||
self.assertFalse(req_input.profile_by_stage)
|
||||
|
||||
def test_profile_req_input_merge_profiles_default_value(self):
|
||||
# Test with minimal data
|
||||
json_data = {"output_dir": "/tmp/test"}
|
||||
|
||||
req_input = ProfileReqInput(**json_data)
|
||||
self.assertFalse(req_input.merge_profiles)
|
||||
|
||||
def test_profile_req_input_merge_profiles_explicit_false(self):
|
||||
json_data = {"output_dir": "/tmp/test", "merge_profiles": False}
|
||||
|
||||
req_input = ProfileReqInput(**json_data)
|
||||
self.assertFalse(req_input.merge_profiles)
|
||||
|
||||
def test_http_api_parameter_flow(self):
|
||||
# Simulate HTTP request data
|
||||
request_data = {
|
||||
"output_dir": "/tmp/test",
|
||||
"num_steps": 5,
|
||||
"activities": ["CPU", "GPU"],
|
||||
"profile_by_stage": True,
|
||||
"merge_profiles": True,
|
||||
}
|
||||
|
||||
# Create ProfileReqInput as HTTP server would
|
||||
obj = ProfileReqInput(**request_data)
|
||||
|
||||
# Verify the parameter is set correctly
|
||||
self.assertTrue(obj.merge_profiles)
|
||||
self.assertEqual(obj.output_dir, "/tmp/test")
|
||||
self.assertEqual(obj.num_steps, 5)
|
||||
self.assertEqual(obj.activities, ["CPU", "GPU"])
|
||||
self.assertTrue(obj.profile_by_stage)
|
||||
|
||||
def test_http_api_parameter_validation(self):
|
||||
# Test with True
|
||||
json_data = {"merge_profiles": True}
|
||||
req_input = ProfileReqInput(**json_data)
|
||||
self.assertTrue(req_input.merge_profiles)
|
||||
|
||||
# Test with False
|
||||
json_data = {"merge_profiles": False}
|
||||
req_input = ProfileReqInput(**json_data)
|
||||
self.assertFalse(req_input.merge_profiles)
|
||||
|
||||
# Test with string "true" (should be converted by JSON parser)
|
||||
json_data = {"merge_profiles": "true"}
|
||||
req_input = ProfileReqInput(**json_data)
|
||||
self.assertEqual(req_input.merge_profiles, "true") # String, not boolean
|
||||
|
||||
def test_http_api_backward_compatibility(self):
|
||||
# Test minimal request (no merge_profiles)
|
||||
json_data = {}
|
||||
req_input = ProfileReqInput(**json_data)
|
||||
self.assertFalse(req_input.merge_profiles) # Should default to False
|
||||
|
||||
# Test with other parameters but no merge_profiles
|
||||
json_data = {
|
||||
"output_dir": "/tmp/test",
|
||||
"num_steps": 5,
|
||||
"activities": ["CPU", "GPU"],
|
||||
}
|
||||
req_input = ProfileReqInput(**json_data)
|
||||
self.assertFalse(req_input.merge_profiles) # Should default to False
|
||||
|
||||
def test_http_api_parameter_combinations(self):
|
||||
test_cases = [
|
||||
{
|
||||
"name": "minimal with merge_profiles",
|
||||
"data": {"merge_profiles": True},
|
||||
"expected_merge": True,
|
||||
},
|
||||
{
|
||||
"name": "full parameters with merge_profiles=True",
|
||||
"data": {
|
||||
"output_dir": "/tmp/test",
|
||||
"num_steps": 10,
|
||||
"activities": ["CPU", "GPU", "MEM"],
|
||||
"profile_by_stage": True,
|
||||
"with_stack": True,
|
||||
"record_shapes": True,
|
||||
"merge_profiles": True,
|
||||
},
|
||||
"expected_merge": True,
|
||||
},
|
||||
{
|
||||
"name": "full parameters with merge_profiles=False",
|
||||
"data": {
|
||||
"output_dir": "/tmp/test",
|
||||
"num_steps": 10,
|
||||
"activities": ["CPU", "GPU", "MEM"],
|
||||
"profile_by_stage": False,
|
||||
"with_stack": False,
|
||||
"record_shapes": False,
|
||||
"merge_profiles": False,
|
||||
},
|
||||
"expected_merge": False,
|
||||
},
|
||||
]
|
||||
|
||||
for test_case in test_cases:
|
||||
with self.subTest(test_case["name"]):
|
||||
req_input = ProfileReqInput(**test_case["data"])
|
||||
self.assertEqual(req_input.merge_profiles, test_case["expected_merge"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user