vec-inf 0.7.0__py3-none-any.whl → 0.7.2__py3-none-any.whl

vec_inf/cli/_cli.py

@@ -69,6 +69,16 @@ def cli() -> None:
     type=int,
     help="Number of GPUs/node to use, default to suggested resource allocation for model",
 )
+@click.option(
+    "--cpus-per-task",
+    type=int,
+    help="Number of CPU cores per task",
+)
+@click.option(
+    "--mem-per-node",
+    type=str,
+    help="Memory allocation per node in GB format (e.g., '32G')",
+)
 @click.option(
     "--account",
     "-A",
@@ -165,6 +175,10 @@ def launch(
             Number of nodes to use
         - gpus_per_node : int, optional
             Number of GPUs per node
+        - cpus_per_task : int, optional
+            Number of CPU cores per task
+        - mem_per_node : str, optional
+            Memory allocation per node in GB format (e.g., '32G')
         - account : str, optional
             Charge resources used by this job to specified account
         - work_dir : str, optional
@@ -447,7 +461,7 @@ def metrics(slurm_job_id: str) -> None:
                     metrics_formatter.format_metrics()
 
                 live.update(metrics_formatter.table)
-                time.sleep(2)
+                time.sleep(1)
     except click.ClickException as e:
         raise e
     except Exception as e:

vec_inf/cli/_helper.py

@@ -36,6 +36,43 @@ class LaunchResponseFormatter:
         self.model_name = model_name
         self.params = params
 
+    def _add_resource_allocation_details(self, table: Table) -> None:
+        """Add resource allocation details to the table."""
+        optional_fields = [
+            ("account", "Account"),
+            ("work_dir", "Working Directory"),
+            ("resource_type", "Resource Type"),
+            ("partition", "Partition"),
+            ("qos", "QoS"),
+        ]
+        for key, label in optional_fields:
+            if self.params.get(key):
+                table.add_row(label, self.params[key])
+
+    def _add_vllm_config(self, table: Table) -> None:
+        """Add vLLM configuration details to the table."""
+        if self.params.get("vllm_args"):
+            table.add_row("vLLM Arguments:", style="magenta")
+            for arg, value in self.params["vllm_args"].items():
+                table.add_row(f"  {arg}:", str(value))
+
+    def _add_env_vars(self, table: Table) -> None:
+        """Add environment variable configuration details to the table."""
+        if self.params.get("env"):
+            table.add_row("Environment Variables", style="magenta")
+            for arg, value in self.params["env"].items():
+                table.add_row(f"  {arg}:", str(value))
+
+    def _add_bind_paths(self, table: Table) -> None:
+        """Add bind path configuration details to the table."""
+        if self.params.get("bind"):
+            table.add_row("Bind Paths", style="magenta")
+            for path in self.params["bind"].split(","):
+                host = target = path
+                if ":" in path:
+                    host, target = path.split(":")
+                table.add_row(f"  {host}:", target)
+
     def format_table_output(self) -> Table:
         """Format output as rich Table.
 
@@ -59,16 +96,7 @@ class LaunchResponseFormatter:
         table.add_row("Vocabulary Size", self.params["vocab_size"])
 
         # Add resource allocation details
-        if self.params.get("account"):
-            table.add_row("Account", self.params["account"])
-        if self.params.get("work_dir"):
-            table.add_row("Working Directory", self.params["work_dir"])
-        if self.params.get("resource_type"):
-            table.add_row("Resource Type", self.params["resource_type"])
-        if self.params.get("partition"):
-            table.add_row("Partition", self.params["partition"])
-        if self.params.get("qos"):
-            table.add_row("QoS", self.params["qos"])
+        self._add_resource_allocation_details(table)
         table.add_row("Time Limit", self.params["time"])
         table.add_row("Num Nodes", self.params["num_nodes"])
         table.add_row("GPUs/Node", self.params["gpus_per_node"])
@@ -76,21 +104,18 @@ class LaunchResponseFormatter:
         table.add_row("Memory/Node", self.params["mem_per_node"])
 
         # Add job config details
+        if self.params.get("venv"):
+            table.add_row("Virtual Environment", self.params["venv"])
         table.add_row(
             "Model Weights Directory",
             str(Path(self.params["model_weights_parent_dir"], self.model_name)),
         )
         table.add_row("Log Directory", self.params["log_dir"])
 
-        # Add vLLM configuration details
-        table.add_row("vLLM Arguments:", style="magenta")
-        for arg, value in self.params["vllm_args"].items():
-            table.add_row(f"  {arg}:", str(value))
-
-        # Add Environment Variable Configuration Details
-        table.add_row("Environment Variables", style="magenta")
-        for arg, value in self.params["env"].items():
-            table.add_row(f"  {arg}:", str(value))
+        # Add configuration details
+        self._add_vllm_config(table)
+        self._add_env_vars(table)
+        self._add_bind_paths(table)
 
         return table
 

vec_inf/client/_client_vars.py

@@ -71,10 +71,3 @@ VLLM_SHORT_TO_LONG_MAP = {
 
 # Required matching arguments for batch mode
 BATCH_MODE_REQUIRED_MATCHING_ARGS = ["venv", "log_dir"]
-
-# Required arguments for launching jobs that don't have a default value and their
-# corresponding environment variables
-REQUIRED_ARGS = {
-    "account": "VEC_INF_ACCOUNT",
-    "work_dir": "VEC_INF_WORK_DIR",
-}

vec_inf/client/_helper.py

@@ -31,6 +31,7 @@ from vec_inf.client._slurm_script_generator import (
     BatchSlurmScriptGenerator,
     SlurmScriptGenerator,
 )
+from vec_inf.client._slurm_vars import CONTAINER_MODULE_NAME, IMAGE_PATH
 from vec_inf.client.config import ModelConfig
 from vec_inf.client.models import (
     BatchLaunchResponse,
@@ -195,23 +196,14 @@ class ModelLauncher:
                         print(f"WARNING: Could not parse env var: {line}")
         return env_vars
 
-    def _get_launch_params(self) -> dict[str, Any]:
-        """Prepare launch parameters, set log dir, and validate required fields.
-
-        Returns
-        -------
-        dict[str, Any]
-            Dictionary of prepared launch parameters
+    def _apply_cli_overrides(self, params: dict[str, Any]) -> None:
+        """Apply CLI argument overrides to params.
 
-        Raises
-        ------
-        MissingRequiredFieldsError
-            If required fields are missing or tensor parallel size is not specified
-            when using multiple GPUs
+        Parameters
+        ----------
+        params : dict[str, Any]
+            Dictionary of launch parameters to override
         """
-        params = self.model_config.model_dump(exclude_none=True)
-
-        # Override config defaults with CLI arguments
         if self.kwargs.get("vllm_args"):
             vllm_args = self._process_vllm_args(self.kwargs["vllm_args"])
             for key, value in vllm_args.items():
@@ -224,13 +216,29 @@ class ModelLauncher:
                 params["env"][key] = str(value)
             del self.kwargs["env"]
 
+        if self.kwargs.get("bind") and params.get("bind"):
+            params["bind"] = f"{params['bind']},{self.kwargs['bind']}"
+            del self.kwargs["bind"]
+
         for key, value in self.kwargs.items():
             params[key] = value
 
-        # Check for required fields without default vals, will raise an error if missing
-        utils.check_required_fields(params)
+    def _validate_resource_allocation(self, params: dict[str, Any]) -> None:
+        """Validate resource allocation and parallelization settings.
 
-        # Validate resource allocation and parallelization settings
+        Parameters
+        ----------
+        params : dict[str, Any]
+            Dictionary of launch parameters to validate
+
+        Raises
+        ------
+        MissingRequiredFieldsError
+            If tensor parallel size is not specified when using multiple GPUs
+        ValueError
+            If total # of GPUs requested is not a power of two
+            If mismatch between total # of GPUs requested and parallelization settings
+        """
         if (
             int(params["gpus_per_node"]) > 1
             and params["vllm_args"].get("--tensor-parallel-size") is None
@@ -251,19 +259,18 @@ class ModelLauncher:
                 "Mismatch between total number of GPUs requested and parallelization settings"
             )
 
-        # Convert gpus_per_node and resource_type to gres
-        resource_type = params.get("resource_type")
-        if resource_type:
-            params["gres"] = f"gpu:{resource_type}:{params['gpus_per_node']}"
-        else:
-            params["gres"] = f"gpu:{params['gpus_per_node']}"
+    def _setup_log_files(self, params: dict[str, Any]) -> None:
+        """Set up log directory and file paths.
 
-        # Create log directory
+        Parameters
+        ----------
+        params : dict[str, Any]
+            Dictionary of launch parameters to set up log files
+        """
         params["log_dir"] = Path(params["log_dir"], params["model_family"]).expanduser()
         params["log_dir"].mkdir(parents=True, exist_ok=True)
         params["src_dir"] = SRC_DIR
 
-        # Construct slurm log file paths
         params["out_file"] = (
             f"{params['log_dir']}/{self.model_name}.%j/{self.model_name}.%j.out"
         )
@@ -274,6 +281,35 @@ class ModelLauncher:
             f"{params['log_dir']}/{self.model_name}.$SLURM_JOB_ID/{self.model_name}.$SLURM_JOB_ID.json"
         )
 
+    def _get_launch_params(self) -> dict[str, Any]:
+        """Prepare launch parameters, set log dir, and validate required fields.
+
+        Returns
+        -------
+        dict[str, Any]
+            Dictionary of prepared launch parameters
+        """
+        params = self.model_config.model_dump(exclude_none=True)
+
+        # Override config defaults with CLI arguments
+        self._apply_cli_overrides(params)
+
+        # Check for required fields without default vals, will raise an error if missing
+        utils.check_required_fields(params)
+
+        # Validate resource allocation and parallelization settings
+        self._validate_resource_allocation(params)
+
+        # Convert gpus_per_node and resource_type to gres
+        resource_type = params.get("resource_type")
+        if resource_type:
+            params["gres"] = f"gpu:{resource_type}:{params['gpus_per_node']}"
+        else:
+            params["gres"] = f"gpu:{params['gpus_per_node']}"
+
+        # Setup log files
+        self._setup_log_files(params)
+
         # Convert path to string for JSON serialization
         for field in params:
             if field in ["vllm_args", "env"]:
@@ -332,6 +368,10 @@ class ModelLauncher:
             job_log_dir / f"{self.model_name}.{self.slurm_job_id}.sbatch"
         )
 
