Add New Models (#22)

* Add New Models

* Add MiniMaxAI/MiniMax-M2

* Change Default Duration of Running of All Models to 2 Hours

Author: AryanAhadinia

src/swiss_ai_model_launch/assets/envs/sglang_kimi.toml

@@ -0,0 +1,36 @@
+image = "/capstor/store/cscs/swissai/infra01/container-images/sglang_kimi_k2.5_cuda13.sqsh"
+
+# "src_path:trg_path" mounts the src_path on the host inside the container at the trg_path.
+mounts = [
+  "/iopsstor/store/cscs/swissai/a09/xyao/bin:/ocfbin",
+  "/capstor",
+  "/iopsstor",
+  "/usr/lib64/libhwloc.so.15:/usr/lib/libhwloc.so.15",
+  "/usr/lib64/libpciaccess.so.0:/usr/lib/libpciaccess.so.0",
+  "/usr/lib64/libxml2.so.2:/usr/lib/libxml2.so.2",
+  "/usr/lib64/libnuma.so.1:/usr/lib/libnuma.so.1",
+]
+
+workdir = "/opt"
+
+[env]
+# NCCL_DEBUG = "INFO"  # uncomment for debugging
+# NCCL_DEBUG_SUBSYS = "INIT,NET"  # uncomment for debugging
+NCCL_NET = "AWS Libfabric"
+NCCL_CROSS_NIC = "1"
+NCCL_NET_GDR_LEVEL = "PHB"
+NCCL_SOCKET_IFNAME = "hsn"
+NCCL_PROTO = "^LL128"
+FI_CXI_COMPAT = "0"
+FI_MR_CACHE_MONITOR = "userfaultfd"
+FI_CXI_RX_MATCH_MODE = "software"
+FI_CXI_DEFAULT_CQ_SIZE = "131072"
+FI_CXI_DEFAULT_TX_SIZE = "32768"
+FI_CXI_DISABLE_HOST_REGISTER = "1"
+OFI_NCCL_DISABLE_DMABUF = "1"
+SGL_ENABLE_JIT_DEEPGEMM = "0"
+
+[annotations]
+com.hooks.aws_ofi_nccl.enabled = "true"
+com.hooks.aws_ofi_nccl.variant = "cuda13"
+com.hooks.cxi.enabled = "true"

src/swiss_ai_model_launch/assets/models.json

@@ -6,7 +6,8 @@
     "environment": null,
     "workers": 1,
     "nodes_per_worker": 1,
-    "time": "00:30:00"
+    "time": "02:00:00",
+    "framework_args": ""
   },
   {
     "vendor": "swiss-ai",
@@ -15,7 +16,8 @@
     "environment": null,
     "workers": 1,
     "nodes_per_worker": 1,
-    "time": "00:30:00"
+    "time": "02:00:00",
+    "framework_args": ""
   },
   {
     "vendor": "swiss-ai",
@@ -24,7 +26,8 @@
     "environment": null,
     "workers": 1,
     "nodes_per_worker": 1,
-    "time": "00:30:00"
+    "time": "02:00:00",
+    "framework_args": ""
   },
   {
     "vendor": "swiss-ai",
@@ -33,24 +36,37 @@
     "environment": null,
     "workers": 1,
     "nodes_per_worker": 1,
-    "time": "00:30:00"
+    "time": "02:00:00",
+    "framework_args": ""
   },
   {
-    "vendor": "zai-org",
-    "model_name": "GLM-4.7-Flash",
+    "vendor": "Qwen",
+    "model_name": "Qwen3-235B-A22B-Instruct-2507",
     "framework": "sglang",
     "environment": null,
     "workers": 1,
-    "nodes_per_worker": 1,
-    "time": "00:30:00"
+    "nodes_per_worker": 2,
+    "time": "02:00:00",
+    "framework_args": "--tp-size 8"
   },
   {
-    "vendor": "zai-org",
-    "model_name": "GLM-4.7-Flash",
-    "framework": "vllm",
+    "vendor": "moonshotai",
+    "model_name": "Kimi-K2.5",
+    "framework": "sglang",
+    "environment": "src/swiss_ai_model_launch/assets/envs/sglang_kimi.toml",
+    "workers": 1,
+    "nodes_per_worker": 4,
+    "time": "02:00:00",
+    "framework_args": "--tp-size 16 --trust-remote-code --tool-call-parser kimi_k2 --reasoning-parser kimi_k2"
+  },
+  {
+    "vendor": "MiniMaxAI",
+    "model_name": "MiniMax-M2",
+    "framework": "sglang",
     "environment": null,
     "workers": 1,
-    "nodes_per_worker": 1,
-    "time": "00:30:00"
+    "nodes_per_worker": 2,
+    "time": "02:00:00",
+    "framework_args": "--tp-size 8 --ep-size 8 --tool-call-parser minimax-m2 --reasoning-parser minimax-append-think --trust-remote-code --mem-fraction-static 0.85"
   }
 ]

