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

vec_inf/README.md

@@ -2,7 +2,7 @@
 
 * `launch`: Specify a model family and other optional parameters to launch an OpenAI compatible inference server.
 * `batch-launch`: Specify a list of models to launch multiple OpenAI compatible inference servers at the same time.
-* `status`: Check the model status by providing its Slurm job ID.
+* `status`: Check the status of all `vec-inf` jobs, or a specific job by providing its job ID.
 * `metrics`: Streams performance metrics to the console.
 * `shutdown`: Shutdown a model by providing its Slurm job ID.
 * `list`: List all available model names, or view the default/cached configuration of a specific model.
@@ -14,6 +14,7 @@ Use `--help` to see all available options
 
 * `launch_model`: Launch an OpenAI compatible inference server.
 * `batch_launch_models`: Launch multiple OpenAI compatible inference servers.
+* `fetch_running_jobs`: Get the running `vec-inf` job IDs.
 * `get_status`: Get the status of a running model.
 * `get_metrics`: Get the performance metrics of a running model.
 * `shutdown_model`: Shutdown a running model.

vec_inf/cli/_cli.py

@@ -30,6 +30,7 @@ from vec_inf.cli._helper import (
     BatchLaunchResponseFormatter,
     LaunchResponseFormatter,
     ListCmdDisplay,
+    ListStatusDisplay,
     MetricsResponseFormatter,
     StatusResponseFormatter,
 )
