45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import re
|
||
|
|
import sys
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api"
|
||
|
|
if str(PACKAGE_DIR) in sys.path:
|
||
|
|
sys.path.remove(str(PACKAGE_DIR))
|
||
|
|
sys.path.insert(0, str(PACKAGE_DIR))
|
||
|
|
|
||
|
|
from template_selector import TemplateSelector # noqa: E402
|
||
|
|
|
||
|
|
|
||
|
|
class PlatformTemplateValidationTests(unittest.TestCase):
|
||
|
|
@classmethod
|
||
|
|
def setUpClass(cls) -> None:
|
||
|
|
cls.selector = TemplateSelector()
|
||
|
|
|
||
|
|
def test_removed_kunlunxin_gpu_is_not_auto_selected(self) -> None:
|
||
|
|
supported = self.selector.supported_target_gpus("text-generation", auto_only=True)
|
||
|
|
self.assertNotIn("Kunlunxin_r-200-8f", supported)
|
||
|
|
self.assertEqual(
|
||
|
|
set(),
|
||
|
|
self.selector.supported_frameworks_for_auto("text-generation", "Kunlunxin_r-200-8f"),
|
||
|
|
)
|
||
|
|
with self.assertRaisesRegex(KeyError, "no longer accepted"):
|
||
|
|
self.selector.normalize_gpu("Kunlunxin_r-200-8f")
|
||
|
|
|
||
|
|
def test_biren_template_is_rendered_with_single_gpu_parallelism(self) -> None:
|
||
|
|
template = self.selector.select_template("text-generation", "vllm", "Biren_166m")
|
||
|
|
config = self.selector.render_config(template)
|
||
|
|
|
||
|
|
gpu_counts = re.findall(r"(?m)^\s*gpu_num:\s*(\d+)", config)
|
||
|
|
tensor_parallel_values = re.findall(r"(?:-tp|--tensor-parallel-size)(?:\s+|=|,\s*['\"]?)(\d+)", config)
|
||
|
|
self.assertEqual(["1", "1"], gpu_counts)
|
||
|
|
self.assertEqual(["1", "1"], tensor_parallel_values)
|
||
|
|
self.assertNotIn("gpu_num: 4", config)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|