+        # Replace venv with image path if using container
+        if self.params["venv"] == CONTAINER_MODULE_NAME:
+            self.params["venv"] = IMAGE_PATH
+
         with job_json.open("w") as file:
             json.dump(self.params, file, indent=4)
 

vec_inf/client/_slurm_script_generator.py

@@ -14,6 +14,7 @@ from vec_inf.client._slurm_templates import (
     BATCH_SLURM_SCRIPT_TEMPLATE,
     SLURM_SCRIPT_TEMPLATE,
 )
+from vec_inf.client._slurm_vars import CONTAINER_MODULE_NAME
 
 
 class SlurmScriptGenerator:
@@ -32,24 +33,35 @@ class SlurmScriptGenerator:
     def __init__(self, params: dict[str, Any]):
         self.params = params
         self.is_multinode = int(self.params["num_nodes"]) > 1
-        self.use_container = (
-            self.params["venv"] == "singularity" or self.params["venv"] == "apptainer"
-        )
+        self.use_container = self.params["venv"] == CONTAINER_MODULE_NAME
         self.additional_binds = self.params.get("bind", "")
         if self.additional_binds:
             self.additional_binds = f" --bind {self.additional_binds}"
         self.model_weights_path = str(
             Path(self.params["model_weights_parent_dir"], self.params["model_name"])
         )
