Move model types to environment config to be dynamic

Author: XkunW

vec_inf/cli/_vars.py

@@ -4,20 +4,40 @@
 used in the CLI display formatting.
 """
 
+from typing import get_args
+
+from vec_inf.client._slurm_vars import MODEL_TYPES
+
+
+# Extract model type values from the Literal type
+_MODEL_TYPES = get_args(MODEL_TYPES)
+
+# Rich color options (prioritizing current colors, with fallbacks for additional types)
+_RICH_COLORS = [
+    "cyan",
+    "bright_blue",
+    "purple",
+    "bright_magenta",
+    "green",
+    "yellow",
+    "bright_green",
+    "bright_yellow",
+    "red",
+    "bright_red",
+    "blue",
+    "magenta",
+    "bright_cyan",
+    "white",
+    "bright_white",
+]
+
 # Mapping of model types to their display priority (lower numbers shown first)
-MODEL_TYPE_PRIORITY = {
-    "LLM": 0,
-    "VLM": 1,
-    "Text_Embedding": 2,
-    "Reward_Modeling": 3,
-}
+MODEL_TYPE_PRIORITY = {model_type: idx for idx, model_type in enumerate(_MODEL_TYPES)}
 
 # Mapping of model types to their display colors in Rich
 MODEL_TYPE_COLORS = {
-    "LLM": "cyan",
-    "VLM": "bright_blue",
-    "Text_Embedding": "purple",
-    "Reward_Modeling": "bright_magenta",
+    model_type: _RICH_COLORS[idx % len(_RICH_COLORS)]
+    for idx, model_type in enumerate(_MODEL_TYPES)
 }
 
 # Inference engine choice and name mapping

vec_inf/client/_slurm_vars.py

@@ -82,11 +82,13 @@ def create_literal_type(values: list[str], fallback: str = "") -> Any:
     _config["allowed_values"]["resource_type"]
 )
 
-# Extract required arguments, for launching jobs that don't have a default value and
-# their corresponding environment variables
+# Model types available derived from the cached model config
+MODEL_TYPES: TypeAlias = create_literal_type(_config["model_types"])  # type: ignore[valid-type]
+
+# Required arguments for launching jobs and corresponding environment variables
 REQUIRED_ARGS: dict[str, str | None] = _config["required_args"]
 
-# Extract python version, running sglang requires python version
+# Running sglang requires python version
 PYTHON_VERSION: str = _config["python_version"]
 
 # Extract default arguments

vec_inf/client/config.py

@@ -8,13 +8,13 @@
 from typing import Any, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, Field
-from typing_extensions import Literal
 
 from vec_inf.client._slurm_vars import (
     DEFAULT_ARGS,
     MAX_CPUS_PER_TASK,
     MAX_GPUS_PER_NODE,
     MAX_NUM_NODES,
+    MODEL_TYPES,
     PARTITION,
     QOS,
     RESOURCE_TYPE,
@@ -88,9 +88,7 @@ class ModelConfig(BaseModel):
     model_variant: Optional[str] = Field(
         default=None, description="Specific variant/version of the model family"
     )
-    model_type: Literal["LLM", "VLM", "Text_Embedding", "Reward_Modeling"] = Field(
-        ..., description="Type of model architecture"
-    )
+    model_type: MODEL_TYPES = Field(..., description="Type of model architecture")
     gpus_per_node: int = Field(
         ..., gt=0, le=MAX_GPUS_PER_NODE, description="GPUs per node"
     )

vec_inf/client/models.py

@@ -25,7 +25,9 @@
 
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, get_args
+
+from vec_inf.client._slurm_vars import MODEL_TYPES
 
 
 class ModelStatus(str, Enum):
@@ -55,25 +57,23 @@ class ModelStatus(str, Enum):
     UNAVAILABLE = "UNAVAILABLE"
 
 
-class ModelType(str, Enum):
-    """Enum representing the possible model types.
+# Extract model type values from the Literal type
+_MODEL_TYPE_VALUES = get_args(MODEL_TYPES)
+
+
+def _model_type_to_enum_name(model_type: str) -> str:
+    """Convert a model type string to a valid enum attribute name."""
+    # Convert to uppercase and replace hyphens with underscores
+    return model_type.upper().replace("-", "_")
 
-    Attributes
-    ----------
-    LLM : str
-        Large Language Model
-    VLM : str
-        Vision Language Model
-    TEXT_EMBEDDING : str
-        Text Embedding Model
-    REWARD_MODELING : str
-        Reward Modeling Model
-    """
 
-    LLM = "LLM"
-    VLM = "VLM"
-    TEXT_EMBEDDING = "Text_Embedding"
-    REWARD_MODELING = "Reward_Modeling"
+# Create ModelType enum dynamically from MODEL_TYPES
+ModelType = Enum(
+    "ModelType",
+    {_model_type_to_enum_name(mt): mt for mt in _MODEL_TYPE_VALUES},
+    type=str,
+    module=__name__,
+)
 
 
 @dataclass

vec_inf/config/environment.yaml

@@ -24,6 +24,8 @@ required_args:
 
 python_version: "python3.12"
 
+model_types: ["LLM", "VLM", "Text_Embedding", "Reward_Modeling", "OCR"] # Derived from models.yaml
+
 default_args:
   cpus_per_task: "16"
   mem_per_node: "64G"