mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2025-12-24 13:28:13 +08:00
6fd3e72da1
* bench command * bench command * bench command * bench command * bench command --------- Co-authored-by: K11OntheBoat <your_email@example.com>
105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from fastdeploy.benchmarks.lib import utils
|
|
|
|
|
|
class TestConvertToPytorchBenchmarkFormat(unittest.TestCase):
|
|
def test_empty_metrics(self):
|
|
args = MagicMock()
|
|
args.model = "test_model"
|
|
metrics = {}
|
|
extra_info = {}
|
|
result = utils.convert_to_pytorch_benchmark_format(args, metrics, extra_info)
|
|
self.assertEqual(result, [])
|
|
|
|
def test_with_metrics_no_save_env(self):
|
|
args = MagicMock()
|
|
args.model = "test_model"
|
|
args.tensor_parallel_size = 2
|
|
metrics = {"latency": [100, 200]}
|
|
extra_info = {"batch_size": 32}
|
|
|
|
with patch.dict(os.environ, {"SAVE_TO_PYTORCH_BENCHMARK_FORMAT": "False"}):
|
|
with patch.object(utils, "os") as mock_os:
|
|
mock_os.environ.get.return_value = False
|
|
result = utils.convert_to_pytorch_benchmark_format(args, metrics, extra_info)
|
|
self.assertEqual(result, [])
|
|
|
|
def test_with_metrics_and_save_env(self):
|
|
args = MagicMock()
|
|
args.model = "test_model"
|
|
args.tensor_parallel_size = 2
|
|
metrics = {"latency": [100, 200]}
|
|
extra_info = {"batch_size": 32}
|
|
|
|
with patch.dict(os.environ, {"SAVE_TO_PYTORCH_BENCHMARK_FORMAT": "True"}):
|
|
result = utils.convert_to_pytorch_benchmark_format(args, metrics, extra_info)
|
|
self.assertEqual(len(result), 1)
|
|
self.assertEqual(result[0]["model"]["name"], "test_model")
|
|
self.assertEqual(result[0]["metric"]["name"], "latency")
|
|
self.assertEqual(result[0]["metric"]["benchmark_values"], [100, 200])
|
|
|
|
|
|
class TestInfEncoder(unittest.TestCase):
|
|
def test_clear_inf_with_dict(self):
|
|
encoder = utils.InfEncoder()
|
|
data = {"a": float("inf"), "b": 1.0}
|
|
result = encoder.clear_inf(data)
|
|
self.assertEqual(result, {"a": "inf", "b": 1.0})
|
|
|
|
def test_clear_inf_with_list(self):
|
|
encoder = utils.InfEncoder()
|
|
data = [float("inf"), 1.0]
|
|
result = encoder.clear_inf(data)
|
|
self.assertEqual(result, ["inf", 1.0])
|
|
|
|
def test_clear_inf_with_other_types(self):
|
|
encoder = utils.InfEncoder()
|
|
self.assertEqual(encoder.clear_inf("test"), "test")
|
|
self.assertEqual(encoder.clear_inf(123), 123)
|
|
self.assertEqual(encoder.clear_inf(None), None)
|
|
|
|
|
|
class TestWriteToJson(unittest.TestCase):
|
|
def test_write_to_json(self):
|
|
test_data = [{"key": "value"}, {"key2": 123}]
|
|
|
|
with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
|
|
temp_file_path = temp_file.name
|
|
|
|
try:
|
|
utils.write_to_json(temp_file_path, test_data)
|
|
|
|
with open(temp_file_path, "r") as f:
|
|
loaded_data = json.load(f)
|
|
|
|
self.assertEqual(loaded_data, test_data)
|
|
finally:
|
|
if os.path.exists(temp_file_path):
|
|
os.remove(temp_file_path)
|
|
|
|
def test_write_to_json_with_inf(self):
|
|
test_data = [{"key": float("inf")}]
|
|
|
|
with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
|
|
temp_file_path = temp_file.name
|
|
|
|
try:
|
|
utils.write_to_json(temp_file_path, test_data)
|
|
|
|
with open(temp_file_path, "r") as f:
|
|
loaded_data = json.load(f)
|
|
|
|
self.assertEqual(loaded_data, [{"key": "inf"}])
|
|
finally:
|
|
if os.path.exists(temp_file_path):
|
|
os.remove(temp_file_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|