+        self.env_str = self._generate_env_str()
+
+    def _generate_env_str(self) -> str:
+        """Generate the environment variables string for the Slurm script.
+
+        Returns
+        -------
+        str
+            Formatted env vars string for container or shell export commands.
+        """
         env_dict: dict[str, str] = self.params.get("env", {})
-        # Create string of environment variables
-        self.env_str = ""
-        for key, val in env_dict.items():
-            if len(self.env_str) == 0:
-                self.env_str = "--env "
-            else:
-                self.env_str += ","
-            self.env_str += key + "=" + val
+
+        if not env_dict:
+            return ""
+
+        if self.use_container:
+            # Format for container: --env KEY1=VAL1,KEY2=VAL2
+            env_pairs = [f"{key}={val}" for key, val in env_dict.items()]
+            return f"--env {','.join(env_pairs)}"
+        # Format for shell: export KEY1=VAL1\nexport KEY2=VAL2
+        export_lines = [f"export {key}={val}" for key, val in env_dict.items()]
+        return "\n".join(export_lines)
 
     def _generate_script_content(self) -> str:
         """Generate the complete Slurm script content.
@@ -95,7 +107,12 @@ class SlurmScriptGenerator:
         server_script = ["\n"]
         if self.use_container:
             server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["container_setup"]))
-        server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["env_vars"]))
+            server_script.append("\n".join(SLURM_SCRIPT_TEMPLATE["container_env_vars"]))
+        else:
+            server_script.append(
+                SLURM_SCRIPT_TEMPLATE["activate_venv"].format(venv=self.params["venv"])
+            )
+            server_script.append(self.env_str)
         server_script.append(
             SLURM_SCRIPT_TEMPLATE["imports"].format(src_dir=self.params["src_dir"])
         )
@@ -112,6 +129,11 @@ class SlurmScriptGenerator:
                         env_str=self.env_str,
                     ),
                 )
+            else:
+                server_setup_str = server_setup_str.replace(
+                    "CONTAINER_PLACEHOLDER",
+                    "\\",
+                )
         else:
             server_setup_str = "\n".join(
                 SLURM_SCRIPT_TEMPLATE["server_setup"]["single_node"]
@@ -145,10 +167,7 @@ class SlurmScriptGenerator:
                     env_str=self.env_str,
                 )
             )
-        else:
-            launcher_script.append(
-                SLURM_SCRIPT_TEMPLATE["activate_venv"].format(venv=self.params["venv"])
-            )
+
         launcher_script.append(
             "\n".join(SLURM_SCRIPT_TEMPLATE["launch_cmd"]).format(
                 model_weights_path=self.model_weights_path,
@@ -194,9 +213,7 @@ class BatchSlurmScriptGenerator:
     def __init__(self, params: dict[str, Any]):
         self.params = params
         self.script_paths: list[Path] = []
-        self.use_container = (
-            self.params["venv"] == "singularity" or self.params["venv"] == "apptainer"
-        )
+        self.use_container = self.params["venv"] == CONTAINER_MODULE_NAME
         for model_name in self.params["models"]:
             self.params["models"][model_name]["additional_binds"] = ""
             if self.params["models"][model_name].get("bind"):

vec_inf/client/_slurm_templates.py

@@ -74,7 +74,7 @@ class SlurmScriptTemplate(TypedDict):
     shebang: ShebangConfig
     container_setup: list[str]
     imports: str
-    env_vars: list[str]
+    container_env_vars: list[str]
     container_command: str
     activate_venv: str
     server_setup: ServerSetupConfig
@@ -96,8 +96,8 @@ SLURM_SCRIPT_TEMPLATE: SlurmScriptTemplate = {
         f"{CONTAINER_MODULE_NAME} exec {IMAGE_PATH} ray stop",
     ],
     "imports": "source {src_dir}/find_port.sh",
-    "env_vars": [
-        f"export {CONTAINER_MODULE_NAME}_BINDPATH=${CONTAINER_MODULE_NAME}_BINDPATH,$(echo /dev/infiniband* | sed -e 's/ /,/g')"
+    "container_env_vars": [
+        f"export {CONTAINER_MODULE_NAME.upper()}_BINDPATH=${CONTAINER_MODULE_NAME.upper()}_BINDPATH,/dev,/tmp"
     ],
     "container_command": f"{CONTAINER_MODULE_NAME} exec --nv {{env_str}} --bind {{model_weights_path}}{{additional_binds}} --containall {IMAGE_PATH} \\",
     "activate_venv": "source {venv}/bin/activate",
@@ -112,6 +112,23 @@ SLURM_SCRIPT_TEMPLATE: SlurmScriptTemplate = {
             "nodes_array=($nodes)",
             "head_node=${{nodes_array[0]}}",
             'head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)',
+            "\n# Check for RDMA devices and set environment variable accordingly",
+            "if ! command -v ibv_devices >/dev/null 2>&1; then",
+            '   echo "ibv_devices not found; forcing TCP. (No RDMA userland on host?)"',
+            "   export NCCL_IB_DISABLE=1",
+            '   export NCCL_ENV_ARG="--env NCCL_IB_DISABLE=1"',
+            "else",
+            "   # Pick GID index based on link layer (IB vs RoCE)",
+            '   if ibv_devinfo 2>/dev/null | grep -q "link_layer:.*Ethernet"; then',
+            "       # RoCEv2 typically needs a nonzero GID index; 3 is common, try 2 if your fabric uses it",
+            "       export NCCL_IB_GID_INDEX={{NCCL_IB_GID_INDEX:-3}}",
+            '       export NCCL_ENV_ARG="--env NCCL_IB_GID_INDEX={{NCCL_IB_GID_INDEX:-3}}"',
+            "   else",
+            "       # Native InfiniBand => GID 0",
+            "       export NCCL_IB_GID_INDEX={{NCCL_IB_GID_INDEX:-0}}",
+            '       export NCCL_ENV_ARG="--env NCCL_IB_GID_INDEX={{NCCL_IB_GID_INDEX:-0}}"',
+            "   fi",
+            "fi",
             "\n# Start Ray head node",
             "head_node_port=$(find_available_port $head_node_ip 8080 65535)",
             "ray_head=$head_node_ip:$head_node_port",

vec_inf/client/_slurm_vars.py

@@ -78,5 +78,9 @@ RESOURCE_TYPE: TypeAlias = create_literal_type(  # type: ignore[valid-type]
     _config["allowed_values"]["resource_type"]
 )
 
+# Extract required arguments, for launching jobs that don't have a default value and
+# their corresponding environment variables
+REQUIRED_ARGS: dict[str, str] = _config["required_args"]
+
 # Extract default arguments
 DEFAULT_ARGS: dict[str, str] = _config["default_args"]

vec_inf/client/_utils.py

@@ -14,9 +14,9 @@ from typing import Any, Optional, Union, cast
 import requests
 import yaml
 
-from vec_inf.client._client_vars import MODEL_READY_SIGNATURE, REQUIRED_ARGS
+from vec_inf.client._client_vars import MODEL_READY_SIGNATURE
 from vec_inf.client._exceptions import MissingRequiredFieldsError
-from vec_inf.client._slurm_vars import CACHED_CONFIG_DIR
+from vec_inf.client._slurm_vars import CACHED_CONFIG_DIR, REQUIRED_ARGS
 from vec_inf.client.config import ModelConfig
 from vec_inf.client.models import ModelStatus
 
@@ -108,15 +108,64 @@ def is_server_running(
     if isinstance(log_content, str):
         return log_content
 
-    status: Union[str, tuple[ModelStatus, str]] = ModelStatus.LAUNCHING
+    # Patterns that indicate fatal errors (not just warnings)
+    fatal_error_patterns = [
+        "traceback",
+        "exception",
+        "fatal error",
+        "critical error",
+        "failed to",
+        "could not",
+        "unable to",
+        "error:",
+    ]
+
+    # Patterns to ignore (non-fatal warnings/info messages)
+    ignore_patterns = [
+        "deprecated",
+        "futurewarning",
+        "userwarning",
+        "deprecationwarning",
+        "slurmstepd: error:",  # SLURM cancellation messages (often after server started)
+    ]
+
+    ready_signature_found = False
+    fatal_error_line = None
 
     for line in log_content:
-        if "error" in line.lower():
-            status = (ModelStatus.FAILED, line.strip("\n"))
+        line_lower = line.lower()
+
+        # Check for ready signature first - if found, server is running
         if MODEL_READY_SIGNATURE in line:
-            status = "RUNNING"
+            ready_signature_found = True
+            # Continue checking to see if there are errors after startup
+
+        # Check for fatal errors (only if we haven't seen ready signature yet)
+        if not ready_signature_found:
+            # Skip lines that match ignore patterns
+            if any(ignore_pattern in line_lower for ignore_pattern in ignore_patterns):
+                continue
 
-    return status
+            # Check for fatal error patterns
+            for pattern in fatal_error_patterns:
+                if pattern in line_lower:
+                    # Additional check: skip if it's part of a warning message
+                    # (warnings often contain "error:" but aren't fatal)
+                    if "warning" in line_lower and "error:" in line_lower:
+                        continue
+                    fatal_error_line = line.strip("\n")
+                    break
+
+    # If we found a fatal error, mark as failed
+    if fatal_error_line:
+        return (ModelStatus.FAILED, fatal_error_line)
+
+    # If ready signature was found and no fatal errors, server is running
+    if ready_signature_found:
+        return "RUNNING"
+
+    # Otherwise, still launching
+    return ModelStatus.LAUNCHING
 
 
 def get_base_url(slurm_job_name: str, slurm_job_id: str, log_dir: str) -> str:

vec_inf/client/api.py

@@ -81,7 +81,7 @@ class VecInfClient:
 
     def __init__(self) -> None:
         """Initialize the Vector Inference client."""
-        pass
+        self._metrics_collectors: dict[str, PerformanceMetricsCollector] = {}
 
     def list_models(self) -> list[ModelInfo]:
         """List all available models.
