From 670c76fe4ef5f9673058a0bbc7640525912fc975 Mon Sep 17 00:00:00 2001 From: CoolBoy Date: Sat, 1 Aug 2026 17:16:03 +0800 Subject: [PATCH] fix: reject invalid GPU submission templates --- modelhub_submmit_api/template_selector.py | 38 ++++++++++++++++++-- tests/test_template_validation.py | 44 +++++++++++++++++++++++ 2 files changed, 79 insertions(+), 3 deletions(-) create mode 100644 tests/test_template_validation.py diff --git a/modelhub_submmit_api/template_selector.py b/modelhub_submmit_api/template_selector.py index e314446e..cc301642 100644 --- a/modelhub_submmit_api/template_selector.py +++ b/modelhub_submmit_api/template_selector.py @@ -52,6 +52,26 @@ GPU_ALIASES = { "pt200x1": "Sunrise_pt-200-x1", } +# The platform's task-add API no longer accepts these historical GPU values, +# even though templates for them are still present in the public template file. +DISABLED_PLATFORM_GPUS = { + "Kunlunxin_r-200-8f", +} + + +def _force_single_gpu_parallelism(config: str) -> str: + config = re.sub( + r"(?m)^(\s*(?:gpu_num|tensor_parallel_size|pipeline_parallel_size)\s*:\s*)\d+", + r"\g<1>1", + config, + ) + return re.sub( + r"(?P(?:-tp|--tensor-parallel-size|--pipeline-parallel-size)(?:\s+|=|,\s*['\"]?))" + r"(?P\d+)(?P['\"]?)", + r"\g1\g", + config, + ) + class TemplateSelector: def __init__(self, template_file: Path | None = None) -> None: @@ -117,7 +137,10 @@ class TemplateSelector: def normalize_gpu(self, gpu: str) -> str: normalized = _normalize_token(gpu) if normalized in GPU_ALIASES: - return GPU_ALIASES[normalized] + target_gpu = GPU_ALIASES[normalized] + if target_gpu in DISABLED_PLATFORM_GPUS: + raise KeyError(f"GPU is no longer accepted by ModelHub: {target_gpu}") + return target_gpu raise KeyError(f"Unsupported GPU alias: {gpu}") def supported_frameworks(self, task_type: str, target_gpu: str) -> set[str]: @@ -131,7 +154,9 @@ class TemplateSelector: gpus = { template.target_gpu for template in self.templates - if template.task_type == task_type and (not auto_only or self.is_auto_template(template)) + if template.task_type == task_type + and template.target_gpu not in DISABLED_PLATFORM_GPUS + and (not auto_only or self.is_auto_template(template)) } return sorted(gpus) @@ -139,7 +164,10 @@ class TemplateSelector: return { template.framework for template in self.templates - if template.task_type == task_type and template.target_gpu == target_gpu and self.is_auto_template(template) + if template.task_type == task_type + and template.target_gpu == target_gpu + and target_gpu not in DISABLED_PLATFORM_GPUS + and self.is_auto_template(template) } def select_template(self, task_type: str, framework: str, target_gpu: str) -> TemplateRecord: @@ -150,6 +178,8 @@ class TemplateSelector: @staticmethod def is_auto_template(template: TemplateRecord) -> bool: + if template.target_gpu in DISABLED_PLATFORM_GPUS: + return False config = template.config_params if "PLACEHOLDER" in config: return False @@ -159,6 +189,8 @@ class TemplateSelector: def render_config(self, template: TemplateRecord, gguf_filename: str | None = None) -> str: config = template.config_params + if template.target_gpu == "Biren_166m": + config = _force_single_gpu_parallelism(config) if template.framework != "llamacpp": return config if not gguf_filename: diff --git a/tests/test_template_validation.py b/tests/test_template_validation.py new file mode 100644 index 00000000..0ebd3076 --- /dev/null +++ b/tests/test_template_validation.py @@ -0,0 +1,44 @@ +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()