fix: reject invalid GPU submission templates

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<prefix>(?:-tp|--tensor-parallel-size|--pipeline-parallel-size)(?:\s+|=|,\s*['\"]?))"
+        r"(?P<value>\d+)(?P<suffix>['\"]?)",
+        r"\g<prefix>1\g<suffix>",
+        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:

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()