src/swiss_ai_model_launch/cli/configuration/models.py

@@ -185,9 +185,11 @@ async def _resolve_options(
         return await cast(Callable[[], Awaitable[OptionsDict]], self.options_factory)()
 
     async def aconfigure(self, get_value: GetValueFn | None = None) -> None:
-        self.value = await self._build_question(
-            await self._resolve_options(get_value)
-        ).ask_async()
+        options = await self._resolve_options(get_value)
+        if len(options) == 1:
+            self.value = next(iter(options))
+        else:
+            self.value = await self._build_question(options).ask_async()
         self._on_answer()
 
 

src/swiss_ai_model_launch/cli/healthcheck/checker.py

@@ -23,6 +23,8 @@ async def check_model_health(served_model_name: str, api_key: str) -> ModelHealt
                 },
                 timeout=_TIMEOUT_SECONDS,
             )
-        return ModelHealth.HEALTHY if response.is_success else ModelHealth.ERROR
+        return (
+            ModelHealth.HEALTHY if response.is_success else ModelHealth.NOT_RESPONDING
+        )
     except (httpx.TransportError, httpx.TimeoutException):
-        return ModelHealth.NOT_RESPONDING
+        return ModelHealth.ERROR

src/swiss_ai_model_launch/cli/main.py

@@ -1,6 +1,5 @@
 import asyncio
 import re
-from collections.abc import Awaitable, Callable
 
 import firecrest as f7t
 
@@ -87,11 +86,18 @@ async def _get_partitions(
     )
 
 
+def _split_vendor_model(combined: str) -> tuple[str, str]:
+    vendor, model_name = combined.split("::", 1)
+    return vendor, model_name
+
+
 async def _get_preconfigured_default(
     get_value_from_context: GetValueFn, preconfigured: list[LaunchRequest], field: str
 ) -> str | None:
