vec-inf 0.7.3__py3-none-any.whl → 0.8.1__py3-none-any.whl

vec_inf/client/models.py

@@ -25,7 +25,9 @@ ModelInfo : datacitten
 
 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(  # type: ignore[misc]
+    "ModelType",
+    {_model_type_to_enum_name(mt): mt for mt in _MODEL_TYPE_VALUES},
+    type=str,
+    module=__name__,
+)
 
 
 @dataclass
@@ -222,8 +222,12 @@ class LaunchOptions:
         Directory for logs
     model_weights_parent_dir : str, optional
         Parent directory containing model weights
+    engine: str, optional
+        Inference engine to use
     vllm_args : str, optional
-        Additional arguments for vLLM
+        vLLM engine arguments
+    sglang_args : str, optional
+        SGLang engine arguments
     env : str, optional
         Environment variables to be set
     config : str, optional
@@ -250,7 +254,9 @@ class LaunchOptions:
     venv: Optional[str] = None
     log_dir: Optional[str] = None
     model_weights_parent_dir: Optional[str] = None
+    engine: Optional[str] = None
     vllm_args: Optional[str] = None
+    sglang_args: Optional[str] = None
     env: Optional[str] = None
     config: Optional[str] = None
 

vec_inf/config/README.md

@@ -1,6 +1,6 @@
 # Configs
 
 * [`environment.yaml`](environment.yaml): Configuration for the Slurm cluster environment, including image paths, resource availabilities, default value, and etc.
-* [`models.yaml`](models.yaml): Configuration for launching model inference servers, including Slurm parameters as well as `vllm serve` arguments.
+* [`models.yaml`](models.yaml): Configuration for launching model inference servers, including Slurm parameters as well as inference engine arguments.
 
 **NOTE**: These configs acts as last resort fallbacks in the `vec-inf` package, they will be updated to match the latest cached config on the Vector Killarney cluster with each new package version release.

vec_inf/config/environment.yaml

@@ -1,5 +1,8 @@
 paths:
-  image_path: "/model-weights/vec-inf-shared/vector-inference_latest.sif"
+  image_path: "/model-weights/vec-inf-shared/vector-inference_latest.sif" # Maintains backwards compatibility
+  vllm_image_path: "/model-weights/vec-inf-shared/vector-inference-vllm_latest.sif"
+  sglang_image_path: "/model-weights/vec-inf-shared/vector-inference-sglang_latest.sif"
+  cached_model_config_path: "/model-weights/vec-inf-shared/models.yaml"
 
 containerization:
   module_load_cmd: "module load apptainer"
@@ -19,13 +22,17 @@ required_args:
   account: "VEC_INF_ACCOUNT"
   work_dir: "VEC_INF_WORK_DIR"
 
+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"
   time: "08:00:00"
   qos: ""
   partition: ""
-  resource_type: ""
+  resource_type: "l40s"
   exclude: ""
   nodelist: ""
   bind: ""