@@ -218,7 +218,13 @@ class VecInfClient:
             - Performance metrics or error message
             - Timestamp of collection
         """
-        performance_metrics_collector = PerformanceMetricsCollector(slurm_job_id)
+        # Use cached collector to preserve state between calls to compute throughput
+        if slurm_job_id not in self._metrics_collectors:
+            self._metrics_collectors[slurm_job_id] = PerformanceMetricsCollector(
+                slurm_job_id
+            )
+
+        performance_metrics_collector = self._metrics_collectors[slurm_job_id]
 
         metrics: Union[dict[str, float], str]
         if not performance_metrics_collector.metrics_url.startswith("http"):

vec_inf/client/models.py

@@ -194,6 +194,10 @@ class LaunchOptions:
         Number of nodes to allocate
     gpus_per_node : int, optional
         Number of GPUs per node
+    cpus_per_task : int, optional
+        Number of CPUs per task
+    mem_per_node : str, optional
+        Memory per node
     account : str, optional
         Account name for job scheduling
     work_dir : str, optional
@@ -232,6 +236,8 @@ class LaunchOptions:
     resource_type: Optional[str] = None
     num_nodes: Optional[int] = None
     gpus_per_node: Optional[int] = None
+    cpus_per_task: Optional[int] = None
+    mem_per_node: Optional[str] = None
     account: Optional[str] = None
     work_dir: Optional[str] = None
     qos: Optional[str] = None

vec_inf/config/environment.yaml

@@ -15,6 +15,10 @@ allowed_values:
   partition: []
   resource_type: ["l40s", "h100"]
 
+required_args:
+  account: "VEC_INF_ACCOUNT"
+  work_dir: "VEC_INF_WORK_DIR"
+
 default_args:
   cpus_per_task: "16"
   mem_per_node: "64G"