-    vendor = get_value_from_context("model_vendor")
-    model_name = get_value_from_context("model_name")
+    combined = get_value_from_context("model_vendor_model")
+    if combined is None:
+        return None
+    vendor, model_name = _split_vendor_model(combined)
     framework = get_value_from_context("framework")
     match = next(
         (
@@ -108,43 +114,36 @@ async def _get_preconfigured_default(
     return str(getattr(match, field))
 
 
-def _make_served_model_name_default(
-    preconfigured: list[LaunchRequest],
-) -> Callable[[GetValueFn], Awaitable[str]]:
-    async def _default(get_value: GetValueFn) -> str:
-        value = await _get_preconfigured_default(
-            get_value, preconfigured, "served_model_name"
-        )
-        if value and value != "None":
-            return value
-        return f"{get_value('model_vendor')}/{get_value('model_name')}-{create_salt(4)}"
-
-    return _default
+async def _get_router_options(get_value: GetValueFn) -> dict[str, tuple[str, str]]:
+    workers = get_value("workers")
+    if workers is not None and int(workers) > 1:
+        return {
+            "yes": ("Yes", "Use router to load balance across workers"),
+            "no": ("No", "Do not use router"),
+        }
+    return {
+        "no": ("No", "Do not use router"),
+    }
 
 
 async def _get_launch_request(launcher: Launcher) -> LaunchRequest:
     preconfigured_launch_requests = await launcher.get_preconfigured_models()
 
-    async def _get_vendors() -> dict[str, tuple[str, str]]:
-        return {
-            lr.vendor: (lr.vendor, lr.vendor) for lr in preconfigured_launch_requests
-        }
-
-    async def _get_models(
-        get_value_from_context: GetValueFn,
-    ) -> dict[str, tuple[str, str]]:
-        vendor = get_value_from_context("model_vendor")
-        return {
-            lr.model_name: (lr.model_name, lr.model_name)
-            for lr in preconfigured_launch_requests
-            if lr.vendor == vendor
-        }
+    async def _get_vendor_models() -> dict[str, tuple[str, str]]:
+        seen: dict[str, tuple[str, str]] = {}
+        for lr in preconfigured_launch_requests:
+            key = f"{lr.vendor}::{lr.model_name}"
+            if key not in seen:
+                seen[key] = (lr.model_name, lr.vendor)
+        return seen
 
     async def _get_frameworks(
         get_value_from_context: GetValueFn,
     ) -> dict[str, tuple[str, str]]:
-        vendor = get_value_from_context("model_vendor")
-        model_name = get_value_from_context("model_name")
+        combined = get_value_from_context("model_vendor_model")
+        if combined is None:
+            return {}
+        vendor, model_name = _split_vendor_model(combined)
         return {
             lr.framework: (lr.framework, lr.framework)
             for lr in preconfigured_launch_requests
@@ -155,14 +154,9 @@ async def _get_frameworks(
         name="launcher_request_configuration",
         chain=[
             OptionsConfiguration(
-                name="model_vendor",
-                prompt="Choose the model vendor.",
-                options_factory=_get_vendors,
-            ),
-            OptionsConfiguration(
-                name="model_name",
+                name="model_vendor_model",
                 prompt="Choose the model to launch.",
-                options_factory=_get_models,
+                options_factory=_get_vendor_models,
             ),
             OptionsConfiguration(
                 name="framework",
@@ -177,13 +171,10 @@ async def _get_frameworks(
                     get_value, preconfigured_launch_requests, "workers"
                 ),
             ),
-            TextConfiguration(
-                name="nodes_per_worker",
-                prompt="Number of nodes to use per worker for running the model.",
-                validator=lambda v: v.isdigit() and int(v) > 0,
-                default_factory=lambda get_value: _get_preconfigured_default(
-                    get_value, preconfigured_launch_requests, "nodes_per_worker"
-                ),
+            OptionsConfiguration(
+                name="use_router",
+                prompt="Use router to load balance across workers.",
+                options_factory=lambda get_value: _get_router_options(get_value),
             ),
             TextConfiguration(
                 name="time",
@@ -195,26 +186,35 @@ async def _get_frameworks(
                     get_value, preconfigured_launch_requests, "time"
                 ),
             ),
-            TextConfiguration(
-                name="served_model_name",
-                prompt="Served model name.",
-                validator=lambda s: len(s) > 0,
-                default_factory=_make_served_model_name_default(
-                    preconfigured_launch_requests
-                ),
-            ),
         ],
     )
     await launch_req_config.aconfigure()
 
+    vendor, model_name = _split_vendor_model(
+        launch_req_config.get_non_none_value("model_vendor_model")
+    )
+    framework = launch_req_config.get_non_none_value("framework")
+    preconfigured = next(
+        (
+            lr
+            for lr in preconfigured_launch_requests
+            if lr.vendor == vendor
+            and lr.model_name == model_name
+            and lr.framework == framework
+        ),
+        None,
+    )
     return LaunchRequest(
-        vendor=launch_req_config.get_non_none_value("model_vendor"),
-        model_name=launch_req_config.get_non_none_value("model_name"),
-        framework=launch_req_config.get_non_none_value("framework"),
+        vendor=vendor,
+        model_name=model_name,
+        framework=framework,
+        environment=preconfigured.environment if preconfigured else None,
         workers=int(launch_req_config.get_non_none_value("workers")),
-        nodes_per_worker=int(launch_req_config.get_non_none_value("nodes_per_worker")),
+        nodes_per_worker=preconfigured.nodes_per_worker if preconfigured else 1,
         time=launch_req_config.get_non_none_value("time"),
-        served_model_name=launch_req_config.get_non_none_value("served_model_name"),
+        served_model_name=f"{vendor}/{model_name}-{create_salt(4)}",
+        framework_args=preconfigured.framework_args if preconfigured else None,
+        use_router=launch_req_config.get_non_none_value("use_router") == "yes",
     )
 
 
@@ -262,3 +262,7 @@ async def _monitor() -> None:
 
 def main() -> None:
     asyncio.run(_main())
+
+
+if __name__ == "__main__":
+    main()

src/swiss_ai_model_launch/launchers/firecrest_launcher.py

@@ -85,6 +85,7 @@ def _get_launch_args_from_request(
                 )
             ),
             telemetry_endpoint=self.telemetry_endpoint,
+            use_router=launch_request.use_router,
         )
 
     def _get_local_env_file_path(self, launch_request: LaunchRequest) -> str:

src/swiss_ai_model_launch/launchers/launch_request.py

@@ -13,3 +13,4 @@ class LaunchRequest(BaseModel):
     time: str
     served_model_name: str | None = None
     framework_args: str | None = None
+    use_router: bool = False