Decouple slurm variables from package and move them into a config, default cached config directory path is still set within the package

Author: XkunW

vec_inf/client/_slurm_templates.py

@@ -6,10 +6,11 @@
 
 from typing import TypedDict
 
-from vec_inf.client.slurm_vars import (
+from vec_inf.client._slurm_vars import (
     LD_LIBRARY_PATH,
     SINGULARITY_IMAGE,
     SINGULARITY_LOAD_CMD,
+    SINGULARITY_MODULE_NAME,
     VLLM_NCCL_SO_PATH,
 )
 
@@ -93,14 +94,14 @@ class SlurmScriptTemplate(TypedDict):
     },
     "singularity_setup": [
         SINGULARITY_LOAD_CMD,
-        f"singularity exec {SINGULARITY_IMAGE} ray stop",
+        f"{SINGULARITY_MODULE_NAME} exec {SINGULARITY_IMAGE} ray stop",
     ],
     "imports": "source {src_dir}/find_port.sh",
     "env_vars": [
         f"export LD_LIBRARY_PATH={LD_LIBRARY_PATH}",
         f"export VLLM_NCCL_SO_PATH={VLLM_NCCL_SO_PATH}",
     ],
-    "singularity_command": f"singularity exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {SINGULARITY_IMAGE} \\",
+    "singularity_command": f"{SINGULARITY_MODULE_NAME} exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {SINGULARITY_IMAGE} \\",
     "activate_venv": "source {venv}/bin/activate",
     "server_setup": {
         "single_node": [
@@ -240,7 +241,7 @@ class BatchModelLaunchScriptTemplate(TypedDict):
         '    "$json_path" > temp_{model_name}.json \\',
         '    && mv temp_{model_name}.json "$json_path"\n',
     ],
-    "singularity_command": f"singularity exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {SINGULARITY_IMAGE} \\",
+    "singularity_command": f"{SINGULARITY_MODULE_NAME} exec --nv --bind {{model_weights_path}}{{additional_binds}} --containall {SINGULARITY_IMAGE} \\",
     "launch_cmd": [
         "vllm serve {model_weights_path} \\",
         "    --served-model-name {model_name} \\",

vec_inf/client/_slurm_vars.py

@@ -0,0 +1,75 @@
+"""Slurm cluster configuration variables."""
+
+import os
+import warnings
+from pathlib import Path
+from typing import Any, TypeAlias
+
+import yaml
+from typing_extensions import Literal
+
+
+CACHED_CONFIG_DIR = Path("/model-weights/vec-inf-shared")
+
+
+def load_env_config() -> dict[str, Any]:
+    """Load the environment configuration."""
+
+    def load_yaml_config(path: Path) -> dict[str, Any]:
+        """Load YAML config with error handling."""
+        try:
+            with path.open() as f:
+                return yaml.safe_load(f) or {}
+        except FileNotFoundError as err:
+            raise FileNotFoundError(f"Could not find config: {path}") from err
+        except yaml.YAMLError as err:
+            raise ValueError(f"Error parsing YAML config at {path}: {err}") from err
+
+    cached_config_path = CACHED_CONFIG_DIR / "environment.yaml"
+    default_path = (
+        cached_config_path
+        if cached_config_path.exists()
+        else Path(__file__).resolve().parent.parent / "config" / "environment.yaml"
+    )
+    config = load_yaml_config(default_path)
+
+    user_path = os.getenv("VEC_INF_CONFIG_DIR")
+    if user_path:
+        user_path_obj = Path(user_path, "environment.yaml")
+        if user_path_obj.exists():
+            user_config = load_yaml_config(user_path_obj)
+            config.update(user_config)
+        else:
+            warnings.warn(
+                f"WARNING: Could not find user config directory: {user_path}, revert to default config located at {default_path}",
+                UserWarning,
+                stacklevel=2,
+            )
+
+    return config
+
+
+_config = load_env_config()
+
+# Extract path values
+LD_LIBRARY_PATH = _config["paths"]["ld_library_path"]
+SINGULARITY_IMAGE = _config["paths"]["image_path"]
+VLLM_NCCL_SO_PATH = _config["paths"]["vllm_nccl_so_path"]
+
+# Extract containerization info
+SINGULARITY_LOAD_CMD = _config["containerization"]["module_load_cmd"]
+SINGULARITY_MODULE_NAME = _config["containerization"]["module_name"]
+
+# Extract limits
+MAX_GPUS_PER_NODE = _config["limits"]["max_gpus_per_node"]
+MAX_NUM_NODES = _config["limits"]["max_num_nodes"]
+MAX_CPUS_PER_TASK = _config["limits"]["max_cpus_per_task"]
+
+# Create dynamic Literal types
+QOS: TypeAlias = str  # Runtime validation will handle the actual values
+PARTITION: TypeAlias = str  # Runtime validation will handle the actual values
+QOS = Literal[tuple(_config["allowed_values"]["qos"])]
+PARTITION = Literal[tuple(_config["allowed_values"]["partition"])]
+
+# Extract default arguments
+DEFAULT_ARGS = _config["default_args"]

vec_inf/client/_utils.py

@@ -15,9 +15,9 @@
 import yaml
 
 from vec_inf.client._client_vars import MODEL_READY_SIGNATURE
+from vec_inf.client._slurm_vars import CACHED_CONFIG_DIR
 from vec_inf.client.config import ModelConfig
 from vec_inf.client.models import ModelStatus
-from vec_inf.client.slurm_vars import CACHED_CONFIG
 
 
 def run_bash_command(command: str) -> tuple[str, str]:
@@ -217,44 +217,54 @@ def load_yaml_config(path: Path) -> dict[str, Any]:
         except yaml.YAMLError as err:
             raise ValueError(f"Error parsing YAML config at {path}: {err}") from err
 
-    # 1. If config_path is given, use only that
-    if config_path:
-        config = load_yaml_config(Path(config_path))
+    def process_config(config: dict[str, Any]) -> list[ModelConfig]:
+        """Process the config based on the config type."""
         return [
             ModelConfig(model_name=name, **model_data)
             for name, model_data in config.get("models", {}).items()
         ]
+        
+
+    def update_config(
+        config: dict[str, Any], user_config: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Update the config with the user config."""
+        for name, data in user_config.get("models", {}).items():
+            if name in config.get("models", {}):
+                config["models"][name].update(data)
+            else:
+                config.setdefault("models", {})[name] = data
+
+        return config
+
+    # 1. If config_path is given, use only that
+    if config_path:
+        config = load_yaml_config(Path(config_path))
+        return process_config(config)
 
     # 2. Otherwise, load default config
     default_path = (
-        CACHED_CONFIG
-        if CACHED_CONFIG.exists()
+        CACHED_CONFIG_DIR / "models_latest.yaml"
+        if CACHED_CONFIG_DIR.exists()
         else Path(__file__).resolve().parent.parent / "config" / "models.yaml"
     )
     config = load_yaml_config(default_path)
 
     # 3. If user config exists, merge it
-    user_path = os.getenv("VEC_INF_CONFIG")
+    user_path = os.getenv("VEC_INF_CONFIG_DIR")
     if user_path:
-        user_path_obj = Path(user_path)
+        user_path_obj = Path(user_path, "models.yaml")
         if user_path_obj.exists():
             user_config = load_yaml_config(user_path_obj)
-            for name, data in user_config.get("models", {}).items():
-                if name in config.get("models", {}):
-                    config["models"][name].update(data)
-                else:
-                    config.setdefault("models", {})[name] = data
+            config = update_config(config, user_config)
         else:
             warnings.warn(
-                f"WARNING: Could not find user config: {user_path}, revert to default config located at {default_path}",
+                f"WARNING: Could not find user config directory: {user_path}, revert to default config located at {default_path}",
                 UserWarning,
                 stacklevel=2,
             )
 
-    return [
-        ModelConfig(model_name=name, **model_data)
-        for name, model_data in config.get("models", {}).items()
-    ]
+    return process_config(config)
 
 
 def parse_launch_output(output: str) -> tuple[str, dict[str, str]]:

vec_inf/client/config.py

@@ -10,7 +10,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import Literal
 
-from vec_inf.client.slurm_vars import (
+from vec_inf.client._slurm_vars import (
     DEFAULT_ARGS,
     MAX_CPUS_PER_TASK,
     MAX_GPUS_PER_NODE,
@@ -132,7 +132,6 @@ class ModelConfig(BaseModel):
     vllm_args: Optional[dict[str, Any]] = Field(
         default={}, description="vLLM engine arguments"
     )
-
     model_config = ConfigDict(
         extra="forbid", str_strip_whitespace=True, validate_default=True, frozen=True
     )

vec_inf/client/slurm_vars.py

@@ -0,0 +1,27 @@
+paths:
+  ld_library_path: "/scratch/ssd001/pkgs/cudnn-11.7-v8.5.0.96/lib/:/scratch/ssd001/pkgs/cuda-11.7/targets/x86_64-linux/lib/"
+  image_path: "/model-weights/vec-inf-shared/vector-inference_latest.sif"
+  vllm_nccl_so_path: "/vec-inf/nccl/libnccl.so.2.18.1"
+
+containerization:
+  module_load_cmd: "module load singularity-ce/3.8.2"
+  module_name: "singularity"
+
+limits:
+  max_gpus_per_node: 8
+  max_num_nodes: 16
+  max_cpus_per_task: 128
+
+allowed_values:
+  qos: ["normal", "m", "m2", "m3", "m4", "m5", "long", "deadline", "high", "scavenger", "llm", "a100"]
+  partition: ["a40", "a100", "t4v1", "t4v2", "rtx6000"]
+
+default_args:
+  cpus_per_task: 16
+  mem_per_node: "64G"
+  qos: "m2"
+  time: "08:00:00"
+  partition: "a40"
+  data_type: "auto"
+  log_dir: "~/.vec-inf-logs"
+  model_weights_parent_dir: "/model-weights"