Default tp_size based on slurm number of GPUs

src/mmirage/core/process/processors/llm/config.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass, field
 
 import logging
+import os
 from typing import Dict, Optional, Sequence, Type, Any, List
 from pydantic import BaseModel, create_model
 
@@ -15,6 +16,39 @@
 env = Environment()
 
 
+def _parse_tp_size_from_env() -> int:
+    """Parse tensor parallelism size from SLURM_GPUS_ON_NODE environment variable.
+
+    Defensively parses the environment variable, handling invalid values:
+    - Returns 1 if the variable is None or empty
+    - Strips whitespace before parsing
+    - Returns 1 for non-integer values
+    - Returns 1 for values <= 0
+
+    Returns:
+        Tensor parallelism size (>= 1), defaults to 1 on any parsing error.
+    """
+    env_value = os.environ.get("SLURM_GPUS_ON_NODE")
+    if not env_value:
+        return 1
+
+    try:
+        tp_size = int(env_value.strip())
+        # Ensure tp_size is positive (must be >= 1)
+        if tp_size <= 0:
+            logger.warning(
+                f"Invalid SLURM_GPUS_ON_NODE value '{env_value}' (must be > 0), defaulting tp_size to 1"
+            )
+            return 1
+        return tp_size
+    except ValueError:
+        # ValueError: invalid integer format
+        logger.warning(
+            f"Invalid SLURM_GPUS_ON_NODE value '{env_value}', defaulting tp_size to 1"
+        )
+        return 1
+
+
 @dataclass
 class SGLangServerArgs:
     """Server arguments for SGLang engine.
@@ -27,7 +61,7 @@ class SGLangServerArgs:
     """
 
     model_path: str = "none"
-    tp_size: int = 1
+    tp_size: int = field(default_factory=_parse_tp_size_from_env)
     trust_remote_code: bool = True
     disable_custom_all_reduce: bool = False