@@ -69,6 +70,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 +176,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
@@ -299,14 +314,14 @@ def batch_launch(
         raise click.ClickException(f"Batch launch failed: {str(e)}") from e
 
 
-@cli.command("status", help="Check the status of a running model on the cluster.")
-@click.argument("slurm_job_id", type=str, nargs=1)
+@cli.command("status", help="Check the status of running vec-inf jobs on the cluster.")
+@click.argument("slurm_job_id", required=False)
 @click.option(
     "--json-mode",
     is_flag=True,
     help="Output in JSON string",
 )
-def status(slurm_job_id: str, json_mode: bool = False) -> None:
+def status(slurm_job_id: Optional[str] = None, json_mode: bool = False) -> None:
     """Get the status of a running model on the cluster.
 
     Parameters
@@ -324,14 +339,28 @@ def status(slurm_job_id: str, json_mode: bool = False) -> None:
     try:
         # Start the client and get model inference server status
         client = VecInfClient()
-        status_response = client.get_status(slurm_job_id)
+        if not slurm_job_id:
+            slurm_job_ids = client.fetch_running_jobs()
+            if not slurm_job_ids:
+                click.echo("No running jobs found.")
+                return
+        else:
+            slurm_job_ids = [slurm_job_id]
+        responses = []
+        for job_id in slurm_job_ids:
+            responses.append(client.get_status(job_id))
+
         # Display status information
-        status_formatter = StatusResponseFormatter(status_response)
-        if json_mode:
-            status_formatter.output_json()
+        if slurm_job_id:
+            status_formatter = StatusResponseFormatter(responses[0])
+            if json_mode:
+                status_formatter.output_json()
+            else:
+                status_info_table = status_formatter.output_table()
+                CONSOLE.print(status_info_table)
         else:
-            status_info_table = status_formatter.output_table()
-            CONSOLE.print(status_info_table)
+            list_status_display = ListStatusDisplay(slurm_job_ids, responses, json_mode)
+            list_status_display.display_multiple_status_output(CONSOLE)
 
     except click.ClickException as e:
         raise e
@@ -447,7 +476,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
 
@@ -226,6 +251,62 @@ class StatusResponseFormatter:
         return table
 
 
+class ListStatusDisplay:
+    """CLI Helper class for formatting a list of StatusResponse.
+
+    A formatter class that handles the presentation of multiple job statuses
+    in a table format.
+
+    Parameters
+    ----------
+    statuses : list[StatusResponse]
+        List of model status information
+    """
+
+    def __init__(
+        self,
+        job_ids: list[str],
+        statuses: list[StatusResponse],
+        json_mode: bool = False,
+    ):
+        self.job_ids = job_ids
+        self.statuses = statuses
+        self.json_mode = json_mode
+
+        self.table = Table(show_header=True, header_style="bold magenta")
+        self.table.add_column("Job ID")
+        self.table.add_column("Model Name")
+        self.table.add_column("Status", style="blue")
+        self.table.add_column("Base URL")
+
+    def display_multiple_status_output(self, console: Console) -> None:
+        """Format and display all model statuses.
+
+        Formats each model's status and adds it to the table.
+        """
+        if self.json_mode:
+            json_data = [
+                {
+                    "job_id": status.model_name,
+                    "model_name": status.model_name,
+                    "model_status": status.server_status,
+                    "base_url": status.base_url,
+                }
+                for status in self.statuses
+            ]
+            click.echo(json.dumps(json_data, indent=4))
+            return
+
+        for i, status in enumerate(self.statuses):
+            self.table.add_row(
+                self.job_ids[i],
+                status.model_name,
+                status.server_status,
+                status.base_url,
+            )
+        console.print(self.table)
+
+
 class MetricsResponseFormatter:
     """CLI Helper class for formatting MetricsResponse.
 

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.
+    def _apply_cli_overrides(self, params: dict[str, Any]) -> None:
+        """Apply CLI argument overrides to params.
 
-        Returns
-        -------
-        dict[str, Any]
-            Dictionary of prepared launch parameters
-
-        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)
 
@@ -429,16 +469,15 @@ class BatchModelLauncher:
             If required fields are missing or tensor parallel size is not specified
             when using multiple GPUs
         """
-        params: dict[str, Any] = {
-            "models": {},
+        common_params: dict[str, Any] = {
             "slurm_job_name": self.slurm_job_name,
             "src_dir": str(SRC_DIR),
             "account": account,
             "work_dir": work_dir,
         }
 
-        # Check for required fields without default vals, will raise an error if missing
-        utils.check_required_fields(params)
+        params: dict[str, Any] = common_params.copy()
+        params["models"] = {}
 
         for i, (model_name, config) in enumerate(self.model_configs.items()):
             params["models"][model_name] = config.model_dump(exclude_none=True)
@@ -515,6 +554,16 @@ class BatchModelLauncher:
                     raise ValueError(
                         f"Mismatch found for {arg}: {params[arg]} != {params['models'][model_name][arg]}, check your configuration"
                     )
+            # Check for required fields and return environment variable overrides
+            env_overrides = utils.check_required_fields(
+                {**params["models"][model_name], **common_params}
+            )
+
+            for arg, value in env_overrides.items():
+                if arg in common_params:
+                    params[arg] = value
+                else:
+                    params["models"][model_name][arg] = value
 
         return params
 
@@ -678,7 +727,7 @@ class ModelStatusMonitor:
             Basic status information for the job
         """
         try:
-            job_name = self.job_status["JobName"]
+            job_name = self.job_status["JobName"].removesuffix("-vec-inf")
             job_state = self.job_status["JobState"]
         except KeyError:
             job_name = "UNAVAILABLE"

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 = (
+            f",{self.params['bind']}" if self.params.get("bind") else ""
         )
-        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.
@@ -77,6 +89,8 @@ class SlurmScriptGenerator:
         for arg, value in SLURM_JOB_CONFIG_ARGS.items():
             if self.params.get(value):
                 shebang.append(f"#SBATCH --{arg}={self.params[value]}")
+            if value == "model_name":
+                shebang[-1] += "-vec-inf"
         if self.is_multinode:
             shebang += SLURM_SCRIPT_TEMPLATE["shebang"]["multinode"]
         return "\n".join(shebang)
@@ -95,7 +109,17 @@ 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(
+                SLURM_SCRIPT_TEMPLATE["bind_path"].format(
+                    model_weights_path=self.model_weights_path,
+                    additional_binds=self.additional_binds,
+                )
+            )
+        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"])
         )
@@ -108,10 +132,14 @@ class SlurmScriptGenerator:
                     "CONTAINER_PLACEHOLDER",
                     SLURM_SCRIPT_TEMPLATE["container_command"].format(
                         model_weights_path=self.model_weights_path,
-                        additional_binds=self.additional_binds,
                         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"]
@@ -141,14 +169,10 @@ class SlurmScriptGenerator:
             launcher_script.append(
                 SLURM_SCRIPT_TEMPLATE["container_command"].format(
                     model_weights_path=self.model_weights_path,
-                    additional_binds=self.additional_binds,
                     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,15 +218,13 @@ 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"):
-                self.params["models"][model_name]["additional_binds"] = (
-                    f" --bind {self.params['models'][model_name]['bind']}"
-                )
+            self.params["models"][model_name]["additional_binds"] = (
+                f",{self.params['models'][model_name]['bind']}"
+                if self.params["models"][model_name].get("bind")
+                else ""
+            )
             self.params["models"][model_name]["model_weights_path"] = str(
                 Path(
                     self.params["models"][model_name]["model_weights_parent_dir"],
@@ -242,7 +264,12 @@ class BatchSlurmScriptGenerator:
         script_content.append(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["shebang"])
         if self.use_container:
             script_content.append(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["container_setup"])
-        script_content.append("\n".join(BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["env_vars"]))
+        script_content.append(
+            BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["bind_path"].format(
+                model_weights_path=model_params["model_weights_path"],
+                additional_binds=model_params["additional_binds"],
+            )
+        )
         script_content.append(
             "\n".join(
                 BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["server_address_setup"]
@@ -260,7 +287,6 @@ class BatchSlurmScriptGenerator:
             script_content.append(
                 BATCH_MODEL_LAUNCH_SCRIPT_TEMPLATE["container_command"].format(
                     model_weights_path=model_params["model_weights_path"],
-                    additional_binds=model_params["additional_binds"],
                 )
             )
         script_content.append(
@@ -304,6 +330,8 @@ class BatchSlurmScriptGenerator:
                 model_params = self.params["models"][model_name]
                 if model_params.get(value) and value not in ["out_file", "err_file"]:
                     shebang.append(f"#SBATCH --{arg}={model_params[value]}")
+                if value == "model_name":
+                    shebang[-1] += "-vec-inf"
             shebang[-1] += "\n"
             shebang.append(BATCH_SLURM_SCRIPT_TEMPLATE["hetjob"])
         # Remove the last hetjob line